Python: add MiddlewareFailure, a first-class fatal signal for function middleware (#7562)
* feat(core): first-class fatal signal (MiddlewareFailure) for function middleware The function-invocation loop converts every exception raised by function middleware into a tool-error result and keeps looping, so middleware that needs fail-closed semantics (enforcement layers, guardrails) had no loud escape: the agent-hooks feature simulated one by mutating shared run state, raising MiddlewareTermination, and re-raising the real failure two hops away at the run boundary. Introduce MiddlewareFailure (a MiddlewareException sibling of MiddlewareTermination) as the loop's explicit fail-closed escape: - _auto_invoke_function re-raises it (both the direct and the pipeline path) instead of absorbing it into a tool-error result; ordinary exceptions keep the absorb-and-continue contract. - A failing call fails the whole parallel batch: in-flight sibling tool tasks are cancelled and awaited before the failure propagates. - Every existing MiddlewareTermination absorb site (agent/chat pipelines, _execute_single_function_call, harness loop, purview) passes it through untouched by construction, and agent/chat middleware exceptions already propagate, so one exception type gives uniform fail-loud semantics across all three categories. Migrate the agent-hooks feature to the new signal: delete the _RunState.halted back-channel and its three run-boundary re-raise checks, drop the halted arm of the termination special case in the function middleware (the approval-request pass-through moves to the single approval check on the normal path), and fail partial installs loudly. Tool-seam host_error blocks keep surfacing as InterceptionBlocked at the run boundary via the exception cause chain (one deny surface at every seam, pinned by tests). Spec 004 gains the middleware-failure invariants and matrix rows. Closes #7522 Signed-off-by: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> * fix(core): harden tool-seam unwrap and pin review findings Review round follow-ups for the MiddlewareFailure feature: - Only agent-hooks' own tagged tool-seam halts (_ToolSeamBlockFailure) authorize re-raising the chained InterceptionBlocked at the run boundary; a third-party MiddlewareFailure with a crafted InterceptionBlocked cause now propagates as raised instead of laundering an attacker-shaped interception record into the feature's deny surface (regression test added, verified by mutation). - Document that middleware must not catch MiddlewareFailure (docstring and spec 004): swallowing it converts a fail-closed abort back into a running, possibly unguarded loop. - Pin the trailing termination re-raise in the agent-hooks function middleware: an inner short-circuit is bracketed and still propagates, skipping outer middleware post-code (test fails with the re-raise removed). Signed-off-by: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> * fix(core): acyclic tool-seam unwrap chain; document cooperative batch cancellation Address two automated-review findings on the MiddlewareFailure PR, both confirmed empirically: - _reraise_tool_seam_block created a two-object exception-chain cycle (block.__cause__ -> wrapper -> block) by re-raising the chained InterceptionBlocked `from` its transport wrapper. Detach the wrapper's back-links and re-raise bare, recording the wrapper as the block's __context__ — acyclic, both exceptions still visible in tracebacks. Regression test walks the chain and pins finiteness (verified to fail against the cyclic re-raise). - Batch cancellation is cooperative: a synchronous tool body already running in a worker thread (asyncio.to_thread) cannot be interrupted by task cancellation and may complete its side effects after the failure reached the caller; its result is discarded either way and propagation is not delayed behind it. Narrow the stated contract (MiddlewareFailure docstring, loop comment, spec 004) and pin it with a blocking-sync-sibling regression test. Signed-off-by: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> * fix(core): settle dangling calls on service-managed conversations on abort Address maintainer review on the MiddlewareFailure PR: - A MiddlewareFailure escaping a tool batch on a service-managed conversation left the hosted thread ending in unresolved function_call items: _update_continuation_state persists session.service_session_id when the model turn completes (before tool execution), and probe-verified the next run sends only the new user message against that conversation — OpenAI-style continuations reject such a request, so a routine policy abort left the session permanently stuck. Both loops now settle the thread before propagating: one error function_result per dangling call, submitted with tool_choice="none" in a single extra request whose response is discarded; a settlement failure never masks the abort, and runs without a service-managed conversation make no extra request. Pinned by three regression tests (non-streaming, streaming, and the no-conversation no-cost case); spec 004 and the MiddlewareFailure docstring updated. - Make the three tool-bracket escape tuples in the agent-hooks function middleware identical (MiddlewareTermination, MiddlewareFailure, CancelledError): a MiddlewareFailure raised inside the post/error-bracket emit bodies is unreachable today, but the uniform tuples remove the need to reason about why they would differ, and preserve the exact exception (including the private tool-seam tag) if the emitter ever surfaces one. Signed-off-by: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> * fix(core): advance settled continuation; settle approved-replay aborts Address maintainer review on the MiddlewareFailure settlement path, both probe-verified (branch rebased onto current main first): - Advance the persisted continuation to the settlement response. For response-ID continuations (OpenAI Responses store=True, where the response id is the continuation handle) the settlement response is the first endpoint whose chain includes the synthetic tool outputs; leaving session.service_session_id on the pre-settlement response made the settlement ineffective — the next run would continue from the still-unresolved turn. The settlement response now runs through _update_function_invocation_continuation_state (a no-op for stable conversation-object ids). Pinned by a regression test that fails with the advance removed. - Cover the approval-resolution phase: a MiddlewareFailure raised while an approved tool is replayed escapes loudly (probe-verified, already the case) but executed before the loops' settlement seams, leaving the original — already service-persisted — call unresolved. _resolve_approval_responses now takes a settle_dangling_calls callback invoked with the approved batch on abort; the settlement helper became a layer method taking explicit calls (approval-response wrappers unwrap to their underlying calls, hosted-tool approvals are left to their provider protocol) and carries its own best-effort containment. Pinned by deny-during- replay regression tests in both response modes, mutation-verified. Spec 004 invariants and matrix rows updated accordingly. Signed-off-by: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> --------- Signed-off-by: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com>
This commit is contained in:
committed by
GitHub
parent
2213ef8493
commit
58da0cc253
@@ -18,7 +18,7 @@ It covers:
|
||||
- approved, rejected, mixed, and replayed approval rounds;
|
||||
- reasoning content and opaque reasoning signatures bound to function calls;
|
||||
- history persistence and service-side continuation;
|
||||
- error, user-input, middleware-termination, and loop-limit paths;
|
||||
- error, user-input, middleware-termination, middleware-failure, and loop-limit paths;
|
||||
- provider and transport serialization of function calls and results.
|
||||
|
||||
The primary implementation is in `python/packages/core/agent_framework/_tools.py`. History replay behavior in
|
||||
@@ -319,7 +319,24 @@ that manually replay messages own the equivalent rule: do not resend an approval
|
||||
### Function calls and results
|
||||
|
||||
- Every actionable local `function_call` produces exactly one terminal `function_result`, unless execution pauses
|
||||
for a new user-input request.
|
||||
for a new user-input request or the run is aborted by `MiddlewareFailure`.
|
||||
- An ordinary exception raised by function middleware or a tool body becomes one terminal error `function_result`
|
||||
and the loop continues; `MiddlewareFailure` is the loop's only fail-closed escape: it is never converted into a
|
||||
tool result, the in-flight parallel batch is cancelled, no further tool call starts, no further model turn is
|
||||
consumed, and the exception propagates to the caller (for streaming runs, when the stream is consumed). On a
|
||||
service-managed conversation the loop first settles the aborted batch — one error `function_result` per dangling
|
||||
call (approval-response wrappers unwrap to their underlying calls; hosted-tool approvals are left to their own
|
||||
provider protocol), submitted with `tool_choice="none"` in a single extra request — so the hosted thread is not
|
||||
left ending in unresolved function calls that the service would reject on the session's next request; the
|
||||
persisted continuation then advances to the settlement response (for response-ID continuations the settled
|
||||
endpoint is the new handle; for conversation-object ids the advance is a no-op) and the settlement response is
|
||||
otherwise discarded. Settlement covers the approval-resolution phase too: a fatal abort while an approved tool is
|
||||
replayed settles the original, already-persisted calls. Without a service-managed conversation no extra request
|
||||
is made. Batch
|
||||
cancellation is cooperative: an async sibling stops at its next suspension point, while a synchronous tool body
|
||||
already executing in a worker thread cannot be interrupted and may complete its side effects — its result is
|
||||
discarded either way and never reaches the transcript, the model, or history. Middleware must not catch
|
||||
`MiddlewareFailure` — swallowing it converts a fail-closed abort back into a running, possibly unguarded loop.
|
||||
- Parallel calls retain model order in the returned transcript.
|
||||
- Reused `call_id` values are correlated by logical occurrence, not one global value per id.
|
||||
- A completed function call/result pair is inert on later turns.
|
||||
@@ -506,6 +523,10 @@ that manually replay messages own the equivalent rule: do not resend an approval
|
||||
| Consecutive error cap | Error threshold stops repeated failures, submits collected results, and makes only the required final no-tool model call. | `test_function_invocation_config_max_consecutive_errors`, `test_streaming_function_invocation_config_max_consecutive_errors`, `test_approval_resume_error_limit_forces_final_no_tool_response` |
|
||||
| Unknown call handling | Configured false returns an error result; configured true raises. | `test_function_invocation_config_terminate_on_unknown_calls_false`, `test_function_invocation_config_terminate_on_unknown_calls_true`, streaming equivalents |
|
||||
| Middleware termination | Normal non-approval loop stops without a second model call. | `test_terminate_loop_single_function_call`, `test_terminate_loop_multiple_function_calls_one_terminates`, `test_terminate_loop_streaming_single_function_call` |
|
||||
| Middleware failure (fatal) | `MiddlewareFailure` from function middleware or a tool body propagates to the caller without becoming a tool result; the tool does not execute (pre-invocation) or its result never feeds another model call (post-invocation); the cause chain is preserved; ordinary exceptions still become tool-error results and the loop continues. | `packages/core/tests/core/test_middleware_with_agent.py::TestMiddlewareFailure::test_failure_before_tool_aborts_run`, `test_failure_after_tool_aborts_run_before_next_model_turn`, `test_failure_cause_chain_reaches_caller`, `test_failure_from_tool_escapes_without_middleware`, `test_failure_streaming_reaches_stream_consumer`, `test_ordinary_exception_still_becomes_tool_error` |
|
||||
| Middleware failure batch cancellation | A fatal signal fails the whole parallel batch: in-flight sibling tool invocations are cancelled and awaited before the failure propagates. Cancellation is cooperative — an async sibling stops at its next suspension point; a synchronous tool body already executing in a worker thread cannot be interrupted and may complete its side effects, but its result is discarded and never reaches the transcript, the model, or history, and failure propagation is not delayed behind it. | `TestMiddlewareFailure::test_failure_cancels_concurrent_sibling_tool`, `test_failure_with_sync_sibling_discards_late_result` |
|
||||
| Middleware failure on a service-managed conversation | The continuation state is already persisted when the batch fails, so before propagating, the loop settles the hosted thread: one error `function_result` per dangling call, sent with `tool_choice="none"` in one extra request; the persisted continuation advances to the settlement response (required for response-ID continuations, a no-op for conversation-object ids) and the settlement response is otherwise discarded; a settlement failure never masks the abort. Without a service-managed conversation no extra request is made. | `TestMiddlewareFailure::test_failure_settles_dangling_calls_on_service_conversation`, `test_failure_settles_service_conversation_streaming`, `test_failure_settlement_advances_response_id_continuation`, `test_failure_without_service_conversation_makes_no_settlement_request` |
|
||||
| Middleware failure during approved-tool replay | A fatal abort while the approval-resolution phase replays an approved tool escapes loudly (never absorbed into a rejection result), the tool's original — already service-persisted — call is settled the same way, and the continuation advances; both response modes. | `TestMiddlewareFailure::test_failure_during_approved_replay_settles_and_escapes`, `test_failure_during_approved_replay_streaming` |
|
||||
| Maximum iterations | No orphan calls; a final no-tool response or deterministic fallback is returned. | `test_max_iterations_limit`, `test_max_iterations_no_orphaned_function_calls`, `test_max_iterations_makes_final_toolchoice_none_call`, `test_max_iterations_blank_final_fallback_synthesizes_message`, streaming equivalents |
|
||||
| Maximum function calls | Parallel overshoot is bounded after the batch; every executed result group counts even without a `function_result`; blank final responses get fallback content. | `test_max_function_calls_limits_parallel_invocations`, `test_max_function_calls_single_calls_per_iteration`, `test_user_input_request_multiple_contents_propagate`, `test_approval_resume_user_input_counts_toward_function_call_budget`, `test_max_function_calls_blank_final_fallback_synthesizes_message`, streaming equivalent |
|
||||
| Provider tool content after an active limit | Locally actionable calls and local approval requests returned despite `tool_choice="none"` are removed in both response modes. Provider-executed informational call/result pairs, hosted approval requests, and metadata-only streaming updates remain visible; fallback text never replaces retained transcript content. | `test_function_invocation_limit_drops_unexecutable_tool_content`, `test_streaming_function_invocation_limit_drops_unexecutable_tool_content`, `test_streaming_function_invocation_limit_preserves_metadata_after_tool_content_is_dropped`, `test_function_invocation_limit_preserves_provider_executed_tool_pair`, `test_streaming_function_invocation_limit_preserves_provider_executed_tool_pair`, `test_function_invocation_limit_appends_fallback_after_provider_executed_tool_pair`, `test_streaming_function_invocation_limit_appends_fallback_after_provider_executed_tool_pair`, `test_function_invocation_limit_preserves_hosted_approval_request`, `test_streaming_function_invocation_limit_preserves_hosted_approval_request` |
|
||||
@@ -593,6 +614,7 @@ Before accepting an update, reviewers must confirm:
|
||||
## Related issues
|
||||
|
||||
- #7241 — approval-resolution result streaming
|
||||
- #7522 — first-class fatal signal (`MiddlewareFailure`) for function middleware
|
||||
- #7267 / #7271 and #7304 — replayed calls and reused ids
|
||||
- #7043 — provider-injected approval execution
|
||||
- #6828 — AG-UI `confirm_changes` snapshot correlation
|
||||
|
||||
@@ -270,6 +270,7 @@ AgentFrameworkException # Base for all AF exceptions
|
||||
│ └── ToolExecutionException # Failure during tool execution
|
||||
│
|
||||
├── MiddlewareException # Middleware failures
|
||||
│ ├── MiddlewareFailure # Control-flow: fatal fail-closed abort of the run
|
||||
│ └── MiddlewareTermination # Control-flow: early middleware termination
|
||||
│
|
||||
└── SettingNotFoundError # Required setting not resolved from any source
|
||||
|
||||
@@ -185,6 +185,7 @@ _LAZY_MODULE_EXPORTS: Final[Mapping[str, tuple[str, ...]]] = {
|
||||
"FunctionMiddleware",
|
||||
"FunctionMiddlewareTypes",
|
||||
"MiddlewareBundle",
|
||||
"MiddlewareFailure",
|
||||
"MiddlewareTermination",
|
||||
"MiddlewareType",
|
||||
"MiddlewareTypes",
|
||||
@@ -510,6 +511,7 @@ __all__ = [
|
||||
"MessageInjectionMiddleware",
|
||||
"MiddlewareBundle",
|
||||
"MiddlewareException",
|
||||
"MiddlewareFailure",
|
||||
"MiddlewareTermination",
|
||||
"MiddlewareType",
|
||||
"MiddlewareTypes",
|
||||
|
||||
@@ -144,6 +144,7 @@ from ._middleware import (
|
||||
FunctionMiddleware,
|
||||
FunctionMiddlewareTypes,
|
||||
MiddlewareBundle,
|
||||
MiddlewareFailure,
|
||||
MiddlewareTermination,
|
||||
MiddlewareType,
|
||||
MiddlewareTypes,
|
||||
@@ -474,6 +475,7 @@ __all__ = [
|
||||
"MessageInjectionMiddleware",
|
||||
"MiddlewareBundle",
|
||||
"MiddlewareException",
|
||||
"MiddlewareFailure",
|
||||
"MiddlewareTermination",
|
||||
"MiddlewareType",
|
||||
"MiddlewareTypes",
|
||||
|
||||
@@ -37,7 +37,13 @@ Enforcement semantics (``mode="enforce"``):
|
||||
not executed (or its result is discarded) and a tool-error payload is surfaced to the
|
||||
model so the agent loop can continue, per the spec's block-propagation rules. A
|
||||
``host_error:*`` deny at the tool seam additionally halts the run (the enforcement
|
||||
layer itself failed, so continuing would be unreliable).
|
||||
layer itself failed, so continuing would be unreliable): the run is aborted through
|
||||
the function-invocation loop's fail-closed escape
|
||||
(:class:`~agent_framework.MiddlewareFailure`) and the
|
||||
:class:`agent_hooks.InterceptionBlocked` propagates to the caller, exactly like a
|
||||
run-level deny. Other unexpected failures inside the enforcement layer (projection
|
||||
bug, emitter fault) abort the run the same way and surface as
|
||||
:class:`~agent_framework.MiddlewareFailure`.
|
||||
- Framework middleware short-circuits (``MiddlewareTermination``) are guarded: a result
|
||||
substituted by another middleware still passes ``output`` / ``post_model_call`` /
|
||||
``post_tool_call`` before it egresses or enters the transcript.
|
||||
@@ -113,6 +119,7 @@ from ._middleware import (
|
||||
FunctionInvocationContext,
|
||||
FunctionMiddleware,
|
||||
MiddlewareBundle,
|
||||
MiddlewareFailure,
|
||||
MiddlewareTermination,
|
||||
)
|
||||
from ._serialization import make_json_safe
|
||||
@@ -231,7 +238,6 @@ class _RunState:
|
||||
builder: AgentContextBuilder
|
||||
session_scoped: bool
|
||||
config: _AgentHooksConfig
|
||||
halted: BaseException | None = None
|
||||
|
||||
|
||||
_RUN_STATE: ContextVar[_RunState | None] = ContextVar("agent_framework_agent_hooks_run_state", default=None)
|
||||
@@ -895,27 +901,63 @@ def _is_approval_request(result: Any) -> bool:
|
||||
return isinstance(result, Content) and result.type == "function_approval_request"
|
||||
|
||||
|
||||
def _halt_on_enforcement_failure(
|
||||
state: _RunState, context: FunctionInvocationContext, exc: BaseException, point: str
|
||||
) -> NoReturn:
|
||||
"""Route an unexpected failure inside the enforcement layer through the fail-closed halt path.
|
||||
def _halt_on_enforcement_failure(exc: BaseException, point: str) -> NoReturn:
|
||||
"""Abort the run fail-closed on an unexpected failure inside the enforcement layer.
|
||||
|
||||
The function-invocation loop converts arbitrary exceptions raised by function
|
||||
middleware into tool-error results and keeps running; for a failure of the
|
||||
enforcement layer itself (projection bug, emitter fault) that would fail open —
|
||||
the failure would vanish from the audit trail and the run would continue unguarded.
|
||||
Instead the loop is stopped via ``MiddlewareTermination`` (its only loud escape) and
|
||||
the agent middleware re-raises the failure to the caller at the run boundary.
|
||||
:class:`MiddlewareFailure` is the loop's explicit fail-closed escape: it propagates
|
||||
to the caller at the run boundary, with the underlying failure chained as its cause.
|
||||
"""
|
||||
message = f"agent-hooks {point} enforcement failed: {type(exc).__name__}"
|
||||
context.result = {"error": message}
|
||||
if isinstance(exc, MiddlewareException):
|
||||
failure: BaseException = exc
|
||||
else:
|
||||
failure = MiddlewareException(message)
|
||||
failure.__cause__ = exc
|
||||
state.halted = failure
|
||||
raise MiddlewareTermination(message) from exc
|
||||
raise MiddlewareFailure(str(exc)) from exc
|
||||
raise MiddlewareFailure(f"agent-hooks {point} enforcement failed: {type(exc).__name__}") from exc
|
||||
|
||||
|
||||
class _ToolSeamBlockFailure(MiddlewareFailure):
|
||||
"""Private tag for this feature's own tool-seam ``host_error`` halts.
|
||||
|
||||
Only failures raised by :meth:`_AgentHooksFunctionMiddleware._maybe_halt` carry
|
||||
this type, which is what authorizes :func:`_reraise_tool_seam_block` to unwrap
|
||||
the chained :class:`agent_hooks.InterceptionBlocked` at the run boundary. A plain
|
||||
``MiddlewareFailure`` raised by third-party middleware is never unwrapped — even
|
||||
when its ``__cause__`` happens to be an ``InterceptionBlocked`` — so untrusted
|
||||
code cannot launder a crafted interception record into this feature's deny
|
||||
surface.
|
||||
"""
|
||||
|
||||
|
||||
def _reraise_tool_seam_block(failure: MiddlewareFailure) -> NoReturn:
|
||||
"""Surface a tool-seam ``host_error`` block as the block itself at the run boundary.
|
||||
|
||||
The tool seam sits behind the function-invocation loop, so its fail-closed halts
|
||||
travel as :class:`MiddlewareFailure` (the loop's loud escape) with the triggering
|
||||
exception chained as the cause. When the failure is this feature's own tagged
|
||||
:class:`_ToolSeamBlockFailure`, re-raise its :class:`agent_hooks.InterceptionBlocked`
|
||||
cause so callers see the same deny surface — ``InterceptionBlocked`` carrying the
|
||||
interception record — at every seam (run-, model-, and tool-level). Any other
|
||||
failure (including a third-party ``MiddlewareFailure`` with a crafted
|
||||
``InterceptionBlocked`` cause) propagates exactly as raised.
|
||||
"""
|
||||
from agent_hooks import InterceptionBlocked
|
||||
|
||||
block = failure.__cause__
|
||||
if isinstance(failure, _ToolSeamBlockFailure) and isinstance(block, InterceptionBlocked):
|
||||
# Detach the transport wrapper's back-links (both were set to the block when
|
||||
# _maybe_halt raised the wrapper from it) before re-raising: re-raising the
|
||||
# block while they are intact would make the two exceptions each other's
|
||||
# cause/context — a chain cycle that every consumer walking
|
||||
# __cause__/__context__ would have to guard against. The bare raise below
|
||||
# records the wrapper as the block's __context__ instead (truthful: the block
|
||||
# is re-raised while the wrapper is being handled), keeping both exceptions
|
||||
# visible in tracebacks, acyclically.
|
||||
failure.__cause__ = None
|
||||
if failure.__context__ is block:
|
||||
failure.__context__ = None
|
||||
raise block
|
||||
raise failure
|
||||
|
||||
|
||||
# endregion
|
||||
@@ -1025,8 +1067,6 @@ class _AgentHooksAgentMiddleware(_AgentHooksMiddlewareBase, AgentMiddleware):
|
||||
# A middleware short-circuited the run; any substituted result
|
||||
# still egresses to the caller, so it passes the output point.
|
||||
termination = exc
|
||||
if state.halted is not None:
|
||||
raise state.halted
|
||||
result = context.result
|
||||
if isinstance(result, AgentResponse):
|
||||
await self._emit_output(state, result)
|
||||
@@ -1046,6 +1086,14 @@ class _AgentHooksAgentMiddleware(_AgentHooksMiddlewareBase, AgentMiddleware):
|
||||
shutdown_reason = "error"
|
||||
context.result = None
|
||||
raise
|
||||
except MiddlewareFailure as failure:
|
||||
# A fail-closed abort from behind the function-invocation loop: the
|
||||
# gated persistence is dropped and, for a tool-seam host_error block,
|
||||
# the block itself is surfaced (one deny surface at every seam).
|
||||
gate.drop()
|
||||
shutdown_reason = "error"
|
||||
context.result = None
|
||||
_reraise_tool_seam_block(failure)
|
||||
except asyncio.CancelledError:
|
||||
shutdown_reason = "cancelled"
|
||||
raise
|
||||
@@ -1093,10 +1141,6 @@ class _AgentHooksAgentMiddleware(_AgentHooksMiddlewareBase, AgentMiddleware):
|
||||
termination = exc
|
||||
inner = context.result
|
||||
if inner is None and termination is not None:
|
||||
if state.halted is not None:
|
||||
# The enforcement layer itself failed: strand the deferred
|
||||
# persistence (fail-closed) and surface the halt.
|
||||
raise state.halted
|
||||
# Terminated without a result: nothing will egress, so the no-egress
|
||||
# termination is a permitted outcome. Release the persistence the
|
||||
# drained in-pipeline work deferred — history of model calls that
|
||||
@@ -1119,6 +1163,12 @@ class _AgentHooksAgentMiddleware(_AgentHooksMiddlewareBase, AgentMiddleware):
|
||||
# A middleware substituted its own stream: it is guarded, and the
|
||||
# termination still short-circuits the rest of the pipeline.
|
||||
raise termination
|
||||
except MiddlewareFailure as failure:
|
||||
# A fail-closed abort during the pipeline descent (e.g. a retry middleware
|
||||
# drained an attempt whose tool seam hit a host_error block): surface the
|
||||
# block itself for tool-seam blocks (one deny surface at every seam). The
|
||||
# deferred persistence is stranded un-flushed — fail-closed.
|
||||
_reraise_tool_seam_block(failure)
|
||||
except asyncio.CancelledError:
|
||||
shutdown_reason = "cancelled"
|
||||
raise
|
||||
@@ -1151,6 +1201,14 @@ class _AgentHooksAgentMiddleware(_AgentHooksMiddlewareBase, AgentMiddleware):
|
||||
except asyncio.CancelledError:
|
||||
await self._emit_shutdown(state, "cancelled")
|
||||
raise
|
||||
except MiddlewareFailure as failure:
|
||||
# A fail-closed abort from behind the function-invocation loop: close
|
||||
# the trail and, for a tool-seam host_error block, surface the block
|
||||
# itself (one deny surface at every seam). The deferred persistence
|
||||
# is dropped — denied content never becomes durable.
|
||||
gate_handle.drop()
|
||||
await self._emit_shutdown(state, "error")
|
||||
_reraise_tool_seam_block(failure)
|
||||
except BaseException:
|
||||
await self._emit_shutdown(state, "error")
|
||||
raise
|
||||
@@ -1164,8 +1222,6 @@ class _AgentHooksAgentMiddleware(_AgentHooksMiddlewareBase, AgentMiddleware):
|
||||
from agent_hooks import InterceptionBlocked
|
||||
|
||||
try:
|
||||
if state.halted is not None:
|
||||
raise state.halted
|
||||
if not isinstance(final, AgentResponse):
|
||||
raise MiddlewareException(
|
||||
f"agent-hooks cannot guard a streamed run result of type {type(final).__name__}; "
|
||||
@@ -1323,21 +1379,20 @@ class _AgentHooksChatMiddleware(_AgentHooksMiddlewareBase, ChatMiddleware):
|
||||
class _AgentHooksFunctionMiddleware(_AgentHooksMiddlewareBase, FunctionMiddleware):
|
||||
"""Tool bracket: ``pre_tool_call`` and ``post_tool_call``."""
|
||||
|
||||
def _block(
|
||||
self, state: _RunState, context: FunctionInvocationContext, exc: InterceptionBlocked, point: str
|
||||
) -> None:
|
||||
def _block(self, context: FunctionInvocationContext, exc: InterceptionBlocked, point: str) -> None:
|
||||
"""Enforce a tool-seam deny: surface a tool error and, on host errors, halt the run."""
|
||||
context.result = _blocked_tool_result(point, exc.result)
|
||||
self._maybe_halt(state, exc, point)
|
||||
self._maybe_halt(exc, point)
|
||||
|
||||
def _maybe_halt(self, state: _RunState, exc: InterceptionBlocked, point: str) -> None:
|
||||
def _maybe_halt(self, exc: InterceptionBlocked, point: str) -> None:
|
||||
record: InterceptionRecord = exc.result
|
||||
if _is_host_error(record):
|
||||
# The enforcement layer itself failed (interceptor crash/timeout, invalid
|
||||
# context): continuing the loop would run unguarded. Halt the run; the
|
||||
# agent middleware re-raises the block to the caller.
|
||||
state.halted = exc
|
||||
raise MiddlewareTermination(
|
||||
# context): continuing the loop would run unguarded. Abort the run through
|
||||
# the loop's fail-closed escape; the agent middleware re-raises the block
|
||||
# itself to the caller (the private tag is what authorizes the unwrap —
|
||||
# see _reraise_tool_seam_block).
|
||||
raise _ToolSeamBlockFailure(
|
||||
f"agent-hooks {point} failed closed: {record.verdict.reason}",
|
||||
) from exc
|
||||
|
||||
@@ -1346,21 +1401,15 @@ class _AgentHooksFunctionMiddleware(_AgentHooksMiddlewareBase, FunctionMiddlewar
|
||||
|
||||
state = _RUN_STATE.get()
|
||||
if state is None:
|
||||
# No run state means the bundle's agent middleware never ran. A plain
|
||||
# exception raised here would be converted into a tool error by the
|
||||
# function-invocation loop and the run would continue unguarded (fail
|
||||
# open); MiddlewareTermination is the only loud escape: the tool is never
|
||||
# dispatched and the loop stops.
|
||||
message = _TRIO_REQUIRED_MESSAGE.format(seam="function")
|
||||
context.result = {"error": message}
|
||||
raise MiddlewareTermination(message)
|
||||
# No run state means the bundle's agent middleware never ran (a partial
|
||||
# install the public API makes impossible). The tool must never be
|
||||
# dispatched: abort the run through the loop's fail-closed escape.
|
||||
raise MiddlewareFailure(_TRIO_REQUIRED_MESSAGE.format(seam="function"))
|
||||
if not self._shares_config(state.config):
|
||||
# A different bundle owns the innermost run state (stacked bundles):
|
||||
# binding to it would silently misroute emissions. Halt the run
|
||||
# fail-closed (the loop swallows plain exceptions, so route through the
|
||||
# halt path).
|
||||
# binding to it would silently misroute emissions. Abort fail-closed.
|
||||
_halt_on_enforcement_failure(
|
||||
state, context, MiddlewareException(_FOREIGN_TRIO_MESSAGE.format(seam="function")), "pre_tool_call"
|
||||
MiddlewareException(_FOREIGN_TRIO_MESSAGE.format(seam="function")), "pre_tool_call"
|
||||
)
|
||||
|
||||
try:
|
||||
@@ -1379,23 +1428,20 @@ class _AgentHooksFunctionMiddleware(_AgentHooksMiddlewareBase, FunctionMiddlewar
|
||||
args = effective
|
||||
except InterceptionBlocked as exc:
|
||||
# §6.2: the tool is not dispatched and no post_tool_call is emitted.
|
||||
self._block(state, context, exc, "pre_tool_call")
|
||||
self._block(context, exc, "pre_tool_call")
|
||||
return
|
||||
except (MiddlewareTermination, asyncio.CancelledError):
|
||||
except (MiddlewareTermination, MiddlewareFailure, asyncio.CancelledError):
|
||||
raise
|
||||
except BaseException as exc:
|
||||
_halt_on_enforcement_failure(state, context, exc, "pre_tool_call")
|
||||
_halt_on_enforcement_failure(exc, "pre_tool_call")
|
||||
|
||||
termination: MiddlewareTermination | None = None
|
||||
try:
|
||||
await call_next()
|
||||
except MiddlewareTermination as exc:
|
||||
if state.halted is not None or _is_approval_request(context.result):
|
||||
# Our own halt path, or framework approval control flow (the tool did
|
||||
# not run; an approved replay re-enters through pre_tool_call).
|
||||
raise
|
||||
# A middleware short-circuited with a substituted result; that result
|
||||
# still enters the transcript, so it is bracketed below.
|
||||
# still enters the transcript, so it is bracketed below (approval-request
|
||||
# control flow is passed through un-emitted further down instead).
|
||||
termination = exc
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
@@ -1411,19 +1457,21 @@ class _AgentHooksFunctionMiddleware(_AgentHooksMiddlewareBase, FunctionMiddlewar
|
||||
except InterceptionBlocked as blocked:
|
||||
# A policy deny over an already-errored call changes nothing (the
|
||||
# result is discarded either way); a host error still halts the run.
|
||||
self._maybe_halt(state, blocked, "post_tool_call")
|
||||
except asyncio.CancelledError:
|
||||
self._maybe_halt(blocked, "post_tool_call")
|
||||
except (MiddlewareTermination, MiddlewareFailure, asyncio.CancelledError):
|
||||
raise
|
||||
except BaseException as emit_exc:
|
||||
_halt_on_enforcement_failure(state, context, emit_exc, "post_tool_call")
|
||||
_halt_on_enforcement_failure(emit_exc, "post_tool_call")
|
||||
raise
|
||||
|
||||
if _is_approval_request(context.result):
|
||||
# Framework approval control flow on the normal return path (a middleware
|
||||
# set an approval request and returned): the tool has not run, so there is
|
||||
# no result to bracket — pass the control object through un-emitted, exactly
|
||||
# like the termination branch above. The approved replay re-enters through
|
||||
# pre_tool_call.
|
||||
# Framework approval control flow (a middleware set an approval request,
|
||||
# whether it returned or short-circuited): the tool has not run, so there
|
||||
# is no result to bracket — pass the control object through un-emitted.
|
||||
# The approved replay re-enters through pre_tool_call. A short-circuit is
|
||||
# re-raised so the loop still stops and surfaces the request.
|
||||
if termination is not None:
|
||||
raise termination
|
||||
return
|
||||
|
||||
try:
|
||||
@@ -1434,11 +1482,11 @@ class _AgentHooksFunctionMiddleware(_AgentHooksMiddlewareBase, FunctionMiddlewar
|
||||
context.result = _ToolResultCodec.write_back(context.result, value, outcome.target)
|
||||
except InterceptionBlocked as exc:
|
||||
# §6.1: the result must be discarded as if the call had errored.
|
||||
self._block(state, context, exc, "post_tool_call")
|
||||
except (MiddlewareTermination, asyncio.CancelledError):
|
||||
self._block(context, exc, "post_tool_call")
|
||||
except (MiddlewareTermination, MiddlewareFailure, asyncio.CancelledError):
|
||||
raise
|
||||
except BaseException as exc:
|
||||
_halt_on_enforcement_failure(state, context, exc, "post_tool_call")
|
||||
_halt_on_enforcement_failure(exc, "post_tool_call")
|
||||
if termination is not None:
|
||||
raise termination
|
||||
|
||||
|
||||
@@ -82,6 +82,64 @@ class MiddlewareTermination(MiddlewareException):
|
||||
self.result = result
|
||||
|
||||
|
||||
class MiddlewareFailure(MiddlewareException):
|
||||
"""Fatal middleware signal that aborts the run instead of being absorbed.
|
||||
|
||||
Ordinary exceptions raised by **function** middleware (or by the tool it wraps) are
|
||||
converted into tool-error results by the function-invocation loop, which then keeps
|
||||
running — appropriate for recoverable tool failures, but fail-open for enforcement
|
||||
layers and guardrails. ``MiddlewareFailure`` is the loop's explicit fail-closed
|
||||
escape: it is never converted into a tool result, the current batch of concurrent
|
||||
tool calls is cancelled, no further tool call starts, and the exception propagates
|
||||
to the caller of :meth:`Agent.run` (for streaming runs, it is raised when the
|
||||
stream is consumed). Cancellation is cooperative: an async sibling stops at its
|
||||
next suspension point, while a synchronous tool body already executing in a worker
|
||||
thread cannot be interrupted and may still complete its side effects — its result
|
||||
is discarded either way and never reaches the transcript, the model, or history.
|
||||
On a service-managed conversation (a persisted conversation id), the loop first
|
||||
settles the aborted batch by submitting one error ``function_result`` per dangling
|
||||
call — one extra request — so the hosted thread is not left with unresolved tool
|
||||
calls that would make the session's next request fail; the persisted continuation
|
||||
advances to the settlement response (the new handle for response-ID continuations)
|
||||
and the settlement response is otherwise discarded. This also covers a failure
|
||||
raised while an approved tool is replayed after an approval pause.
|
||||
|
||||
Agent and chat middleware do not need a dedicated signal — every exception they
|
||||
raise already propagates to the caller — and ``MiddlewareFailure`` behaves the same
|
||||
there, so one exception type gives uniform fail-loud semantics across all three
|
||||
middleware categories.
|
||||
|
||||
Contrast with :class:`MiddlewareTermination`, which stops the loop *gracefully*
|
||||
(optionally substituting a result that still flows back to the caller):
|
||||
``MiddlewareFailure`` produces no result at all.
|
||||
|
||||
Middleware must not catch ``MiddlewareFailure`` (let it propagate through
|
||||
``call_next()``): swallowing it converts a fail-closed abort back into a running —
|
||||
and possibly unguarded — loop.
|
||||
|
||||
Chain the underlying error so it reaches the caller intact:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework import FunctionMiddleware, FunctionInvocationContext, MiddlewareFailure
|
||||
|
||||
|
||||
class EnforcementMiddleware(FunctionMiddleware):
|
||||
async def process(self, context: FunctionInvocationContext, call_next):
|
||||
try:
|
||||
verdict = await self.check(context.arguments)
|
||||
except Exception as exc:
|
||||
# The enforcement layer itself failed: abort instead of running unguarded.
|
||||
raise MiddlewareFailure("policy check failed") from exc
|
||||
if verdict.deny:
|
||||
raise MiddlewareFailure(f"denied: {verdict.reason}")
|
||||
await call_next()
|
||||
"""
|
||||
|
||||
def __init__(self, message: str = "Middleware failed.") -> None:
|
||||
super().__init__(message, log_level=None)
|
||||
|
||||
|
||||
class MiddlewareType(str, Enum):
|
||||
"""Enum representing the type of middleware.
|
||||
|
||||
@@ -544,6 +602,13 @@ class FunctionMiddleware(ABC):
|
||||
FunctionMiddleware is an abstract base class. You must subclass it and implement
|
||||
the ``process()`` method to create custom function middleware.
|
||||
|
||||
Note:
|
||||
Exception semantics inside the function-invocation loop: an ordinary exception
|
||||
raised from function middleware is converted into a tool-error result and the
|
||||
loop keeps running; raise :class:`MiddlewareTermination` to stop the loop
|
||||
gracefully (optionally substituting a result), or :class:`MiddlewareFailure` to
|
||||
abort the run fail-closed and propagate the failure to the caller.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
|
||||
@@ -1459,6 +1459,10 @@ async def _auto_invoke_function(
|
||||
Raises:
|
||||
KeyError: If the requested function is not found in the tool map.
|
||||
MiddlewareTermination: If middleware requests loop termination.
|
||||
MiddlewareFailure: If middleware (or the tool) aborts the run fail-closed.
|
||||
Unlike ordinary exceptions, which are converted into tool-error results,
|
||||
this explicit signal is re-raised so it propagates to the run's caller.
|
||||
UserInputRequiredException: If the tool requires user input to proceed.
|
||||
"""
|
||||
from ._types import Content
|
||||
|
||||
@@ -1533,7 +1537,7 @@ async def _auto_invoke_function(
|
||||
additional_properties=function_call_content.additional_properties,
|
||||
)
|
||||
|
||||
from ._middleware import FunctionInvocationContext
|
||||
from ._middleware import FunctionInvocationContext, MiddlewareFailure
|
||||
|
||||
if middleware_pipeline is None or not middleware_pipeline.has_middlewares:
|
||||
# No middleware - execute directly
|
||||
@@ -1557,7 +1561,9 @@ async def _auto_invoke_function(
|
||||
result=function_result,
|
||||
additional_properties=function_call_content.additional_properties,
|
||||
)
|
||||
except UserInputRequiredException:
|
||||
except (MiddlewareFailure, UserInputRequiredException):
|
||||
# Explicit control-flow signals escape the loop; only ordinary exceptions
|
||||
# are absorbed into tool-error results below.
|
||||
raise
|
||||
except Exception as exc:
|
||||
return _function_execution_error_result(function_call_content, tool.name, exc, config)
|
||||
@@ -1621,7 +1627,10 @@ async def _auto_invoke_function(
|
||||
additional_properties=function_call_content.additional_properties,
|
||||
)
|
||||
raise
|
||||
except UserInputRequiredException:
|
||||
except (MiddlewareFailure, UserInputRequiredException):
|
||||
# MiddlewareFailure is the loop's explicit fail-closed escape: middleware that
|
||||
# must abort the run (enforcement layers, guardrails) raises it instead of
|
||||
# relying on the tool-error conversion below, and it propagates to the caller.
|
||||
raise
|
||||
except Exception as exc:
|
||||
return _function_execution_error_result(function_call_content, tool.name, exc, config)
|
||||
@@ -1846,7 +1855,20 @@ async def _try_execute_function_call_groups(
|
||||
)
|
||||
for function_call in function_calls
|
||||
]
|
||||
execution_results = await asyncio.gather(*execution_tasks)
|
||||
try:
|
||||
execution_results = await asyncio.gather(*execution_tasks)
|
||||
except BaseException:
|
||||
# A loud escape from one call (e.g. MiddlewareFailure aborting the run
|
||||
# fail-closed) fails the whole batch: cancel in-flight siblings and wait for
|
||||
# them so no new tool work starts after the loop is abandoned. Cancellation
|
||||
# is cooperative — a synchronous tool body already running in a worker thread
|
||||
# (asyncio.to_thread) cannot be interrupted and may complete its side effects,
|
||||
# but its result is discarded with the batch and never reaches the transcript,
|
||||
# the model, or history.
|
||||
for task in execution_tasks:
|
||||
task.cancel()
|
||||
await asyncio.gather(*execution_tasks, return_exceptions=True)
|
||||
raise
|
||||
|
||||
should_terminate = any(terminate for _, terminate in execution_results)
|
||||
return [result_contents for result_contents, _ in execution_results], should_terminate
|
||||
@@ -2851,8 +2873,16 @@ async def _resolve_approval_responses(
|
||||
max_errors: int,
|
||||
execute_function_calls: _FunctionCallExecutor,
|
||||
invocation_session: AgentSession | None = None,
|
||||
settle_dangling_calls: Callable[[Sequence[Content]], Awaitable[None]] | None = None,
|
||||
) -> _FunctionProcessingResult:
|
||||
"""Resolve inbound approval responses before the next model call."""
|
||||
"""Resolve inbound approval responses before the next model call.
|
||||
|
||||
``settle_dangling_calls``, when provided, is invoked with the approved batch if
|
||||
executing it aborts with ``MiddlewareFailure``, so a service-managed conversation
|
||||
can be settled before the abort propagates (the replay's original calls belong to
|
||||
an earlier, already-persisted model turn).
|
||||
"""
|
||||
from ._middleware import MiddlewareFailure
|
||||
from ._types import Message
|
||||
|
||||
_bind_approval_responses_to_pending_requests(prepared_messages, invocation_session)
|
||||
@@ -2884,10 +2914,18 @@ async def _resolve_approval_responses(
|
||||
should_terminate = False
|
||||
reached_error_limit = False
|
||||
if responses_to_execute:
|
||||
execution = await execute_function_calls(
|
||||
function_calls=responses_to_execute,
|
||||
options=options,
|
||||
)
|
||||
try:
|
||||
execution = await execute_function_calls(
|
||||
function_calls=responses_to_execute,
|
||||
options=options,
|
||||
)
|
||||
except MiddlewareFailure:
|
||||
# Fail-closed abort during an approved replay: the original calls belong
|
||||
# to an already-persisted model turn, so settle them on a service-managed
|
||||
# conversation before propagating (best-effort inside the callback).
|
||||
if settle_dangling_calls is not None:
|
||||
await settle_dangling_calls(responses_to_execute)
|
||||
raise
|
||||
execution_result_groups = execution.result_groups
|
||||
should_terminate = execution.should_terminate
|
||||
errors_in_a_row, reached_error_limit = _update_consecutive_error_count(
|
||||
@@ -3036,6 +3074,83 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
|
||||
):
|
||||
session.service_session_id = conversation_id
|
||||
|
||||
async def _settle_dangling_service_function_calls(
|
||||
self,
|
||||
*,
|
||||
super_get_response: Callable[..., Any],
|
||||
function_calls: Sequence[Content],
|
||||
options: dict[str, Any],
|
||||
request_kwargs: dict[str, Any],
|
||||
compaction_strategy: CompactionStrategy | None,
|
||||
tokenizer: TokenizerProtocol | None,
|
||||
invocation_session: AgentSession | None,
|
||||
response_conversation_id: str | None = None,
|
||||
) -> None:
|
||||
"""Resolve an aborted batch's function calls on a service-managed conversation.
|
||||
|
||||
When ``MiddlewareFailure`` aborts a tool batch, the local run raises before
|
||||
any result exists — but on a service-managed conversation the continuation
|
||||
state (``session.service_session_id``) was already persisted, so the hosted
|
||||
thread ends with unresolved ``function_call`` items and OpenAI-style
|
||||
continuations reject the session's next request (missing tool output). Settle
|
||||
the thread by submitting one error ``function_result`` per dangling call
|
||||
(approval-response wrappers are unwrapped to their underlying calls;
|
||||
hosted-tool approvals are left to their own provider protocol) with
|
||||
``tool_choice="none"`` so no new calls are requested, then advance the
|
||||
persisted continuation to the settlement response: for response-ID
|
||||
continuations the settlement response is the first endpoint whose chain
|
||||
includes the synthetic outputs, so the next run must start from it (for
|
||||
conversation-object ids the advance is a no-op). The settlement response is
|
||||
otherwise discarded and the run still fails with the original
|
||||
``MiddlewareFailure``. Everything here is best-effort — a settlement failure
|
||||
is logged and never masks the abort. Costs one extra request, only on the
|
||||
failure path and only when a service-managed conversation is in play.
|
||||
"""
|
||||
from ._types import ChatResponse, Content, Message
|
||||
|
||||
if response_conversation_id is None and not options.get("conversation_id"):
|
||||
return
|
||||
try:
|
||||
error_results: list[Content] = []
|
||||
for function_call in function_calls:
|
||||
if _is_hosted_tool_approval(function_call):
|
||||
continue
|
||||
underlying_call = _underlying_function_call(function_call)
|
||||
if underlying_call.type != "function_call" or underlying_call.call_id is None:
|
||||
continue
|
||||
error_results.append(
|
||||
Content.from_function_result(
|
||||
call_id=underlying_call.call_id,
|
||||
result="Error: Tool execution was aborted by middleware before a result was produced.",
|
||||
exception="MiddlewareFailure",
|
||||
additional_properties=underlying_call.additional_properties,
|
||||
)
|
||||
)
|
||||
if not error_results:
|
||||
return
|
||||
options["tool_choice"] = "none"
|
||||
settlement_response = await super_get_response(
|
||||
messages=[Message(role="tool", contents=error_results)],
|
||||
stream=False,
|
||||
options=options,
|
||||
compaction_strategy=compaction_strategy,
|
||||
tokenizer=tokenizer,
|
||||
client_kwargs=request_kwargs,
|
||||
)
|
||||
if isinstance(settlement_response, ChatResponse):
|
||||
self._update_function_invocation_continuation_state(
|
||||
request_kwargs,
|
||||
cast("ChatResponse[Any]", settlement_response),
|
||||
session=invocation_session,
|
||||
options=options,
|
||||
)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Failed to settle dangling function calls on the service-managed conversation; "
|
||||
"the next request over this conversation may be rejected by the service.",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
def _get_function_middleware_pipeline(
|
||||
self,
|
||||
runtime_middleware: Sequence[FunctionMiddlewareTypes],
|
||||
@@ -3066,6 +3181,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
|
||||
max_errors: int,
|
||||
) -> ChatResponse[Any]:
|
||||
"""Run the non-streaming function invocation loop."""
|
||||
from ._middleware import MiddlewareFailure
|
||||
from ._types import ChatResponse, add_usage_details
|
||||
|
||||
errors_in_a_row = 0
|
||||
@@ -3078,6 +3194,17 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
|
||||
max_iterations = self.function_invocation_configuration.get("max_iterations", DEFAULT_MAX_ITERATIONS)
|
||||
attempt_start = int(budget_state.get("attempt_count", 0) or 0)
|
||||
|
||||
async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> None:
|
||||
await self._settle_dangling_service_function_calls(
|
||||
super_get_response=super_get_response,
|
||||
function_calls=function_calls,
|
||||
options=options,
|
||||
request_kwargs=request_kwargs,
|
||||
compaction_strategy=compaction_strategy,
|
||||
tokenizer=tokenizer,
|
||||
invocation_session=invocation_session,
|
||||
)
|
||||
|
||||
# Phase 1: resolve inbound approvals before consuming another model iteration.
|
||||
approval_processing = await _resolve_approval_responses(
|
||||
prepared_messages=prepared_messages,
|
||||
@@ -3086,6 +3213,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
|
||||
max_errors=max_errors,
|
||||
execute_function_calls=execute_function_calls,
|
||||
invocation_session=invocation_session,
|
||||
settle_dangling_calls=settle_approval_replay_calls,
|
||||
)
|
||||
function_call_messages.extend(approval_processing.response_messages)
|
||||
errors_in_a_row = approval_processing.errors_in_a_row
|
||||
@@ -3129,15 +3257,32 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
|
||||
options=options,
|
||||
)
|
||||
|
||||
function_processing = await _process_model_function_calls(
|
||||
response=response,
|
||||
options=options,
|
||||
function_call_messages=function_call_messages,
|
||||
errors_in_a_row=errors_in_a_row,
|
||||
max_errors=max_errors,
|
||||
execute_function_calls=execute_function_calls,
|
||||
invocation_session=invocation_session,
|
||||
)
|
||||
try:
|
||||
function_processing = await _process_model_function_calls(
|
||||
response=response,
|
||||
options=options,
|
||||
function_call_messages=function_call_messages,
|
||||
errors_in_a_row=errors_in_a_row,
|
||||
max_errors=max_errors,
|
||||
execute_function_calls=execute_function_calls,
|
||||
invocation_session=invocation_session,
|
||||
)
|
||||
except MiddlewareFailure:
|
||||
# Fail-closed abort: before propagating, settle the batch's calls on a
|
||||
# service-managed conversation and advance the persisted continuation
|
||||
# to the settled endpoint (best-effort — a settlement failure never
|
||||
# masks the abort).
|
||||
await self._settle_dangling_service_function_calls(
|
||||
super_get_response=super_get_response,
|
||||
function_calls=_extract_function_calls(response),
|
||||
options=options,
|
||||
request_kwargs=request_kwargs,
|
||||
compaction_strategy=compaction_strategy,
|
||||
tokenizer=tokenizer,
|
||||
invocation_session=invocation_session,
|
||||
response_conversation_id=response.conversation_id,
|
||||
)
|
||||
raise
|
||||
total_function_calls = _record_function_calls(
|
||||
budget_state,
|
||||
total_function_calls,
|
||||
@@ -3199,6 +3344,8 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
|
||||
max_errors: int,
|
||||
) -> AsyncIterable[ChatResponseUpdate]:
|
||||
"""Run the streaming function invocation loop."""
|
||||
from ._middleware import MiddlewareFailure
|
||||
|
||||
errors_in_a_row = 0
|
||||
total_function_calls = int(budget_state.get("total_function_calls", 0) or 0)
|
||||
max_function_calls = self.function_invocation_configuration.get("max_function_calls")
|
||||
@@ -3207,6 +3354,17 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
|
||||
max_iterations = self.function_invocation_configuration.get("max_iterations", DEFAULT_MAX_ITERATIONS)
|
||||
attempt_start = int(budget_state.get("attempt_count", 0) or 0)
|
||||
|
||||
async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> None:
|
||||
await self._settle_dangling_service_function_calls(
|
||||
super_get_response=super_get_response,
|
||||
function_calls=function_calls,
|
||||
options=options,
|
||||
request_kwargs=request_kwargs,
|
||||
compaction_strategy=compaction_strategy,
|
||||
tokenizer=tokenizer,
|
||||
invocation_session=invocation_session,
|
||||
)
|
||||
|
||||
# Phase 1: resolve and emit inbound approval outcomes before opening another provider stream.
|
||||
approval_processing = await _resolve_approval_responses(
|
||||
prepared_messages=prepared_messages,
|
||||
@@ -3215,6 +3373,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
|
||||
max_errors=max_errors,
|
||||
execute_function_calls=execute_function_calls,
|
||||
invocation_session=invocation_session,
|
||||
settle_dangling_calls=settle_approval_replay_calls,
|
||||
)
|
||||
errors_in_a_row = approval_processing.errors_in_a_row
|
||||
total_function_calls = _record_function_calls(
|
||||
@@ -3280,15 +3439,31 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
|
||||
yield _function_invocation_limit_fallback_update()
|
||||
return
|
||||
|
||||
function_processing = await _process_model_function_calls(
|
||||
response=response,
|
||||
options=options,
|
||||
function_call_messages=None,
|
||||
errors_in_a_row=errors_in_a_row,
|
||||
max_errors=max_errors,
|
||||
execute_function_calls=execute_function_calls,
|
||||
invocation_session=invocation_session,
|
||||
)
|
||||
try:
|
||||
function_processing = await _process_model_function_calls(
|
||||
response=response,
|
||||
options=options,
|
||||
function_call_messages=None,
|
||||
errors_in_a_row=errors_in_a_row,
|
||||
max_errors=max_errors,
|
||||
execute_function_calls=execute_function_calls,
|
||||
invocation_session=invocation_session,
|
||||
)
|
||||
except MiddlewareFailure:
|
||||
# See the non-streaming loop: settle a service-managed conversation's
|
||||
# dangling calls and advance the persisted continuation before
|
||||
# propagating the fail-closed abort (best-effort).
|
||||
await self._settle_dangling_service_function_calls(
|
||||
super_get_response=super_get_response,
|
||||
function_calls=_extract_function_calls(response),
|
||||
options=options,
|
||||
request_kwargs=request_kwargs,
|
||||
compaction_strategy=compaction_strategy,
|
||||
tokenizer=tokenizer,
|
||||
invocation_session=invocation_session,
|
||||
response_conversation_id=response.conversation_id,
|
||||
)
|
||||
raise
|
||||
errors_in_a_row = function_processing.errors_in_a_row
|
||||
total_function_calls = _record_function_calls(
|
||||
budget_state,
|
||||
|
||||
@@ -26,6 +26,7 @@ from agent_framework import (
|
||||
Message,
|
||||
MiddlewareBundle,
|
||||
MiddlewareException,
|
||||
MiddlewareFailure,
|
||||
MiddlewareTermination,
|
||||
ResponseStream,
|
||||
create_agent_hooks_middleware,
|
||||
@@ -708,6 +709,184 @@ async def test_interceptor_crash_fails_closed_and_halts_run(chat_client_base: Mo
|
||||
assert points(records)[-1] == "agent_shutdown"
|
||||
|
||||
|
||||
@requires_sdk
|
||||
async def test_interceptor_crash_at_tool_seam_fails_closed_streaming(chat_client_base: MockBaseChatClient) -> None:
|
||||
# The streaming twin of the halt above: the tool-seam host_error block travels the
|
||||
# function-invocation loop as MiddlewareFailure and is surfaced to the stream
|
||||
# consumer as the InterceptionBlocked itself — one deny surface at every seam.
|
||||
records: list[InterceptionRecord] = []
|
||||
chat_client_base.streaming_responses = [
|
||||
[
|
||||
ChatResponseUpdate(
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
call_id="call_1", name="weather_tool", arguments='{"location": "Seattle"}'
|
||||
)
|
||||
],
|
||||
role="assistant",
|
||||
finish_reason="tool_calls",
|
||||
)
|
||||
],
|
||||
]
|
||||
agent = Agent(
|
||||
client=chat_client_base,
|
||||
tools=[weather_tool],
|
||||
middleware=[create_agent_hooks_middleware([CrashingGuard("post_tool_call")], record_sink=records.append)],
|
||||
)
|
||||
|
||||
updates: list[AgentResponseUpdate] = []
|
||||
with pytest.raises(InterceptionBlocked) as exc_info:
|
||||
async for update in agent.run("get the weather", stream=True):
|
||||
updates.append(update)
|
||||
|
||||
assert exc_info.value.result.verdict.reason == "host_error:interceptor_failed"
|
||||
assert updates == [] # nothing egressed before the halt
|
||||
assert points(records)[-1] == "agent_shutdown"
|
||||
|
||||
|
||||
@requires_sdk
|
||||
async def test_tool_seam_block_exception_chain_is_acyclic(chat_client_base: MockBaseChatClient) -> None:
|
||||
# Re-raising the InterceptionBlocked with the transport wrapper's back-links
|
||||
# intact would make the two exceptions each other's cause/context — a chain
|
||||
# cycle every __cause__/__context__ walker would have to guard against. The
|
||||
# unwrap must detach the wrapper first, keeping both exceptions visible in a
|
||||
# finite traceback.
|
||||
import traceback
|
||||
|
||||
chat_client_base.run_responses = [tool_call_response(), final_response()]
|
||||
agent = Agent(
|
||||
client=chat_client_base,
|
||||
tools=[weather_tool],
|
||||
middleware=[create_agent_hooks_middleware([CrashingGuard("post_tool_call")])],
|
||||
)
|
||||
|
||||
with pytest.raises(InterceptionBlocked) as exc_info:
|
||||
await agent.run("get the weather")
|
||||
|
||||
block = exc_info.value
|
||||
seen: set[int] = set()
|
||||
node: BaseException | None = block
|
||||
while node is not None:
|
||||
assert id(node) not in seen, "exception chain contains a cycle"
|
||||
seen.add(id(node))
|
||||
# Follow the chain the way traceback rendering does.
|
||||
node = node.__cause__ if (node.__cause__ is not None or node.__suppress_context__) else node.__context__
|
||||
|
||||
formatted = "".join(traceback.format_exception(type(block), block, block.__traceback__))
|
||||
assert "InterceptionBlocked" in formatted
|
||||
# The loop-transport wrapper stays visible as context, acyclically.
|
||||
assert "failed closed" in formatted
|
||||
|
||||
|
||||
@requires_sdk
|
||||
async def test_third_party_middleware_failure_is_bracketed_and_propagates(
|
||||
chat_client_base: MockBaseChatClient,
|
||||
) -> None:
|
||||
# A MiddlewareFailure raised by another (inner) function middleware is not
|
||||
# agent-hooks' own halt: the tool bracket still closes with is_error=True (only
|
||||
# the exception type name crosses the boundary) and the failure itself propagates
|
||||
# to the caller un-unwrapped.
|
||||
class InnerEnforcement(FunctionMiddleware):
|
||||
async def process(self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
raise MiddlewareFailure("inner enforcement denied")
|
||||
|
||||
guard = AllowGuard()
|
||||
records: list[InterceptionRecord] = []
|
||||
chat_client_base.run_responses = [tool_call_response(), final_response()]
|
||||
agent = Agent(
|
||||
client=chat_client_base,
|
||||
tools=[weather_tool],
|
||||
middleware=[create_agent_hooks_middleware([guard], record_sink=records.append), InnerEnforcement()],
|
||||
)
|
||||
|
||||
with pytest.raises(MiddlewareFailure, match="inner enforcement denied"):
|
||||
await agent.run("get the weather")
|
||||
|
||||
assert weather_tool_calls == []
|
||||
post_tool = guard.contexts_for("post_tool_call")[0]
|
||||
assert post_tool["tool_result"]["is_error"] is True
|
||||
assert post_tool["tool_result"]["value"] == "MiddlewareFailure"
|
||||
assert points(records)[-1] == "agent_shutdown"
|
||||
|
||||
|
||||
@requires_sdk
|
||||
async def test_third_party_crafted_interception_cause_is_not_unwrapped(chat_client_base: MockBaseChatClient) -> None:
|
||||
# Adversarial probe: only this feature's own tagged tool-seam halts authorize
|
||||
# unwrapping the chained InterceptionBlocked at the run boundary. A third-party
|
||||
# MiddlewareFailure whose __cause__ is a crafted InterceptionBlocked must surface
|
||||
# AS the MiddlewareFailure — otherwise untrusted middleware could launder an
|
||||
# attacker-shaped interception record into this feature's audit-bearing deny
|
||||
# surface.
|
||||
from agent_hooks import EnforcementMode, InterceptionPoint
|
||||
|
||||
crafted = InterceptionBlocked(
|
||||
InterceptionRecord(
|
||||
interception_point=InterceptionPoint.PRE_TOOL_CALL,
|
||||
mode=EnforcementMode.ENFORCE,
|
||||
verdict=Verdict.deny(reason="forged_deny"),
|
||||
input_identity=None,
|
||||
enforced_identity=None,
|
||||
)
|
||||
)
|
||||
|
||||
class Laundering(FunctionMiddleware):
|
||||
async def process(self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
raise MiddlewareFailure("third-party failure") from crafted
|
||||
|
||||
chat_client_base.run_responses = [tool_call_response(), final_response()]
|
||||
agent = Agent(
|
||||
client=chat_client_base,
|
||||
tools=[weather_tool],
|
||||
middleware=[create_agent_hooks_middleware([AllowGuard()]), Laundering()],
|
||||
)
|
||||
|
||||
with pytest.raises(MiddlewareFailure, match="third-party failure"):
|
||||
await agent.run("get the weather")
|
||||
|
||||
assert weather_tool_calls == []
|
||||
|
||||
|
||||
@requires_sdk
|
||||
async def test_inner_termination_re_raises_through_after_bracketing(chat_client_base: MockBaseChatClient) -> None:
|
||||
# Pins the trailing `raise termination` in the function middleware: an inner
|
||||
# short-circuit is bracketed (post_tool_call over the substituted result) and then
|
||||
# still propagates as a short-circuit — outer middleware post-call_next code is
|
||||
# skipped and the loop stops without another model call.
|
||||
outer_events: list[str] = []
|
||||
|
||||
class Outer(FunctionMiddleware):
|
||||
async def process(self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
outer_events.append("before")
|
||||
await call_next()
|
||||
outer_events.append("after") # must be skipped by the re-raised termination
|
||||
|
||||
class InnerShortCircuit(FunctionMiddleware):
|
||||
async def process(self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
context.result = "substituted"
|
||||
raise MiddlewareTermination("stop")
|
||||
|
||||
guard = AllowGuard()
|
||||
chat_client_base.run_responses = [tool_call_response(), final_response()]
|
||||
agent = Agent(
|
||||
client=chat_client_base,
|
||||
tools=[weather_tool],
|
||||
middleware=[Outer(), create_agent_hooks_middleware([guard]), InnerShortCircuit()],
|
||||
)
|
||||
|
||||
response = await agent.run("get the weather")
|
||||
|
||||
assert outer_events == ["before"]
|
||||
assert weather_tool_calls == []
|
||||
assert chat_client_base.call_count == 1
|
||||
# The substituted result was bracketed before the short-circuit propagated.
|
||||
post_tool = guard.contexts_for("post_tool_call")[0]
|
||||
assert post_tool["tool_result"]["value"] == "substituted"
|
||||
results = [
|
||||
content for message in response.messages for content in message.contents if content.type == "function_result"
|
||||
]
|
||||
assert len(results) == 1
|
||||
|
||||
|
||||
@requires_sdk
|
||||
async def test_interceptor_crash_at_input_fails_closed(chat_client_base: MockBaseChatClient) -> None:
|
||||
agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware([CrashingGuard("input")])])
|
||||
@@ -1185,7 +1364,8 @@ async def test_enforcement_failure_at_function_seam_halts_run(chat_client_base:
|
||||
async def test_function_seam_without_run_state_blocks_tool(chat_client_base: MockBaseChatClient) -> None:
|
||||
# The bundle makes a partial install impossible through the public API; this
|
||||
# exercises the internal defense directly: the private function middleware invoked
|
||||
# without an active agent-hooks run must never dispatch the tool.
|
||||
# without an active agent-hooks run must never dispatch the tool. The run aborts
|
||||
# loudly through the loop's fail-closed escape (MiddlewareFailure).
|
||||
bundle = create_agent_hooks_middleware([AllowGuard()])
|
||||
chat_client_base.run_responses = [tool_call_response(), final_response()]
|
||||
agent = Agent(
|
||||
@@ -1194,13 +1374,12 @@ async def test_function_seam_without_run_state_blocks_tool(chat_client_base: Moc
|
||||
middleware=[_bundle_member(bundle, "FunctionMiddleware")],
|
||||
)
|
||||
|
||||
response = await agent.run("get the weather")
|
||||
with pytest.raises(MiddlewareFailure, match="without an active agent-hooks run"):
|
||||
await agent.run("get the weather")
|
||||
|
||||
# The tool is never dispatched and the loop terminates instead of continuing.
|
||||
# The tool is never dispatched and the loop stops instead of continuing.
|
||||
assert weather_tool_calls == []
|
||||
assert chat_client_base.call_count == 1
|
||||
transcript = str([content.result for message in response.messages for content in message.contents])
|
||||
assert "without an active agent-hooks run" in transcript
|
||||
|
||||
|
||||
@requires_sdk
|
||||
@@ -1208,7 +1387,7 @@ async def test_bundle_passed_to_chat_client_call_raises_instead_of_dropping_the_
|
||||
chat_client_base: MockBaseChatClient,
|
||||
) -> None:
|
||||
# The chat-client middleware seam installs only chat and function middleware; the
|
||||
# bundle's agent member carries the output gate and the halt re-raise, so silently
|
||||
# bundle's agent member carries the output gate and the deny surface, so silently
|
||||
# dropping it would install partial enforcement. The seam must raise instead.
|
||||
bundle = create_agent_hooks_middleware([AllowGuard()])
|
||||
with pytest.raises(MiddlewareException, match="cannot be partially installed"):
|
||||
@@ -2249,6 +2428,45 @@ async def test_approval_request_on_normal_return_path_passes_through(chat_client
|
||||
assert len(approval_requests) == 1
|
||||
|
||||
|
||||
@requires_sdk
|
||||
async def test_approval_request_on_termination_path_passes_through(chat_client_base: MockBaseChatClient) -> None:
|
||||
class ApprovalGate(FunctionMiddleware):
|
||||
"""Framework pattern: request human approval and short-circuit the pipeline."""
|
||||
|
||||
async def process(self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
context.result = Content.from_function_approval_request(
|
||||
id=str(context.metadata.get("call_id")),
|
||||
function_call=Content.from_function_call(
|
||||
str(context.metadata.get("call_id")), context.function.name, arguments={}
|
||||
),
|
||||
)
|
||||
raise MiddlewareTermination("needs approval")
|
||||
|
||||
records: list[InterceptionRecord] = []
|
||||
chat_client_base.run_responses = [tool_call_response(), final_response()]
|
||||
agent = Agent(
|
||||
client=chat_client_base,
|
||||
tools=[weather_tool],
|
||||
middleware=[create_agent_hooks_middleware([AllowGuard()], record_sink=records.append), ApprovalGate()],
|
||||
)
|
||||
|
||||
response = await agent.run("get the weather")
|
||||
|
||||
# The tool never ran and the short-circuit still stopped the loop: the control
|
||||
# object is passed through un-bracketed (no post_tool_call reporting a value for
|
||||
# a tool that never executed) and surfaces to the caller.
|
||||
assert weather_tool_calls == []
|
||||
assert "post_tool_call" not in points(records)
|
||||
assert chat_client_base.call_count == 1
|
||||
approval_requests = [
|
||||
content
|
||||
for message in response.messages
|
||||
for content in message.contents
|
||||
if content.type == "function_approval_request"
|
||||
]
|
||||
assert len(approval_requests) == 1
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
# region Tool-call transforms at post_model_call
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any, cast
|
||||
|
||||
@@ -21,6 +23,7 @@ from agent_framework import (
|
||||
FunctionTool,
|
||||
Message,
|
||||
MiddlewareException,
|
||||
MiddlewareFailure,
|
||||
MiddlewareTermination,
|
||||
MiddlewareType,
|
||||
SupportsChatGetResponse,
|
||||
@@ -2398,3 +2401,641 @@ class TestCallableClassMiddlewareErrorHandling:
|
||||
|
||||
assert "UndeterminedCallableMiddleware" in str(exc_info.value)
|
||||
assert "Cannot determine middleware type" in str(exc_info.value)
|
||||
|
||||
|
||||
# region MiddlewareFailure fail-closed escape
|
||||
|
||||
|
||||
def _tool_call_response(name: str = "sample_tool_function", call_id: str = "call_1") -> ChatResponse:
|
||||
return ChatResponse(
|
||||
messages=[
|
||||
Message(
|
||||
role="assistant",
|
||||
contents=[Content.from_function_call(call_id=call_id, name=name, arguments='{"location": "Seattle"}')],
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
class TestMiddlewareFailure:
|
||||
"""The explicit fail-closed escape from the function-invocation loop.
|
||||
|
||||
Ordinary exceptions raised by function middleware are converted into tool-error
|
||||
results and the loop keeps running (a public behavior, pinned below); only the
|
||||
explicit ``MiddlewareFailure`` signal escapes the loop and propagates to the
|
||||
``run()`` caller.
|
||||
"""
|
||||
|
||||
async def test_failure_before_tool_aborts_run(self, chat_client_base: "MockBaseChatClient") -> None:
|
||||
executed: list[str] = []
|
||||
|
||||
def tool_impl(location: str) -> str:
|
||||
executed.append(location)
|
||||
return "ran"
|
||||
|
||||
tracked_tool = FunctionTool(func=tool_impl, name="sample_tool_function", approval_mode="never_require")
|
||||
|
||||
class Enforcement(FunctionMiddleware):
|
||||
async def process(
|
||||
self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]
|
||||
) -> None:
|
||||
raise MiddlewareFailure("policy denied")
|
||||
|
||||
chat_client_base.run_responses = [
|
||||
_tool_call_response(),
|
||||
ChatResponse(messages=[Message(role="assistant", contents=["Final"])]),
|
||||
]
|
||||
agent = Agent(client=chat_client_base, middleware=[Enforcement()], tools=[tracked_tool])
|
||||
|
||||
with pytest.raises(MiddlewareFailure, match="policy denied"):
|
||||
await agent.run("get the weather")
|
||||
|
||||
# Fail-closed: the tool never executed and no further model turn was consumed.
|
||||
assert executed == []
|
||||
assert chat_client_base.call_count == 1
|
||||
|
||||
async def test_failure_after_tool_aborts_run_before_next_model_turn(
|
||||
self, chat_client_base: "MockBaseChatClient"
|
||||
) -> None:
|
||||
executed: list[str] = []
|
||||
|
||||
def tool_impl(location: str) -> str:
|
||||
executed.append(location)
|
||||
return "ran"
|
||||
|
||||
tracked_tool = FunctionTool(func=tool_impl, name="sample_tool_function", approval_mode="never_require")
|
||||
|
||||
class PostCheck(FunctionMiddleware):
|
||||
async def process(
|
||||
self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]
|
||||
) -> None:
|
||||
await call_next()
|
||||
raise MiddlewareFailure("result rejected")
|
||||
|
||||
chat_client_base.run_responses = [
|
||||
_tool_call_response(),
|
||||
ChatResponse(messages=[Message(role="assistant", contents=["Final"])]),
|
||||
]
|
||||
agent = Agent(client=chat_client_base, middleware=[PostCheck()], tools=[tracked_tool])
|
||||
|
||||
with pytest.raises(MiddlewareFailure, match="result rejected"):
|
||||
await agent.run("get the weather")
|
||||
|
||||
# The tool ran once, but its result never fed another model iteration.
|
||||
assert executed == ["Seattle"]
|
||||
assert chat_client_base.call_count == 1
|
||||
|
||||
async def test_failure_cause_chain_reaches_caller(self, chat_client_base: "MockBaseChatClient") -> None:
|
||||
class Enforcement(FunctionMiddleware):
|
||||
async def process(
|
||||
self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]
|
||||
) -> None:
|
||||
try:
|
||||
raise ValueError("enforcement backend down")
|
||||
except ValueError as exc:
|
||||
raise MiddlewareFailure("enforcement failed") from exc
|
||||
|
||||
chat_client_base.run_responses = [_tool_call_response()]
|
||||
agent = Agent(client=chat_client_base, middleware=[Enforcement()], tools=[sample_tool_function])
|
||||
|
||||
with pytest.raises(MiddlewareFailure, match="enforcement failed") as exc_info:
|
||||
await agent.run("get the weather")
|
||||
|
||||
assert isinstance(exc_info.value.__cause__, ValueError)
|
||||
|
||||
async def test_ordinary_exception_still_becomes_tool_error(self, chat_client_base: "MockBaseChatClient") -> None:
|
||||
"""The loop's absorb-into-tool-error contract for ordinary exceptions is unchanged."""
|
||||
|
||||
class Broken(FunctionMiddleware):
|
||||
async def process(
|
||||
self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]
|
||||
) -> None:
|
||||
raise RuntimeError("middleware bug")
|
||||
|
||||
chat_client_base.run_responses = [
|
||||
_tool_call_response(),
|
||||
ChatResponse(messages=[Message(role="assistant", contents=["Final"])]),
|
||||
]
|
||||
agent = Agent(client=chat_client_base, middleware=[Broken()], tools=[sample_tool_function])
|
||||
|
||||
response = await agent.run("get the weather")
|
||||
|
||||
# The exception was converted into a tool-error result and the loop continued.
|
||||
assert chat_client_base.call_count == 2
|
||||
error_results = [
|
||||
content
|
||||
for message in response.messages
|
||||
for content in message.contents
|
||||
if content.type == "function_result" and content.exception is not None
|
||||
]
|
||||
assert len(error_results) == 1
|
||||
|
||||
async def test_failure_from_tool_escapes_without_middleware(self, chat_client_base: "MockBaseChatClient") -> None:
|
||||
"""The direct (no-middleware) execution path honors the same explicit signal."""
|
||||
|
||||
def tool_impl(location: str) -> str:
|
||||
raise MiddlewareFailure("tool aborted the run")
|
||||
|
||||
failing_tool = FunctionTool(func=tool_impl, name="sample_tool_function", approval_mode="never_require")
|
||||
chat_client_base.run_responses = [_tool_call_response()]
|
||||
agent = Agent(client=chat_client_base, tools=[failing_tool])
|
||||
|
||||
with pytest.raises(MiddlewareFailure, match="tool aborted the run"):
|
||||
await agent.run("get the weather")
|
||||
|
||||
assert chat_client_base.call_count == 1
|
||||
|
||||
async def test_failure_cancels_concurrent_sibling_tool(self, chat_client_base: "MockBaseChatClient") -> None:
|
||||
"""A fatal signal fails the whole batch: in-flight siblings are cancelled."""
|
||||
sibling_started = asyncio.Event()
|
||||
sibling_cancelled: list[bool] = []
|
||||
|
||||
async def slow_impl(location: str) -> str:
|
||||
sibling_started.set()
|
||||
try:
|
||||
await asyncio.Event().wait() # blocks until cancelled
|
||||
except asyncio.CancelledError:
|
||||
sibling_cancelled.append(True)
|
||||
raise
|
||||
return "never"
|
||||
|
||||
slow_tool = FunctionTool(func=slow_impl, name="slow_tool", approval_mode="never_require")
|
||||
|
||||
class FailFast(FunctionMiddleware):
|
||||
async def process(
|
||||
self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]
|
||||
) -> None:
|
||||
if context.function.name == "sample_tool_function":
|
||||
# Fail only after the sibling is genuinely in flight.
|
||||
await sibling_started.wait()
|
||||
raise MiddlewareFailure("abort the batch")
|
||||
await call_next()
|
||||
|
||||
batch_response = ChatResponse(
|
||||
messages=[
|
||||
Message(
|
||||
role="assistant",
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
call_id="call_1", name="sample_tool_function", arguments='{"location": "Seattle"}'
|
||||
),
|
||||
Content.from_function_call(call_id="call_2", name="slow_tool", arguments='{"location": "x"}'),
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
chat_client_base.run_responses = [batch_response]
|
||||
agent = Agent(client=chat_client_base, middleware=[FailFast()], tools=[sample_tool_function, slow_tool])
|
||||
|
||||
with pytest.raises(MiddlewareFailure, match="abort the batch"):
|
||||
await agent.run("run both tools")
|
||||
|
||||
assert sibling_cancelled == [True]
|
||||
assert chat_client_base.call_count == 1
|
||||
|
||||
async def test_failure_with_sync_sibling_discards_late_result(self, chat_client_base: "MockBaseChatClient") -> None:
|
||||
"""Batch cancellation is cooperative: a synchronous sibling cannot be interrupted.
|
||||
|
||||
A synchronous tool body runs in a worker thread (``asyncio.to_thread``);
|
||||
cancelling its wrapping task cannot stop the thread, so the body may complete
|
||||
its side effects after the failure has already reached the caller. Its result
|
||||
is discarded either way — the loop stops at one model call — and failure
|
||||
propagation is not delayed behind the still-running thread.
|
||||
"""
|
||||
sync_started = threading.Event()
|
||||
sync_release = threading.Event()
|
||||
sync_completed: list[str] = []
|
||||
|
||||
def sync_slow(location: str) -> str:
|
||||
sync_started.set()
|
||||
sync_release.wait(10)
|
||||
sync_completed.append("side effect")
|
||||
return "late sync result"
|
||||
|
||||
slow_tool = FunctionTool(func=sync_slow, name="slow_tool", approval_mode="never_require")
|
||||
|
||||
class FailFast(FunctionMiddleware):
|
||||
async def process(
|
||||
self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]
|
||||
) -> None:
|
||||
if context.function.name == "sample_tool_function":
|
||||
# Fail only once the synchronous sibling body is genuinely running.
|
||||
await asyncio.to_thread(sync_started.wait, 5)
|
||||
raise MiddlewareFailure("abort the batch")
|
||||
await call_next()
|
||||
|
||||
batch_response = ChatResponse(
|
||||
messages=[
|
||||
Message(
|
||||
role="assistant",
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
call_id="call_1", name="sample_tool_function", arguments='{"location": "Seattle"}'
|
||||
),
|
||||
Content.from_function_call(call_id="call_2", name="slow_tool", arguments='{"location": "x"}'),
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
chat_client_base.run_responses = [batch_response]
|
||||
agent = Agent(client=chat_client_base, middleware=[FailFast()], tools=[sample_tool_function, slow_tool])
|
||||
|
||||
try:
|
||||
with pytest.raises(MiddlewareFailure, match="abort the batch"):
|
||||
await agent.run("run both tools")
|
||||
# The failure reached the caller while the synchronous body was still running.
|
||||
assert sync_completed == []
|
||||
finally:
|
||||
sync_release.set()
|
||||
|
||||
for _ in range(500):
|
||||
if sync_completed:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
# The worker thread survived cancellation and completed its side effect
|
||||
# (the documented cooperative-cancellation limitation) ...
|
||||
assert sync_completed == ["side effect"]
|
||||
# ... but its result went nowhere: the loop never made another model call.
|
||||
assert chat_client_base.call_count == 1
|
||||
|
||||
async def test_failure_streaming_reaches_stream_consumer(self, chat_client_base: "MockBaseChatClient") -> None:
|
||||
executed: list[str] = []
|
||||
|
||||
def tool_impl(location: str) -> str:
|
||||
executed.append(location)
|
||||
return "ran"
|
||||
|
||||
tracked_tool = FunctionTool(func=tool_impl, name="sample_tool_function", approval_mode="never_require")
|
||||
|
||||
class Enforcement(FunctionMiddleware):
|
||||
async def process(
|
||||
self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]
|
||||
) -> None:
|
||||
raise MiddlewareFailure("policy denied")
|
||||
|
||||
chat_client_base.streaming_responses = [
|
||||
[
|
||||
ChatResponseUpdate(
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
call_id="call_1", name="sample_tool_function", arguments='{"location": "Seattle"}'
|
||||
)
|
||||
],
|
||||
role="assistant",
|
||||
)
|
||||
]
|
||||
]
|
||||
agent = Agent(client=chat_client_base, middleware=[Enforcement()], tools=[tracked_tool])
|
||||
|
||||
with pytest.raises(MiddlewareFailure, match="policy denied"):
|
||||
async for _ in agent.run("get the weather", stream=True):
|
||||
pass
|
||||
|
||||
assert executed == []
|
||||
assert chat_client_base.call_count == 1
|
||||
|
||||
async def test_failure_from_agent_middleware_propagates(self, chat_client_base: "MockBaseChatClient") -> None:
|
||||
"""Agent (and chat) middleware exceptions already propagate; the explicit signal behaves the same."""
|
||||
|
||||
class Guard(AgentMiddleware):
|
||||
async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
raise MiddlewareFailure("run denied")
|
||||
|
||||
agent = Agent(client=chat_client_base, middleware=[Guard()])
|
||||
|
||||
with pytest.raises(MiddlewareFailure, match="run denied"):
|
||||
await agent.run("hello")
|
||||
|
||||
assert chat_client_base.call_count == 0
|
||||
|
||||
async def test_failure_settles_dangling_calls_on_service_conversation(
|
||||
self, chat_client_base: "MockBaseChatClient"
|
||||
) -> None:
|
||||
"""A service-managed conversation is settled before the failure propagates.
|
||||
|
||||
The continuation state (``session.service_session_id``) is persisted when the
|
||||
model turn completes — before tool execution — so an aborted batch would leave
|
||||
the hosted thread ending in unresolved function calls, and OpenAI-style
|
||||
continuations reject the next request over such a thread. The loop submits one
|
||||
error ``function_result`` per dangling call (``tool_choice="none"``) and
|
||||
discards the settlement response; the run still fails.
|
||||
"""
|
||||
from agent_framework import AgentSession
|
||||
|
||||
requests: list[dict[str, Any]] = []
|
||||
|
||||
class Recorder(ChatMiddleware):
|
||||
async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
requests.append({
|
||||
"contents": [(m.role, [c for c in m.contents]) for m in context.messages],
|
||||
"conversation_id": (context.options or {}).get("conversation_id"),
|
||||
"tool_choice": (context.options or {}).get("tool_choice"),
|
||||
})
|
||||
await call_next()
|
||||
|
||||
class Enforcement(FunctionMiddleware):
|
||||
async def process(
|
||||
self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]
|
||||
) -> None:
|
||||
raise MiddlewareFailure("policy denied")
|
||||
|
||||
tool_turn = _tool_call_response()
|
||||
tool_turn.conversation_id = "conv_123"
|
||||
chat_client_base.run_responses = [tool_turn]
|
||||
session = AgentSession()
|
||||
agent = Agent(client=chat_client_base, middleware=[Enforcement(), Recorder()], tools=[sample_tool_function])
|
||||
|
||||
with pytest.raises(MiddlewareFailure, match="policy denied"):
|
||||
await agent.run("get the weather", session=session)
|
||||
|
||||
# The continuation state was already durable when the batch failed ...
|
||||
assert session.service_session_id == "conv_123"
|
||||
# ... so the loop settled the thread: one extra request carrying an error
|
||||
# function_result for the dangling call, with tool calling disabled.
|
||||
assert len(requests) == 2
|
||||
settlement = requests[1]
|
||||
assert settlement["conversation_id"] == "conv_123"
|
||||
assert settlement["tool_choice"] == "none"
|
||||
settlement_results = [
|
||||
content
|
||||
for _, contents in settlement["contents"]
|
||||
for content in contents
|
||||
if content.type == "function_result"
|
||||
]
|
||||
assert [result.call_id for result in settlement_results] == ["call_1"]
|
||||
assert settlement_results[0].exception == "MiddlewareFailure"
|
||||
|
||||
async def test_failure_settles_service_conversation_streaming(self, chat_client_base: "MockBaseChatClient") -> None:
|
||||
"""The streaming loop settles a service-managed conversation the same way."""
|
||||
requests: list[dict[str, Any]] = []
|
||||
|
||||
class Recorder(ChatMiddleware):
|
||||
async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
requests.append({
|
||||
"tool_choice": (context.options or {}).get("tool_choice"),
|
||||
"messages": list(context.messages),
|
||||
})
|
||||
await call_next()
|
||||
|
||||
class Enforcement(FunctionMiddleware):
|
||||
async def process(
|
||||
self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]
|
||||
) -> None:
|
||||
raise MiddlewareFailure("policy denied")
|
||||
|
||||
chat_client_base.streaming_responses = [
|
||||
[
|
||||
ChatResponseUpdate(
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
call_id="call_1", name="sample_tool_function", arguments='{"location": "Seattle"}'
|
||||
)
|
||||
],
|
||||
role="assistant",
|
||||
conversation_id="conv_123",
|
||||
)
|
||||
]
|
||||
]
|
||||
agent = Agent(client=chat_client_base, middleware=[Enforcement(), Recorder()], tools=[sample_tool_function])
|
||||
|
||||
with pytest.raises(MiddlewareFailure, match="policy denied"):
|
||||
async for _ in agent.run("get the weather", stream=True):
|
||||
pass
|
||||
|
||||
assert len(requests) == 2
|
||||
assert requests[1]["tool_choice"] == "none"
|
||||
settlement_results = [
|
||||
content
|
||||
for message in requests[1]["messages"]
|
||||
for content in message.contents
|
||||
if content.type == "function_result"
|
||||
]
|
||||
assert [result.call_id for result in settlement_results] == ["call_1"]
|
||||
|
||||
async def test_failure_without_service_conversation_makes_no_settlement_request(
|
||||
self, chat_client_base: "MockBaseChatClient"
|
||||
) -> None:
|
||||
"""No service-managed conversation, no settlement cost: the abort stays at one model call."""
|
||||
|
||||
class Enforcement(FunctionMiddleware):
|
||||
async def process(
|
||||
self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]
|
||||
) -> None:
|
||||
raise MiddlewareFailure("policy denied")
|
||||
|
||||
chat_client_base.run_responses = [_tool_call_response()]
|
||||
agent = Agent(client=chat_client_base, middleware=[Enforcement()], tools=[sample_tool_function])
|
||||
|
||||
with pytest.raises(MiddlewareFailure, match="policy denied"):
|
||||
await agent.run("get the weather")
|
||||
|
||||
assert chat_client_base.call_count == 1
|
||||
|
||||
async def test_failure_settlement_advances_response_id_continuation(
|
||||
self, chat_client_base: "MockBaseChatClient"
|
||||
) -> None:
|
||||
"""The persisted continuation advances to the settlement response.
|
||||
|
||||
For response-ID continuations (OpenAI Responses with ``store=True``, where
|
||||
each response id is the continuation handle) the settlement response is the
|
||||
first endpoint whose chain includes the synthetic tool outputs; leaving
|
||||
``session.service_session_id`` on the pre-settlement response would make the
|
||||
next run continue from the still-unresolved turn.
|
||||
"""
|
||||
from agent_framework import AgentSession
|
||||
|
||||
class Enforcement(FunctionMiddleware):
|
||||
async def process(
|
||||
self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]
|
||||
) -> None:
|
||||
raise MiddlewareFailure("policy denied")
|
||||
|
||||
tool_turn = _tool_call_response()
|
||||
tool_turn.conversation_id = "resp_1"
|
||||
settlement_turn = ChatResponse(
|
||||
messages=[Message(role="assistant", contents=["settled"])], conversation_id="resp_2"
|
||||
)
|
||||
chat_client_base.run_responses = [tool_turn, settlement_turn]
|
||||
session = AgentSession()
|
||||
agent = Agent(client=chat_client_base, middleware=[Enforcement()], tools=[sample_tool_function])
|
||||
|
||||
with pytest.raises(MiddlewareFailure, match="policy denied"):
|
||||
await agent.run("get the weather", session=session)
|
||||
|
||||
# The stored continuation points at the settled endpoint, not the aborted turn.
|
||||
assert session.service_session_id == "resp_2"
|
||||
assert chat_client_base.call_count == 2
|
||||
|
||||
async def test_failure_during_approved_replay_settles_and_escapes(
|
||||
self, chat_client_base: "MockBaseChatClient"
|
||||
) -> None:
|
||||
"""A fatal signal during an approved-tool replay escapes loudly and settles.
|
||||
|
||||
The replayed call belongs to an earlier, already-persisted model turn, so the
|
||||
service-managed conversation must be settled from the approval-resolution
|
||||
phase too (which runs before any model call of the resumed run).
|
||||
"""
|
||||
from agent_framework import AgentSession
|
||||
|
||||
executed: list[str] = []
|
||||
|
||||
def tool_impl(location: str) -> str:
|
||||
executed.append(location)
|
||||
return "ran"
|
||||
|
||||
guarded_tool = FunctionTool(func=tool_impl, name="guarded_tool", approval_mode="always_require")
|
||||
|
||||
class FailReplay(FunctionMiddleware):
|
||||
async def process(
|
||||
self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]
|
||||
) -> None:
|
||||
raise MiddlewareFailure("replay denied")
|
||||
|
||||
requests: list[dict[str, Any]] = []
|
||||
|
||||
class Recorder(ChatMiddleware):
|
||||
async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
requests.append({
|
||||
"tool_choice": (context.options or {}).get("tool_choice"),
|
||||
"conversation_id": (context.options or {}).get("conversation_id"),
|
||||
"results": [
|
||||
content
|
||||
for message in context.messages
|
||||
for content in message.contents
|
||||
if content.type == "function_result"
|
||||
],
|
||||
})
|
||||
await call_next()
|
||||
|
||||
tool_turn = ChatResponse(
|
||||
messages=[
|
||||
Message(
|
||||
role="assistant",
|
||||
contents=[
|
||||
Content.from_function_call(call_id="call_1", name="guarded_tool", arguments='{"location": "x"}')
|
||||
],
|
||||
)
|
||||
],
|
||||
conversation_id="resp_1",
|
||||
)
|
||||
chat_client_base.run_responses = [tool_turn]
|
||||
session = AgentSession()
|
||||
agent = Agent(client=chat_client_base, middleware=[FailReplay(), Recorder()], tools=[guarded_tool])
|
||||
|
||||
paused = await agent.run("go", session=session)
|
||||
approvals = [
|
||||
content
|
||||
for message in paused.messages
|
||||
for content in message.contents
|
||||
if content.type == "function_approval_request"
|
||||
]
|
||||
assert len(approvals) == 1
|
||||
assert session.service_session_id == "resp_1"
|
||||
|
||||
settlement_turn = ChatResponse(
|
||||
messages=[Message(role="assistant", contents=["settled"])], conversation_id="resp_2"
|
||||
)
|
||||
chat_client_base.run_responses = [settlement_turn]
|
||||
approval_request = approvals[0]
|
||||
assert approval_request.id is not None
|
||||
assert approval_request.function_call is not None
|
||||
approval_message = Message(
|
||||
role="user",
|
||||
contents=[
|
||||
Content.from_function_approval_response(True, approval_request.id, approval_request.function_call)
|
||||
],
|
||||
)
|
||||
|
||||
with pytest.raises(MiddlewareFailure, match="replay denied"):
|
||||
await agent.run([approval_message], session=session)
|
||||
|
||||
# The tool never ran, and the settlement request resolved the original call
|
||||
# on the persisted conversation before the abort propagated.
|
||||
assert executed == []
|
||||
settlement = requests[-1]
|
||||
assert settlement["tool_choice"] == "none"
|
||||
assert settlement["conversation_id"] == "resp_1"
|
||||
assert [result.call_id for result in settlement["results"]] == ["call_1"]
|
||||
assert settlement["results"][0].exception == "MiddlewareFailure"
|
||||
# The continuation advanced to the settled endpoint.
|
||||
assert session.service_session_id == "resp_2"
|
||||
|
||||
async def test_failure_during_approved_replay_streaming(self, chat_client_base: "MockBaseChatClient") -> None:
|
||||
"""The streaming loop's approval-resolution phase settles and escapes the same way."""
|
||||
from agent_framework import AgentSession
|
||||
|
||||
executed: list[str] = []
|
||||
|
||||
def tool_impl(location: str) -> str:
|
||||
executed.append(location)
|
||||
return "ran"
|
||||
|
||||
guarded_tool = FunctionTool(func=tool_impl, name="guarded_tool", approval_mode="always_require")
|
||||
|
||||
class FailReplay(FunctionMiddleware):
|
||||
async def process(
|
||||
self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]
|
||||
) -> None:
|
||||
raise MiddlewareFailure("replay denied")
|
||||
|
||||
requests: list[dict[str, Any]] = []
|
||||
|
||||
class Recorder(ChatMiddleware):
|
||||
async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
requests.append({
|
||||
"tool_choice": (context.options or {}).get("tool_choice"),
|
||||
"results": [
|
||||
content
|
||||
for message in context.messages
|
||||
for content in message.contents
|
||||
if content.type == "function_result"
|
||||
],
|
||||
})
|
||||
await call_next()
|
||||
|
||||
tool_turn = ChatResponse(
|
||||
messages=[
|
||||
Message(
|
||||
role="assistant",
|
||||
contents=[
|
||||
Content.from_function_call(call_id="call_1", name="guarded_tool", arguments='{"location": "x"}')
|
||||
],
|
||||
)
|
||||
],
|
||||
conversation_id="resp_1",
|
||||
)
|
||||
chat_client_base.run_responses = [tool_turn]
|
||||
session = AgentSession()
|
||||
agent = Agent(client=chat_client_base, middleware=[FailReplay(), Recorder()], tools=[guarded_tool])
|
||||
|
||||
paused = await agent.run("go", session=session)
|
||||
approvals = [
|
||||
content
|
||||
for message in paused.messages
|
||||
for content in message.contents
|
||||
if content.type == "function_approval_request"
|
||||
]
|
||||
assert len(approvals) == 1
|
||||
|
||||
chat_client_base.run_responses = [ChatResponse(messages=[Message(role="assistant", contents=["settled"])])]
|
||||
approval_request = approvals[0]
|
||||
assert approval_request.id is not None
|
||||
assert approval_request.function_call is not None
|
||||
approval_message = Message(
|
||||
role="user",
|
||||
contents=[
|
||||
Content.from_function_approval_response(True, approval_request.id, approval_request.function_call)
|
||||
],
|
||||
)
|
||||
|
||||
with pytest.raises(MiddlewareFailure, match="replay denied"):
|
||||
async for _ in agent.run([approval_message], session=session, stream=True):
|
||||
pass
|
||||
|
||||
assert executed == []
|
||||
settlement = requests[-1]
|
||||
assert settlement["tool_choice"] == "none"
|
||||
assert [result.call_id for result in settlement["results"]] == ["call_1"]
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
Generated
+4203
-4203
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user