From 0a85dd339a4f4984e27b2bda5947a3f0df96ca3a Mon Sep 17 00:00:00 2001 From: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> Date: Mon, 3 Aug 2026 19:15:33 +0000 Subject: [PATCH 1/8] feat(python): add agent-hooks middleware as experimental core feature Implement the AGENT-HOOKS-0.1 interception contract as a first-class experimental feature in agent_framework core. - Single public factory agent_hooks_middleware() returning a private agent/chat/function middleware trio (one object per middleware category); partial or stacked installs fail closed with loud errors. - All eight interception points: input/output at the agent seam, pre/post_model_call at the chat seam, pre/post_tool_call at the function seam, agent_startup/agent_shutdown bracketing each run. - Fail-closed enforcement throughout: transforms write back into the native contexts (messages, arguments, results) or raise; content is preserved as Content objects; MiddlewareTermination short-circuits are guarded at every seam; enforcement-layer failures halt the run; interceptor crashes surface as host_error denies. - Streaming is fully buffered per spec buffered_output semantics: no update egresses before the post_model_call/output verdicts; a deny at pull time releases zero updates; run state stays active across lazy pulls with cleanup on every exit path. - Session scoping: per-run by default (startup/shutdown bracket each run) or host-owned via emitter/builder parameters for one session spanning multiple runs. - agent-hooks-sdk is an opt-in agent-hooks extra (not in all), lazy-imported per the _mcp.py pattern; core imports cleanly without it and the factory raises a clear ModuleNotFoundError. - ExperimentalFeature.AGENT_HOOKS + @experimental decorator, lazy root export, typing surface, PACKAGE_STATUS.md entry. - 55 tests built on real Agent/mock-client flows covering deny-before- execution, transform write-back, rich-content preservation, complete streaming ordering, error cleanup, concurrency isolation, nested agents, and importability without the optional SDK. Signed-off-by: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> --- python/PACKAGE_STATUS.md | 7 + .../packages/core/agent_framework/__init__.py | 2 + .../core/agent_framework/__init__.pyi | 2 + .../core/agent_framework/_agent_hooks.py | 1375 +++++++++++++++++ .../core/agent_framework/_feature_stage.py | 1 + python/packages/core/pyproject.toml | 5 + .../core/pyrightconfig.dependency.json | 1 + .../core/tests/core/test_agent_hooks.py | 1334 ++++++++++++++++ python/uv.lock | 15 +- 9 files changed, 2741 insertions(+), 1 deletion(-) create mode 100644 python/packages/core/agent_framework/_agent_hooks.py create mode 100644 python/packages/core/tests/core/test_agent_hooks.py diff --git a/python/PACKAGE_STATUS.md b/python/PACKAGE_STATUS.md index 37053838277..98ff7e5c75b 100644 --- a/python/PACKAGE_STATUS.md +++ b/python/PACKAGE_STATUS.md @@ -64,6 +64,13 @@ listed below. ### Experimental features +#### `AGENT_HOOKS` + +- `agent-framework-core`: `agent_hooks_middleware` from `agent_framework/_agent_hooks.py`, + the AGENT-HOOKS-0.1 enforcement middleware trio. Requires the opt-in + `agent-framework-core[agent-hooks]` extra (`agent-hooks-sdk`), which is deliberately + not part of `agent-framework-core[all]`. + #### `DECLARATIVE_AGENTS` - `agent-framework-declarative`: declarative agent loading APIs from diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index ff0aa20e1a0..055fc6fa28c 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -41,6 +41,7 @@ ) _LAZY_MODULE_EXPORTS: Final[Mapping[str, tuple[str, ...]]] = { + "._agent_hooks": ("agent_hooks_middleware",), "._agents": ("Agent", "BaseAgent", "RawAgent", "SupportsAgentRun"), "._clients": ( "BaseChatClient", @@ -593,6 +594,7 @@ "WorkflowViz", "__version__", "add_usage_details", + "agent_hooks_middleware", "agent_middleware", "annotate_message_groups", "apply_compaction", diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi index 31fa53a56ff..d9662f7ff74 100644 --- a/python/packages/core/agent_framework/__init__.pyi +++ b/python/packages/core/agent_framework/__init__.pyi @@ -4,6 +4,7 @@ from typing import Final __version__: Final[str] +from ._agent_hooks import agent_hooks_middleware from ._agents import Agent, BaseAgent, RawAgent, SupportsAgentRun from ._clients import ( BaseChatClient, @@ -557,6 +558,7 @@ __all__ = [ "WorkflowViz", "__version__", "add_usage_details", + "agent_hooks_middleware", "agent_middleware", "annotate_message_groups", "apply_compaction", diff --git a/python/packages/core/agent_framework/_agent_hooks.py b/python/packages/core/agent_framework/_agent_hooks.py new file mode 100644 index 00000000000..6e7ede73c5d --- /dev/null +++ b/python/packages/core/agent_framework/_agent_hooks.py @@ -0,0 +1,1375 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""AGENT-HOOKS-0.1 enforcement middleware for Agent Framework (experimental). + +This module implements the `agent-hooks `_ +control contract as one coherent feature on the framework's native middleware seams: + +- ``agent_startup`` / ``input`` / ``output`` / ``agent_shutdown`` ride the agent seam, +- ``pre_model_call`` / ``post_model_call`` ride the chat seam, +- ``pre_tool_call`` / ``post_tool_call`` ride the function seam. + +The single public entry point is :func:`agent_hooks_middleware`, which returns the three +middleware objects as one unit. The implementations are deliberately private: installing +only part of the trio would enforce only part of the control contract. Install exactly +one trio per agent, placed first in the middleware list: middleware listed before the +trio runs outside the enforcement boundary (outer position is outer trust) — e.g. a +function middleware placed before the trio can substitute a tool result that the tool +seam never brackets; the final ``output`` point still guards whatever egresses. + +Enforcement semantics (``mode="enforce"``): + +- Every interception point is emitted **before** the guarded action runs (pre points) or + before its result is incorporated (post points). Emission failures inside the SDK + (interceptor crash/timeout, invalid context) synthesize ``host_error:*`` denies and are + treated as blocks — the feature never fails open. +- ``transform`` verdicts are written back into the native middleware contexts + (``messages`` / ``arguments`` / ``result``), so the framework executes exactly the + value the interceptors approved. Content objects are preserved: rich (non-text) + message content is projected as content dictionaries, never flattened to text. +- A ``deny`` at ``input``, ``pre_model_call``, ``post_model_call``, or ``output`` + terminates the run: :class:`agent_hooks.InterceptionBlocked` propagates to the caller + of :meth:`Agent.run` (for streaming runs, it is raised when the stream is consumed). +- A ``deny`` at ``pre_tool_call`` / ``post_tool_call`` blocks the tool call: the tool is + 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). +- 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. +- The agent middleware verifies at run start that its chat and function siblings are + installed with it (the trio is created together by the factory), results the feature + cannot project are rejected loudly, and unexpected failures inside the enforcement + layer itself halt the run instead of degrading into tool errors. + +Streaming is supported **fail-closed by buffering**: the model/agent stream is fully +consumed internally, ``post_model_call`` / ``output`` verdicts are applied to the +assembled response, and only then are the (possibly transformed) updates released to the +consumer. No partial content ever egresses ahead of a verdict (spec §12.1/§12.1a +``buffered_output: true`` behaviour). + +Session scoping: by default each agent run is one agent-hooks session (fresh emitter and +sequence, ``agent_startup``/``agent_shutdown`` bracket the run). A host that owns a +longer-lived session can pass its own ``emitter`` and ``builder``; the middleware then +emits only the per-run points and the host owns the session boundaries. + +The ``agent-hooks-sdk`` dependency is optional: importing this module (and the lazy root +export) works without it, and :func:`agent_hooks_middleware` raises a descriptive +``ModuleNotFoundError`` when the SDK is missing. Install it via +``pip install agent-framework-core[agent-hooks]``. +""" + +from __future__ import annotations + +import asyncio +import base64 +import contextlib +import json +import uuid +from collections.abc import Awaitable, Callable, Mapping, Sequence +from contextvars import ContextVar +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Generic, NoReturn, TypeVar, cast + +from pydantic import BaseModel + +from ._feature_stage import ExperimentalFeature, experimental +from ._middleware import ( + AgentContext, + AgentMiddleware, + ChatContext, + ChatMiddleware, + FunctionInvocationContext, + FunctionMiddleware, + MiddlewareTermination, + MiddlewareTypes, +) +from ._types import ( + AgentResponse, + AgentResponseUpdate, + ChatResponse, + ChatResponseUpdate, + Content, + Message, + ResponseStream, +) +from .exceptions import MiddlewareException + +if TYPE_CHECKING: + from agent_hooks import ( + AgentContextBuilder, + ApprovalResolver, + CompositionConfig, + EmitOutcome, + EnforcementMode, + IdentityProvider, + InterceptionBlocked, + InterceptionEmitter, + InterceptionRecord, + Interceptor, + ) + +_UpdateT = TypeVar("_UpdateT") + +_FRAMEWORK_NAME = "agent-framework" +_HOST_ERROR_PREFIX = "host_error:" +_JCS_SHA256 = "jcs-sha256" +_DEFAULT_TIMEOUT = 5.0 + +_SDK_MISSING_MESSAGE = ( + "agent_hooks_middleware requires the optional `agent-hooks-sdk` package. " + "Please install `agent-framework-core[agent-hooks]` (or `agent-hooks-sdk`)." +) + +_TRIO_REQUIRED_MESSAGE = ( + "agent-hooks {seam} middleware was invoked without an active agent-hooks run. " + "The middleware returned by agent_hooks_middleware() must be installed together " + "on an Agent, e.g. Agent(client=..., middleware=agent_hooks_middleware([...]))." +) + +_SIBLINGS_REQUIRED_MESSAGE = ( + "agent-hooks agent middleware could not find its chat and function middleware " + "siblings in this run's effective middleware. The middleware returned by " + "agent_hooks_middleware() must be installed together on an Agent, e.g. " + "Agent(client=..., middleware=agent_hooks_middleware([...]))." +) + +_FOREIGN_TRIO_MESSAGE = ( + "agent-hooks {seam} middleware found an active agent-hooks run owned by a different " + "agent_hooks_middleware() trio. Stacking multiple trios on one agent (or splitting " + "trios across agent- and client-level middleware) is not supported: emissions would " + "silently bind to the wrong emitter. Install exactly one trio per agent." +) + + +class _AgentHooksWriteBackError(MiddlewareException): + """A transform verdict could not be converted back into the native context. + + Raised (and deliberately never caught by this module) so an unappliable transform + fails the run closed instead of silently proceeding with the untransformed value. + """ + + +# region Run state + + +@dataclass +class _RunState: + """Per-run enforcement state shared by the middleware trio via a ContextVar.""" + + emitter: InterceptionEmitter + 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) + + +@dataclass +class _AgentHooksConfig: + """Configuration shared by one middleware trio (one factory call).""" + + interceptors: tuple[tuple[str | None, Interceptor], ...] + resolver: ApprovalResolver | None + mode: EnforcementMode | None + composition: CompositionConfig | None + identity_provider: str | IdentityProvider | None + timeout: float | None + record_sink: Callable[[InterceptionRecord], None] | None + emitter: InterceptionEmitter | None + builder: AgentContextBuilder | None + + +# endregion + +# region Wire projections (framework objects -> AGENT-HOOKS wire JSON) + + +def _json_safe(value: Any) -> Any: + """Best-effort projection of an arbitrary Python value into JSON-native values. + + The SDK's own marshalling guard still fails closed on anything this misses + (``host_error:context_invalid``), so this is a fidelity aid, not a safety gate. + """ + if value is None or isinstance(value, (str, bool, int, float)): + return value + if isinstance(value, Content): + return _json_safe(value.to_dict()) + if isinstance(value, Mapping): + return {str(key): _json_safe(item) for key, item in cast("Mapping[Any, Any]", value).items()} + if isinstance(value, (bytes, bytearray)): + return base64.b64encode(bytes(value)).decode("ascii") + if isinstance(value, Sequence): + return [_json_safe(item) for item in cast("Sequence[Any]", value)] + if isinstance(value, BaseModel): + return _json_safe(value.model_dump()) + to_dict = getattr(value, "to_dict", None) + if callable(to_dict): + with contextlib.suppress(Exception): + return _json_safe(to_dict()) + return str(value) + + +def _role_str(role: Any) -> str: + return str(getattr(role, "value", role) or "user") + + +def _input_role(role: Any) -> str: + """Map a framework role onto the spec's input role enum (user | system | external).""" + value = _role_str(role) + return value if value in ("user", "system") else "external" + + +def _finish_reason_str(finish_reason: Any) -> str: + return str(getattr(finish_reason, "value", finish_reason) or "stop") + + +def _contents_to_wire(contents: Sequence[Content]) -> str | list[dict[str, Any]]: + """Project message contents faithfully: plain text as a string, rich content as dicts.""" + if len(contents) == 1 and contents[0].type == "text": + return contents[0].text or "" + return [cast("dict[str, Any]", _json_safe(content.to_dict())) for content in contents] + + +def _message_to_wire(message: Message) -> dict[str, Any]: + return {"role": _role_str(message.role), "content": _contents_to_wire(message.contents)} + + +def _messages_to_wire(messages: Sequence[Message]) -> list[dict[str, Any]]: + return [_message_to_wire(message) for message in messages] + + +def _wire_to_contents(value: Any, *, point: str) -> list[Content]: + """Decode a transformed wire content value back into framework Content objects.""" + if value is None: + return [] + if isinstance(value, str): + return [Content.from_text(value)] + if isinstance(value, Mapping): + items: list[Any] = [value] + elif isinstance(value, Sequence): + items = list(cast("Sequence[Any]", value)) + else: + raise _AgentHooksWriteBackError(f"agent-hooks {point} transform produced an unsupported content value type.") + contents: list[Content] = [] + for item in items: + if isinstance(item, str): + contents.append(Content.from_text(item)) + continue + if isinstance(item, Mapping) and "type" in item: + try: + contents.append(Content.from_dict(cast("Mapping[str, Any]", item))) + continue + except Exception as exc: + raise _AgentHooksWriteBackError( + f"agent-hooks {point} transform produced an undecodable content item." + ) from exc + raise _AgentHooksWriteBackError(f"agent-hooks {point} transform produced an unsupported content item.") + return contents + + +def _write_back_message_list( + originals: Sequence[Message], + before: Sequence[Mapping[str, Any]], + after: Any, + *, + point: str, +) -> list[Message]: + """Convert a transformed wire message list back into framework messages. + + Unchanged messages keep their original ``Message`` object (identity, ids, and + non-projected metadata). Changed messages are mutated in place when the role + matches (so shared conversation history adopts the transform) and rebuilt + otherwise. The transformed list is authoritative: added or removed messages are + reflected as-is. + """ + if not isinstance(after, list): + raise _AgentHooksWriteBackError(f"agent-hooks {point} transform must produce a list of messages.") + result: list[Message] = [] + for index, item in enumerate(cast("list[Any]", after)): + if not isinstance(item, Mapping) or "content" not in item: + raise _AgentHooksWriteBackError(f"agent-hooks {point} transform produced a message without role/content.") + wire_message = cast("Mapping[str, Any]", item) + role = str(wire_message.get("role") or "user") + if index < len(before) and dict(wire_message) == dict(before[index]): + result.append(originals[index]) + continue + if index < len(originals) and role == str(before[index].get("role")): + message = originals[index] + message.contents = _wire_to_contents(wire_message.get("content"), point=point) + result.append(message) + continue + result.append(Message(role, _wire_to_contents(wire_message.get("content"), point=point))) + return result + + +def _input_message_to_wire(message: Message) -> dict[str, Any]: + """Project one input message with the spec's input role mapping.""" + return {"role": _input_role(message.role), "content": _contents_to_wire(message.contents)} + + +def _input_messages_to_wire(messages: Sequence[Message]) -> list[dict[str, Any]]: + return [_input_message_to_wire(message) for message in messages] + + +def _input_to_wire(messages: Sequence[Message]) -> dict[str, Any]: + """Project run input per the spec's ``input`` payload schema. + + A single plain-text message projects as its content string (so string-matching + perimeter guards fire); multi-message or rich input projects as a list of + per-message ``{"role", "content"}`` objects (roles mapped onto the spec's input + role enum) whose contents are strings when plain. + """ + if len(messages) == 1: + return {"content": _contents_to_wire(messages[0].contents), "role": _input_role(messages[0].role)} + return {"content": _input_messages_to_wire(messages), "role": "user"} + + +def _looks_like_message_dicts(value: Any) -> bool: + if not isinstance(value, list): + return False + items = cast("list[Any]", value) + return bool(items) and all(isinstance(item, Mapping) and "content" in item for item in items) + + +def _apply_input_write_back(messages: list[Message], before: Mapping[str, Any], after: Any) -> None: + """Write a transformed ``input`` target back into the run's message list.""" + if after is None or after == before: + return + if not isinstance(after, Mapping): + raise _AgentHooksWriteBackError("agent-hooks input transform must produce an input object target.") + after_map = cast("Mapping[str, Any]", after) + after_role = after_map.get("role") + if after_role != before.get("role"): + # The role field is per-message only for single-message input; for multi-message + # input the top-level role is synthetic and a transform against it is ambiguous. + if len(messages) != 1 or not isinstance(after_role, str): + raise _AgentHooksWriteBackError( + "agent-hooks input transform changed the input role in a way that cannot be written back." + ) + messages[0].role = after_role + after_content = after_map.get("content") + if after_content == before.get("content"): + return + if len(messages) == 1 and not _looks_like_message_dicts(after_content): + messages[0].contents = _wire_to_contents(after_content, point="input") + return + before_list = _input_messages_to_wire(messages) + messages[:] = _write_back_message_list(list(messages), before_list, after_content, point="input") + + +def _arguments_to_wire(arguments: Any) -> dict[str, Any]: + """Project tool-call arguments as the spec's ``args`` object.""" + if arguments is None: + return {} + if isinstance(arguments, BaseModel): + return {str(key): _json_safe(item) for key, item in arguments.model_dump().items()} + if isinstance(arguments, Mapping): + return {str(key): _json_safe(item) for key, item in cast("Mapping[Any, Any]", arguments).items()} + if isinstance(arguments, str): + with contextlib.suppress(ValueError): + parsed = json.loads(arguments) + if isinstance(parsed, dict): + return cast("dict[str, Any]", parsed) + return {"raw_arguments": arguments} + return {"raw_arguments": str(arguments)} + + +def _tool_calls_to_wire(messages: Sequence[Message]) -> list[dict[str, Any]]: + calls: list[dict[str, Any]] = [] + for message in messages: + for content in message.contents: + if content.type == "function_call" and not content.informational_only: + calls.append({ + "id": str(content.call_id or ""), + "name": str(content.name or ""), + "args": _arguments_to_wire(content.arguments), + }) + return calls + + +def _response_content_to_wire(messages: Sequence[Message]) -> str | list[dict[str, Any]] | None: + """Project a model response's non-tool-call content (tool calls ride ``tool_calls``).""" + parts: list[dict[str, Any]] = [] + for message in messages: + visible = [content for content in message.contents if content.type != "function_call"] + if not visible: + continue + parts.append({"role": _role_str(message.role), "content": _contents_to_wire(visible)}) + if not parts: + return None + if len(parts) == 1 and isinstance(parts[0]["content"], str): + return parts[0]["content"] + return parts + + +def _usage_to_wire(usage_details: Any) -> dict[str, int] | None: + if not isinstance(usage_details, Mapping): + return None + usage = { + str(key): item + for key, item in cast("Mapping[Any, Any]", usage_details).items() + if isinstance(item, int) and not isinstance(item, bool) + } + return usage or None + + +def _tool_result_to_wire(value: Any) -> Any: + """Project a tool result faithfully, unwrapping framework ``Content`` containers. + + Text content projects as its text, ``function_result`` content projects as its + canonical ``result`` value, and any other content projects as its full content + dictionary — never as ``str(Content)`` reprs. + """ + if value is None or isinstance(value, (str, bool, int, float)): + return value + if isinstance(value, list) and len(cast("list[Any]", value)) == 1 and isinstance(value[0], Content): + # The canonical single-content result (e.g. the default parser's wrapped text) + # projects as the content's value itself, matching what the model sees. + return _tool_result_to_wire(value[0]) + if isinstance(value, Content): + if value.type == "text": + return value.text or "" + if value.type == "function_result": + if value.result is not None: + return _tool_result_to_wire(value.result) + if value.items is not None: + return [_tool_result_to_wire(item) for item in value.items] + return None + return _json_safe(value.to_dict()) + if isinstance(value, Mapping): + return {str(key): _tool_result_to_wire(item) for key, item in cast("Mapping[Any, Any]", value).items()} + if isinstance(value, Sequence) and not isinstance(value, (bytes, bytearray)): + return [_tool_result_to_wire(item) for item in cast("Sequence[Any]", value)] + return _json_safe(value) + + +def _apply_tool_result_write_back(original: Any, after: Any) -> Any: + """Convert a transformed ``post_tool_call`` value back into the native result shape. + + When the original result is the framework's canonical ``list[Content]`` and the + transformed value is shape-compatible, the Content wrappers are preserved; + otherwise the transformed wire value becomes the result as-is (the function + invocation layer serializes arbitrary JSON-native results faithfully). + """ + if ( + isinstance(original, list) + and original + and all(isinstance(item, Content) for item in cast("list[Any]", original)) + ): + original_contents = cast("list[Content]", original) + if isinstance(after, str) and len(original_contents) == 1 and original_contents[0].type == "text": + return [Content.from_text(after)] + after_items = cast("list[Any]", after) if isinstance(after, list) else None + if after_items is not None and len(after_items) == len(original_contents): + rebuilt: list[Content] = [] + for content, item in zip(original_contents, after_items): + if _tool_result_to_wire(content) == item: + rebuilt.append(content) + elif content.type == "text" and isinstance(item, str): + rebuilt.append(Content.from_text(item)) + else: + return after_items + return rebuilt + return cast(Any, after) + + +def _agent_output_to_wire(response: AgentResponse) -> str | list[dict[str, Any]]: + """Project the run output: a single plain-text message as a string, else per-message objects.""" + parts = _messages_to_wire(response.messages) + if len(parts) == 1 and isinstance(parts[0]["content"], str): + return parts[0]["content"] + return parts + + +def _apply_output_write_back(response: AgentResponse, before_content: Any, after: Any) -> bool: + """Write a transformed ``output`` target back into the agent response. Returns whether it changed.""" + if after is None: + return False + if not isinstance(after, Mapping): + raise _AgentHooksWriteBackError("agent-hooks output transform must produce an output object target.") + after_content = cast("Mapping[str, Any]", after).get("content") + if after_content == before_content: + return False + originals = list(response.messages) + if isinstance(after_content, str): + if len(originals) == 1: + originals[0].contents = _wire_to_contents(after_content, point="output") + else: + response.messages = [Message("assistant", [after_content])] + return True + if after_content is None: + response.messages = [] + return True + before_list = _messages_to_wire(originals) + response.messages = _write_back_message_list(originals, before_list, after_content, point="output") + return True + + +def _write_back_tool_calls(response: ChatResponse[Any], after_calls: Any) -> bool: + """Reconcile transformed ``tool_calls`` with the response's function-call contents.""" + if not isinstance(after_calls, list): + raise _AgentHooksWriteBackError("agent-hooks post_model_call transform must keep tool_calls a list.") + wire_calls: list[Mapping[str, Any]] = [] + for item in cast("list[Any]", after_calls): + if not isinstance(item, Mapping) or "id" not in item or "name" not in item: + raise _AgentHooksWriteBackError( + "agent-hooks post_model_call transform produced a tool call without id/name." + ) + wire_calls.append(cast("Mapping[str, Any]", item)) + calls_by_id = {str(call["id"]): call for call in wire_calls} + consumed: set[str] = set() + changed = False + for message in response.messages: + kept: list[Content] = [] + for content in message.contents: + if content.type != "function_call" or content.informational_only: + kept.append(content) + continue + wire = calls_by_id.get(str(content.call_id)) + if wire is None: + changed = True # the transform dropped this tool call + continue + consumed.add(str(content.call_id)) + wire_args = wire.get("args") + if isinstance(wire_args, Mapping) and _arguments_to_wire(content.arguments) != dict( + cast("Mapping[str, Any]", wire_args) + ): + content.arguments = {str(key): item for key, item in cast("Mapping[Any, Any]", wire_args).items()} + changed = True + kept.append(content) + if len(kept) != len(message.contents): + message.contents = kept + added = [call for call in wire_calls if str(call["id"]) not in consumed] + if added: + contents = [ + Content.from_function_call( + str(call["id"]), + str(call["name"]), + arguments={str(key): item for key, item in cast("Mapping[Any, Any]", call.get("args") or {}).items()}, + ) + for call in added + ] + target = next((m for m in reversed(response.messages) if _role_str(m.role) == "assistant"), None) + if target is not None: + target.contents = [*target.contents, *contents] + else: + response.messages.append(Message("assistant", contents)) + changed = True + return changed + + +def _write_back_response_content(response: ChatResponse[Any], after_content: Any) -> None: + """Rebuild the response's visible content from a transformed ``response.content`` value.""" + calls = [ + content for message in response.messages for content in message.contents if content.type == "function_call" + ] + base: list[Message] + if after_content is None: + base = [] + elif isinstance(after_content, str): + base = [Message("assistant", [after_content])] + elif isinstance(after_content, list): + base = [] + for item in cast("list[Any]", after_content): + if not isinstance(item, Mapping) or "content" not in item: + raise _AgentHooksWriteBackError( + "agent-hooks post_model_call transform produced content without role/content." + ) + wire_message = cast("Mapping[str, Any]", item) + base.append( + Message( + str(wire_message.get("role") or "assistant"), + _wire_to_contents(wire_message.get("content"), point="post_model_call"), + ) + ) + else: + raise _AgentHooksWriteBackError("agent-hooks post_model_call transform produced unsupported content.") + if calls: + if base and _role_str(base[-1].role) == "assistant": + base[-1].contents = [*base[-1].contents, *calls] + else: + base.append(Message("assistant", calls)) + response.messages = base + + +def _apply_response_write_back(response: ChatResponse[Any], before: Mapping[str, Any], after: Any) -> bool: + """Write a transformed ``post_model_call`` target back into the chat response.""" + if after is None or after == before: + return False + if not isinstance(after, Mapping): + raise _AgentHooksWriteBackError("agent-hooks post_model_call transform must produce a response object.") + after_map = cast("Mapping[str, Any]", after) + changed = False + after_finish = after_map.get("finish_reason") + if after_finish != before.get("finish_reason"): + if not isinstance(after_finish, str): + raise _AgentHooksWriteBackError("agent-hooks post_model_call transform must keep finish_reason a string.") + response.finish_reason = cast(Any, after_finish) + changed = True + after_calls = after_map.get("tool_calls") + if after_calls != before.get("tool_calls"): + changed = _write_back_tool_calls(response, after_calls) or changed + after_content = after_map.get("content") + if after_content != before.get("content"): + _write_back_response_content(response, after_content) + changed = True + return changed + + +def _chat_updates_from_response(response: ChatResponse[Any]) -> list[ChatResponseUpdate]: + """Re-derive stream updates from a (transformed) assembled chat response.""" + updates = [ + ChatResponseUpdate( + contents=list(message.contents), + role=cast(Any, message.role), + author_name=message.author_name, + message_id=message.message_id, + response_id=response.response_id, + model=response.model, + ) + for message in response.messages + ] + if not updates: + updates = [ChatResponseUpdate(role="assistant", response_id=response.response_id)] + updates[-1].finish_reason = response.finish_reason + return updates + + +def _agent_updates_from_response(response: AgentResponse[Any]) -> list[AgentResponseUpdate]: + """Re-derive stream updates from a (transformed) assembled agent response.""" + updates = [ + AgentResponseUpdate( + contents=list(message.contents), + role=message.role, + author_name=message.author_name, + message_id=message.message_id, + response_id=response.response_id, + ) + for message in response.messages + ] + if not updates: + updates = [AgentResponseUpdate(role="assistant", response_id=response.response_id)] + return updates + + +def _tool_names(context: AgentContext) -> list[str]: + tools: Any = context.tools if context.tools is not None else getattr(context.agent, "tools", None) + if tools is None: + return [] + items: Sequence[Any] = ( + cast("Sequence[Any]", tools) if isinstance(tools, Sequence) and not isinstance(tools, (str, bytes)) else [tools] + ) + names: list[str] = [] + for item in items: + name = getattr(item, "name", None) or getattr(item, "__name__", None) + names.append(str(name) if name else type(item).__name__) + return names + + +def _is_host_error(record: InterceptionRecord) -> bool: + return bool(record.verdict.reason and record.verdict.reason.startswith(_HOST_ERROR_PREFIX)) + + +def _blocked_tool_result(point: str, record: InterceptionRecord) -> dict[str, Any]: + """Tool-error payload surfaced to the model for a blocked tool call (no target content).""" + payload: dict[str, Any] = { + "error": f"Tool call blocked by agent-hooks at {point}.", + "reason": record.verdict.reason or "deny", + } + if record.verdict.message: + payload["message"] = record.verdict.message + return payload + + +def _is_approval_request(result: Any) -> bool: + """Whether a function result is the framework's approval-request control object.""" + 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 halt path. + + 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 while the run continued. 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. + """ + 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 + + +# eq=False keeps identity semantics (and hashability): the middleware pipeline caches +# compare middleware tuples with ==, and value-equality would let a pipeline cached for +# one trio be reused for a field-equal fresh trio, breaking the identity-based sibling +# verification on the second run. +@dataclass(eq=False) +class _AgentHooksMiddlewareBase: + """Shared base carrying the trio's config (also enables sibling identity checks).""" + + _config: _AgentHooksConfig + + def _shares_config(self, config: _AgentHooksConfig) -> bool: + """Whether this middleware was created by the same factory call as ``config``.""" + return self._config is config + + +class _ReplayAsyncIterator(Generic[_UpdateT]): + """Async iterator that replays a buffered list of updates.""" + + def __init__(self, updates: Sequence[_UpdateT]) -> None: + self._iterator = iter(updates) + + def __aiter__(self) -> _ReplayAsyncIterator[_UpdateT]: + return self + + async def __anext__(self) -> _UpdateT: + try: + return next(self._iterator) + except StopIteration: + raise StopAsyncIteration from None + + +# endregion + +# region Middleware implementations (private: the trio is one coherent feature) + + +class _AgentHooksAgentMiddleware(_AgentHooksMiddlewareBase, AgentMiddleware): + """Run bracket: ``agent_startup``, ``input``, ``output``, ``agent_shutdown``.""" + + def _new_run_state(self, context: AgentContext) -> _RunState: + config = self._config + if config.emitter is not None and config.builder is not None: + return _RunState(emitter=config.emitter, builder=config.builder, session_scoped=True, config=config) + from agent_hooks import AgentContextBuilder, InterceptionEmitter + + agent = context.agent + agent_name = getattr(agent, "name", None) + agent_id = str(getattr(agent, "id", None) or agent_name or "agent") + builder = AgentContextBuilder( + agent_id=agent_id, + framework=_FRAMEWORK_NAME, + session_id=uuid.uuid4().hex, + agent_name=str(agent_name) if agent_name else None, + ) + kwargs: dict[str, Any] = { + "resolver": config.resolver, + "timeout": config.timeout, + "composition": config.composition, + "identity_provider": config.identity_provider, + } + if config.mode is not None: + kwargs["mode"] = config.mode + emitter = InterceptionEmitter(**kwargs) + for name, interceptor in config.interceptors: + emitter.register(interceptor, name) + if config.record_sink is not None: + emitter.set_record_sink(config.record_sink) + return _RunState(emitter=emitter, builder=builder, session_scoped=False, config=config) + + async def _emit_run_start(self, context: AgentContext, state: _RunState) -> None: + """Emit ``agent_startup`` (per-run sessions) and ``input``; apply input transforms.""" + if not state.session_scoped: + await state.emitter.emit(state.builder.agent_startup(tools_registered=_tool_names(context))) + before = _input_to_wire(context.messages) + outcome: EmitOutcome = await state.emitter.emit( + state.builder.input(content=before["content"], role=before["role"]) + ) + _apply_input_write_back(context.messages, before, outcome.target) + + async def _emit_output(self, state: _RunState, response: AgentResponse[Any]) -> bool: + """Emit ``output`` over the assembled response; apply output transforms.""" + before_content = _agent_output_to_wire(response) + outcome: EmitOutcome = await state.emitter.emit(state.builder.output(content=before_content)) + return _apply_output_write_back(response, before_content, outcome.target) + + async def _emit_shutdown(self, state: _RunState, reason: str) -> None: + """Best-effort ``agent_shutdown`` (per-run sessions only; blocks there are record-only).""" + if state.session_scoped: + return + with contextlib.suppress(Exception): + await state.emitter.emit_unchecked(state.builder.agent_shutdown(reason=reason)) + + def _verify_siblings(self, context: AgentContext) -> None: + """Fail closed unless the chat and function siblings from this trio are installed. + + The trio is created together by the factory, so the siblings are located by + identity of the shared config. The agent seam sees the combined chat/function + middleware for the run in ``context.client_kwargs["middleware"]`` (the agent + layer assembles it there before building the context). + """ + installed = context.client_kwargs.get("middleware") + candidates = cast("Sequence[Any]", installed) if isinstance(installed, Sequence) else () + has_chat = any( + isinstance(item, _AgentHooksChatMiddleware) and item._shares_config(self._config) for item in candidates + ) + has_function = any( + isinstance(item, _AgentHooksFunctionMiddleware) and item._shares_config(self._config) for item in candidates + ) + if not (has_chat and has_function): + raise MiddlewareException(_SIBLINGS_REQUIRED_MESSAGE) + + async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + from agent_hooks import InterceptionBlocked + + self._verify_siblings(context) + state = self._new_run_state(context) + if context.stream: + await self._process_streaming(context, call_next, state) + return + + token = _RUN_STATE.set(state) + shutdown_reason = "completed" + try: + try: + await self._emit_run_start(context, state) + termination: MiddlewareTermination | None = None + try: + await call_next() + except MiddlewareTermination as exc: + # 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) + elif result is not None: + raise MiddlewareException( + f"agent-hooks cannot guard a run result of type {type(result).__name__}; " + "the output interception point was not emitted." + ) + if termination is not None: + raise termination + except InterceptionBlocked: + shutdown_reason = "error" + context.result = None + raise + except asyncio.CancelledError: + shutdown_reason = "cancelled" + raise + except MiddlewareTermination: + # Deliberate short-circuit: the result (if any) was guarded above. + raise + except BaseException: + shutdown_reason = "error" + raise + finally: + await self._emit_shutdown(state, shutdown_reason) + _RUN_STATE.reset(token) + + async def _process_streaming( + self, context: AgentContext, call_next: Callable[[], Awaitable[None]], state: _RunState + ) -> None: + """Streaming run setup (executed lazily on the first pull of the outer stream). + + Ownership of the session trail passes to the guarded stream only once it is + installed as the run result; every earlier exit (pre-run deny, middleware + exception, unguardable result) still closes the trail with ``agent_shutdown``. + """ + token = _RUN_STATE.set(state) + # Pessimistic default: any exit before the guarded stream takes ownership + # closes the trail as an error; ``None`` means ownership was handed off. + shutdown_reason: str | None = "error" + try: + await self._emit_run_start(context, state) + termination: MiddlewareTermination | None = None + try: + await call_next() + except MiddlewareTermination as exc: + termination = exc + inner = context.result + if inner is None and termination is not None: + # Terminated without a result: nothing will egress. + shutdown_reason = "completed" + raise termination + if not isinstance(inner, ResponseStream): + raise MiddlewareException( + "agent-hooks streaming enforcement requires a ResponseStream agent result; " + f"got {type(inner).__name__}." + ) + context.result = cast( + "ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", + cast(Any, ResponseStream).from_awaitable( + self._guarded_agent_stream( + state, cast("ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", inner) + ) + ), + ) + # The guarded stream now owns the shutdown emission. + shutdown_reason = None + if termination is not None: + # A middleware substituted its own stream: it is guarded, and the + # termination still short-circuits the rest of the pipeline. + raise termination + except asyncio.CancelledError: + shutdown_reason = "cancelled" + raise + finally: + if shutdown_reason is not None: + await self._emit_shutdown(state, shutdown_reason) + _RUN_STATE.reset(token) + + async def _guarded_agent_stream( + self, + state: _RunState, + inner: ResponseStream[AgentResponseUpdate, AgentResponse[Any]], + ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: + """Consume the run fully, apply the ``output`` verdict, then release buffered updates. + + Fail-closed streamed egress: no update is released to the consumer before the + ``output`` emission permits, and a transformed output replaces the buffered + updates entirely. + """ + from agent_hooks import InterceptionBlocked + + token = _RUN_STATE.set(state) + shutdown_reason = "completed" + transformed = False + final: AgentResponse[Any] + try: + try: + final = await inner.get_final_response() + 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__}; " + "the output interception point was not emitted." + ) + transformed = await self._emit_output(state, final) + except InterceptionBlocked: + shutdown_reason = "error" + raise + except asyncio.CancelledError: + shutdown_reason = "cancelled" + raise + except BaseException: + shutdown_reason = "error" + raise + finally: + await self._emit_shutdown(state, shutdown_reason) + _RUN_STATE.reset(token) + + updates = _agent_updates_from_response(final) if transformed else list(inner.updates) + + def _finalize(_: Sequence[AgentResponseUpdate]) -> AgentResponse[Any]: + return final + + return ResponseStream(_ReplayAsyncIterator(updates), finalizer=_finalize) + + +class _AgentHooksChatMiddleware(_AgentHooksMiddlewareBase, ChatMiddleware): + """Model bracket: ``pre_model_call`` and ``post_model_call``.""" + + async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None: + from agent_hooks import InterceptionBlocked + + state = _RUN_STATE.get() + if state is None: + raise MiddlewareException(_TRIO_REQUIRED_MESSAGE.format(seam="chat")) + if not self._shares_config(state.config): + # A different trio's agent middleware owns the innermost run state (stacked + # trios, or a trio split across agent- and client-level middleware): binding + # to it would silently misroute emissions. Fail closed instead. + raise MiddlewareException(_FOREIGN_TRIO_MESSAGE.format(seam="chat")) + + options = context.options or {} + model_id = str(options.get("model") or type(context.client).__name__) + before = _messages_to_wire(context.messages) + outcome: EmitOutcome = await state.emitter.emit( + state.builder.pre_model_call(model_id=model_id, messages=before) + ) + if outcome.target != before: + context.messages = _write_back_message_list( + list(context.messages), before, outcome.target, point="pre_model_call" + ) + + termination: MiddlewareTermination | None = None + try: + await call_next() + except MiddlewareTermination as exc: + # A chat middleware short-circuited with a substituted result; whatever + # was substituted still flows into the agent loop, so it is guarded below. + termination = exc + + result = context.result + if isinstance(result, ResponseStream): + context.result = cast( + "ResponseStream[ChatResponseUpdate, ChatResponse[Any]]", + cast(Any, ResponseStream).from_awaitable( + self._guarded_chat_stream( + state, model_id, cast("ResponseStream[ChatResponseUpdate, ChatResponse[Any]]", result) + ) + ), + ) + elif isinstance(result, ChatResponse): + try: + await self._emit_post_model_call(state, model_id, result) + except InterceptionBlocked: + # §6.1: the denied response must not be incorporated. + context.result = None + raise + elif result is not None: + raise MiddlewareException( + f"agent-hooks cannot guard a chat result of type {type(result).__name__}; " + "the post_model_call interception point was not emitted." + ) + elif context.stream and termination is None: + raise MiddlewareException("agent-hooks streaming enforcement requires a ResponseStream chat result.") + if termination is not None: + raise termination + + async def _emit_post_model_call(self, state: _RunState, model_id: str, response: ChatResponse[Any]) -> bool: + """Emit ``post_model_call`` over the assembled response; apply transforms. Returns whether changed.""" + before: dict[str, Any] = { + "content": _response_content_to_wire(response.messages), + "tool_calls": _tool_calls_to_wire(response.messages), + "finish_reason": _finish_reason_str(response.finish_reason), + } + outcome: EmitOutcome = await state.emitter.emit( + state.builder.post_model_call( + model_id=str(response.model or model_id), + content=before["content"], + tool_calls=before["tool_calls"], + finish_reason=before["finish_reason"], + usage=_usage_to_wire(response.usage_details), + request_id=response.response_id, + ) + ) + return _apply_response_write_back(response, before, outcome.target) + + async def _guarded_chat_stream( + self, + state: _RunState, + model_id: str, + inner: ResponseStream[ChatResponseUpdate, ChatResponse[Any]], + ) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: + """Assemble the streamed response, apply the ``post_model_call`` verdict, then replay. + + Spec §12.1: the complete response is assembled before ``post_model_call`` is + emitted, and nothing (updates or tool calls) is released beforehand. A deny + raises before any update egresses. + """ + response = await inner.get_final_response() + if not isinstance(response, ChatResponse): + raise MiddlewareException( + f"agent-hooks cannot guard a streamed chat result of type {type(response).__name__}; " + "the post_model_call interception point was not emitted." + ) + changed = await self._emit_post_model_call(state, model_id, response) + updates = _chat_updates_from_response(response) if changed else list(inner.updates) + + def _finalize(_: Sequence[ChatResponseUpdate]) -> ChatResponse[Any]: + return response + + return ResponseStream(_ReplayAsyncIterator(updates), finalizer=_finalize) + + +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: + """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) + + def _maybe_halt(self, state: _RunState, 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( + f"agent-hooks {point} failed closed: {record.verdict.reason}", + ) from exc + + async def process(self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]) -> None: + from agent_hooks import InterceptionBlocked + + state = _RUN_STATE.get() + if state is None: + # No run state means the trio was split. 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) + if not self._shares_config(state.config): + # A different trio owns the innermost run state (stacked trios): binding to + # it would silently misroute emissions. Halt the run fail-closed (the loop + # swallows plain exceptions, so route through the halt path). + _halt_on_enforcement_failure( + state, context, MiddlewareException(_FOREIGN_TRIO_MESSAGE.format(seam="function")), "pre_tool_call" + ) + + try: + raw_call_id = context.metadata.get("call_id") + call_id = str(raw_call_id) if raw_call_id else uuid.uuid4().hex + name = str(getattr(context.function, "name", context.function)) + args = _arguments_to_wire(context.arguments) + outcome: EmitOutcome = await state.emitter.emit( + state.builder.pre_tool_call(call_id=call_id, name=name, args=args) + ) + target = outcome.target + if not isinstance(target, Mapping): + raise _AgentHooksWriteBackError("agent-hooks pre_tool_call transform must produce an arguments object.") + effective = {str(key): item for key, item in cast("Mapping[Any, Any]", target).items()} + if effective != args: + # The transform rewrote the arguments: execute the approved value. + context.arguments = dict(effective) + 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") + return + except (MiddlewareTermination, asyncio.CancelledError): + raise + except BaseException as exc: + _halt_on_enforcement_failure(state, context, 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. + termination = exc + except asyncio.CancelledError: + raise + except BaseException as exc: + # The invocation errored: the contract still brackets it (is_error=True). + # Only the exception type name crosses the boundary (spec §6.3/§14). + try: + await state.emitter.emit( + state.builder.post_tool_call( + call_id=call_id, name=name, args=args, value=type(exc).__name__, is_error=True + ) + ) + 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: + raise + except BaseException as emit_exc: + _halt_on_enforcement_failure(state, context, emit_exc, "post_tool_call") + raise + + try: + value = _tool_result_to_wire(context.result) + outcome = await state.emitter.emit( + state.builder.post_tool_call(call_id=call_id, name=name, args=args, value=value) + ) + if outcome.target != value: + context.result = _apply_tool_result_write_back(context.result, 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): + raise + except BaseException as exc: + _halt_on_enforcement_failure(state, context, exc, "post_tool_call") + if termination is not None: + raise termination + + +# endregion + +# region Public factory + + +@experimental(feature_id=ExperimentalFeature.AGENT_HOOKS) +def agent_hooks_middleware( + interceptors: Sequence[Interceptor] | Mapping[str, Interceptor] | None = None, + *, + resolver: ApprovalResolver | None = None, + mode: EnforcementMode | str = "enforce", + composition: CompositionConfig | None = None, + identity_provider: str | IdentityProvider | None = _JCS_SHA256, + timeout: float | None = _DEFAULT_TIMEOUT, + record_sink: Callable[[InterceptionRecord], None] | None = None, + emitter: InterceptionEmitter | None = None, + builder: AgentContextBuilder | None = None, +) -> list[MiddlewareTypes]: + """Build the AGENT-HOOKS-0.1 middleware trio for an :class:`~agent_framework.Agent`. + + The returned middleware emit every applicable interception point of the agent-hooks + control contract and enforce the combined verdicts fail-closed: denies block the + guarded action (a run-level deny raises :class:`agent_hooks.InterceptionBlocked` to + the caller; a tool-seam deny surfaces a tool error to the model), and transforms are + written back into the framework's messages, arguments, and results so execution uses + exactly the values the interceptors approved. Streaming runs are buffered and only + released after the ``output`` verdict permits. + + The three middleware form one coherent feature and must be installed together; + always pass the returned list as a whole to ``Agent(middleware=...)``, and install + exactly one trio per agent (stacked trios are rejected fail-closed). + + Composition order: + Place the trio first (outermost) in the middleware list. Middleware listed + before the trio runs outside the enforcement boundary: a function middleware + placed before the trio, for example, can substitute a tool result that the + tool seam never brackets. Outer position is outer trust — the final + ``output`` point still guards whatever egresses to the caller. + + Session scoping: + By default every agent run is one agent-hooks session: a fresh + ``InterceptionEmitter``/``AgentContextBuilder`` pair is created per run and + ``agent_startup``/``agent_shutdown`` bracket the run. To scope one session + across multiple runs (shared ``sequence``, stateful interceptors, one approval + ledger), construct and own the emitter and builder yourself and pass them via + ``emitter=``/``builder=``; the middleware then emits only the per-run points + (``input`` through ``output``) and your host emits the session boundaries. + Use ``record_sink`` (or your own emitter) to observe interception records. + + Args: + interceptors: The agent-hooks interceptors to register, either as a sequence or + as a mapping of registration name to interceptor (names appear on the + records' verdict summaries). Required unless ``emitter`` is supplied. + + Keyword Args: + resolver: Optional approval resolver consulted for liftable denies. + mode: ``"enforce"`` (default) honours verdicts; ``"evaluate_only"`` records + them without acting. + composition: Composition profile and knobs; ``None`` uses the SDK default + (``sequential/first_deny``, ``on_approval: stop``). + identity_provider: ``"jcs-sha256"`` (default), a custom + :class:`agent_hooks.IdentityProvider`, or ``None`` for identity-unbound + records. + timeout: Per-interceptor/resolver timeout in seconds (spec RECOMMENDED 5.0). + record_sink: Optional callable receiving every interception record. + emitter: A host-owned ``InterceptionEmitter`` for session-scoped enforcement. + Must be provided together with ``builder``; the emitter's own + configuration (interceptors, resolver, mode, ...) then governs and the + per-run configuration parameters must be left at their defaults. + builder: The host-owned ``AgentContextBuilder`` matching ``emitter``. + + Returns: + The middleware list to pass to ``Agent(middleware=...)``. + + Raises: + ModuleNotFoundError: If the optional ``agent-hooks-sdk`` package is not + installed. + ValueError: If the configuration is inconsistent (no interceptors in per-run + mode, ``emitter``/``builder`` not provided together, or per-run parameters + combined with a host-owned emitter). + + Examples: + .. code-block:: python + + from agent_framework import Agent + from agent_hooks import ALLOW, Verdict + + + class EgressGuard: + def intercept(self, context): + if "secret" in str(context.get("target")): + return Verdict.deny(reason="egress_blocked") + return ALLOW + + + agent = Agent( + client=client, + name="assistant", + middleware=agent_hooks_middleware([EgressGuard()]), + ) + """ + try: + from agent_hooks import EnforcementMode + except ImportError as exc: # pragma: no cover - exercised via tests with import hooks + raise ModuleNotFoundError(_SDK_MISSING_MESSAGE) from exc + + if (emitter is None) != (builder is None): + raise ValueError("agent_hooks_middleware requires `emitter` and `builder` to be provided together.") + + if emitter is not None: + conflicting = [ + param_name + for param_name, value in ( + ("interceptors", interceptors), + ("resolver", resolver), + ("composition", composition), + ("record_sink", record_sink), + ) + if value is not None + ] + if (mode.value if isinstance(mode, EnforcementMode) else str(mode)) != "enforce": + conflicting.append("mode") + if identity_provider != _JCS_SHA256: + conflicting.append("identity_provider") + if timeout != _DEFAULT_TIMEOUT: + conflicting.append("timeout") + if conflicting: + raise ValueError( + "agent_hooks_middleware received a host-owned `emitter`; the emitter's own configuration " + f"governs, so these parameters must be left at their defaults: {', '.join(sorted(conflicting))}." + ) + # Host-owned session: only the emitter/builder pair is retained — the emitter's + # own configuration governs, so the per-run knobs are not stored at all. + config = _AgentHooksConfig( + interceptors=(), + resolver=None, + mode=None, + composition=None, + identity_provider=None, + timeout=None, + record_sink=None, + emitter=emitter, + builder=builder, + ) + else: + named: list[tuple[str | None, Interceptor]] + if isinstance(interceptors, Mapping): + named = [(str(name), interceptor) for name, interceptor in interceptors.items()] + else: + named = [(None, interceptor) for interceptor in interceptors or []] + if not named: + raise ValueError( + "agent_hooks_middleware requires at least one interceptor (an emitter with zero " + "interceptors fails closed on every emission)." + ) + config = _AgentHooksConfig( + interceptors=tuple(named), + resolver=resolver, + mode=EnforcementMode(mode) if isinstance(mode, str) else mode, + composition=composition, + identity_provider=identity_provider, + timeout=timeout, + record_sink=record_sink, + emitter=None, + builder=None, + ) + + return [ + _AgentHooksAgentMiddleware(config), + _AgentHooksChatMiddleware(config), + _AgentHooksFunctionMiddleware(config), + ] + + +# endregion diff --git a/python/packages/core/agent_framework/_feature_stage.py b/python/packages/core/agent_framework/_feature_stage.py index 7f6c26ffc08..53893e397f3 100644 --- a/python/packages/core/agent_framework/_feature_stage.py +++ b/python/packages/core/agent_framework/_feature_stage.py @@ -50,6 +50,7 @@ class ExperimentalFeature(str, Enum): on enum membership or attribute presence over time. """ + AGENT_HOOKS = "AGENT_HOOKS" DECLARATIVE_AGENTS = "DECLARATIVE_AGENTS" EVALS = "EVALS" FILE_HISTORY = "FILE_HISTORY" diff --git a/python/packages/core/pyproject.toml b/python/packages/core/pyproject.toml index b5a1a36a49d..0edf08d69b6 100644 --- a/python/packages/core/pyproject.toml +++ b/python/packages/core/pyproject.toml @@ -31,6 +31,11 @@ dependencies = [ ] [project.optional-dependencies] +# Deliberately NOT part of `all`: the agent-hooks enforcement middleware is an +# explicitly opt-in experimental feature. +agent-hooks = [ + "agent-hooks-sdk>=0.1.0a4,<0.2", +] all = [ "mcp>=1.24.0,<2", "agent-framework-a2a", diff --git a/python/packages/core/pyrightconfig.dependency.json b/python/packages/core/pyrightconfig.dependency.json index 2cc3d95b21f..f0be44f3d6e 100644 --- a/python/packages/core/pyrightconfig.dependency.json +++ b/python/packages/core/pyrightconfig.dependency.json @@ -6,6 +6,7 @@ "agent_framework/_workflows" ], "exclude": [ + "agent_framework/_agent_hooks.py", "agent_framework/_harness/_agent.py" ] } diff --git a/python/packages/core/tests/core/test_agent_hooks.py b/python/packages/core/tests/core/test_agent_hooks.py new file mode 100644 index 00000000000..49e13225e67 --- /dev/null +++ b/python/packages/core/tests/core/test_agent_hooks.py @@ -0,0 +1,1334 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterable, Awaitable, Callable +from typing import Any, cast + +import pytest + +import agent_framework +from agent_framework import ( + SKIP_PARSING, + Agent, + AgentContext, + AgentMiddleware, + AgentResponse, + AgentResponseUpdate, + ChatContext, + ChatMiddleware, + ChatResponse, + ChatResponseUpdate, + Content, + FunctionInvocationContext, + FunctionMiddleware, + Message, + MiddlewareException, + MiddlewareTermination, + ResponseStream, + agent_hooks_middleware, + tool, +) + +try: + from agent_hooks import ( + ALLOW, + AgentContextBuilder, + Decision, + InterceptionBlocked, + InterceptionEmitter, + InterceptionRecord, + Transform, + Verdict, + ) + + AGENT_HOOKS_AVAILABLE = True +except ImportError: # pragma: no cover - exercised on envs without the extra + AGENT_HOOKS_AVAILABLE = False + +from .conftest import MockBaseChatClient + +requires_sdk = pytest.mark.skipif(not AGENT_HOOKS_AVAILABLE, reason="agent-hooks-sdk is not installed") + +pytestmark = pytest.mark.filterwarnings("ignore::agent_framework._feature_stage.ExperimentalWarning") + + +# region Helpers + + +class AllowGuard: + """Records every context it sees (deep copies) and allows everything.""" + + def __init__(self) -> None: + self.contexts: list[dict[str, Any]] = [] + + def intercept(self, context: dict[str, Any]) -> Any: + self.contexts.append(context) + return ALLOW + + def contexts_for(self, point: str) -> list[dict[str, Any]]: + return [ctx for ctx in self.contexts if ctx["interception_point"] == point] + + +class PointGuard: + """Returns a configured verdict at one interception point, allows elsewhere.""" + + def __init__(self, point: str, verdict: Any) -> None: + self.point = point + self.verdict = verdict + + def intercept(self, context: dict[str, Any]) -> Any: + if context["interception_point"] == self.point: + verdict = self.verdict + return verdict(context) if callable(verdict) else verdict + return ALLOW + + +class CrashingGuard: + """Raises at one interception point (SDK synthesizes host_error:interceptor_failed).""" + + def __init__(self, point: str) -> None: + self.point = point + + def intercept(self, context: dict[str, Any]) -> Any: + if context["interception_point"] == self.point: + raise RuntimeError("guard crashed") + return ALLOW + + +@tool(approval_mode="never_require") +def weather_tool(location: str) -> str: + """Get the weather for a location.""" + weather_tool_calls.append(location) + return f"weather in {location}" + + +weather_tool_calls: list[str] = [] + + +@pytest.fixture(autouse=True) +def _reset_tool_calls() -> None: + weather_tool_calls.clear() + + +def tool_call_response(location: str = "Seattle", call_id: str = "call_1") -> ChatResponse: + return ChatResponse( + messages=[ + Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id=call_id, name="weather_tool", arguments=f'{{"location": "{location}"}}' + ) + ], + ) + ] + ) + + +def final_response(text: str = "Final response") -> ChatResponse: + return ChatResponse(messages=[Message(role="assistant", contents=[text])]) + + +def points(records: list[Any]) -> list[str]: + return [record.interception_point.value for record in records] + + +FULL_TOOL_RUN_POINTS = [ + "agent_startup", + "input", + "pre_model_call", + "post_model_call", + "pre_tool_call", + "post_tool_call", + "pre_model_call", + "post_model_call", + "output", + "agent_shutdown", +] + + +# endregion + +# region Factory validation + + +@requires_sdk +async def test_factory_requires_interceptors() -> None: + with pytest.raises(ValueError, match="at least one interceptor"): + agent_hooks_middleware([]) + with pytest.raises(ValueError, match="at least one interceptor"): + agent_hooks_middleware() + + +@requires_sdk +async def test_factory_requires_emitter_and_builder_together() -> None: + emitter = InterceptionEmitter() + with pytest.raises(ValueError, match="together"): + agent_hooks_middleware(emitter=emitter) + + +@requires_sdk +async def test_factory_rejects_per_run_config_with_host_emitter() -> None: + emitter = InterceptionEmitter().register(AllowGuard()) + builder = AgentContextBuilder(agent_id="a", framework="agent-framework", session_id="s") + with pytest.raises(ValueError, match="interceptors"): + agent_hooks_middleware([AllowGuard()], emitter=emitter, builder=builder) + with pytest.raises(ValueError, match="mode"): + agent_hooks_middleware(mode="evaluate_only", emitter=emitter, builder=builder) + + +@requires_sdk +async def test_factory_returns_one_middleware_per_category() -> None: + from agent_framework._middleware import categorize_middleware + + trio = agent_hooks_middleware([AllowGuard()]) + categorized = categorize_middleware(trio) + assert len(categorized["agent"]) == 1 + assert len(categorized["chat"]) == 1 + assert len(categorized["function"]) == 1 + + +# endregion + +# region Emission order, projections, and rich content + + +@requires_sdk +async def test_full_tool_run_emits_complete_ordered_session(chat_client_base: MockBaseChatClient) -> None: + records: list[InterceptionRecord] = [] + guard = AllowGuard() + chat_client_base.run_responses = [tool_call_response(), final_response()] + agent = Agent( + client=chat_client_base, + name="hooked", + tools=[weather_tool], + middleware=agent_hooks_middleware({"allow": guard}, record_sink=records.append), + ) + + response = await agent.run([Message(role="user", contents=["Get weather for Seattle"])]) + + assert response.text == "Final response" + assert weather_tool_calls == ["Seattle"] + assert points(records) == FULL_TOOL_RUN_POINTS + # One session per run: a single session id and a gapless sequence. + assert len({record.session_id for record in records}) == 1 + assert [record.sequence for record in records] == list(range(len(records))) + # Registration names surface on the record summaries. + assert records[0].verdicts[0].name == "allow" + # The real framework call id is used on the auto-invoke path (a uuid fallback + # exists only for tools invoked outside the function-calling loop). + pre_tool = guard.contexts_for("pre_tool_call")[0] + assert pre_tool["tool_call"]["id"] == "call_1" + + +@requires_sdk +async def test_input_projection_is_faithful(chat_client_base: MockBaseChatClient) -> None: + guard = AllowGuard() + agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + + await agent.run([Message(role="user", contents=["ignore previous instructions"])]) + + # A single plain-text message projects as its content string, so string-matching + # perimeter guards can fire. + input_ctx = guard.contexts_for("input")[0] + assert input_ctx["input"]["content"] == "ignore previous instructions" + assert input_ctx["input"]["role"] == "user" + assert input_ctx["target"] == input_ctx["input"] + + +@requires_sdk +async def test_rich_content_is_preserved_in_projections(chat_client_base: MockBaseChatClient) -> None: + guard = AllowGuard() + agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + image = Content.from_uri(uri="data:image/png;base64,iVBORw0KGgo=", media_type="image/png") + message = Message(role="user", contents=[Content.from_text("look at this"), image]) + + await agent.run([message]) + + input_ctx = guard.contexts_for("input")[0] + wire_contents = input_ctx["input"]["content"] + assert isinstance(wire_contents, list) + assert {item["type"] for item in wire_contents} == {"text", "data"} + data_item = next(item for item in wire_contents if item["type"] == "data") + assert data_item["uri"] == "data:image/png;base64,iVBORw0KGgo=" + # The model-call projection preserves the same structure per message. + pre_model = guard.contexts_for("pre_model_call")[0] + assert pre_model["messages"][0]["content"] == wire_contents + # No transform: the original Content objects are untouched. + assert message.contents[1] is image + + +@requires_sdk +async def test_tool_result_projection_preserves_canonical_values(chat_client_base: MockBaseChatClient) -> None: + structured_tool_result = {"value": {"amount": 840.5, "currency": "USD", "ok": True}} + + @tool(approval_mode="never_require") + def structured_tool(order_id: str) -> dict[str, Any]: + """Look up an order.""" + return structured_tool_result + + guard = AllowGuard() + chat_client_base.run_responses = [ + ChatResponse( + messages=[ + Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="c9", name="structured_tool", arguments='{"order_id": "1"}') + ], + ) + ] + ), + final_response(), + ] + agent = Agent( + client=chat_client_base, + tools=[structured_tool], + middleware=agent_hooks_middleware([guard]), + ) + + await agent.run("look up order 1") + + post_tool = guard.contexts_for("post_tool_call")[0] + # The default result parser wraps dict results as JSON text content; the wire value + # is the canonical JSON string the model sees — never a str(Content) repr. + value = post_tool["tool_result"]["value"] + assert "Content(" not in str(value) + assert "840.5" in str(value) + assert post_tool["target"] == value + + +# endregion + +# region Deny-before-execution + + +@requires_sdk +async def test_input_deny_blocks_run_before_model_call(chat_client_base: MockBaseChatClient) -> None: + records: list[InterceptionRecord] = [] + guard = PointGuard("input", Verdict.deny(reason="injection_blocked")) + agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard], record_sink=records.append)) + + with pytest.raises(InterceptionBlocked) as exc_info: + await agent.run("evil prompt") + + assert exc_info.value.result.verdict.reason == "injection_blocked" + assert chat_client_base.call_count == 0 + # §6.1a: the record trail is still closed with agent_shutdown. + assert points(records) == ["agent_startup", "input", "agent_shutdown"] + + +@requires_sdk +async def test_pre_model_call_deny_blocks_model_dispatch(chat_client_base: MockBaseChatClient) -> None: + guard = PointGuard("pre_model_call", Verdict.deny(reason="model_blocked")) + agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + + with pytest.raises(InterceptionBlocked): + await agent.run("hello") + + assert chat_client_base.call_count == 0 + + +@requires_sdk +async def test_post_model_call_deny_discards_response(chat_client_base: MockBaseChatClient) -> None: + guard = PointGuard("post_model_call", Verdict.deny(reason="response_blocked")) + agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + + with pytest.raises(InterceptionBlocked) as exc_info: + await agent.run("hello") + + assert exc_info.value.result.verdict.reason == "response_blocked" + assert chat_client_base.call_count == 1 # the action ran; its result was discarded + + +@requires_sdk +async def test_pre_tool_call_deny_blocks_tool_and_continues_loop(chat_client_base: MockBaseChatClient) -> None: + records: list[InterceptionRecord] = [] + guard = PointGuard("pre_tool_call", Verdict.deny(reason="tool_forbidden")) + chat_client_base.run_responses = [tool_call_response(), final_response("Understood.")] + agent = Agent( + client=chat_client_base, + tools=[weather_tool], + middleware=agent_hooks_middleware([guard], record_sink=records.append), + ) + + response = await agent.run("get the weather") + + # Deny before execution: the tool never ran, and no post_tool_call was emitted (§6.2). + assert weather_tool_calls == [] + assert "post_tool_call" not in points(records) + # A tool error surfaced to the model and the loop continued. + assert response.text == "Understood." + transcript = str([content.result for message in response.messages for content in message.contents]) + assert "tool_forbidden" in transcript + + +@requires_sdk +async def test_post_tool_call_deny_discards_result(chat_client_base: MockBaseChatClient) -> None: + guard = PointGuard("post_tool_call", Verdict.deny(reason="result_blocked")) + chat_client_base.run_responses = [tool_call_response(), final_response()] + agent = Agent(client=chat_client_base, tools=[weather_tool], middleware=agent_hooks_middleware([guard])) + + response = await agent.run("get the weather") + + assert weather_tool_calls == ["Seattle"] # the tool ran; its result must be discarded + transcript = str([content.result for message in response.messages for content in message.contents]) + assert "weather in Seattle" not in transcript + assert "result_blocked" in transcript + + +@requires_sdk +async def test_output_deny_blocks_response(chat_client_base: MockBaseChatClient) -> None: + guard = PointGuard("output", Verdict.deny(reason="egress_blocked")) + agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + + with pytest.raises(InterceptionBlocked) as exc_info: + await agent.run("hello") + + assert exc_info.value.result.verdict.reason == "egress_blocked" + + +# endregion + +# region Transform write-back + + +@requires_sdk +async def test_input_transform_writes_back_into_run_messages(chat_client_base: MockBaseChatClient) -> None: + guard = PointGuard( + "input", + Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target.content", value="[redacted]")), + ) + agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + message = Message(role="user", contents=["my SSN is 123-45-6789"]) + + response = await agent.run([message]) + + # The mock echoes the last request message, proving the model saw the redaction. + assert response.text == "test response - [redacted]" + # The caller-held Message object adopted the transform (shared history is redacted). + assert message.text == "[redacted]" + + +@requires_sdk +async def test_pre_model_call_transform_writes_back_into_request(chat_client_base: MockBaseChatClient) -> None: + guard = PointGuard( + "pre_model_call", + Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target[0].content", value="[masked]")), + ) + agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + + response = await agent.run("raw PII 123-45-6789") + + assert response.text == "test response - [masked]" + + +@requires_sdk +async def test_pre_tool_call_transform_writes_back_into_arguments(chat_client_base: MockBaseChatClient) -> None: + guard = AllowGuard() + transformer = PointGuard( + "pre_tool_call", + Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target.location", value="Redmond")), + ) + chat_client_base.run_responses = [tool_call_response("Seattle"), final_response()] + agent = Agent( + client=chat_client_base, + tools=[weather_tool], + middleware=agent_hooks_middleware([transformer, guard]), + ) + + await agent.run("get the weather") + + # The tool executed with exactly the approved arguments. + assert weather_tool_calls == ["Redmond"] + # §4.2: post_tool_call args reflect the post-transform arguments. + post_tool = guard.contexts_for("post_tool_call")[0] + assert post_tool["tool_call"]["args"] == {"location": "Redmond"} + assert post_tool["tool_result"]["value"] == "weather in Redmond" + + +@requires_sdk +async def test_post_tool_call_transform_writes_back_into_result(chat_client_base: MockBaseChatClient) -> None: + guard = PointGuard( + "post_tool_call", + Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target", value="[scrubbed]")), + ) + chat_client_base.run_responses = [tool_call_response(), final_response()] + agent = Agent(client=chat_client_base, tools=[weather_tool], middleware=agent_hooks_middleware([guard])) + + response = await agent.run("get the weather") + + results = [ + content for message in response.messages for content in message.contents if content.type == "function_result" + ] + assert len(results) == 1 + transcript = str(results[0].result) + assert "weather in Seattle" not in transcript + assert "[scrubbed]" in transcript + + +@requires_sdk +async def test_output_transform_writes_back_into_response(chat_client_base: MockBaseChatClient) -> None: + guard = PointGuard( + "output", + Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target.content", value="[card:redacted]")), + ) + agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + + response = await agent.run("what is my card number?") + + assert response.text == "[card:redacted]" + + +# endregion + +# region Streaming + + +@requires_sdk +async def test_streaming_buffers_until_all_verdicts_permit(chat_client_base: MockBaseChatClient) -> None: + 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", + ) + ], + [ + ChatResponseUpdate(contents=[Content.from_text("Final ")], role="assistant"), + ChatResponseUpdate(contents=[Content.from_text("answer")], role="assistant", finish_reason="stop"), + ], + ] + agent = Agent( + client=chat_client_base, + tools=[weather_tool], + middleware=agent_hooks_middleware([AllowGuard()], record_sink=records.append), + ) + + updates: list[AgentResponseUpdate] = [] + points_at_first_update: list[str] = [] + stream = agent.run("get the weather", stream=True) + async for update in stream: + if not updates: + points_at_first_update = points(records) + updates.append(update) + final = await stream.get_final_response() + + # Complete ordered session, and every emission (including output/shutdown) happened + # BEFORE the first update egressed: fail-closed buffered streaming. + assert points(records) == FULL_TOOL_RUN_POINTS + assert points_at_first_update == FULL_TOOL_RUN_POINTS + assert "".join(update.text for update in updates) == "Final answer" + assert final.text == "Final answer" + assert weather_tool_calls == ["Seattle"] + + +@requires_sdk +async def test_streaming_output_deny_releases_nothing(chat_client_base: MockBaseChatClient) -> None: + records: list[InterceptionRecord] = [] + guard = PointGuard("output", Verdict.deny(reason="egress_blocked")) + agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard], record_sink=records.append)) + + updates: list[AgentResponseUpdate] = [] + with pytest.raises(InterceptionBlocked): + async for update in agent.run("hello", stream=True): + updates.append(update) + + assert updates == [] # nothing egressed before the deny + assert points(records)[-1] == "agent_shutdown" # the record trail is closed + + +@requires_sdk +async def test_streaming_post_model_call_deny_releases_nothing(chat_client_base: MockBaseChatClient) -> None: + guard = PointGuard("post_model_call", Verdict.deny(reason="response_blocked")) + agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + + updates: list[AgentResponseUpdate] = [] + with pytest.raises(InterceptionBlocked): + async for update in agent.run("hello", stream=True): + updates.append(update) + + assert updates == [] + + +@requires_sdk +async def test_streaming_output_transform_rewrites_updates(chat_client_base: MockBaseChatClient) -> None: + guard = PointGuard( + "output", + Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target.content", value="[masked]")), + ) + agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + + updates: list[AgentResponseUpdate] = [] + stream = agent.run("hello", stream=True) + async for update in stream: + updates.append(update) + final = await stream.get_final_response() + + assert "".join(update.text for update in updates) == "[masked]" + assert final.text == "[masked]" + + +# endregion + +# region Error cleanup and fail-closed host errors + + +@requires_sdk +async def test_tool_exception_is_bracketed_with_error_post_tool_call(chat_client_base: MockBaseChatClient) -> None: + @tool(approval_mode="never_require") + def broken_tool(location: str) -> str: + """Always fails.""" + raise RuntimeError("boom with secret data") + + guard = AllowGuard() + chat_client_base.run_responses = [ + ChatResponse( + messages=[ + Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="c2", name="broken_tool", arguments='{"location": "x"}') + ], + ) + ] + ), + final_response("Sorry."), + ] + agent = Agent(client=chat_client_base, tools=[broken_tool], middleware=agent_hooks_middleware([guard])) + + response = await agent.run("run the broken tool") + + assert response.text == "Sorry." + post_tool = guard.contexts_for("post_tool_call")[0] + assert post_tool["tool_result"]["is_error"] is True + # Only the exception type name crosses the boundary (§6.3/§14). + assert post_tool["tool_result"]["value"] == "RuntimeError" + assert "secret" not in str(post_tool) + + +@requires_sdk +async def test_interceptor_crash_fails_closed_and_halts_run(chat_client_base: MockBaseChatClient) -> None: + records: list[InterceptionRecord] = [] + chat_client_base.run_responses = [tool_call_response(), final_response()] + agent = Agent( + client=chat_client_base, + tools=[weather_tool], + middleware=agent_hooks_middleware([CrashingGuard("post_tool_call")], record_sink=records.append), + ) + + with pytest.raises(InterceptionBlocked) as exc_info: + await agent.run("get the weather") + + assert exc_info.value.result.verdict.reason == "host_error:interceptor_failed" + # The tool ran, but the enforcement layer failed: the run halted fail-closed and + # the shutdown record still closed the trail. + assert points(records)[-1] == "agent_shutdown" + + +@requires_sdk +async def test_interceptor_crash_at_input_fails_closed(chat_client_base: MockBaseChatClient) -> None: + agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([CrashingGuard("input")])) + + with pytest.raises(InterceptionBlocked) as exc_info: + await agent.run("hello") + + assert exc_info.value.result.verdict.reason == "host_error:interceptor_failed" + assert chat_client_base.call_count == 0 + + +@requires_sdk +async def test_streaming_requires_no_partial_enforcement_on_error(chat_client_base: MockBaseChatClient) -> None: + records: list[InterceptionRecord] = [] + agent = Agent( + client=chat_client_base, + middleware=agent_hooks_middleware([CrashingGuard("output")], record_sink=records.append), + ) + + updates: list[AgentResponseUpdate] = [] + with pytest.raises(InterceptionBlocked): + async for update in agent.run("hello", stream=True): + updates.append(update) + + assert updates == [] + assert points(records)[-1] == "agent_shutdown" + + +# endregion + +# region Concurrency isolation + + +@requires_sdk +async def test_concurrent_runs_are_isolated(chat_client_base: MockBaseChatClient) -> None: + records: list[InterceptionRecord] = [] + agent = Agent( + client=chat_client_base, middleware=agent_hooks_middleware([AllowGuard()], record_sink=records.append) + ) + + await asyncio.gather(agent.run("first"), agent.run("second")) + + sessions: dict[str, list[InterceptionRecord]] = {} + for record in records: + sessions.setdefault(record.session_id, []).append(record) + # Two runs, two independent sessions with complete, gapless sequences each. + assert len(sessions) == 2 + for session_records in sessions.values(): + assert points(session_records) == [ + "agent_startup", + "input", + "pre_model_call", + "post_model_call", + "output", + "agent_shutdown", + ] + assert [record.sequence for record in session_records] == list(range(len(session_records))) + + +# endregion + +# region Session scoping and modes + + +@requires_sdk +async def test_host_owned_session_spans_runs(chat_client_base: MockBaseChatClient) -> None: + records: list[InterceptionRecord] = [] + emitter = InterceptionEmitter() + emitter.register(AllowGuard()) + emitter.set_record_sink(records.append) + builder = AgentContextBuilder(agent_id="host-agent", framework="agent-framework", session_id="session-1") + agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware(emitter=emitter, builder=builder)) + + # The host owns the session boundaries. + await emitter.emit(builder.agent_startup(tools_registered=[])) + await agent.run("first turn") + await agent.run("second turn") + await emitter.emit(builder.agent_shutdown(reason="completed")) + + assert points(records) == [ + "agent_startup", + "input", + "pre_model_call", + "post_model_call", + "output", + "input", + "pre_model_call", + "post_model_call", + "output", + "agent_shutdown", + ] + # One session: a single id and one monotonically increasing sequence across runs. + assert {record.session_id for record in records} == {"session-1"} + assert [record.sequence for record in records] == list(range(len(records))) + + +@requires_sdk +async def test_liftable_deny_is_resolved_through_the_approval_seam(chat_client_base: MockBaseChatClient) -> None: + from agent_hooks import ApprovalOutcome, ApprovalRequest, ApprovalResolution + + class ApproveAll: + def resolve(self, request: ApprovalRequest) -> ApprovalResolution: + return ApprovalResolution( + outcome=ApprovalOutcome.APPROVE, + context_identity=request.context_identity, + verdict=ALLOW, + ) + + records: list[InterceptionRecord] = [] + guard = PointGuard("pre_tool_call", Verdict.escalate(reason="needs_approval")) + chat_client_base.run_responses = [tool_call_response(), final_response()] + agent = Agent( + client=chat_client_base, + tools=[weather_tool], + middleware=agent_hooks_middleware([guard], resolver=ApproveAll(), record_sink=records.append), + ) + + response = await agent.run("get the weather") + + assert response.text == "Final response" + assert weather_tool_calls == ["Seattle"] # the lifted deny let the tool run + lifted = [record for record in records if record.resolved_by == "approval"] + assert len(lifted) == 1 + assert lifted[0].interception_point.value == "pre_tool_call" + + +@requires_sdk +async def test_evaluate_only_records_but_never_blocks(chat_client_base: MockBaseChatClient) -> None: + records: list[InterceptionRecord] = [] + guard = PointGuard("pre_tool_call", Verdict.deny(reason="would_block")) + chat_client_base.run_responses = [tool_call_response(), final_response()] + agent = Agent( + client=chat_client_base, + tools=[weather_tool], + middleware=agent_hooks_middleware([guard], mode="evaluate_only", record_sink=records.append), + ) + + response = await agent.run("get the weather") + + assert response.text == "Final response" + assert weather_tool_calls == ["Seattle"] # evaluate_only never blocks + deny_records = [record for record in records if record.verdict.decision.value == "deny"] + assert len(deny_records) == 1 + assert deny_records[0].mode.value == "evaluate_only" + + +@requires_sdk +async def test_partial_trio_fails_closed(chat_client_base: MockBaseChatClient) -> None: + # Installing only part of the trio (here: on a bare chat client, where the agent + # middleware never runs) must fail closed, not silently skip enforcement. + trio = agent_hooks_middleware([AllowGuard()]) + chat_client_base.chat_middleware = cast( + "list[Any]", + [middleware for middleware in trio if type(middleware).__name__.endswith("ChatMiddleware")], + ) + + with pytest.raises(MiddlewareException, match="installed together"): + await chat_client_base.get_response([Message(role="user", contents=["hi"])]) + + +# endregion + +# region Middleware short-circuits (MiddlewareTermination) are guarded + + +class AgentShortCircuit(AgentMiddleware): + """Framework-documented short-circuit pattern: substitute a result and terminate.""" + + def __init__(self, result: Any) -> None: + self._result = result + + async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + context.result = self._result + raise MiddlewareTermination("cached") + + +class ChatShortCircuit(ChatMiddleware): + def __init__(self, result: Any) -> None: + self._result = result + + async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None: + context.result = self._result + raise MiddlewareTermination("cached") + + +class FunctionShortCircuit(FunctionMiddleware): + def __init__(self, result: Any) -> None: + self._result = result + + async def process(self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]) -> None: + context.result = self._result + raise MiddlewareTermination("cached tool result") + + +def _cached_agent_stream(text: str) -> "ResponseStream[AgentResponseUpdate, AgentResponse[Any]]": + async def _updates() -> AsyncIterable[AgentResponseUpdate]: + yield AgentResponseUpdate(contents=[Content.from_text(text)], role="assistant") + + return ResponseStream(_updates(), finalizer=AgentResponse.from_updates) + + +@requires_sdk +async def test_agent_seam_short_circuit_result_is_guarded(chat_client_base: MockBaseChatClient) -> None: + records: list[InterceptionRecord] = [] + substituted = AgentResponse(messages=[Message(role="assistant", contents=["cached payload"])]) + agent = Agent( + client=chat_client_base, + middleware=[ + *agent_hooks_middleware([AllowGuard()], record_sink=records.append), + AgentShortCircuit(substituted), + ], + ) + + response = await agent.run("hello") + + assert response.text == "cached payload" + assert chat_client_base.call_count == 0 + # The substituted result egressed, so it passed the output point. + assert points(records) == ["agent_startup", "input", "output", "agent_shutdown"] + + +@requires_sdk +async def test_agent_seam_short_circuit_result_can_be_denied(chat_client_base: MockBaseChatClient) -> None: + guard = PointGuard("output", Verdict.deny(reason="egress_blocked")) + substituted = AgentResponse(messages=[Message(role="assistant", contents=["cached payload"])]) + agent = Agent( + client=chat_client_base, + middleware=[*agent_hooks_middleware([guard]), AgentShortCircuit(substituted)], + ) + + with pytest.raises(InterceptionBlocked) as exc_info: + await agent.run("hello") + + assert exc_info.value.result.verdict.reason == "egress_blocked" + + +@requires_sdk +async def test_chat_seam_short_circuit_result_is_guarded(chat_client_base: MockBaseChatClient) -> None: + records: list[InterceptionRecord] = [] + substituted = ChatResponse(messages=[Message(role="assistant", contents=["cached model reply"])]) + agent = Agent( + client=chat_client_base, + middleware=[*agent_hooks_middleware([AllowGuard()], record_sink=records.append), ChatShortCircuit(substituted)], + ) + + response = await agent.run("hello") + + assert response.text == "cached model reply" + assert chat_client_base.call_count == 0 + # The substituted model reply still passed post_model_call (and output). + assert points(records) == [ + "agent_startup", + "input", + "pre_model_call", + "post_model_call", + "output", + "agent_shutdown", + ] + + +@requires_sdk +async def test_chat_seam_short_circuit_result_can_be_denied(chat_client_base: MockBaseChatClient) -> None: + guard = PointGuard("post_model_call", Verdict.deny(reason="response_blocked")) + substituted = ChatResponse(messages=[Message(role="assistant", contents=["cached model reply"])]) + agent = Agent( + client=chat_client_base, + middleware=[*agent_hooks_middleware([guard]), ChatShortCircuit(substituted)], + ) + + with pytest.raises(InterceptionBlocked) as exc_info: + await agent.run("hello") + + assert exc_info.value.result.verdict.reason == "response_blocked" + assert chat_client_base.call_count == 0 + + +@requires_sdk +async def test_streaming_short_circuit_stream_is_guarded(chat_client_base: MockBaseChatClient) -> None: + records: list[InterceptionRecord] = [] + agent = Agent( + client=chat_client_base, + middleware=[ + *agent_hooks_middleware([AllowGuard()], record_sink=records.append), + AgentShortCircuit(_cached_agent_stream("cached stream")), + ], + ) + + updates: list[AgentResponseUpdate] = [] + async for update in agent.run("hello", stream=True): + updates.append(update) + + assert "".join(update.text for update in updates) == "cached stream" + assert chat_client_base.call_count == 0 + assert points(records) == ["agent_startup", "input", "output", "agent_shutdown"] + + +@requires_sdk +async def test_streaming_short_circuit_deny_releases_nothing(chat_client_base: MockBaseChatClient) -> None: + guard = PointGuard("output", Verdict.deny(reason="egress_blocked")) + agent = Agent( + client=chat_client_base, + middleware=[*agent_hooks_middleware([guard]), AgentShortCircuit(_cached_agent_stream("cached stream"))], + ) + + updates: list[AgentResponseUpdate] = [] + with pytest.raises(InterceptionBlocked): + async for update in agent.run("hello", stream=True): + updates.append(update) + + assert updates == [] + + +@requires_sdk +async def test_streaming_short_circuit_without_result_closes_trail(chat_client_base: MockBaseChatClient) -> None: + records: list[InterceptionRecord] = [] + agent = Agent( + client=chat_client_base, + middleware=[*agent_hooks_middleware([AllowGuard()], record_sink=records.append), AgentShortCircuit(None)], + ) + + updates: list[AgentResponseUpdate] = [] + async for update in agent.run("hello", stream=True): + updates.append(update) + + assert updates == [] # nothing egressed + assert points(records) == ["agent_startup", "input", "agent_shutdown"] + + +@requires_sdk +async def test_function_seam_foreign_termination_is_bracketed(chat_client_base: MockBaseChatClient) -> None: + records: list[InterceptionRecord] = [] + guard = AllowGuard() + chat_client_base.run_responses = [tool_call_response(), final_response()] + agent = Agent( + client=chat_client_base, + tools=[weather_tool], + middleware=[ + *agent_hooks_middleware({"allow": guard}, record_sink=records.append), + FunctionShortCircuit({"substituted": "tool result"}), + ], + ) + + response = await agent.run("get the weather") + + assert weather_tool_calls == [] # the terminator pre-empted the tool + # The substituted result entered the transcript, so it was bracketed. + post_tool = guard.contexts_for("post_tool_call")[0] + assert post_tool["tool_result"]["value"] == {"substituted": "tool result"} + assert "post_tool_call" in points(records) + transcript = str([content.result for message in response.messages for content in message.contents]) + assert "substituted" in transcript + + +@requires_sdk +async def test_function_seam_foreign_termination_result_can_be_denied(chat_client_base: MockBaseChatClient) -> None: + guard = PointGuard("post_tool_call", Verdict.deny(reason="result_blocked")) + chat_client_base.run_responses = [tool_call_response(), final_response()] + agent = Agent( + client=chat_client_base, + tools=[weather_tool], + middleware=[*agent_hooks_middleware([guard]), FunctionShortCircuit({"substituted": "tool result"})], + ) + + response = await agent.run("get the weather") + + transcript = str([content.result for message in response.messages for content in message.contents]) + assert "substituted" not in transcript # the denied substitution never entered the transcript + assert "result_blocked" in transcript + + +# endregion + +# region Fail-closed write-back and enforcement failures + + +@requires_sdk +async def test_pre_tool_call_transform_to_non_object_fails_closed(chat_client_base: MockBaseChatClient) -> None: + guard = PointGuard( + "pre_tool_call", + Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target", value="oops")), + ) + chat_client_base.run_responses = [tool_call_response(), final_response()] + agent = Agent(client=chat_client_base, tools=[weather_tool], middleware=agent_hooks_middleware([guard])) + + with pytest.raises(MiddlewareException, match="arguments object"): + await agent.run("get the weather") + + assert weather_tool_calls == [] # the tool never ran with unapproved arguments + + +@requires_sdk +async def test_input_role_transform_is_written_back(chat_client_base: MockBaseChatClient) -> None: + guard = PointGuard( + "input", + Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target.role", value="system")), + ) + agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + message = Message(role="user", contents=["hello"]) + + await agent.run([message]) + + assert message.role == "system" + + +@requires_sdk +async def test_multi_message_input_role_transform_fails_closed(chat_client_base: MockBaseChatClient) -> None: + guard = PointGuard( + "input", + Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target.role", value="system")), + ) + agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + + with pytest.raises(MiddlewareException, match="cannot be written back"): + await agent.run([Message(role="user", contents=["one"]), Message(role="user", contents=["two"])]) + + assert chat_client_base.call_count == 0 + + +@requires_sdk +async def test_non_string_finish_reason_transform_fails_closed(chat_client_base: MockBaseChatClient) -> None: + guard = PointGuard( + "post_model_call", + Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target.finish_reason", value=42)), + ) + agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + + with pytest.raises(MiddlewareException, match="finish_reason a string"): + await agent.run("hello") + + +@requires_sdk +async def test_enforcement_failure_at_function_seam_halts_run(chat_client_base: MockBaseChatClient) -> None: + class Unprojectable: + def __str__(self) -> str: + raise RuntimeError("no projection") + + @tool(approval_mode="never_require", result_parser=SKIP_PARSING) + def opaque_tool(location: str) -> Any: + """Return a value the enforcement layer cannot project.""" + return Unprojectable() + + records: list[InterceptionRecord] = [] + chat_client_base.run_responses = [ + ChatResponse( + messages=[ + Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="c7", name="opaque_tool", arguments='{"location": "x"}') + ], + ) + ] + ), + final_response(), + ] + agent = Agent( + client=chat_client_base, + tools=[opaque_tool], + middleware=agent_hooks_middleware([AllowGuard()], record_sink=records.append), + ) + + # An unexpected failure inside the enforcement layer must halt the run, not + # degrade into a tool error that lets the run continue unaudited. + with pytest.raises(MiddlewareException, match="post_tool_call enforcement failed"): + await agent.run("run the opaque tool") + + assert points(records)[-1] == "agent_shutdown" # the record trail is closed + + +# endregion + +# region Partial installs fail closed + + +@requires_sdk +async def test_agent_seam_verifies_sibling_middleware(chat_client_base: MockBaseChatClient) -> None: + trio = agent_hooks_middleware([AllowGuard()]) + + # agent + chat, function middleware missing + agent = Agent(client=chat_client_base, middleware=[trio[0], trio[1]]) + with pytest.raises(MiddlewareException, match="installed together"): + await agent.run("hello") + assert chat_client_base.call_count == 0 + + # agent + function, chat middleware missing + agent = Agent(client=chat_client_base, middleware=[trio[0], trio[2]]) + with pytest.raises(MiddlewareException, match="installed together"): + await agent.run("hello") + assert chat_client_base.call_count == 0 + + # agent middleware alone + agent = Agent(client=chat_client_base, middleware=[trio[0]]) + with pytest.raises(MiddlewareException, match="installed together"): + await agent.run("hello") + assert chat_client_base.call_count == 0 + + +@requires_sdk +async def test_function_seam_without_run_state_blocks_tool(chat_client_base: MockBaseChatClient) -> None: + trio = agent_hooks_middleware([AllowGuard()]) + chat_client_base.run_responses = [tool_call_response(), final_response()] + agent = Agent(client=chat_client_base, tools=[weather_tool], middleware=[trio[2]]) + + response = await agent.run("get the weather") + + # The tool is never dispatched and the loop terminates 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 "installed together" in transcript + + +@requires_sdk +async def test_fresh_trio_per_run_is_not_conflated_by_pipeline_caching(chat_client_base: MockBaseChatClient) -> None: + # The agent middleware pipeline is cached and compared with ==; the trio must keep + # identity semantics so a field-equal fresh trio on the next run is not conflated + # with the cached one (which would trip the sibling verification on run 2). + guard = AllowGuard() + agent = Agent(client=chat_client_base) + + first = await agent.run("one", middleware=agent_hooks_middleware([guard])) + second = await agent.run("two", middleware=agent_hooks_middleware([guard])) + + assert first.text == "test response - one" + assert second.text == "test response - two" + # Identity semantics: fresh trios never compare equal, and the objects stay hashable. + trio_a = agent_hooks_middleware([guard]) + trio_b = agent_hooks_middleware([guard]) + assert all(a != b for a, b in zip(trio_a, trio_b)) + assert all(isinstance(hash(middleware), int) for middleware in trio_a) + + +@requires_sdk +async def test_stacked_trios_fail_loudly(chat_client_base: MockBaseChatClient) -> None: + records_a: list[InterceptionRecord] = [] + records_b: list[InterceptionRecord] = [] + agent = Agent( + client=chat_client_base, + middleware=[ + *agent_hooks_middleware([AllowGuard()], record_sink=records_a.append), + *agent_hooks_middleware([AllowGuard()], record_sink=records_b.append), + ], + ) + + # Two trios on one agent would silently bind half the seams to the wrong emitter; + # that must be a loud failure, never a silent partial enforcement. + with pytest.raises(MiddlewareException, match="owned by a different"): + await agent.run("hello") + + assert chat_client_base.call_count == 0 + # Neither trio recorded any model/tool points (nothing misbound before the halt), + # and both trails were closed. + for records in (records_a, records_b): + assert set(points(records)) <= {"agent_startup", "input", "agent_shutdown"} + assert points(records)[-1] == "agent_shutdown" + + +@requires_sdk +async def test_nested_agents_with_their_own_trios_stay_isolated(chat_client_base: MockBaseChatClient) -> None: + records_outer: list[InterceptionRecord] = [] + records_inner: list[InterceptionRecord] = [] + inner_client = MockBaseChatClient() + inner_agent = Agent( + client=inner_client, + name="inner", + middleware=agent_hooks_middleware([AllowGuard()], record_sink=records_inner.append), + ) + + @tool(approval_mode="never_require") + async def ask_inner(question: str) -> str: + """Delegate a question to the inner agent.""" + response = await inner_agent.run(question) + return response.text + + chat_client_base.run_responses = [ + ChatResponse( + messages=[ + Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="c8", name="ask_inner", arguments='{"question": "hi"}') + ], + ) + ] + ), + final_response(), + ] + outer_agent = Agent( + client=chat_client_base, + name="outer", + tools=[ask_inner], + middleware=agent_hooks_middleware([AllowGuard()], record_sink=records_outer.append), + ) + + response = await outer_agent.run("go ask the inner agent") + + # Each agent's own trio guards its own run: the inner run (executed inside the + # outer tool call) binds to the inner emitter and restores the outer state after. + assert response.text == "Final response" + assert points(records_outer) == FULL_TOOL_RUN_POINTS + assert points(records_inner) == [ + "agent_startup", + "input", + "pre_model_call", + "post_model_call", + "output", + "agent_shutdown", + ] + assert {record.session_id for record in records_outer}.isdisjoint(record.session_id for record in records_inner) + + +# endregion + +# region Streaming setup failures and unguardable results + + +@requires_sdk +async def test_streaming_setup_failure_still_closes_trail(chat_client_base: MockBaseChatClient) -> None: + class Boom(AgentMiddleware): + async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + raise RuntimeError("boom") + + records: list[InterceptionRecord] = [] + agent = Agent( + client=chat_client_base, + middleware=[*agent_hooks_middleware([AllowGuard()], record_sink=records.append), Boom()], + ) + + updates: list[AgentResponseUpdate] = [] + with pytest.raises(RuntimeError, match="boom"): + async for update in agent.run("hello", stream=True): + updates.append(update) + + assert updates == [] + assert points(records) == ["agent_startup", "input", "agent_shutdown"] + + +@requires_sdk +async def test_unguardable_run_result_fails_closed(chat_client_base: MockBaseChatClient) -> None: + class BadResult(AgentMiddleware): + async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + await call_next() + context.result = cast(Any, "plain string") + + agent = Agent(client=chat_client_base, middleware=[*agent_hooks_middleware([AllowGuard()]), BadResult()]) + + with pytest.raises(MiddlewareException, match="cannot guard a run result"): + await agent.run("hello") + + +@requires_sdk +async def test_unguardable_chat_result_fails_closed(chat_client_base: MockBaseChatClient) -> None: + agent = Agent( + client=chat_client_base, + middleware=[*agent_hooks_middleware([AllowGuard()]), ChatShortCircuit(cast(Any, "plain string"))], + ) + + with pytest.raises(MiddlewareException, match="cannot guard a chat result"): + await agent.run("hello") + + +# endregion + +# region Optional dependency + + +def test_agent_hooks_middleware_importable_without_sdk(monkeypatch: pytest.MonkeyPatch) -> None: + import builtins + import sys + + import agent_framework._agent_hooks as agent_hooks_module + + real_import = builtins.__import__ + + def _import_without_agent_hooks( + name: str, + globals_: dict[str, object] | None = None, + locals_: dict[str, object] | None = None, + fromlist: tuple[str, ...] = (), + level: int = 0, + ) -> object: + if name == "agent_hooks" or name.startswith("agent_hooks."): + raise ModuleNotFoundError(f"No module named '{name}'", name="agent_hooks") + return real_import(name, globals_, locals_, fromlist, level) + + for module_name in list(sys.modules): + if module_name == "agent_hooks" or module_name.startswith("agent_hooks."): + monkeypatch.delitem(sys.modules, module_name) + monkeypatch.setattr(builtins, "__import__", _import_without_agent_hooks) + + # The lazy root export and the module itself stay importable without the SDK. + assert agent_framework.agent_hooks_middleware is agent_hooks_module.agent_hooks_middleware + + with pytest.raises(ModuleNotFoundError, match=r"agent-framework-core\[agent-hooks\]"): + agent_framework.agent_hooks_middleware([cast("Any", object())]) + + +# endregion diff --git a/python/uv.lock b/python/uv.lock index f810d2aa56d..91733357ef0 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -419,6 +419,9 @@ dependencies = [ ] [package.optional-dependencies] +agent-hooks = [ + { name = "agent-hooks-sdk" }, +] all = [ { name = "agent-framework-a2a" }, { name = "agent-framework-ag-ui" }, @@ -490,6 +493,7 @@ requires-dist = [ { name = "agent-framework-purview", marker = "extra == 'all'", editable = "packages/purview" }, { name = "agent-framework-redis", marker = "extra == 'all'", editable = "packages/redis" }, { name = "agent-framework-tools", marker = "extra == 'all'", editable = "packages/tools" }, + { name = "agent-hooks-sdk", marker = "extra == 'agent-hooks'", specifier = ">=0.1.0a4,<0.2" }, { name = "mcp", marker = "extra == 'all'", specifier = ">=1.24.0,<2" }, { name = "msgspec", specifier = ">=0.20.0,<0.22" }, { name = "opentelemetry-api", specifier = ">=1.39.0,<2" }, @@ -497,7 +501,7 @@ requires-dist = [ { name = "python-dotenv", specifier = ">=1,<2" }, { name = "typing-extensions", specifier = ">=4.15.0,<5" }, ] -provides-extras = ["all"] +provides-extras = ["agent-hooks", "all"] [package.metadata.requires-dev] dev = [{ name = "azure-ai-agentserver-core", specifier = ">=2.0.0b7,<3" }] @@ -1004,6 +1008,15 @@ requires-dist = [ { name = "psutil", specifier = ">=5.9" }, ] +[[package]] +name = "agent-hooks-sdk" +version = "0.1.0a4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c4/9b/07e9e54941ce60991ce8ba00fe9dc38ea96d885b38a86cd944a409cd7772/agent_hooks_sdk-0.1.0a4.tar.gz", hash = "sha256:e3affd44d7779b8254fd5142f7d41648e48d3d4646acc94b07a28d13c5c07b1d", size = 173745, upload-time = "2026-07-30T02:11:22.235Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d6/07/dc7f4e7a53c208ba35fcf0bd0208cf77a1e84ffd7a42d2a61c6e9faf274b/agent_hooks_sdk-0.1.0a4-cp310-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:86f5ff8e0c1201a61b113026c688fc98ea590b10f8bbc5ccfa6ac289fdec037d", size = 634247, upload-time = "2026-07-30T02:11:20.497Z" }, +] + [[package]] name = "agentlightning" version = "0.3.0" From 2cd37fcc66ee7f2aa29474e680ed5ea4bf172897 Mon Sep 17 00:00:00 2001 From: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> Date: Tue, 4 Aug 2026 19:09:56 +0000 Subject: [PATCH 2/8] style(python): unquote ResponseStream annotation per pyupgrade The pre-commit pyupgrade hook rewrites the quoted forward reference; ResponseStream is imported at runtime in this module, so the quotes were unnecessary. Signed-off-by: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> --- python/packages/core/tests/core/test_agent_hooks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/packages/core/tests/core/test_agent_hooks.py b/python/packages/core/tests/core/test_agent_hooks.py index 49e13225e67..2babc5f7bb4 100644 --- a/python/packages/core/tests/core/test_agent_hooks.py +++ b/python/packages/core/tests/core/test_agent_hooks.py @@ -828,7 +828,7 @@ async def process(self, context: FunctionInvocationContext, call_next: Callable[ raise MiddlewareTermination("cached tool result") -def _cached_agent_stream(text: str) -> "ResponseStream[AgentResponseUpdate, AgentResponse[Any]]": +def _cached_agent_stream(text: str) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: async def _updates() -> AsyncIterable[AgentResponseUpdate]: yield AgentResponseUpdate(contents=[Content.from_text(text)], role="assistant") From a80b5f3b2e8112a70998824bf1fd7228eabe23a4 Mon Sep 17 00:00:00 2001 From: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> Date: Wed, 5 Aug 2026 06:10:38 +0000 Subject: [PATCH 3/8] refactor(python): address agent-hooks review feedback Reworks the agent-hooks feature per PR review: - Verdicts now precede durability: a run-scoped persistence gate (_sessions.py) defers per-service-call history persistence and after-run provider work until the covering post_model_call/output verdict permits; denied content never persists, transforms persist post-write-back. Unhooked runs are unchanged (verified against an instrumented baseline). - ResponseStream.buffered_and_gated: a buffered-gate combinator that applies the run's pending stream hooks before the gate, then seals the stream, so no middleware can rewrite egress after the output verdict. Replaces the hand-rolled replay iterator. - MiddlewareBundle (public, _middleware.py): the factory returns an indivisible bundle categorize_middleware splits, making partial installs impossible by construction; members are validated at construction. Bare (non-sequence) middleware at agent construction is now normalized instead of silently dropped, and unrecognized middleware logs a warning instead of vanishing. - Factory split and rename: create_agent_hooks_middleware (per-run sessions) and create_agent_hooks_middleware_from_emitter (host-owned); the sentinel parameter-diffing is gone. - Wire conversions live in per-point codec classes owning to_wire and write_back. Fixes in that code: tool-call name transforms apply or raise; non-object args transforms raise; argument write-back merges only changed keys (original values, including bytes, preserved by identity); message-list write-back matches by identity, not index. - function_approval_request objects on the normal return path pass through un-emitted, preserving the human approval pause. - Hosted (service-executed) tool calls surface in the post_model_call content projection; the tool-seam limitation is documented. - Import probe covers the full SDK surface and re-raises as missing-extra only for the agent_hooks module; module logger added; _json_safe replaced by make_json_safe (which gained bytes support); tools_registered uses normalize_tools; dependency-pyright analyzes the module again via the test dependency-group. - Tests: 75 in the feature suite (persistence gating, stream-hook sealing, approval passthrough, codec units, bundle validation, bare-bundle installs), full core suite green. Signed-off-by: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> --- python/PACKAGE_STATUS.md | 10 +- .../packages/core/agent_framework/__init__.py | 7 +- .../core/agent_framework/__init__.pyi | 7 +- .../core/agent_framework/_agent_hooks.py | 1437 +++++++++-------- .../packages/core/agent_framework/_agents.py | 15 + .../core/agent_framework/_middleware.py | 109 +- .../core/agent_framework/_serialization.py | 8 +- .../core/agent_framework/_sessions.py | 34 +- .../packages/core/agent_framework/_types.py | 135 ++ .../core/pyrightconfig.dependency.json | 1 - .../core/tests/core/test_agent_hooks.py | 789 +++++++-- .../tests/core/test_middleware_with_agent.py | 22 + python/pyproject.toml | 3 + python/uv.lock | 7 + 14 files changed, 1792 insertions(+), 792 deletions(-) diff --git a/python/PACKAGE_STATUS.md b/python/PACKAGE_STATUS.md index 98ff7e5c75b..ebb16c62ff0 100644 --- a/python/PACKAGE_STATUS.md +++ b/python/PACKAGE_STATUS.md @@ -66,10 +66,14 @@ listed below. #### `AGENT_HOOKS` -- `agent-framework-core`: `agent_hooks_middleware` from `agent_framework/_agent_hooks.py`, - the AGENT-HOOKS-0.1 enforcement middleware trio. Requires the opt-in +- `agent-framework-core`: `create_agent_hooks_middleware` and + `create_agent_hooks_middleware_from_emitter` from `agent_framework/_agent_hooks.py`, + the AGENT-HOOKS-0.1 enforcement middleware bundle. Requires the opt-in `agent-framework-core[agent-hooks]` extra (`agent-hooks-sdk`), which is deliberately - not part of `agent-framework-core[all]`. + not part of `agent-framework-core[all]`. Known limitation: service-side (hosted) tool + execution never passes through the framework's function-invocation seam, so the + `pre_tool_call`/`post_tool_call` points cannot intercept it; hosted tool calls and + outputs are surfaced in the `post_model_call` content projection instead. #### `DECLARATIVE_AGENTS` diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index 055fc6fa28c..d7ab9168755 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -41,7 +41,7 @@ ) _LAZY_MODULE_EXPORTS: Final[Mapping[str, tuple[str, ...]]] = { - "._agent_hooks": ("agent_hooks_middleware",), + "._agent_hooks": ("create_agent_hooks_middleware", "create_agent_hooks_middleware_from_emitter"), "._agents": ("Agent", "BaseAgent", "RawAgent", "SupportsAgentRun"), "._clients": ( "BaseChatClient", @@ -184,6 +184,7 @@ "FunctionInvocationContext", "FunctionMiddleware", "FunctionMiddlewareTypes", + "MiddlewareBundle", "MiddlewareTermination", "MiddlewareType", "MiddlewareTypes", @@ -504,6 +505,7 @@ "MemoryTopicRecord", "Message", "MessageInjectionMiddleware", + "MiddlewareBundle", "MiddlewareException", "MiddlewareTermination", "MiddlewareType", @@ -594,13 +596,14 @@ "WorkflowViz", "__version__", "add_usage_details", - "agent_hooks_middleware", "agent_middleware", "annotate_message_groups", "apply_compaction", "background_tasks_running", "background_tasks_running_message", "chat_middleware", + "create_agent_hooks_middleware", + "create_agent_hooks_middleware_from_emitter", "create_always_approve_tool_response", "create_always_approve_tool_with_arguments_response", "create_edge_runner", diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi index d9662f7ff74..f72269c02d0 100644 --- a/python/packages/core/agent_framework/__init__.pyi +++ b/python/packages/core/agent_framework/__init__.pyi @@ -4,7 +4,7 @@ from typing import Final __version__: Final[str] -from ._agent_hooks import agent_hooks_middleware +from ._agent_hooks import create_agent_hooks_middleware, create_agent_hooks_middleware_from_emitter from ._agents import Agent, BaseAgent, RawAgent, SupportsAgentRun from ._clients import ( BaseChatClient, @@ -143,6 +143,7 @@ from ._middleware import ( FunctionInvocationContext, FunctionMiddleware, FunctionMiddlewareTypes, + MiddlewareBundle, MiddlewareTermination, MiddlewareType, MiddlewareTypes, @@ -468,6 +469,7 @@ __all__ = [ "MemoryTopicRecord", "Message", "MessageInjectionMiddleware", + "MiddlewareBundle", "MiddlewareException", "MiddlewareTermination", "MiddlewareType", @@ -558,13 +560,14 @@ __all__ = [ "WorkflowViz", "__version__", "add_usage_details", - "agent_hooks_middleware", "agent_middleware", "annotate_message_groups", "apply_compaction", "background_tasks_running", "background_tasks_running_message", "chat_middleware", + "create_agent_hooks_middleware", + "create_agent_hooks_middleware_from_emitter", "create_always_approve_tool_response", "create_always_approve_tool_with_arguments_response", "create_edge_runner", diff --git a/python/packages/core/agent_framework/_agent_hooks.py b/python/packages/core/agent_framework/_agent_hooks.py index 6e7ede73c5d..2a61031c36e 100644 --- a/python/packages/core/agent_framework/_agent_hooks.py +++ b/python/packages/core/agent_framework/_agent_hooks.py @@ -9,13 +9,15 @@ - ``pre_model_call`` / ``post_model_call`` ride the chat seam, - ``pre_tool_call`` / ``post_tool_call`` ride the function seam. -The single public entry point is :func:`agent_hooks_middleware`, which returns the three -middleware objects as one unit. The implementations are deliberately private: installing -only part of the trio would enforce only part of the control contract. Install exactly -one trio per agent, placed first in the middleware list: middleware listed before the -trio runs outside the enforcement boundary (outer position is outer trust) — e.g. a -function middleware placed before the trio can substitute a tool result that the tool -seam never brackets; the final ``output`` point still guards whatever egresses. +The public entry points are :func:`create_agent_hooks_middleware` (one agent-hooks +session per run) and :func:`create_agent_hooks_middleware_from_emitter` (host-owned +session). Both return a :class:`~agent_framework.MiddlewareBundle`: the three middleware +implementations are deliberately private and travel as one indivisible unit, because +installing only part of them would enforce only part of the control contract. Install +exactly one bundle per agent, placed first in the middleware list: middleware listed +before the bundle runs outside the enforcement boundary (outer position is outer trust) +— e.g. a function middleware placed before the bundle can substitute a tool result that +the tool seam never brackets; the final ``output`` point still guards whatever egresses. Enforcement semantics (``mode="enforce"``): @@ -24,9 +26,10 @@ (interceptor crash/timeout, invalid context) synthesize ``host_error:*`` denies and are treated as blocks — the feature never fails open. - ``transform`` verdicts are written back into the native middleware contexts - (``messages`` / ``arguments`` / ``result``), so the framework executes exactly the - value the interceptors approved. Content objects are preserved: rich (non-text) - message content is projected as content dictionaries, never flattened to text. + (``messages`` / ``arguments`` / ``results``) through per-point codecs, so the framework + executes exactly the value the interceptors approved. Content objects are preserved: + rich (non-text) message content is projected as content dictionaries, never flattened + to text. - A ``deny`` at ``input``, ``pre_model_call``, ``post_model_call``, or ``output`` terminates the run: :class:`agent_hooks.InterceptionBlocked` propagates to the caller of :meth:`Agent.run` (for streaming runs, it is raised when the stream is consumed). @@ -38,39 +41,52 @@ - 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. -- The agent middleware verifies at run start that its chat and function siblings are - installed with it (the trio is created together by the factory), results the feature - cannot project are rejected loudly, and unexpected failures inside the enforcement - layer itself halt the run instead of degrading into tool errors. - -Streaming is supported **fail-closed by buffering**: the model/agent stream is fully -consumed internally, ``post_model_call`` / ``output`` verdicts are applied to the -assembled response, and only then are the (possibly transformed) updates released to the -consumer. No partial content ever egresses ahead of a verdict (spec §12.1/§12.1a -``buffered_output: true`` behaviour). +- Durable history persistence is gated behind the verdicts: run-end context-provider + persistence and per-service-call history persistence are deferred (via the run + persistence gate in ``_sessions.py``) until the ``output`` / ``post_model_call`` + emission permits the content, so denied content never becomes durable and transformed + content is persisted post-transform. Each persist is gated by its own covering + verdict: per-service-call history persisted under a permitted ``post_model_call`` + verdict remains durable even if the run's ``output`` is later denied. + +Streaming is supported **fail-closed by buffering** via +:meth:`ResponseStream.buffered_and_gated`: the model/agent stream is fully consumed +internally, middleware stream hooks are applied to the buffered content, the +``post_model_call`` / ``output`` verdict is applied to the finalized result, and only +then are the (possibly transformed) updates released to the consumer. No partial content +ever egresses ahead of a verdict (spec §12.1/§12.1a ``buffered_output: true`` +behaviour), and nothing can rewrite content past the gate. + +Known limitation — service-side (hosted) tool execution: tools executed by the model +provider itself (surfaced as informational-only function calls, e.g. hosted MCP or +web-search tools) never pass through the framework's function-invocation seam, so +``pre_tool_call`` / ``post_tool_call`` cannot intercept them. Their calls and outputs +are surfaced faithfully in the ``post_model_call`` content projection (they are part of +the model response), where interceptors can observe and deny/transform the response +that carries them. Session scoping: by default each agent run is one agent-hooks session (fresh emitter and sequence, ``agent_startup``/``agent_shutdown`` bracket the run). A host that owns a -longer-lived session can pass its own ``emitter`` and ``builder``; the middleware then -emits only the per-run points and the host owns the session boundaries. +longer-lived session constructs its own emitter and builder and installs them via +:func:`create_agent_hooks_middleware_from_emitter`; the middleware then emits only the +per-run points and the host owns the session boundaries. The ``agent-hooks-sdk`` dependency is optional: importing this module (and the lazy root -export) works without it, and :func:`agent_hooks_middleware` raises a descriptive -``ModuleNotFoundError`` when the SDK is missing. Install it via -``pip install agent-framework-core[agent-hooks]``. +exports) works without it, and the factories raise a descriptive ``ModuleNotFoundError`` +when the SDK is missing. Install it via ``pip install agent-framework-core[agent-hooks]``. """ from __future__ import annotations import asyncio -import base64 import contextlib import json +import logging import uuid from collections.abc import Awaitable, Callable, Mapping, Sequence from contextvars import ContextVar from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Generic, NoReturn, TypeVar, cast +from typing import TYPE_CHECKING, Any, NoReturn, cast from pydantic import BaseModel @@ -82,9 +98,11 @@ ChatMiddleware, FunctionInvocationContext, FunctionMiddleware, + MiddlewareBundle, MiddlewareTermination, - MiddlewareTypes, ) +from ._serialization import make_json_safe +from ._sessions import _DEFERRED_RUN_PERSISTENCE # pyright: ignore[reportPrivateUsage] from ._types import ( AgentResponse, AgentResponseUpdate, @@ -110,7 +128,7 @@ Interceptor, ) -_UpdateT = TypeVar("_UpdateT") +logger = logging.getLogger(__name__) _FRAMEWORK_NAME = "agent-framework" _HOST_ERROR_PREFIX = "host_error:" @@ -118,28 +136,21 @@ _DEFAULT_TIMEOUT = 5.0 _SDK_MISSING_MESSAGE = ( - "agent_hooks_middleware requires the optional `agent-hooks-sdk` package. " + "The agent-hooks middleware requires the optional `agent-hooks-sdk` package. " "Please install `agent-framework-core[agent-hooks]` (or `agent-hooks-sdk`)." ) _TRIO_REQUIRED_MESSAGE = ( "agent-hooks {seam} middleware was invoked without an active agent-hooks run. " - "The middleware returned by agent_hooks_middleware() must be installed together " - "on an Agent, e.g. Agent(client=..., middleware=agent_hooks_middleware([...]))." -) - -_SIBLINGS_REQUIRED_MESSAGE = ( - "agent-hooks agent middleware could not find its chat and function middleware " - "siblings in this run's effective middleware. The middleware returned by " - "agent_hooks_middleware() must be installed together on an Agent, e.g. " - "Agent(client=..., middleware=agent_hooks_middleware([...]))." + "The middleware bundle returned by create_agent_hooks_middleware() must be installed " + "as one unit on an Agent, e.g. Agent(client=..., middleware=[create_agent_hooks_middleware([...])])." ) _FOREIGN_TRIO_MESSAGE = ( "agent-hooks {seam} middleware found an active agent-hooks run owned by a different " - "agent_hooks_middleware() trio. Stacking multiple trios on one agent (or splitting " - "trios across agent- and client-level middleware) is not supported: emissions would " - "silently bind to the wrong emitter. Install exactly one trio per agent." + "agent-hooks middleware bundle. Stacking multiple bundles on one agent (or splitting " + "bundles across agent- and client-level middleware) is not supported: emissions would " + "silently bind to the wrong emitter. Install exactly one bundle per agent." ) @@ -151,26 +162,37 @@ class _AgentHooksWriteBackError(MiddlewareException): """ -# region Run state - - -@dataclass -class _RunState: - """Per-run enforcement state shared by the middleware trio via a ContextVar.""" +def _require_sdk() -> None: + """Import the SDK surface this module uses at runtime, with a helpful install hint. - emitter: InterceptionEmitter - builder: AgentContextBuilder - session_scoped: bool - config: _AgentHooksConfig - halted: BaseException | None = None + Only a genuinely missing ``agent_hooks`` package is translated into the + install-the-extra message; anything else (a broken installation, an incompatible + SDK version missing symbols, a failing transitive import) propagates unchanged so + real breakage is not masked as a missing extra. + """ + try: + from agent_hooks import ( + AgentContextBuilder, + EnforcementMode, + InterceptionBlocked, + InterceptionEmitter, + ) + except ModuleNotFoundError as exc: + if exc.name == "agent_hooks" or (exc.name or "").startswith("agent_hooks."): + raise ModuleNotFoundError(_SDK_MISSING_MESSAGE) from exc + raise + else: + # Reference the probed surface so the imports are not "unused": the probe's + # purpose is exactly to verify these symbols resolve. + del AgentContextBuilder, EnforcementMode, InterceptionBlocked, InterceptionEmitter -_RUN_STATE: ContextVar[_RunState | None] = ContextVar("agent_framework_agent_hooks_run_state", default=None) +# region Run state @dataclass class _AgentHooksConfig: - """Configuration shared by one middleware trio (one factory call).""" + """Configuration shared by one middleware bundle (one factory call).""" interceptors: tuple[tuple[str | None, Interceptor], ...] resolver: ApprovalResolver | None @@ -183,34 +205,35 @@ class _AgentHooksConfig: builder: AgentContextBuilder | None -# endregion +@dataclass +class _RunState: + """Per-run enforcement state shared by the middleware trio via a ContextVar.""" + + emitter: InterceptionEmitter + builder: AgentContextBuilder + session_scoped: bool + config: _AgentHooksConfig + halted: BaseException | None = None -# region Wire projections (framework objects -> AGENT-HOOKS wire JSON) +_RUN_STATE: ContextVar[_RunState | None] = ContextVar("agent_framework_agent_hooks_run_state", default=None) + +_PersistCallback = Callable[[], Awaitable[None]] -def _json_safe(value: Any) -> Any: - """Best-effort projection of an arbitrary Python value into JSON-native values. - The SDK's own marshalling guard still fails closed on anything this misses - (``host_error:context_invalid``), so this is a fidelity aid, not a safety gate. +async def _run_pending_persistence(pending: list[_PersistCallback]) -> None: + """Execute persistence work that was deferred behind a permitting verdict. + + The caller must have reset the persistence gate first, so re-entrant persistence + calls made by these callables run inline. """ - if value is None or isinstance(value, (str, bool, int, float)): - return value - if isinstance(value, Content): - return _json_safe(value.to_dict()) - if isinstance(value, Mapping): - return {str(key): _json_safe(item) for key, item in cast("Mapping[Any, Any]", value).items()} - if isinstance(value, (bytes, bytearray)): - return base64.b64encode(bytes(value)).decode("ascii") - if isinstance(value, Sequence): - return [_json_safe(item) for item in cast("Sequence[Any]", value)] - if isinstance(value, BaseModel): - return _json_safe(value.model_dump()) - to_dict = getattr(value, "to_dict", None) - if callable(to_dict): - with contextlib.suppress(Exception): - return _json_safe(to_dict()) - return str(value) + for persist in pending: + await persist() + + +# endregion + +# region Wire building blocks (framework values <-> AGENT-HOOKS wire JSON) def _role_str(role: Any) -> str: @@ -231,7 +254,7 @@ def _contents_to_wire(contents: Sequence[Content]) -> str | list[dict[str, Any]] """Project message contents faithfully: plain text as a string, rich content as dicts.""" if len(contents) == 1 and contents[0].type == "text": return contents[0].text or "" - return [cast("dict[str, Any]", _json_safe(content.to_dict())) for content in contents] + return [cast("dict[str, Any]", make_json_safe(content.to_dict())) for content in contents] def _message_to_wire(message: Message) -> dict[str, Any]: @@ -271,6 +294,13 @@ def _wire_to_contents(value: Any, *, point: str) -> list[Content]: return contents +def _looks_like_message_dicts(value: Any) -> bool: + if not isinstance(value, list): + return False + items = cast("list[Any]", value) + return bool(items) and all(isinstance(item, Mapping) and "content" in item for item in items) + + def _write_back_message_list( originals: Sequence[Message], before: Sequence[Mapping[str, Any]], @@ -280,95 +310,63 @@ def _write_back_message_list( ) -> list[Message]: """Convert a transformed wire message list back into framework messages. - Unchanged messages keep their original ``Message`` object (identity, ids, and - non-projected metadata). Changed messages are mutated in place when the role - matches (so shared conversation history adopts the transform) and rebuilt - otherwise. The transformed list is authoritative: added or removed messages are - reflected as-is. + The transformed list is authoritative. Entries are matched to original messages by + projection identity rather than list position, so a removal or insertion in the + middle does not shift content onto the wrong original: + + - An entry equal to an (unconsumed) original's projection reuses that original + untouched; originals skipped over were removed by the transform. + - A changed entry mutates the next unconsumed original in place only when that + original's projection is not preserved later in the transformed list (i.e. it was + modified, not shifted) and its role is unchanged. At the ``input`` seam this + in-place mutation lets caller-held message objects adopt the transform; at the + chat seam the outgoing messages are framework-rebuilt copies, so the mutation + affects only the guarded request (enforcement is unaffected either way). + - Anything else (insertions, role changes) becomes a new ``Message``. """ if not isinstance(after, list): raise _AgentHooksWriteBackError(f"agent-hooks {point} transform must produce a list of messages.") - result: list[Message] = [] - for index, item in enumerate(cast("list[Any]", after)): + after_items: list[Mapping[str, Any]] = [] + for item in cast("list[Any]", after): if not isinstance(item, Mapping) or "content" not in item: raise _AgentHooksWriteBackError(f"agent-hooks {point} transform produced a message without role/content.") - wire_message = cast("Mapping[str, Any]", item) - role = str(wire_message.get("role") or "user") - if index < len(before) and dict(wire_message) == dict(before[index]): - result.append(originals[index]) - continue - if index < len(originals) and role == str(before[index].get("role")): - message = originals[index] - message.contents = _wire_to_contents(wire_message.get("content"), point=point) - result.append(message) + after_items.append(cast("Mapping[str, Any]", item)) + + before_dicts = [dict(projection) for projection in before] + result: list[Message] = [] + cursor = 0 + for index, item in enumerate(after_items): + item_dict = dict(item) + match_index = next( + (position for position in range(cursor, len(originals)) if before_dicts[position] == item_dict), + None, + ) + if match_index is not None: + result.append(originals[match_index]) + cursor = match_index + 1 continue - result.append(Message(role, _wire_to_contents(wire_message.get("content"), point=point))) + if cursor < len(originals): + candidate_projection = before_dicts[cursor] + preserved_later = any(dict(later) == candidate_projection for later in after_items[index + 1 :]) + role = str(item.get("role") or "user") + if not preserved_later and role == str(candidate_projection.get("role")): + message = originals[cursor] + cursor += 1 + message.contents = _wire_to_contents(item.get("content"), point=point) + result.append(message) + continue + result.append(Message(str(item.get("role") or "user"), _wire_to_contents(item.get("content"), point=point))) return result -def _input_message_to_wire(message: Message) -> dict[str, Any]: - """Project one input message with the spec's input role mapping.""" - return {"role": _input_role(message.role), "content": _contents_to_wire(message.contents)} - - -def _input_messages_to_wire(messages: Sequence[Message]) -> list[dict[str, Any]]: - return [_input_message_to_wire(message) for message in messages] - - -def _input_to_wire(messages: Sequence[Message]) -> dict[str, Any]: - """Project run input per the spec's ``input`` payload schema. - - A single plain-text message projects as its content string (so string-matching - perimeter guards fire); multi-message or rich input projects as a list of - per-message ``{"role", "content"}`` objects (roles mapped onto the spec's input - role enum) whose contents are strings when plain. - """ - if len(messages) == 1: - return {"content": _contents_to_wire(messages[0].contents), "role": _input_role(messages[0].role)} - return {"content": _input_messages_to_wire(messages), "role": "user"} - - -def _looks_like_message_dicts(value: Any) -> bool: - if not isinstance(value, list): - return False - items = cast("list[Any]", value) - return bool(items) and all(isinstance(item, Mapping) and "content" in item for item in items) - - -def _apply_input_write_back(messages: list[Message], before: Mapping[str, Any], after: Any) -> None: - """Write a transformed ``input`` target back into the run's message list.""" - if after is None or after == before: - return - if not isinstance(after, Mapping): - raise _AgentHooksWriteBackError("agent-hooks input transform must produce an input object target.") - after_map = cast("Mapping[str, Any]", after) - after_role = after_map.get("role") - if after_role != before.get("role"): - # The role field is per-message only for single-message input; for multi-message - # input the top-level role is synthetic and a transform against it is ambiguous. - if len(messages) != 1 or not isinstance(after_role, str): - raise _AgentHooksWriteBackError( - "agent-hooks input transform changed the input role in a way that cannot be written back." - ) - messages[0].role = after_role - after_content = after_map.get("content") - if after_content == before.get("content"): - return - if len(messages) == 1 and not _looks_like_message_dicts(after_content): - messages[0].contents = _wire_to_contents(after_content, point="input") - return - before_list = _input_messages_to_wire(messages) - messages[:] = _write_back_message_list(list(messages), before_list, after_content, point="input") - - def _arguments_to_wire(arguments: Any) -> dict[str, Any]: """Project tool-call arguments as the spec's ``args`` object.""" if arguments is None: return {} if isinstance(arguments, BaseModel): - return {str(key): _json_safe(item) for key, item in arguments.model_dump().items()} + return {str(key): make_json_safe(item) for key, item in arguments.model_dump().items()} if isinstance(arguments, Mapping): - return {str(key): _json_safe(item) for key, item in cast("Mapping[Any, Any]", arguments).items()} + return {str(key): make_json_safe(item) for key, item in cast("Mapping[Any, Any]", arguments).items()} if isinstance(arguments, str): with contextlib.suppress(ValueError): parsed = json.loads(arguments) @@ -378,34 +376,6 @@ def _arguments_to_wire(arguments: Any) -> dict[str, Any]: return {"raw_arguments": str(arguments)} -def _tool_calls_to_wire(messages: Sequence[Message]) -> list[dict[str, Any]]: - calls: list[dict[str, Any]] = [] - for message in messages: - for content in message.contents: - if content.type == "function_call" and not content.informational_only: - calls.append({ - "id": str(content.call_id or ""), - "name": str(content.name or ""), - "args": _arguments_to_wire(content.arguments), - }) - return calls - - -def _response_content_to_wire(messages: Sequence[Message]) -> str | list[dict[str, Any]] | None: - """Project a model response's non-tool-call content (tool calls ride ``tool_calls``).""" - parts: list[dict[str, Any]] = [] - for message in messages: - visible = [content for content in message.contents if content.type != "function_call"] - if not visible: - continue - parts.append({"role": _role_str(message.role), "content": _contents_to_wire(visible)}) - if not parts: - return None - if len(parts) == 1 and isinstance(parts[0]["content"], str): - return parts[0]["content"] - return parts - - def _usage_to_wire(usage_details: Any) -> dict[str, int] | None: if not isinstance(usage_details, Mapping): return None @@ -417,207 +387,399 @@ def _usage_to_wire(usage_details: Any) -> dict[str, int] | None: return usage or None -def _tool_result_to_wire(value: Any) -> Any: - """Project a tool result faithfully, unwrapping framework ``Content`` containers. +# endregion - Text content projects as its text, ``function_result`` content projects as its - canonical ``result`` value, and any other content projects as its full content - dictionary — never as ``str(Content)`` reprs. - """ - if value is None or isinstance(value, (str, bool, int, float)): - return value - if isinstance(value, list) and len(cast("list[Any]", value)) == 1 and isinstance(value[0], Content): - # The canonical single-content result (e.g. the default parser's wrapped text) - # projects as the content's value itself, matching what the model sees. - return _tool_result_to_wire(value[0]) - if isinstance(value, Content): - if value.type == "text": - return value.text or "" - if value.type == "function_result": - if value.result is not None: - return _tool_result_to_wire(value.result) - if value.items is not None: - return [_tool_result_to_wire(item) for item in value.items] +# region Interception-point codecs +# +# One codec per interception point owns both directions of the wire conversion: +# ``to_wire`` projects the native framework value into the spec's payload, and +# ``write_back`` converts the (possibly transformed) wire target back into the native +# value. Every ``write_back`` implements the same rule exactly once per point: a wire +# value the interceptors left untouched maps back to the untouched native value — only +# genuine transforms modify native state, and an untranslatable transform raises +# ``_AgentHooksWriteBackError`` (fail closed) rather than being dropped. + + +class _InputCodec: + """``input``: the run's input messages <-> the spec's input payload.""" + + @staticmethod + def message_to_wire(message: Message) -> dict[str, Any]: + """Project one input message with the spec's input role mapping.""" + return {"role": _input_role(message.role), "content": _contents_to_wire(message.contents)} + + @classmethod + def to_wire(cls, messages: Sequence[Message]) -> dict[str, Any]: + """Project run input per the spec's ``input`` payload schema. + + A single plain-text message projects as its content string (so string-matching + perimeter guards fire); multi-message or rich input projects as a list of + per-message ``{"role", "content"}`` objects (roles mapped onto the spec's input + role enum) whose contents are strings when plain. + """ + if len(messages) == 1: + return {"content": _contents_to_wire(messages[0].contents), "role": _input_role(messages[0].role)} + return {"content": [cls.message_to_wire(message) for message in messages], "role": "user"} + + @classmethod + def write_back(cls, messages: list[Message], before: Mapping[str, Any], after: Any) -> None: + """Write a transformed ``input`` target back into the run's message list.""" + if after is None or after == before: + return + if not isinstance(after, Mapping): + raise _AgentHooksWriteBackError("agent-hooks input transform must produce an input object target.") + after_map = cast("Mapping[str, Any]", after) + after_role = after_map.get("role") + if after_role != before.get("role"): + # The role field is per-message only for single-message input; for + # multi-message input the top-level role is synthetic and a transform + # against it is ambiguous. + if len(messages) != 1 or not isinstance(after_role, str): + raise _AgentHooksWriteBackError( + "agent-hooks input transform changed the input role in a way that cannot be written back." + ) + messages[0].role = after_role + after_content = after_map.get("content") + if after_content == before.get("content"): + return + if len(messages) == 1 and not _looks_like_message_dicts(after_content): + messages[0].contents = _wire_to_contents(after_content, point="input") + return + before_list = [cls.message_to_wire(message) for message in messages] + messages[:] = _write_back_message_list(list(messages), before_list, after_content, point="input") + + +class _ModelRequestCodec: + """``pre_model_call``: the outgoing request messages <-> the spec's messages list.""" + + @staticmethod + def to_wire(messages: Sequence[Message]) -> list[dict[str, Any]]: + return _messages_to_wire(messages) + + @staticmethod + def write_back( + messages: Sequence[Message], before: Sequence[Mapping[str, Any]], after: Any + ) -> list[Message] | None: + """Return the transformed message list, or None when the target is untouched.""" + if after == list(before): return None - return _json_safe(value.to_dict()) - if isinstance(value, Mapping): - return {str(key): _tool_result_to_wire(item) for key, item in cast("Mapping[Any, Any]", value).items()} - if isinstance(value, Sequence) and not isinstance(value, (bytes, bytearray)): - return [_tool_result_to_wire(item) for item in cast("Sequence[Any]", value)] - return _json_safe(value) + return _write_back_message_list(list(messages), before, after, point="pre_model_call") -def _apply_tool_result_write_back(original: Any, after: Any) -> Any: - """Convert a transformed ``post_tool_call`` value back into the native result shape. +class _ModelResponseCodec: + """``post_model_call``: the assembled chat response <-> the spec's response payload. - When the original result is the framework's canonical ``list[Content]`` and the - transformed value is shape-compatible, the Content wrappers are preserved; - otherwise the transformed wire value becomes the result as-is (the function - invocation layer serializes arbitrary JSON-native results faithfully). + Host-executed tool calls ride ``tool_calls`` (they drive the function seam); + service-executed (informational-only) tool calls are part of the model response + itself and are surfaced in ``content`` so hosted tool activity is interceptable + here even though the function seam never sees it. """ - if ( - isinstance(original, list) - and original - and all(isinstance(item, Content) for item in cast("list[Any]", original)) - ): - original_contents = cast("list[Content]", original) - if isinstance(after, str) and len(original_contents) == 1 and original_contents[0].type == "text": - return [Content.from_text(after)] - after_items = cast("list[Any]", after) if isinstance(after, list) else None - if after_items is not None and len(after_items) == len(original_contents): - rebuilt: list[Content] = [] - for content, item in zip(original_contents, after_items): - if _tool_result_to_wire(content) == item: - rebuilt.append(content) - elif content.type == "text" and isinstance(item, str): - rebuilt.append(Content.from_text(item)) - else: - return after_items - return rebuilt - return cast(Any, after) - - -def _agent_output_to_wire(response: AgentResponse) -> str | list[dict[str, Any]]: - """Project the run output: a single plain-text message as a string, else per-message objects.""" - parts = _messages_to_wire(response.messages) - if len(parts) == 1 and isinstance(parts[0]["content"], str): - return parts[0]["content"] - return parts - - -def _apply_output_write_back(response: AgentResponse, before_content: Any, after: Any) -> bool: - """Write a transformed ``output`` target back into the agent response. Returns whether it changed.""" - if after is None: - return False - if not isinstance(after, Mapping): - raise _AgentHooksWriteBackError("agent-hooks output transform must produce an output object target.") - after_content = cast("Mapping[str, Any]", after).get("content") - if after_content == before_content: - return False - originals = list(response.messages) - if isinstance(after_content, str): - if len(originals) == 1: - originals[0].contents = _wire_to_contents(after_content, point="output") - else: - response.messages = [Message("assistant", [after_content])] - return True - if after_content is None: - response.messages = [] - return True - before_list = _messages_to_wire(originals) - response.messages = _write_back_message_list(originals, before_list, after_content, point="output") - return True - - -def _write_back_tool_calls(response: ChatResponse[Any], after_calls: Any) -> bool: - """Reconcile transformed ``tool_calls`` with the response's function-call contents.""" - if not isinstance(after_calls, list): - raise _AgentHooksWriteBackError("agent-hooks post_model_call transform must keep tool_calls a list.") - wire_calls: list[Mapping[str, Any]] = [] - for item in cast("list[Any]", after_calls): - if not isinstance(item, Mapping) or "id" not in item or "name" not in item: - raise _AgentHooksWriteBackError( - "agent-hooks post_model_call transform produced a tool call without id/name." - ) - wire_calls.append(cast("Mapping[str, Any]", item)) - calls_by_id = {str(call["id"]): call for call in wire_calls} - consumed: set[str] = set() - changed = False - for message in response.messages: - kept: list[Content] = [] - for content in message.contents: - if content.type != "function_call" or content.informational_only: - kept.append(content) - continue - wire = calls_by_id.get(str(content.call_id)) - if wire is None: - changed = True # the transform dropped this tool call - continue - consumed.add(str(content.call_id)) - wire_args = wire.get("args") - if isinstance(wire_args, Mapping) and _arguments_to_wire(content.arguments) != dict( - cast("Mapping[str, Any]", wire_args) - ): - content.arguments = {str(key): item for key, item in cast("Mapping[Any, Any]", wire_args).items()} - changed = True - kept.append(content) - if len(kept) != len(message.contents): - message.contents = kept - added = [call for call in wire_calls if str(call["id"]) not in consumed] - if added: - contents = [ - Content.from_function_call( - str(call["id"]), - str(call["name"]), - arguments={str(key): item for key, item in cast("Mapping[Any, Any]", call.get("args") or {}).items()}, - ) - for call in added - ] - target = next((m for m in reversed(response.messages) if _role_str(m.role) == "assistant"), None) - if target is not None: - target.contents = [*target.contents, *contents] - else: - response.messages.append(Message("assistant", contents)) - changed = True - return changed + @staticmethod + def content_to_wire(messages: Sequence[Message]) -> str | list[dict[str, Any]] | None: + """Project the response content (everything except host-executed tool calls).""" + parts: list[dict[str, Any]] = [] + for message in messages: + visible = [ + content for content in message.contents if content.type != "function_call" or content.informational_only + ] + if not visible: + continue + parts.append({"role": _role_str(message.role), "content": _contents_to_wire(visible)}) + if not parts: + return None + if len(parts) == 1 and isinstance(parts[0]["content"], str): + return parts[0]["content"] + return parts + + @staticmethod + def tool_calls_to_wire(messages: Sequence[Message]) -> list[dict[str, Any]]: + """Project the host-executed tool calls (the ones the function seam will bracket).""" + calls: list[dict[str, Any]] = [] + for message in messages: + for content in message.contents: + if content.type == "function_call" and not content.informational_only: + calls.append({ + "id": str(content.call_id or ""), + "name": str(content.name or ""), + "args": _arguments_to_wire(content.arguments), + }) + return calls + + @classmethod + def to_wire(cls, response: ChatResponse[Any]) -> dict[str, Any]: + return { + "content": cls.content_to_wire(response.messages), + "tool_calls": cls.tool_calls_to_wire(response.messages), + "finish_reason": _finish_reason_str(response.finish_reason), + } -def _write_back_response_content(response: ChatResponse[Any], after_content: Any) -> None: - """Rebuild the response's visible content from a transformed ``response.content`` value.""" - calls = [ - content for message in response.messages for content in message.contents if content.type == "function_call" - ] - base: list[Message] - if after_content is None: - base = [] - elif isinstance(after_content, str): - base = [Message("assistant", [after_content])] - elif isinstance(after_content, list): - base = [] - for item in cast("list[Any]", after_content): - if not isinstance(item, Mapping) or "content" not in item: + @classmethod + def write_back(cls, response: ChatResponse[Any], before: Mapping[str, Any], after: Any) -> bool: + """Write a transformed ``post_model_call`` target back into the chat response.""" + if after is None or after == before: + return False + if not isinstance(after, Mapping): + raise _AgentHooksWriteBackError("agent-hooks post_model_call transform must produce a response object.") + after_map = cast("Mapping[str, Any]", after) + changed = False + after_finish = after_map.get("finish_reason") + if after_finish != before.get("finish_reason"): + if not isinstance(after_finish, str): raise _AgentHooksWriteBackError( - "agent-hooks post_model_call transform produced content without role/content." + "agent-hooks post_model_call transform must keep finish_reason a string." ) - wire_message = cast("Mapping[str, Any]", item) - base.append( - Message( - str(wire_message.get("role") or "assistant"), - _wire_to_contents(wire_message.get("content"), point="post_model_call"), + response.finish_reason = cast(Any, after_finish) + changed = True + after_calls = after_map.get("tool_calls") + if after_calls != before.get("tool_calls"): + changed = cls._write_back_tool_calls(response, after_calls) or changed + after_content = after_map.get("content") + if after_content != before.get("content"): + cls._write_back_content(response, after_content) + changed = True + return changed + + @staticmethod + def _write_back_tool_calls(response: ChatResponse[Any], after_calls: Any) -> bool: + """Reconcile transformed ``tool_calls`` with the response's function-call contents.""" + if not isinstance(after_calls, list): + raise _AgentHooksWriteBackError("agent-hooks post_model_call transform must keep tool_calls a list.") + wire_calls: list[Mapping[str, Any]] = [] + for item in cast("list[Any]", after_calls): + if not isinstance(item, Mapping) or "id" not in item or "name" not in item: + raise _AgentHooksWriteBackError( + "agent-hooks post_model_call transform produced a tool call without id/name." + ) + wire_calls.append(cast("Mapping[str, Any]", item)) + calls_by_id = {str(call["id"]): call for call in wire_calls} + consumed: set[str] = set() + changed = False + for message in response.messages: + kept: list[Content] = [] + for content in message.contents: + if content.type != "function_call" or content.informational_only: + kept.append(content) + continue + wire = calls_by_id.get(str(content.call_id)) + if wire is None: + changed = True # the transform dropped this tool call + continue + consumed.add(str(content.call_id)) + wire_name = wire.get("name") + if not isinstance(wire_name, str) or not wire_name: + raise _AgentHooksWriteBackError( + "agent-hooks post_model_call transform must keep each tool call's name a non-empty string." + ) + if wire_name != str(content.name): + content.name = wire_name + changed = True + wire_args = wire.get("args") + if not isinstance(wire_args, Mapping): + raise _AgentHooksWriteBackError( + "agent-hooks post_model_call transform must keep each tool call's args an object." + ) + if _arguments_to_wire(content.arguments) != dict(cast("Mapping[str, Any]", wire_args)): + content.arguments = {str(key): item for key, item in cast("Mapping[Any, Any]", wire_args).items()} + changed = True + kept.append(content) + if len(kept) != len(message.contents): + message.contents = kept + added = [call for call in wire_calls if str(call["id"]) not in consumed] + if added: + contents = [ + Content.from_function_call( + str(call["id"]), + str(call["name"]), + arguments={ + str(key): item for key, item in cast("Mapping[Any, Any]", call.get("args") or {}).items() + }, + ) + for call in added + ] + target = next((m for m in reversed(response.messages) if _role_str(m.role) == "assistant"), None) + if target is not None: + target.contents = [*target.contents, *contents] + else: + response.messages.append(Message("assistant", contents)) + changed = True + return changed + + @staticmethod + def _write_back_content(response: ChatResponse[Any], after_content: Any) -> None: + """Rebuild the response's visible content from a transformed ``response.content`` value.""" + calls = [ + content + for message in response.messages + for content in message.contents + if content.type == "function_call" and not content.informational_only + ] + base: list[Message] + if after_content is None: + base = [] + elif isinstance(after_content, str): + base = [Message("assistant", [after_content])] + elif isinstance(after_content, list): + base = [] + for item in cast("list[Any]", after_content): + if not isinstance(item, Mapping) or "content" not in item: + raise _AgentHooksWriteBackError( + "agent-hooks post_model_call transform produced content without role/content." + ) + wire_message = cast("Mapping[str, Any]", item) + base.append( + Message( + str(wire_message.get("role") or "assistant"), + _wire_to_contents(wire_message.get("content"), point="post_model_call"), + ) ) - ) - else: - raise _AgentHooksWriteBackError("agent-hooks post_model_call transform produced unsupported content.") - if calls: - if base and _role_str(base[-1].role) == "assistant": - base[-1].contents = [*base[-1].contents, *calls] else: - base.append(Message("assistant", calls)) - response.messages = base + raise _AgentHooksWriteBackError("agent-hooks post_model_call transform produced unsupported content.") + if calls: + if base and _role_str(base[-1].role) == "assistant": + base[-1].contents = [*base[-1].contents, *calls] + else: + base.append(Message("assistant", calls)) + response.messages = base + + +class _ToolArgumentsCodec: + """``pre_tool_call``: the native tool arguments <-> the spec's args object.""" + + @staticmethod + def to_wire(arguments: Any) -> dict[str, Any]: + return _arguments_to_wire(arguments) + + @staticmethod + def write_back(arguments: Any, before: Mapping[str, Any], after: Any) -> tuple[Any, dict[str, Any]]: + """Merge a transformed ``args`` target back onto the native arguments. + + Returns ``(native_arguments, effective_wire_args)``. Only the keys the + transform actually changed (or added/removed) are taken from the wire value; + untouched keys keep their original native values, so non-JSON-native argument + values (bytes, rich objects) survive a transform that did not touch them. + """ + if not isinstance(after, Mapping): + raise _AgentHooksWriteBackError("agent-hooks pre_tool_call transform must produce an arguments object.") + effective = {str(key): item for key, item in cast("Mapping[Any, Any]", after).items()} + if effective == dict(before): + return arguments, effective + if isinstance(arguments, BaseModel): + native: dict[str, Any] = {name: getattr(arguments, name) for name in type(arguments).model_fields} + elif isinstance(arguments, Mapping): + native = {str(key): item for key, item in cast("Mapping[Any, Any]", arguments).items()} + else: + native = {} + merged = {key: item for key, item in native.items() if key in effective} + for key, item in effective.items(): + if key not in before or before[key] != item: + merged[key] = item + return merged, effective -def _apply_response_write_back(response: ChatResponse[Any], before: Mapping[str, Any], after: Any) -> bool: - """Write a transformed ``post_model_call`` target back into the chat response.""" - if after is None or after == before: - return False - if not isinstance(after, Mapping): - raise _AgentHooksWriteBackError("agent-hooks post_model_call transform must produce a response object.") - after_map = cast("Mapping[str, Any]", after) - changed = False - after_finish = after_map.get("finish_reason") - if after_finish != before.get("finish_reason"): - if not isinstance(after_finish, str): - raise _AgentHooksWriteBackError("agent-hooks post_model_call transform must keep finish_reason a string.") - response.finish_reason = cast(Any, after_finish) - changed = True - after_calls = after_map.get("tool_calls") - if after_calls != before.get("tool_calls"): - changed = _write_back_tool_calls(response, after_calls) or changed - after_content = after_map.get("content") - if after_content != before.get("content"): - _write_back_response_content(response, after_content) - changed = True - return changed +class _ToolResultCodec: + """``post_tool_call``: the native tool result <-> the spec's result value.""" + + @staticmethod + def to_wire(value: Any) -> Any: + """Project a tool result faithfully, unwrapping framework ``Content`` containers. + + Text content projects as its text, ``function_result`` content projects as its + canonical ``result`` value, and any other content projects as its full content + dictionary — never as ``str(Content)`` reprs. + """ + if value is None or isinstance(value, (str, bool, int, float)): + return value + if isinstance(value, list) and len(cast("list[Any]", value)) == 1 and isinstance(value[0], Content): + # The canonical single-content result (e.g. the default parser's wrapped + # text) projects as the content's value itself, matching what the model sees. + return _ToolResultCodec.to_wire(value[0]) + if isinstance(value, Content): + if value.type == "text": + return value.text or "" + if value.type == "function_result": + if value.result is not None: + return _ToolResultCodec.to_wire(value.result) + if value.items is not None: + return [_ToolResultCodec.to_wire(item) for item in value.items] + return None + return make_json_safe(value.to_dict()) + if isinstance(value, Mapping): + return {str(key): _ToolResultCodec.to_wire(item) for key, item in cast("Mapping[Any, Any]", value).items()} + if isinstance(value, Sequence) and not isinstance(value, (bytes, bytearray)): + return [_ToolResultCodec.to_wire(item) for item in cast("Sequence[Any]", value)] + return make_json_safe(value) + + @staticmethod + def write_back(original: Any, after: Any) -> Any: + """Convert a transformed ``post_tool_call`` value back into the native result shape. + + When the original result is the framework's canonical ``list[Content]`` and the + transformed value is shape-compatible, the Content wrappers are preserved; + otherwise the transformed wire value becomes the result as-is (the function + invocation layer serializes arbitrary JSON-native results faithfully). + """ + if ( + isinstance(original, list) + and original + and all(isinstance(item, Content) for item in cast("list[Any]", original)) + ): + original_contents = cast("list[Content]", original) + if isinstance(after, str) and len(original_contents) == 1 and original_contents[0].type == "text": + return [Content.from_text(after)] + after_items = cast("list[Any]", after) if isinstance(after, list) else None + if after_items is not None and len(after_items) == len(original_contents): + rebuilt: list[Content] = [] + for content, item in zip(original_contents, after_items): + if _ToolResultCodec.to_wire(content) == item: + rebuilt.append(content) + elif content.type == "text" and isinstance(item, str): + rebuilt.append(Content.from_text(item)) + else: + return after_items + return rebuilt + return cast(Any, after) + + +class _OutputCodec: + """``output``: the final agent response <-> the spec's output payload.""" + + @staticmethod + def to_wire(response: AgentResponse[Any]) -> str | list[dict[str, Any]]: + """Project the run output: a single plain-text message as a string, else per-message objects.""" + parts = _messages_to_wire(response.messages) + if len(parts) == 1 and isinstance(parts[0]["content"], str): + return parts[0]["content"] + return parts + + @staticmethod + def write_back(response: AgentResponse[Any], before_content: Any, after: Any) -> bool: + """Write a transformed ``output`` target back into the agent response. Returns whether it changed.""" + if after is None: + return False + if not isinstance(after, Mapping): + raise _AgentHooksWriteBackError("agent-hooks output transform must produce an output object target.") + after_content = cast("Mapping[str, Any]", after).get("content") + if after_content == before_content: + return False + originals = list(response.messages) + if isinstance(after_content, str): + if len(originals) == 1: + originals[0].contents = _wire_to_contents(after_content, point="output") + else: + response.messages = [Message("assistant", [after_content])] + return True + if after_content is None: + response.messages = [] + return True + before_list = _messages_to_wire(originals) + response.messages = _write_back_message_list(originals, before_list, after_content, point="output") + return True + + +# endregion + +# region Enforcement helpers def _chat_updates_from_response(response: ChatResponse[Any]) -> list[ChatResponseUpdate]: @@ -657,16 +819,21 @@ def _agent_updates_from_response(response: AgentResponse[Any]) -> list[AgentResp def _tool_names(context: AgentContext) -> list[str]: + """Project the registered tool names for ``agent_startup`` (spec ``tools_registered``).""" + from ._tools import _get_tool_name, normalize_tools # type: ignore[reportPrivateUsage] + tools: Any = context.tools if context.tools is not None else getattr(context.agent, "tools", None) if tools is None: return [] - items: Sequence[Any] = ( - cast("Sequence[Any]", tools) if isinstance(tools, Sequence) and not isinstance(tools, (str, bytes)) else [tools] - ) + try: + normalized = normalize_tools(tools) + except Exception: + logger.warning("agent-hooks could not normalize the run's tools for the agent_startup projection.") + return [] names: list[str] = [] - for item in items: - name = getattr(item, "name", None) or getattr(item, "__name__", None) - names.append(str(name) if name else type(item).__name__) + for item in normalized: + name = _get_tool_name(item) + names.append(name if name else type(item).__name__) return names @@ -693,14 +860,14 @@ def _is_approval_request(result: Any) -> bool: def _halt_on_enforcement_failure( state: _RunState, context: FunctionInvocationContext, exc: BaseException, point: str ) -> NoReturn: - """Route an unexpected failure inside the enforcement layer through the halt path. + """Route an unexpected failure inside the enforcement layer through the fail-closed halt path. 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 while the run continued. 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. + 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. """ message = f"agent-hooks {point} enforcement failed: {type(exc).__name__}" context.result = {"error": message} @@ -713,13 +880,17 @@ def _halt_on_enforcement_failure( raise MiddlewareTermination(message) from exc +# endregion + +# region Middleware implementations (private: the bundle is one coherent feature) + + # eq=False keeps identity semantics (and hashability): the middleware pipeline caches # compare middleware tuples with ==, and value-equality would let a pipeline cached for -# one trio be reused for a field-equal fresh trio, breaking the identity-based sibling -# verification on the second run. +# one bundle be reused for a field-equal fresh bundle, conflating their run states. @dataclass(eq=False) class _AgentHooksMiddlewareBase: - """Shared base carrying the trio's config (also enables sibling identity checks).""" + """Shared base carrying the bundle's config (also enables ownership checks).""" _config: _AgentHooksConfig @@ -728,27 +899,6 @@ def _shares_config(self, config: _AgentHooksConfig) -> bool: return self._config is config -class _ReplayAsyncIterator(Generic[_UpdateT]): - """Async iterator that replays a buffered list of updates.""" - - def __init__(self, updates: Sequence[_UpdateT]) -> None: - self._iterator = iter(updates) - - def __aiter__(self) -> _ReplayAsyncIterator[_UpdateT]: - return self - - async def __anext__(self) -> _UpdateT: - try: - return next(self._iterator) - except StopIteration: - raise StopAsyncIteration from None - - -# endregion - -# region Middleware implementations (private: the trio is one coherent feature) - - class _AgentHooksAgentMiddleware(_AgentHooksMiddlewareBase, AgentMiddleware): """Run bracket: ``agent_startup``, ``input``, ``output``, ``agent_shutdown``.""" @@ -786,48 +936,34 @@ async def _emit_run_start(self, context: AgentContext, state: _RunState) -> None """Emit ``agent_startup`` (per-run sessions) and ``input``; apply input transforms.""" if not state.session_scoped: await state.emitter.emit(state.builder.agent_startup(tools_registered=_tool_names(context))) - before = _input_to_wire(context.messages) + before = _InputCodec.to_wire(context.messages) outcome: EmitOutcome = await state.emitter.emit( state.builder.input(content=before["content"], role=before["role"]) ) - _apply_input_write_back(context.messages, before, outcome.target) + _InputCodec.write_back(context.messages, before, outcome.target) async def _emit_output(self, state: _RunState, response: AgentResponse[Any]) -> bool: """Emit ``output`` over the assembled response; apply output transforms.""" - before_content = _agent_output_to_wire(response) + before_content = _OutputCodec.to_wire(response) outcome: EmitOutcome = await state.emitter.emit(state.builder.output(content=before_content)) - return _apply_output_write_back(response, before_content, outcome.target) + return _OutputCodec.write_back(response, before_content, outcome.target) async def _emit_shutdown(self, state: _RunState, reason: str) -> None: """Best-effort ``agent_shutdown`` (per-run sessions only; blocks there are record-only).""" if state.session_scoped: return - with contextlib.suppress(Exception): + try: await state.emitter.emit_unchecked(state.builder.agent_shutdown(reason=reason)) - - def _verify_siblings(self, context: AgentContext) -> None: - """Fail closed unless the chat and function siblings from this trio are installed. - - The trio is created together by the factory, so the siblings are located by - identity of the shared config. The agent seam sees the combined chat/function - middleware for the run in ``context.client_kwargs["middleware"]`` (the agent - layer assembles it there before building the context). - """ - installed = context.client_kwargs.get("middleware") - candidates = cast("Sequence[Any]", installed) if isinstance(installed, Sequence) else () - has_chat = any( - isinstance(item, _AgentHooksChatMiddleware) and item._shares_config(self._config) for item in candidates - ) - has_function = any( - isinstance(item, _AgentHooksFunctionMiddleware) and item._shares_config(self._config) for item in candidates - ) - if not (has_chat and has_function): - raise MiddlewareException(_SIBLINGS_REQUIRED_MESSAGE) + except Exception: + logger.warning( + "agent-hooks failed to emit agent_shutdown (reason=%r); the session trail is incomplete.", + reason, + exc_info=True, + ) async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: from agent_hooks import InterceptionBlocked - self._verify_siblings(context) state = self._new_run_state(context) if context.stream: await self._process_streaming(context, call_next, state) @@ -839,12 +975,17 @@ async def process(self, context: AgentContext, call_next: Callable[[], Awaitable try: await self._emit_run_start(context, state) termination: MiddlewareTermination | None = None + pending: list[_PersistCallback] = [] + persistence_token = _DEFERRED_RUN_PERSISTENCE.set(pending) try: - await call_next() - except MiddlewareTermination as exc: - # A middleware short-circuited the run; any substituted result - # still egresses to the caller, so it passes the output point. - termination = exc + try: + await call_next() + except MiddlewareTermination as exc: + # A middleware short-circuited the run; any substituted result + # still egresses to the caller, so it passes the output point. + termination = exc + finally: + _DEFERRED_RUN_PERSISTENCE.reset(persistence_token) if state.halted is not None: raise state.halted result = context.result @@ -855,6 +996,10 @@ async def process(self, context: AgentContext, call_next: Callable[[], Awaitable f"agent-hooks cannot guard a run result of type {type(result).__name__}; " "the output interception point was not emitted." ) + # The verdict permitted the content (or nothing egresses): run the + # persistence the run deferred behind the gate. A deny above skips + # this entirely, so denied content never becomes durable. + await _run_pending_persistence(pending) if termination is not None: raise termination except InterceptionBlocked: @@ -879,13 +1024,13 @@ async def _process_streaming( ) -> None: """Streaming run setup (executed lazily on the first pull of the outer stream). - Ownership of the session trail passes to the guarded stream only once it is + Ownership of the session trail passes to the gated stream only once it is installed as the run result; every earlier exit (pre-run deny, middleware exception, unguardable result) still closes the trail with ``agent_shutdown``. """ token = _RUN_STATE.set(state) - # Pessimistic default: any exit before the guarded stream takes ownership - # closes the trail as an error; ``None`` means ownership was handed off. + # Pessimistic default: any exit before the gated stream takes ownership closes + # the trail as an error; ``None`` means ownership was handed off. shutdown_reason: str | None = "error" try: await self._emit_run_start(context, state) @@ -904,15 +1049,10 @@ async def _process_streaming( "agent-hooks streaming enforcement requires a ResponseStream agent result; " f"got {type(inner).__name__}." ) - context.result = cast( - "ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", - cast(Any, ResponseStream).from_awaitable( - self._guarded_agent_stream( - state, cast("ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", inner) - ) - ), + context.result = self._gated_agent_stream( + state, cast("ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", inner) ) - # The guarded stream now owns the shutdown emission. + # The gated stream now owns the shutdown emission. shutdown_reason = None if termination is not None: # A middleware substituted its own stream: it is guarded, and the @@ -926,26 +1066,44 @@ async def _process_streaming( await self._emit_shutdown(state, shutdown_reason) _RUN_STATE.reset(token) - async def _guarded_agent_stream( + def _gated_agent_stream( self, state: _RunState, inner: ResponseStream[AgentResponseUpdate, AgentResponse[Any]], ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: - """Consume the run fully, apply the ``output`` verdict, then release buffered updates. - - Fail-closed streamed egress: no update is released to the consumer before the - ``output`` emission permits, and a transformed output replaces the buffered - updates entirely. + """Guard a streaming run with a fail-closed buffered gate. + + The run is fully consumed (with the run state and the persistence gate active), + middleware stream hooks are applied by the combinator before the verdict, the + ``output`` verdict is applied to the finalized response, deferred persistence + runs only after the verdict permits, and only then are updates released. A + transformed output (or any applied stream hooks) re-derives the released + updates from the verdicted response so streamed egress cannot diverge from it. """ - from agent_hooks import InterceptionBlocked + pending: list[_PersistCallback] = [] - token = _RUN_STATE.set(state) - shutdown_reason = "completed" - transformed = False - final: AgentResponse[Any] - try: + async def _consume() -> tuple[Sequence[AgentResponseUpdate], AgentResponse[Any]]: + run_token = _RUN_STATE.set(state) + persistence_token = _DEFERRED_RUN_PERSISTENCE.set(pending) try: final = await inner.get_final_response() + except asyncio.CancelledError: + await self._emit_shutdown(state, "cancelled") + raise + except BaseException: + await self._emit_shutdown(state, "error") + raise + finally: + _DEFERRED_RUN_PERSISTENCE.reset(persistence_token) + _RUN_STATE.reset(run_token) + return list(inner.updates), final + + async def _gate( + updates: list[AgentResponseUpdate], final: AgentResponse[Any], hooks_applied: bool + ) -> tuple[Sequence[AgentResponseUpdate], AgentResponse[Any]]: + from agent_hooks import InterceptionBlocked + + try: if state.halted is not None: raise state.halted if not isinstance(final, AgentResponse): @@ -954,25 +1112,24 @@ async def _guarded_agent_stream( "the output interception point was not emitted." ) transformed = await self._emit_output(state, final) + await _run_pending_persistence(pending) + await self._emit_shutdown(state, "completed") + released = _agent_updates_from_response(final) if (transformed or hooks_applied) else updates + return released, final except InterceptionBlocked: - shutdown_reason = "error" + await self._emit_shutdown(state, "error") raise except asyncio.CancelledError: - shutdown_reason = "cancelled" + await self._emit_shutdown(state, "cancelled") raise except BaseException: - shutdown_reason = "error" + await self._emit_shutdown(state, "error") raise - finally: - await self._emit_shutdown(state, shutdown_reason) - _RUN_STATE.reset(token) - updates = _agent_updates_from_response(final) if transformed else list(inner.updates) - - def _finalize(_: Sequence[AgentResponseUpdate]) -> AgentResponse[Any]: - return final - - return ResponseStream(_ReplayAsyncIterator(updates), finalizer=_finalize) + return cast( + "ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", + cast(Any, ResponseStream).buffered_and_gated(consume=_consume, gate=_gate), + ) class _AgentHooksChatMiddleware(_AgentHooksMiddlewareBase, ChatMiddleware): @@ -985,47 +1142,48 @@ async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[ if state is None: raise MiddlewareException(_TRIO_REQUIRED_MESSAGE.format(seam="chat")) if not self._shares_config(state.config): - # A different trio's agent middleware owns the innermost run state (stacked - # trios, or a trio split across agent- and client-level middleware): binding - # to it would silently misroute emissions. Fail closed instead. + # A different bundle's agent middleware owns the innermost run state + # (stacked bundles, or a bundle split across agent- and client-level + # middleware): binding to it would silently misroute emissions. raise MiddlewareException(_FOREIGN_TRIO_MESSAGE.format(seam="chat")) options = context.options or {} model_id = str(options.get("model") or type(context.client).__name__) - before = _messages_to_wire(context.messages) + before = _ModelRequestCodec.to_wire(context.messages) outcome: EmitOutcome = await state.emitter.emit( state.builder.pre_model_call(model_id=model_id, messages=before) ) - if outcome.target != before: - context.messages = _write_back_message_list( - list(context.messages), before, outcome.target, point="pre_model_call" - ) + transformed_messages = _ModelRequestCodec.write_back(context.messages, before, outcome.target) + if transformed_messages is not None: + context.messages = transformed_messages termination: MiddlewareTermination | None = None + pending: list[_PersistCallback] = [] + persistence_token = _DEFERRED_RUN_PERSISTENCE.set(pending) try: - await call_next() - except MiddlewareTermination as exc: - # A chat middleware short-circuited with a substituted result; whatever - # was substituted still flows into the agent loop, so it is guarded below. - termination = exc + try: + await call_next() + except MiddlewareTermination as exc: + # A chat middleware short-circuited with a substituted result; whatever + # was substituted still flows into the agent loop, so it is guarded below. + termination = exc + finally: + _DEFERRED_RUN_PERSISTENCE.reset(persistence_token) result = context.result if isinstance(result, ResponseStream): - context.result = cast( - "ResponseStream[ChatResponseUpdate, ChatResponse[Any]]", - cast(Any, ResponseStream).from_awaitable( - self._guarded_chat_stream( - state, model_id, cast("ResponseStream[ChatResponseUpdate, ChatResponse[Any]]", result) - ) - ), + context.result = self._gated_chat_stream( + state, model_id, cast("ResponseStream[ChatResponseUpdate, ChatResponse[Any]]", result), pending ) elif isinstance(result, ChatResponse): try: await self._emit_post_model_call(state, model_id, result) except InterceptionBlocked: - # §6.1: the denied response must not be incorporated. + # §6.1: the denied response must not be incorporated (and the deferred + # per-service-call persistence for it is dropped, never executed). context.result = None raise + await _run_pending_persistence(pending) elif result is not None: raise MiddlewareException( f"agent-hooks cannot guard a chat result of type {type(result).__name__}; " @@ -1038,11 +1196,7 @@ async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[ async def _emit_post_model_call(self, state: _RunState, model_id: str, response: ChatResponse[Any]) -> bool: """Emit ``post_model_call`` over the assembled response; apply transforms. Returns whether changed.""" - before: dict[str, Any] = { - "content": _response_content_to_wire(response.messages), - "tool_calls": _tool_calls_to_wire(response.messages), - "finish_reason": _finish_reason_str(response.finish_reason), - } + before = _ModelResponseCodec.to_wire(response) outcome: EmitOutcome = await state.emitter.emit( state.builder.post_model_call( model_id=str(response.model or model_id), @@ -1053,33 +1207,48 @@ async def _emit_post_model_call(self, state: _RunState, model_id: str, response: request_id=response.response_id, ) ) - return _apply_response_write_back(response, before, outcome.target) + return _ModelResponseCodec.write_back(response, before, outcome.target) - async def _guarded_chat_stream( + def _gated_chat_stream( self, state: _RunState, model_id: str, inner: ResponseStream[ChatResponseUpdate, ChatResponse[Any]], + pending: list[_PersistCallback], ) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: - """Assemble the streamed response, apply the ``post_model_call`` verdict, then replay. + """Guard a streaming model call with a fail-closed buffered gate. Spec §12.1: the complete response is assembled before ``post_model_call`` is emitted, and nothing (updates or tool calls) is released beforehand. A deny - raises before any update egresses. + raises before any update egresses, and per-service-call history persistence + deferred by the run persistence gate executes only after the verdict permits. """ - response = await inner.get_final_response() - if not isinstance(response, ChatResponse): - raise MiddlewareException( - f"agent-hooks cannot guard a streamed chat result of type {type(response).__name__}; " - "the post_model_call interception point was not emitted." - ) - changed = await self._emit_post_model_call(state, model_id, response) - updates = _chat_updates_from_response(response) if changed else list(inner.updates) - - def _finalize(_: Sequence[ChatResponseUpdate]) -> ChatResponse[Any]: - return response - return ResponseStream(_ReplayAsyncIterator(updates), finalizer=_finalize) + async def _consume() -> tuple[Sequence[ChatResponseUpdate], ChatResponse[Any]]: + persistence_token = _DEFERRED_RUN_PERSISTENCE.set(pending) + try: + response = await inner.get_final_response() + finally: + _DEFERRED_RUN_PERSISTENCE.reset(persistence_token) + return list(inner.updates), response + + async def _gate( + updates: list[ChatResponseUpdate], final: ChatResponse[Any], hooks_applied: bool + ) -> tuple[Sequence[ChatResponseUpdate], ChatResponse[Any]]: + if not isinstance(final, ChatResponse): + raise MiddlewareException( + f"agent-hooks cannot guard a streamed chat result of type {type(final).__name__}; " + "the post_model_call interception point was not emitted." + ) + changed = await self._emit_post_model_call(state, model_id, final) + await _run_pending_persistence(pending) + released = _chat_updates_from_response(final) if (changed or hooks_applied) else updates + return released, final + + return cast( + "ResponseStream[ChatResponseUpdate, ChatResponse[Any]]", + cast(Any, ResponseStream).buffered_and_gated(consume=_consume, gate=_gate), + ) class _AgentHooksFunctionMiddleware(_AgentHooksMiddlewareBase, FunctionMiddleware): @@ -1108,17 +1277,19 @@ async def process(self, context: FunctionInvocationContext, call_next: Callable[ state = _RUN_STATE.get() if state is None: - # No run state means the trio was split. 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. + # 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) if not self._shares_config(state.config): - # A different trio owns the innermost run state (stacked trios): binding to - # it would silently misroute emissions. Halt the run fail-closed (the loop - # swallows plain exceptions, so route through the halt path). + # 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). _halt_on_enforcement_failure( state, context, MiddlewareException(_FOREIGN_TRIO_MESSAGE.format(seam="function")), "pre_tool_call" ) @@ -1127,18 +1298,16 @@ async def process(self, context: FunctionInvocationContext, call_next: Callable[ raw_call_id = context.metadata.get("call_id") call_id = str(raw_call_id) if raw_call_id else uuid.uuid4().hex name = str(getattr(context.function, "name", context.function)) - args = _arguments_to_wire(context.arguments) + args = _ToolArgumentsCodec.to_wire(context.arguments) outcome: EmitOutcome = await state.emitter.emit( state.builder.pre_tool_call(call_id=call_id, name=name, args=args) ) - target = outcome.target - if not isinstance(target, Mapping): - raise _AgentHooksWriteBackError("agent-hooks pre_tool_call transform must produce an arguments object.") - effective = {str(key): item for key, item in cast("Mapping[Any, Any]", target).items()} - if effective != args: - # The transform rewrote the arguments: execute the approved value. - context.arguments = dict(effective) - args = effective + arguments, effective = _ToolArgumentsCodec.write_back(context.arguments, args, outcome.target) + if arguments is not context.arguments: + # The transform rewrote (some of) the arguments: execute the approved + # values, keeping untouched keys' original native values. + context.arguments = arguments + 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") @@ -1180,13 +1349,21 @@ async def process(self, context: FunctionInvocationContext, call_next: Callable[ _halt_on_enforcement_failure(state, context, 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. + return + try: - value = _tool_result_to_wire(context.result) + value = _ToolResultCodec.to_wire(context.result) outcome = await state.emitter.emit( state.builder.post_tool_call(call_id=call_id, name=name, args=args, value=value) ) if outcome.target != value: - context.result = _apply_tool_result_write_back(context.result, outcome.target) + context.result = _ToolResultCodec.write_back(context.result, 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") @@ -1200,12 +1377,12 @@ async def process(self, context: FunctionInvocationContext, call_next: Callable[ # endregion -# region Public factory +# region Public factories @experimental(feature_id=ExperimentalFeature.AGENT_HOOKS) -def agent_hooks_middleware( - interceptors: Sequence[Interceptor] | Mapping[str, Interceptor] | None = None, +def create_agent_hooks_middleware( + interceptors: Sequence[Interceptor] | Mapping[str, Interceptor], *, resolver: ApprovalResolver | None = None, mode: EnforcementMode | str = "enforce", @@ -1213,44 +1390,36 @@ def agent_hooks_middleware( identity_provider: str | IdentityProvider | None = _JCS_SHA256, timeout: float | None = _DEFAULT_TIMEOUT, record_sink: Callable[[InterceptionRecord], None] | None = None, - emitter: InterceptionEmitter | None = None, - builder: AgentContextBuilder | None = None, -) -> list[MiddlewareTypes]: - """Build the AGENT-HOOKS-0.1 middleware trio for an :class:`~agent_framework.Agent`. +) -> MiddlewareBundle: + """Build the AGENT-HOOKS-0.1 enforcement middleware for an :class:`~agent_framework.Agent`. - The returned middleware emit every applicable interception point of the agent-hooks - control contract and enforce the combined verdicts fail-closed: denies block the + The returned bundle emits every applicable interception point of the agent-hooks + control contract and enforces the combined verdicts fail-closed: denies block the guarded action (a run-level deny raises :class:`agent_hooks.InterceptionBlocked` to - the caller; a tool-seam deny surfaces a tool error to the model), and transforms are + the caller; a tool-seam deny surfaces a tool error to the model), transforms are written back into the framework's messages, arguments, and results so execution uses - exactly the values the interceptors approved. Streaming runs are buffered and only - released after the ``output`` verdict permits. + exactly the values the interceptors approved, streaming runs are buffered and only + released after the ``output`` verdict permits, and durable history persistence is + deferred until the covering verdict permits the content. - The three middleware form one coherent feature and must be installed together; - always pass the returned list as a whole to ``Agent(middleware=...)``, and install - exactly one trio per agent (stacked trios are rejected fail-closed). + Every agent run is one agent-hooks session: a fresh ``InterceptionEmitter`` and + ``AgentContextBuilder`` pair is created per run and ``agent_startup`` / + ``agent_shutdown`` bracket the run. To scope one session across multiple runs, see + :func:`create_agent_hooks_middleware_from_emitter`. Use ``record_sink`` to observe + interception records. Composition order: - Place the trio first (outermost) in the middleware list. Middleware listed - before the trio runs outside the enforcement boundary: a function middleware - placed before the trio, for example, can substitute a tool result that the - tool seam never brackets. Outer position is outer trust — the final - ``output`` point still guards whatever egresses to the caller. - - Session scoping: - By default every agent run is one agent-hooks session: a fresh - ``InterceptionEmitter``/``AgentContextBuilder`` pair is created per run and - ``agent_startup``/``agent_shutdown`` bracket the run. To scope one session - across multiple runs (shared ``sequence``, stateful interceptors, one approval - ledger), construct and own the emitter and builder yourself and pass them via - ``emitter=``/``builder=``; the middleware then emits only the per-run points - (``input`` through ``output``) and your host emits the session boundaries. - Use ``record_sink`` (or your own emitter) to observe interception records. + Pass the bundle as one element of ``Agent(middleware=[...])``, placed first + (outermost), and install exactly one bundle per agent (stacked bundles are + rejected fail-closed). Middleware listed before the bundle runs outside the + enforcement boundary: a function middleware placed before it, for example, can + substitute a tool result that the tool seam never brackets. Outer position is + outer trust — the final ``output`` point still guards whatever egresses. Args: interceptors: The agent-hooks interceptors to register, either as a sequence or as a mapping of registration name to interceptor (names appear on the - records' verdict summaries). Required unless ``emitter`` is supplied. + records' verdict summaries). At least one interceptor is required. Keyword Args: resolver: Optional approval resolver consulted for liftable denies. @@ -1263,21 +1432,14 @@ def agent_hooks_middleware( records. timeout: Per-interceptor/resolver timeout in seconds (spec RECOMMENDED 5.0). record_sink: Optional callable receiving every interception record. - emitter: A host-owned ``InterceptionEmitter`` for session-scoped enforcement. - Must be provided together with ``builder``; the emitter's own - configuration (interceptors, resolver, mode, ...) then governs and the - per-run configuration parameters must be left at their defaults. - builder: The host-owned ``AgentContextBuilder`` matching ``emitter``. Returns: - The middleware list to pass to ``Agent(middleware=...)``. + The middleware bundle to pass (as one element) to ``Agent(middleware=...)``. Raises: ModuleNotFoundError: If the optional ``agent-hooks-sdk`` package is not installed. - ValueError: If the configuration is inconsistent (no interceptors in per-run - mode, ``emitter``/``builder`` not provided together, or per-run parameters - combined with a host-owned emitter). + ValueError: If no interceptors are provided. Examples: .. code-block:: python @@ -1296,80 +1458,89 @@ def intercept(self, context): agent = Agent( client=client, name="assistant", - middleware=agent_hooks_middleware([EgressGuard()]), + middleware=[create_agent_hooks_middleware([EgressGuard()])], ) """ - try: - from agent_hooks import EnforcementMode - except ImportError as exc: # pragma: no cover - exercised via tests with import hooks - raise ModuleNotFoundError(_SDK_MISSING_MESSAGE) from exc - - if (emitter is None) != (builder is None): - raise ValueError("agent_hooks_middleware requires `emitter` and `builder` to be provided together.") - - if emitter is not None: - conflicting = [ - param_name - for param_name, value in ( - ("interceptors", interceptors), - ("resolver", resolver), - ("composition", composition), - ("record_sink", record_sink), - ) - if value is not None - ] - if (mode.value if isinstance(mode, EnforcementMode) else str(mode)) != "enforce": - conflicting.append("mode") - if identity_provider != _JCS_SHA256: - conflicting.append("identity_provider") - if timeout != _DEFAULT_TIMEOUT: - conflicting.append("timeout") - if conflicting: - raise ValueError( - "agent_hooks_middleware received a host-owned `emitter`; the emitter's own configuration " - f"governs, so these parameters must be left at their defaults: {', '.join(sorted(conflicting))}." - ) - # Host-owned session: only the emitter/builder pair is retained — the emitter's - # own configuration governs, so the per-run knobs are not stored at all. - config = _AgentHooksConfig( - interceptors=(), - resolver=None, - mode=None, - composition=None, - identity_provider=None, - timeout=None, - record_sink=None, - emitter=emitter, - builder=builder, - ) + _require_sdk() + from agent_hooks import EnforcementMode + + named: list[tuple[str | None, Interceptor]] + if isinstance(interceptors, Mapping): + named = [(str(name), interceptor) for name, interceptor in interceptors.items()] else: - named: list[tuple[str | None, Interceptor]] - if isinstance(interceptors, Mapping): - named = [(str(name), interceptor) for name, interceptor in interceptors.items()] - else: - named = [(None, interceptor) for interceptor in interceptors or []] - if not named: - raise ValueError( - "agent_hooks_middleware requires at least one interceptor (an emitter with zero " - "interceptors fails closed on every emission)." - ) - config = _AgentHooksConfig( - interceptors=tuple(named), - resolver=resolver, - mode=EnforcementMode(mode) if isinstance(mode, str) else mode, - composition=composition, - identity_provider=identity_provider, - timeout=timeout, - record_sink=record_sink, - emitter=None, - builder=None, + named = [(None, interceptor) for interceptor in interceptors] + if not named: + raise ValueError( + "create_agent_hooks_middleware requires at least one interceptor (an emitter with " + "zero interceptors fails closed on every emission)." ) + config = _AgentHooksConfig( + interceptors=tuple(named), + resolver=resolver, + mode=EnforcementMode(mode) if isinstance(mode, str) else mode, + composition=composition, + identity_provider=identity_provider, + timeout=timeout, + record_sink=record_sink, + emitter=None, + builder=None, + ) + return _build_bundle(config) + + +@experimental(feature_id=ExperimentalFeature.AGENT_HOOKS) +def create_agent_hooks_middleware_from_emitter( + emitter: InterceptionEmitter, + builder: AgentContextBuilder, +) -> MiddlewareBundle: + """Build the agent-hooks middleware bundle around a host-owned session. + + Use this form when your host scopes one agent-hooks session across multiple agent + runs (shared ``sequence``, stateful interceptors, one approval ledger): construct + and configure the ``InterceptionEmitter`` (interceptors, resolver, mode, + composition, identity provider, timeout, record sink) and the matching + ``AgentContextBuilder`` yourself and pass them here. The middleware then emits only + the per-run points (``input`` through ``output``) on your emitter, and your host + owns the ``agent_startup`` / ``agent_shutdown`` session boundaries. + + Enforcement semantics and composition-order requirements are identical to + :func:`create_agent_hooks_middleware`. - return [ + Args: + emitter: The host-owned, fully configured ``InterceptionEmitter``. + builder: The host-owned ``AgentContextBuilder`` matching ``emitter``. + + Returns: + The middleware bundle to pass (as one element) to ``Agent(middleware=...)``. + + Raises: + ModuleNotFoundError: If the optional ``agent-hooks-sdk`` package is not + installed. + ValueError: If ``emitter`` or ``builder`` is missing. + """ + _require_sdk() + if not emitter or not builder: + raise ValueError("create_agent_hooks_middleware_from_emitter requires both an emitter and a builder.") + config = _AgentHooksConfig( + interceptors=(), + resolver=None, + mode=None, + composition=None, + identity_provider=None, + timeout=None, + record_sink=None, + emitter=emitter, + builder=builder, + ) + return _build_bundle(config) + + +def _build_bundle(config: _AgentHooksConfig) -> MiddlewareBundle: + return MiddlewareBundle([ _AgentHooksAgentMiddleware(config), _AgentHooksChatMiddleware(config), _AgentHooksFunctionMiddleware(config), - ] + ]) # endregion diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index 3ded2f67639..5f1c9574e6f 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -35,6 +35,7 @@ PerServiceCallHistoryPersistingMiddleware, ServiceSessionId, SessionContext, + _defer_run_persistence, # pyright: ignore[reportPrivateUsage] is_local_history_conversation_id, ) from ._telemetry import FeatureIndex, mark_feature_used @@ -439,6 +440,12 @@ def __init__( self.name = name self.description = description self.context_providers: list[ContextProvider] = list(context_providers or []) + # Normalize a bare middleware source (a single middleware object or a + # MiddlewareBundle) into a single-element list, mirroring how + # categorize_middleware treats non-sequence sources, so construction-time + # and run-time middleware arguments behave identically. + if middleware is not None and not isinstance(middleware, Sequence): + middleware = [middleware] self.middleware: list[MiddlewareTypes] | None = ( cast(list[MiddlewareTypes], middleware) if middleware is not None else None ) @@ -532,10 +539,18 @@ async def _run_after_providers( ) -> None: """Run after_run on all context providers in reverse order. + When an egress-enforcement gate is active for this run (see + ``_sessions._defer_run_persistence``), the provider work is deferred to the gate + owner so denied or transformed content never becomes durable ahead of its + verdict. The gate owner resets the gate before executing deferred callables, so + the re-entrant call below runs inline. + Keyword Args: session: The conversation session. context: The invocation context with response populated. """ + if _defer_run_persistence(partial(self._run_after_providers, session=session, context=context)): + return provider_session = session if provider_session is None and self.context_providers: provider_session = AgentSession() diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index c72e1a7277f..4098ca8ee08 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -4,6 +4,7 @@ import contextlib import inspect +import logging import sys from abc import ABC, abstractmethod from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequence @@ -24,6 +25,8 @@ ) from .exceptions import MiddlewareException +logger = logging.getLogger(__name__) + if sys.version_info >= (3, 13): from typing import TypeVar # pragma: no cover else: @@ -668,6 +671,75 @@ async def process( FunctionMiddleware | FunctionMiddlewareCallable | ChatMiddleware | ChatMiddlewareCallable ) + +class MiddlewareBundle: + """An indivisible group of middleware that forms one coherent feature. + + Some features (for example the agent-hooks enforcement middleware) consist of + several middleware objects that only uphold their contract when installed + together. A bundle carries those objects as one opaque unit: pass the bundle + itself anywhere middleware is accepted, and :func:`categorize_middleware` + splits its members into their agent/function/chat categories while the bundle + guarantees the members cannot be installed partially — it is deliberately not + a sequence, so it cannot be unpacked or sliced. + + Examples: + .. code-block:: python + + from agent_framework import Agent + + bundle = create_some_feature_middleware(...) + agent = Agent(client=client, middleware=[bundle, my_other_middleware]) + """ + + def __init__( + self, + middleware: Sequence[ + AgentMiddleware + | AgentMiddlewareCallable + | FunctionMiddleware + | FunctionMiddlewareCallable + | ChatMiddleware + | ChatMiddlewareCallable + ], + ) -> None: + """Initialize the bundle. + + Args: + middleware: The middleware objects that belong together. Order is + preserved when the bundle is expanded into the run's pipelines. + Every member must be categorizable by the framework's own rules + (an agent/function/chat middleware instance, or a callable with a + recognizable middleware signature); nested bundles are rejected. + + Raises: + MiddlewareException: If a member is a nested bundle or cannot be + categorized as agent, function, or chat middleware. + """ + members = tuple(middleware) + for member in members: + if isinstance(member, MiddlewareBundle): + raise MiddlewareException( + "MiddlewareBundle members must be middleware objects; nesting a " + "MiddlewareBundle inside another bundle is not supported." + ) + if isinstance(member, (AgentMiddleware, FunctionMiddleware, ChatMiddleware)): + continue + if callable(member): + # Raises MiddlewareException when the callable's category cannot be + # determined — the same validation categorize_middleware applies. + _determine_middleware_type(member) + continue + raise MiddlewareException( + f"MiddlewareBundle members must be agent, function, or chat middleware; got {type(member).__name__}." + ) + self._middleware = members + + def __repr__(self) -> str: + members = ", ".join(type(middleware).__name__ for middleware in self._middleware) + return f"{type(self).__name__}({members})" + + # Type alias for all middleware types MiddlewareTypes: TypeAlias = ( AgentMiddleware @@ -676,6 +748,7 @@ async def process( | FunctionMiddlewareCallable | ChatMiddleware | ChatMiddlewareCallable + | MiddlewareBundle ) @@ -842,6 +915,16 @@ def _register_middleware_with_wrapper( self._middleware.append(middleware) elif callable(middleware): self._middleware.append(MiddlewareWrapper(middleware)) # type: ignore[arg-type] + else: + # Preserve the long-standing lenient behavior (do not fail the run), but + # never skip silently: an unrecognized object here means middleware the + # caller supplied will not execute. + logger.warning( + "Ignoring unrecognized middleware of type %s: it is neither a %s nor a callable " + "and will not be executed.", + type(middleware).__name__, + expected_type.__name__, + ) class AgentMiddlewarePipeline(BaseMiddlewarePipeline): @@ -1353,11 +1436,18 @@ def run( client_kwargs: Mapping[str, Any] | None = None, ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: """MiddlewareTypes-enabled unified run method.""" - # Re-categorize self.middleware at runtime to support dynamic changes + # Re-categorize self.middleware at runtime to support dynamic changes. A bare + # single source (one middleware object or a MiddlewareBundle assigned directly + # to the attribute) is treated as a one-element list, mirroring how + # categorize_middleware handles non-sequence sources — never silently dropped. base_middleware_attr = getattr(self, "middleware", None) - base_middleware: Sequence[MiddlewareTypes] = ( - cast(Sequence[MiddlewareTypes], base_middleware_attr) if isinstance(base_middleware_attr, Sequence) else [] - ) + base_middleware: Sequence[MiddlewareTypes] + if isinstance(base_middleware_attr, Sequence) and not isinstance(base_middleware_attr, (str, bytes)): + base_middleware = cast(Sequence[MiddlewareTypes], base_middleware_attr) + elif base_middleware_attr is not None: + base_middleware = [cast(MiddlewareTypes, base_middleware_attr)] + else: + base_middleware = [] base_middleware_list = categorize_middleware(base_middleware) run_middleware_list = categorize_middleware(middleware) pipeline = self._get_agent_middleware_pipeline([*base_middleware_list["agent"], *run_middleware_list["agent"]]) @@ -1543,6 +1633,17 @@ def categorize_middleware( else: all_middleware.append(source) + # Expand bundles first: a bundle's members are categorized individually (in + # order) but travel as one unit, so a feature spanning several categories can + # never be partially installed. + expanded_middleware: list[Any] = [] + for middleware in all_middleware: + if isinstance(middleware, MiddlewareBundle): + expanded_middleware.extend(middleware._middleware) # pyright: ignore[reportPrivateUsage] + else: + expanded_middleware.append(middleware) + all_middleware = expanded_middleware + # Categorize each middleware item for middleware in all_middleware: if isinstance(middleware, AgentMiddleware): diff --git a/python/packages/core/agent_framework/_serialization.py b/python/packages/core/agent_framework/_serialization.py index a4d00fb4b70..cb7f3a0db89 100644 --- a/python/packages/core/agent_framework/_serialization.py +++ b/python/packages/core/agent_framework/_serialization.py @@ -2,6 +2,7 @@ from __future__ import annotations +import base64 import copy import json import logging @@ -649,8 +650,9 @@ def make_json_safe(obj: Any) -> Any: """Recursively convert an object to a JSON-serializable form. Handles dataclasses, Pydantic models, objects with ``to_dict``/``dict``/``__dict__``, - datetimes, lists, dicts, and primitives. Falls back to ``str()`` for any remaining - non-serializable value so that ``json.dumps`` never raises a ``TypeError``. + datetimes, bytes (base64), lists, dicts, and primitives. Falls back to ``str()`` for + any remaining non-serializable value so that ``json.dumps`` never raises a + ``TypeError``. Args: obj: Object to make JSON safe. @@ -662,6 +664,8 @@ def make_json_safe(obj: Any) -> Any: return obj if isinstance(obj, (datetime, date)): return obj.isoformat() + if isinstance(obj, (bytes, bytearray)): + return base64.b64encode(bytes(obj)).decode("ascii") if is_dataclass(obj) and not isinstance(obj, type): return make_json_safe(asdict(obj)) if type(obj) is dict: diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index fc6e395a20a..1b4a155eebd 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -30,7 +30,9 @@ from base64 import urlsafe_b64encode from collections import deque from collections.abc import AsyncIterable, Awaitable, Callable, Iterable, Mapping, Sequence +from contextvars import ContextVar from dataclasses import dataclass +from functools import partial from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias, TypeGuard, TypeVar, cast @@ -1025,6 +1027,30 @@ async def after_run( await self.save_messages(context.session_id, messages_to_store, state=state) +# Run-scoped gate for durable persistence side effects. Egress-enforcement middleware +# (agent-hooks) sets this to a collector list around the guarded section of a run; the +# persistence call sites below consult it via ``_defer_run_persistence`` so that history +# only becomes durable after the run's egress verdict permits the content. When no gate +# is active (the default), persistence runs inline exactly as before. +_DEFERRED_RUN_PERSISTENCE: ContextVar[list[Callable[[], Awaitable[None]]] | None] = ContextVar( + "agent_framework_deferred_run_persistence", default=None +) + + +def _defer_run_persistence(persist: Callable[[], Awaitable[None]]) -> bool: + """Queue a persistence side effect behind the active run-persistence gate, if any. + + Returns True when ``persist`` was deferred to the gate owner (which MUST reset the + context variable before executing the collected callables, so re-entrant calls run + inline), and False when no gate is active and the caller must persist inline. + """ + pending = _DEFERRED_RUN_PERSISTENCE.get() + if pending is None: + return False + pending.append(persist) + return True + + LOCAL_HISTORY_CONVERSATION_ID = "agent_framework_local_history_persistence" @@ -1332,10 +1358,16 @@ async def _finalize_response( "instead." ) - await self._persist_service_call_response( + # Durability is gated: when an egress-enforcement gate is active for this run, + # the persist is deferred until the model-call verdict permits the content + # (the sentinel/validation logic above stays inline — it drives control flow). + persist = partial( + self._persist_service_call_response, service_call_context=service_call_context, response=response, ) + if not _defer_run_persistence(persist): + await persist() # The local sentinel only applies when the service does not store history; when it does, # the real conversation id already drives function-loop continuation. if not self._service_stores_history and _response_contains_follow_up_request(response): diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 7a0a982ba8e..1d3d797d618 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -10,6 +10,7 @@ import sys from asyncio import iscoroutine from collections.abc import ( + AsyncGenerator, AsyncIterable, AsyncIterator, Awaitable, @@ -3223,6 +3224,49 @@ def from_awaitable( stream._wrap_inner = True return stream + @classmethod + def buffered_and_gated( + cls, + consume: Callable[[], Awaitable[tuple[Sequence[UpdateT], FinalT]]], + gate: Callable[[list[UpdateT], FinalT, bool], Awaitable[tuple[Sequence[UpdateT], FinalT]]], + ) -> ResponseStream[UpdateT, FinalT]: + """Create a fully buffered stream whose content is finalized by a gate. + + This combinator exists for egress-gating middleware (e.g. policy enforcement) + that must apply a verdict to a run's *complete* content before anything is + released, and it states the hook-ordering contract in one place: + + 1. On the first pull, ``consume`` runs and produces the buffered updates and + the finalized result. Nothing has egressed yet. + 2. Every transform/result/cleanup hook registered on the *returned* stream so + far (for example hooks attached by middleware pipelines after they unwind) + is applied to the buffered updates and result now — **before** the gate — + so the gate's verdict covers their effect. The hooks are consumed: they are + not applied again during replay. + 3. ``gate`` receives the post-hook updates, the post-hook result, and whether + any pending transform/result hooks were applied. It returns the updates and + result to release (it may raise to block egress entirely). A gate that + issues its verdict over the final result should re-derive the released + updates from that result whenever hooks were applied, so streamed egress + cannot diverge from the verdicted content. + 4. The stream is then sealed: the gated updates and result are replayed + verbatim, and registering further transform or result hooks raises + ``RuntimeError`` — nothing can rewrite content past the gate. + + Args: + consume: Produces the buffered updates and finalized result. Runs inside + the caller's context (the caller owns any context-variable scoping). + gate: Receives ``(updates, final, hooks_applied)`` after pending hooks are + applied; returns the ``(updates, final)`` to release. + + Returns: + A sealed, fully buffered ResponseStream. + """ + return cast( + "ResponseStream[UpdateT, FinalT]", + cast(Any, _GatedResponseStream).create_buffered_and_gated(consume, gate), + ) + async def _get_stream(self) -> AsyncIterable[UpdateT]: if self._stream is None: if hasattr(self._stream_source, "__aiter__"): @@ -3481,6 +3525,97 @@ def updates(self) -> Sequence[UpdateT]: return self._updates +class _GatedResponseStream(ResponseStream[UpdateT, FinalT]): + """ResponseStream whose content is sealed once its gate has run. + + Created by :meth:`ResponseStream.buffered_and_gated`. Hooks registered before + the gate runs are applied to the buffered content ahead of the gate; once the + gate has run, registering transform or result hooks raises so nothing can + rewrite content past the gate. (Cleanup hooks remain allowed: they cannot + influence content.) + """ + + _gate_sealed: bool = False + + def with_transform_hook( + self, + hook: Callable[[UpdateT], UpdateT | Awaitable[UpdateT | None] | None], + ) -> ResponseStream[UpdateT, FinalT]: + """Register a transform hook; rejected once the stream's gate has run.""" + if self._gate_sealed: + raise RuntimeError( + "Cannot register a transform hook on a gated ResponseStream after its gate has " + "run: content is sealed by the gate's verdict." + ) + return super().with_transform_hook(hook) + + def with_result_hook( + self, + hook: Callable[[FinalT], FinalT | Awaitable[FinalT | None] | None], + ) -> ResponseStream[UpdateT, FinalT]: + """Register a result hook; rejected once the stream's gate has run.""" + if self._gate_sealed: + raise RuntimeError( + "Cannot register a result hook on a gated ResponseStream after its gate has " + "run: content is sealed by the gate's verdict." + ) + return super().with_result_hook(hook) + + @classmethod + def create_buffered_and_gated( + cls, + consume: Callable[[], Awaitable[tuple[Sequence[UpdateT], FinalT]]], + gate: Callable[[list[UpdateT], FinalT, bool], Awaitable[tuple[Sequence[UpdateT], FinalT]]], + ) -> _GatedResponseStream[UpdateT, FinalT]: + """Build the gated stream for :meth:`ResponseStream.buffered_and_gated`.""" + holder: dict[str, Any] = {} + + async def _materialize() -> AsyncGenerator[UpdateT]: + stream = cast(_GatedResponseStream[UpdateT, FinalT], holder["stream"]) + updates, final = await consume() + # Drain the hooks registered on the gated stream so far and apply them to + # the buffered content before the gate (contract step 2). Draining also + # means _record_update applies nothing during replay. + transform_hooks = list(stream._transform_hooks) + stream._transform_hooks.clear() + result_hooks = list(stream._result_hooks) + stream._result_hooks.clear() + cleanup_hooks = list(stream._cleanup_hooks) + stream._cleanup_hooks.clear() + hooked_updates: list[UpdateT] = [] + for update in updates: + hooked_update = update + for hook in transform_hooks: + hooked = hook(hooked_update) + if isawaitable(hooked): + hooked = await hooked + if hooked is not None: + hooked_update = cast(UpdateT, hooked) + hooked_updates.append(hooked_update) + for result_hook in result_hooks: + hooked_final = result_hook(final) + if isawaitable(hooked_final): + hooked_final = await hooked_final + if hooked_final is not None: + final = cast(FinalT, hooked_final) + for cleanup_hook in cleanup_hooks: + cleanup_result = cleanup_hook() + if isawaitable(cleanup_result): + await cleanup_result + gated_updates, gated_final = await gate(hooked_updates, final, bool(transform_hooks or result_hooks)) + holder["final"] = gated_final + stream._gate_sealed = True + for update in gated_updates: + yield update + + def _finalizer(_: Sequence[UpdateT]) -> FinalT: + return cast(FinalT, holder["final"]) + + stream: _GatedResponseStream[UpdateT, FinalT] = cls(_materialize(), finalizer=_finalizer) + holder["stream"] = stream + return stream + + # region ChatOptions diff --git a/python/packages/core/pyrightconfig.dependency.json b/python/packages/core/pyrightconfig.dependency.json index f0be44f3d6e..2cc3d95b21f 100644 --- a/python/packages/core/pyrightconfig.dependency.json +++ b/python/packages/core/pyrightconfig.dependency.json @@ -6,7 +6,6 @@ "agent_framework/_workflows" ], "exclude": [ - "agent_framework/_agent_hooks.py", "agent_framework/_harness/_agent.py" ] } diff --git a/python/packages/core/tests/core/test_agent_hooks.py b/python/packages/core/tests/core/test_agent_hooks.py index 2babc5f7bb4..31811a38272 100644 --- a/python/packages/core/tests/core/test_agent_hooks.py +++ b/python/packages/core/tests/core/test_agent_hooks.py @@ -10,7 +10,6 @@ import agent_framework from agent_framework import ( - SKIP_PARSING, Agent, AgentContext, AgentMiddleware, @@ -24,10 +23,12 @@ FunctionInvocationContext, FunctionMiddleware, Message, + MiddlewareBundle, MiddlewareException, MiddlewareTermination, ResponseStream, - agent_hooks_middleware, + create_agent_hooks_middleware, + create_agent_hooks_middleware_from_emitter, tool, ) @@ -157,37 +158,76 @@ def points(records: list[Any]) -> list[str]: @requires_sdk async def test_factory_requires_interceptors() -> None: with pytest.raises(ValueError, match="at least one interceptor"): - agent_hooks_middleware([]) - with pytest.raises(ValueError, match="at least one interceptor"): - agent_hooks_middleware() + create_agent_hooks_middleware([]) @requires_sdk -async def test_factory_requires_emitter_and_builder_together() -> None: - emitter = InterceptionEmitter() - with pytest.raises(ValueError, match="together"): - agent_hooks_middleware(emitter=emitter) +async def test_from_emitter_factory_requires_both_arguments() -> None: + emitter = InterceptionEmitter().register(AllowGuard()) + builder = AgentContextBuilder(agent_id="a", framework="agent-framework", session_id="s") + with pytest.raises(ValueError, match="both an emitter and a builder"): + create_agent_hooks_middleware_from_emitter(emitter, cast("Any", None)) + with pytest.raises(ValueError, match="both an emitter and a builder"): + create_agent_hooks_middleware_from_emitter(cast("Any", None), builder) @requires_sdk -async def test_factory_rejects_per_run_config_with_host_emitter() -> None: - emitter = InterceptionEmitter().register(AllowGuard()) - builder = AgentContextBuilder(agent_id="a", framework="agent-framework", session_id="s") - with pytest.raises(ValueError, match="interceptors"): - agent_hooks_middleware([AllowGuard()], emitter=emitter, builder=builder) - with pytest.raises(ValueError, match="mode"): - agent_hooks_middleware(mode="evaluate_only", emitter=emitter, builder=builder) +async def test_bare_bundle_at_construction_is_fully_enforced(chat_client_base: MockBaseChatClient) -> None: + # Passing the bundle bare (instead of inside a list) at construction must install + # it exactly like `middleware=[bundle]` — previously it was silently dropped and + # the run executed fully unhooked. + records: list[InterceptionRecord] = [] + guard = PointGuard("output", Verdict.deny(reason="egress_blocked")) + agent = Agent( + client=chat_client_base, + middleware=cast("Any", create_agent_hooks_middleware([guard], record_sink=records.append)), + ) + + with pytest.raises(InterceptionBlocked) as exc_info: + await agent.run("hello") + + assert exc_info.value.result.verdict.reason == "egress_blocked" + # The full session was emitted: enforcement was installed, not silently skipped. + assert points(records) == [ + "agent_startup", + "input", + "pre_model_call", + "post_model_call", + "output", + "agent_shutdown", + ] + + +def test_middleware_bundle_rejects_invalid_members() -> None: + # A nested bundle (or any uncategorizable member) would previously fall through + # categorization and be skipped silently at pipeline registration. + inner = MiddlewareBundle([AgentShortCircuit(None)]) + with pytest.raises(MiddlewareException, match="nesting"): + MiddlewareBundle([cast("Any", inner)]) + with pytest.raises(MiddlewareException, match="must be agent, function, or chat middleware"): + MiddlewareBundle([cast("Any", object())]) + with pytest.raises(MiddlewareException): + # A callable whose middleware category cannot be determined is rejected by the + # same validation categorize_middleware applies. + MiddlewareBundle([cast("Any", lambda context, call_next: None)]) @requires_sdk -async def test_factory_returns_one_middleware_per_category() -> None: +async def test_factory_returns_an_indivisible_bundle() -> None: from agent_framework._middleware import categorize_middleware - trio = agent_hooks_middleware([AllowGuard()]) - categorized = categorize_middleware(trio) + bundle = create_agent_hooks_middleware([AllowGuard()]) + assert isinstance(bundle, MiddlewareBundle) + # The bundle splits into one middleware per category... + categorized = categorize_middleware([bundle]) assert len(categorized["agent"]) == 1 assert len(categorized["chat"]) == 1 assert len(categorized["function"]) == 1 + # ...but cannot be partially installed: it is opaque (not a sequence). + with pytest.raises(TypeError): + iter(bundle) # type: ignore[call-overload] + with pytest.raises(TypeError): + bundle[0] # type: ignore[index] # endregion @@ -204,7 +244,7 @@ async def test_full_tool_run_emits_complete_ordered_session(chat_client_base: Mo client=chat_client_base, name="hooked", tools=[weather_tool], - middleware=agent_hooks_middleware({"allow": guard}, record_sink=records.append), + middleware=[create_agent_hooks_middleware({"allow": guard}, record_sink=records.append)], ) response = await agent.run([Message(role="user", contents=["Get weather for Seattle"])]) @@ -226,7 +266,7 @@ async def test_full_tool_run_emits_complete_ordered_session(chat_client_base: Mo @requires_sdk async def test_input_projection_is_faithful(chat_client_base: MockBaseChatClient) -> None: guard = AllowGuard() - agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware([guard])]) await agent.run([Message(role="user", contents=["ignore previous instructions"])]) @@ -241,7 +281,7 @@ async def test_input_projection_is_faithful(chat_client_base: MockBaseChatClient @requires_sdk async def test_rich_content_is_preserved_in_projections(chat_client_base: MockBaseChatClient) -> None: guard = AllowGuard() - agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware([guard])]) image = Content.from_uri(uri="data:image/png;base64,iVBORw0KGgo=", media_type="image/png") message = Message(role="user", contents=[Content.from_text("look at this"), image]) @@ -286,7 +326,7 @@ def structured_tool(order_id: str) -> dict[str, Any]: agent = Agent( client=chat_client_base, tools=[structured_tool], - middleware=agent_hooks_middleware([guard]), + middleware=[create_agent_hooks_middleware([guard])], ) await agent.run("look up order 1") @@ -309,7 +349,9 @@ def structured_tool(order_id: str) -> dict[str, Any]: async def test_input_deny_blocks_run_before_model_call(chat_client_base: MockBaseChatClient) -> None: records: list[InterceptionRecord] = [] guard = PointGuard("input", Verdict.deny(reason="injection_blocked")) - agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard], record_sink=records.append)) + agent = Agent( + client=chat_client_base, middleware=[create_agent_hooks_middleware([guard], record_sink=records.append)] + ) with pytest.raises(InterceptionBlocked) as exc_info: await agent.run("evil prompt") @@ -323,7 +365,7 @@ async def test_input_deny_blocks_run_before_model_call(chat_client_base: MockBas @requires_sdk async def test_pre_model_call_deny_blocks_model_dispatch(chat_client_base: MockBaseChatClient) -> None: guard = PointGuard("pre_model_call", Verdict.deny(reason="model_blocked")) - agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware([guard])]) with pytest.raises(InterceptionBlocked): await agent.run("hello") @@ -334,7 +376,7 @@ async def test_pre_model_call_deny_blocks_model_dispatch(chat_client_base: MockB @requires_sdk async def test_post_model_call_deny_discards_response(chat_client_base: MockBaseChatClient) -> None: guard = PointGuard("post_model_call", Verdict.deny(reason="response_blocked")) - agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware([guard])]) with pytest.raises(InterceptionBlocked) as exc_info: await agent.run("hello") @@ -351,7 +393,7 @@ async def test_pre_tool_call_deny_blocks_tool_and_continues_loop(chat_client_bas agent = Agent( client=chat_client_base, tools=[weather_tool], - middleware=agent_hooks_middleware([guard], record_sink=records.append), + middleware=[create_agent_hooks_middleware([guard], record_sink=records.append)], ) response = await agent.run("get the weather") @@ -369,7 +411,7 @@ async def test_pre_tool_call_deny_blocks_tool_and_continues_loop(chat_client_bas async def test_post_tool_call_deny_discards_result(chat_client_base: MockBaseChatClient) -> None: guard = PointGuard("post_tool_call", Verdict.deny(reason="result_blocked")) chat_client_base.run_responses = [tool_call_response(), final_response()] - agent = Agent(client=chat_client_base, tools=[weather_tool], middleware=agent_hooks_middleware([guard])) + agent = Agent(client=chat_client_base, tools=[weather_tool], middleware=[create_agent_hooks_middleware([guard])]) response = await agent.run("get the weather") @@ -382,7 +424,7 @@ async def test_post_tool_call_deny_discards_result(chat_client_base: MockBaseCha @requires_sdk async def test_output_deny_blocks_response(chat_client_base: MockBaseChatClient) -> None: guard = PointGuard("output", Verdict.deny(reason="egress_blocked")) - agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware([guard])]) with pytest.raises(InterceptionBlocked) as exc_info: await agent.run("hello") @@ -401,7 +443,7 @@ async def test_input_transform_writes_back_into_run_messages(chat_client_base: M "input", Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target.content", value="[redacted]")), ) - agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware([guard])]) message = Message(role="user", contents=["my SSN is 123-45-6789"]) response = await agent.run([message]) @@ -418,7 +460,7 @@ async def test_pre_model_call_transform_writes_back_into_request(chat_client_bas "pre_model_call", Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target[0].content", value="[masked]")), ) - agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware([guard])]) response = await agent.run("raw PII 123-45-6789") @@ -436,7 +478,7 @@ async def test_pre_tool_call_transform_writes_back_into_arguments(chat_client_ba agent = Agent( client=chat_client_base, tools=[weather_tool], - middleware=agent_hooks_middleware([transformer, guard]), + middleware=[create_agent_hooks_middleware([transformer, guard])], ) await agent.run("get the weather") @@ -456,7 +498,7 @@ async def test_post_tool_call_transform_writes_back_into_result(chat_client_base Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target", value="[scrubbed]")), ) chat_client_base.run_responses = [tool_call_response(), final_response()] - agent = Agent(client=chat_client_base, tools=[weather_tool], middleware=agent_hooks_middleware([guard])) + agent = Agent(client=chat_client_base, tools=[weather_tool], middleware=[create_agent_hooks_middleware([guard])]) response = await agent.run("get the weather") @@ -475,7 +517,7 @@ async def test_output_transform_writes_back_into_response(chat_client_base: Mock "output", Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target.content", value="[card:redacted]")), ) - agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware([guard])]) response = await agent.run("what is my card number?") @@ -510,7 +552,7 @@ async def test_streaming_buffers_until_all_verdicts_permit(chat_client_base: Moc agent = Agent( client=chat_client_base, tools=[weather_tool], - middleware=agent_hooks_middleware([AllowGuard()], record_sink=records.append), + middleware=[create_agent_hooks_middleware([AllowGuard()], record_sink=records.append)], ) updates: list[AgentResponseUpdate] = [] @@ -535,7 +577,9 @@ async def test_streaming_buffers_until_all_verdicts_permit(chat_client_base: Moc async def test_streaming_output_deny_releases_nothing(chat_client_base: MockBaseChatClient) -> None: records: list[InterceptionRecord] = [] guard = PointGuard("output", Verdict.deny(reason="egress_blocked")) - agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard], record_sink=records.append)) + agent = Agent( + client=chat_client_base, middleware=[create_agent_hooks_middleware([guard], record_sink=records.append)] + ) updates: list[AgentResponseUpdate] = [] with pytest.raises(InterceptionBlocked): @@ -549,7 +593,7 @@ async def test_streaming_output_deny_releases_nothing(chat_client_base: MockBase @requires_sdk async def test_streaming_post_model_call_deny_releases_nothing(chat_client_base: MockBaseChatClient) -> None: guard = PointGuard("post_model_call", Verdict.deny(reason="response_blocked")) - agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware([guard])]) updates: list[AgentResponseUpdate] = [] with pytest.raises(InterceptionBlocked): @@ -565,7 +609,7 @@ async def test_streaming_output_transform_rewrites_updates(chat_client_base: Moc "output", Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target.content", value="[masked]")), ) - agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware([guard])]) updates: list[AgentResponseUpdate] = [] stream = agent.run("hello", stream=True) @@ -603,7 +647,7 @@ def broken_tool(location: str) -> str: ), final_response("Sorry."), ] - agent = Agent(client=chat_client_base, tools=[broken_tool], middleware=agent_hooks_middleware([guard])) + agent = Agent(client=chat_client_base, tools=[broken_tool], middleware=[create_agent_hooks_middleware([guard])]) response = await agent.run("run the broken tool") @@ -622,7 +666,7 @@ async def test_interceptor_crash_fails_closed_and_halts_run(chat_client_base: Mo agent = Agent( client=chat_client_base, tools=[weather_tool], - middleware=agent_hooks_middleware([CrashingGuard("post_tool_call")], record_sink=records.append), + middleware=[create_agent_hooks_middleware([CrashingGuard("post_tool_call")], record_sink=records.append)], ) with pytest.raises(InterceptionBlocked) as exc_info: @@ -636,7 +680,7 @@ async def test_interceptor_crash_fails_closed_and_halts_run(chat_client_base: Mo @requires_sdk async def test_interceptor_crash_at_input_fails_closed(chat_client_base: MockBaseChatClient) -> None: - agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([CrashingGuard("input")])) + agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware([CrashingGuard("input")])]) with pytest.raises(InterceptionBlocked) as exc_info: await agent.run("hello") @@ -650,7 +694,7 @@ async def test_streaming_requires_no_partial_enforcement_on_error(chat_client_ba records: list[InterceptionRecord] = [] agent = Agent( client=chat_client_base, - middleware=agent_hooks_middleware([CrashingGuard("output")], record_sink=records.append), + middleware=[create_agent_hooks_middleware([CrashingGuard("output")], record_sink=records.append)], ) updates: list[AgentResponseUpdate] = [] @@ -671,7 +715,7 @@ async def test_streaming_requires_no_partial_enforcement_on_error(chat_client_ba async def test_concurrent_runs_are_isolated(chat_client_base: MockBaseChatClient) -> None: records: list[InterceptionRecord] = [] agent = Agent( - client=chat_client_base, middleware=agent_hooks_middleware([AllowGuard()], record_sink=records.append) + client=chat_client_base, middleware=[create_agent_hooks_middleware([AllowGuard()], record_sink=records.append)] ) await asyncio.gather(agent.run("first"), agent.run("second")) @@ -705,7 +749,7 @@ async def test_host_owned_session_spans_runs(chat_client_base: MockBaseChatClien emitter.register(AllowGuard()) emitter.set_record_sink(records.append) builder = AgentContextBuilder(agent_id="host-agent", framework="agent-framework", session_id="session-1") - agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware(emitter=emitter, builder=builder)) + agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware_from_emitter(emitter, builder)]) # The host owns the session boundaries. await emitter.emit(builder.agent_startup(tools_registered=[])) @@ -748,7 +792,7 @@ def resolve(self, request: ApprovalRequest) -> ApprovalResolution: agent = Agent( client=chat_client_base, tools=[weather_tool], - middleware=agent_hooks_middleware([guard], resolver=ApproveAll(), record_sink=records.append), + middleware=[create_agent_hooks_middleware([guard], resolver=ApproveAll(), record_sink=records.append)], ) response = await agent.run("get the weather") @@ -768,7 +812,7 @@ async def test_evaluate_only_records_but_never_blocks(chat_client_base: MockBase agent = Agent( client=chat_client_base, tools=[weather_tool], - middleware=agent_hooks_middleware([guard], mode="evaluate_only", record_sink=records.append), + middleware=[create_agent_hooks_middleware([guard], mode="evaluate_only", record_sink=records.append)], ) response = await agent.run("get the weather") @@ -780,17 +824,25 @@ async def test_evaluate_only_records_but_never_blocks(chat_client_base: MockBase assert deny_records[0].mode.value == "evaluate_only" -@requires_sdk -async def test_partial_trio_fails_closed(chat_client_base: MockBaseChatClient) -> None: - # Installing only part of the trio (here: on a bare chat client, where the agent - # middleware never runs) must fail closed, not silently skip enforcement. - trio = agent_hooks_middleware([AllowGuard()]) - chat_client_base.chat_middleware = cast( - "list[Any]", - [middleware for middleware in trio if type(middleware).__name__.endswith("ChatMiddleware")], +def _bundle_member(bundle: MiddlewareBundle, suffix: str) -> Any: + """Reach into a bundle's private members (tests only) to simulate a broken install.""" + return next( + member + for member in bundle._middleware # pyright: ignore[reportPrivateUsage] + if type(member).__name__.endswith(suffix) ) - with pytest.raises(MiddlewareException, match="installed together"): + +@requires_sdk +async def test_chat_seam_without_run_state_fails_closed(chat_client_base: MockBaseChatClient) -> None: + # The bundle makes a partial install impossible through the public API; this + # exercises the internal defense directly: the private chat middleware invoked + # without an active agent-hooks run (its agent sibling never ran) must fail + # closed, not silently skip enforcement. + bundle = create_agent_hooks_middleware([AllowGuard()]) + chat_client_base.chat_middleware = cast("list[Any]", [_bundle_member(bundle, "ChatMiddleware")]) + + with pytest.raises(MiddlewareException, match="without an active agent-hooks run"): await chat_client_base.get_response([Message(role="user", contents=["hi"])]) @@ -842,7 +894,7 @@ async def test_agent_seam_short_circuit_result_is_guarded(chat_client_base: Mock agent = Agent( client=chat_client_base, middleware=[ - *agent_hooks_middleware([AllowGuard()], record_sink=records.append), + create_agent_hooks_middleware([AllowGuard()], record_sink=records.append), AgentShortCircuit(substituted), ], ) @@ -861,7 +913,7 @@ async def test_agent_seam_short_circuit_result_can_be_denied(chat_client_base: M substituted = AgentResponse(messages=[Message(role="assistant", contents=["cached payload"])]) agent = Agent( client=chat_client_base, - middleware=[*agent_hooks_middleware([guard]), AgentShortCircuit(substituted)], + middleware=[create_agent_hooks_middleware([guard]), AgentShortCircuit(substituted)], ) with pytest.raises(InterceptionBlocked) as exc_info: @@ -876,7 +928,10 @@ async def test_chat_seam_short_circuit_result_is_guarded(chat_client_base: MockB substituted = ChatResponse(messages=[Message(role="assistant", contents=["cached model reply"])]) agent = Agent( client=chat_client_base, - middleware=[*agent_hooks_middleware([AllowGuard()], record_sink=records.append), ChatShortCircuit(substituted)], + middleware=[ + create_agent_hooks_middleware([AllowGuard()], record_sink=records.append), + ChatShortCircuit(substituted), + ], ) response = await agent.run("hello") @@ -900,7 +955,7 @@ async def test_chat_seam_short_circuit_result_can_be_denied(chat_client_base: Mo substituted = ChatResponse(messages=[Message(role="assistant", contents=["cached model reply"])]) agent = Agent( client=chat_client_base, - middleware=[*agent_hooks_middleware([guard]), ChatShortCircuit(substituted)], + middleware=[create_agent_hooks_middleware([guard]), ChatShortCircuit(substituted)], ) with pytest.raises(InterceptionBlocked) as exc_info: @@ -916,7 +971,7 @@ async def test_streaming_short_circuit_stream_is_guarded(chat_client_base: MockB agent = Agent( client=chat_client_base, middleware=[ - *agent_hooks_middleware([AllowGuard()], record_sink=records.append), + create_agent_hooks_middleware([AllowGuard()], record_sink=records.append), AgentShortCircuit(_cached_agent_stream("cached stream")), ], ) @@ -935,7 +990,7 @@ async def test_streaming_short_circuit_deny_releases_nothing(chat_client_base: M guard = PointGuard("output", Verdict.deny(reason="egress_blocked")) agent = Agent( client=chat_client_base, - middleware=[*agent_hooks_middleware([guard]), AgentShortCircuit(_cached_agent_stream("cached stream"))], + middleware=[create_agent_hooks_middleware([guard]), AgentShortCircuit(_cached_agent_stream("cached stream"))], ) updates: list[AgentResponseUpdate] = [] @@ -951,7 +1006,7 @@ async def test_streaming_short_circuit_without_result_closes_trail(chat_client_b records: list[InterceptionRecord] = [] agent = Agent( client=chat_client_base, - middleware=[*agent_hooks_middleware([AllowGuard()], record_sink=records.append), AgentShortCircuit(None)], + middleware=[create_agent_hooks_middleware([AllowGuard()], record_sink=records.append), AgentShortCircuit(None)], ) updates: list[AgentResponseUpdate] = [] @@ -971,7 +1026,7 @@ async def test_function_seam_foreign_termination_is_bracketed(chat_client_base: client=chat_client_base, tools=[weather_tool], middleware=[ - *agent_hooks_middleware({"allow": guard}, record_sink=records.append), + create_agent_hooks_middleware({"allow": guard}, record_sink=records.append), FunctionShortCircuit({"substituted": "tool result"}), ], ) @@ -994,7 +1049,7 @@ async def test_function_seam_foreign_termination_result_can_be_denied(chat_clien agent = Agent( client=chat_client_base, tools=[weather_tool], - middleware=[*agent_hooks_middleware([guard]), FunctionShortCircuit({"substituted": "tool result"})], + middleware=[create_agent_hooks_middleware([guard]), FunctionShortCircuit({"substituted": "tool result"})], ) response = await agent.run("get the weather") @@ -1016,7 +1071,7 @@ async def test_pre_tool_call_transform_to_non_object_fails_closed(chat_client_ba Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target", value="oops")), ) chat_client_base.run_responses = [tool_call_response(), final_response()] - agent = Agent(client=chat_client_base, tools=[weather_tool], middleware=agent_hooks_middleware([guard])) + agent = Agent(client=chat_client_base, tools=[weather_tool], middleware=[create_agent_hooks_middleware([guard])]) with pytest.raises(MiddlewareException, match="arguments object"): await agent.run("get the weather") @@ -1030,7 +1085,7 @@ async def test_input_role_transform_is_written_back(chat_client_base: MockBaseCh "input", Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target.role", value="system")), ) - agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware([guard])]) message = Message(role="user", contents=["hello"]) await agent.run([message]) @@ -1044,7 +1099,7 @@ async def test_multi_message_input_role_transform_fails_closed(chat_client_base: "input", Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target.role", value="system")), ) - agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware([guard])]) with pytest.raises(MiddlewareException, match="cannot be written back"): await agent.run([Message(role="user", contents=["one"]), Message(role="user", contents=["two"])]) @@ -1058,7 +1113,7 @@ async def test_non_string_finish_reason_transform_fails_closed(chat_client_base: "post_model_call", Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target.finish_reason", value=42)), ) - agent = Agent(client=chat_client_base, middleware=agent_hooks_middleware([guard])) + agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware([guard])]) with pytest.raises(MiddlewareException, match="finish_reason a string"): await agent.run("hello") @@ -1066,40 +1121,28 @@ async def test_non_string_finish_reason_transform_fails_closed(chat_client_base: @requires_sdk async def test_enforcement_failure_at_function_seam_halts_run(chat_client_base: MockBaseChatClient) -> None: - class Unprojectable: - def __str__(self) -> str: - raise RuntimeError("no projection") + from unittest.mock import patch - @tool(approval_mode="never_require", result_parser=SKIP_PARSING) - def opaque_tool(location: str) -> Any: - """Return a value the enforcement layer cannot project.""" - return Unprojectable() + from agent_framework._agent_hooks import _ToolResultCodec records: list[InterceptionRecord] = [] - chat_client_base.run_responses = [ - ChatResponse( - messages=[ - Message( - role="assistant", - contents=[ - Content.from_function_call(call_id="c7", name="opaque_tool", arguments='{"location": "x"}') - ], - ) - ] - ), - final_response(), - ] + chat_client_base.run_responses = [tool_call_response(), final_response()] agent = Agent( client=chat_client_base, - tools=[opaque_tool], - middleware=agent_hooks_middleware([AllowGuard()], record_sink=records.append), + tools=[weather_tool], + middleware=[create_agent_hooks_middleware([AllowGuard()], record_sink=records.append)], ) - # An unexpected failure inside the enforcement layer must halt the run, not - # degrade into a tool error that lets the run continue unaudited. - with pytest.raises(MiddlewareException, match="post_tool_call enforcement failed"): - await agent.run("run the opaque tool") + # An unexpected failure inside the enforcement layer itself (here: a projection + # bug simulated by patching the tool-result codec) must halt the run, not degrade + # into a tool error that lets the run continue unaudited. + with ( + patch.object(_ToolResultCodec, "to_wire", side_effect=RuntimeError("projection bug")), + pytest.raises(MiddlewareException, match="post_tool_call enforcement failed"), + ): + await agent.run("get the weather") + assert weather_tool_calls == ["Seattle"] # the tool ran; the enforcement layer failed after assert points(records)[-1] == "agent_shutdown" # the record trail is closed @@ -1108,34 +1151,18 @@ def opaque_tool(location: str) -> Any: # region Partial installs fail closed -@requires_sdk -async def test_agent_seam_verifies_sibling_middleware(chat_client_base: MockBaseChatClient) -> None: - trio = agent_hooks_middleware([AllowGuard()]) - - # agent + chat, function middleware missing - agent = Agent(client=chat_client_base, middleware=[trio[0], trio[1]]) - with pytest.raises(MiddlewareException, match="installed together"): - await agent.run("hello") - assert chat_client_base.call_count == 0 - - # agent + function, chat middleware missing - agent = Agent(client=chat_client_base, middleware=[trio[0], trio[2]]) - with pytest.raises(MiddlewareException, match="installed together"): - await agent.run("hello") - assert chat_client_base.call_count == 0 - - # agent middleware alone - agent = Agent(client=chat_client_base, middleware=[trio[0]]) - with pytest.raises(MiddlewareException, match="installed together"): - await agent.run("hello") - assert chat_client_base.call_count == 0 - - @requires_sdk async def test_function_seam_without_run_state_blocks_tool(chat_client_base: MockBaseChatClient) -> None: - trio = agent_hooks_middleware([AllowGuard()]) + # 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. + bundle = create_agent_hooks_middleware([AllowGuard()]) chat_client_base.run_responses = [tool_call_response(), final_response()] - agent = Agent(client=chat_client_base, tools=[weather_tool], middleware=[trio[2]]) + agent = Agent( + client=chat_client_base, + tools=[weather_tool], + middleware=[cast("Any", _bundle_member(bundle, "FunctionMiddleware"))], + ) response = await agent.run("get the weather") @@ -1143,27 +1170,30 @@ async def test_function_seam_without_run_state_blocks_tool(chat_client_base: Moc 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 "installed together" in transcript + assert "without an active agent-hooks run" in transcript @requires_sdk -async def test_fresh_trio_per_run_is_not_conflated_by_pipeline_caching(chat_client_base: MockBaseChatClient) -> None: - # The agent middleware pipeline is cached and compared with ==; the trio must keep - # identity semantics so a field-equal fresh trio on the next run is not conflated - # with the cached one (which would trip the sibling verification on run 2). +async def test_fresh_bundle_per_run_is_not_conflated_by_pipeline_caching(chat_client_base: MockBaseChatClient) -> None: + # The agent middleware pipeline is cached and compared with ==; the bundle members + # must keep identity semantics so a field-equal fresh bundle on the next run is not + # conflated with the cached one (which would bind run state to the wrong bundle). + from agent_framework._middleware import categorize_middleware + guard = AllowGuard() agent = Agent(client=chat_client_base) - first = await agent.run("one", middleware=agent_hooks_middleware([guard])) - second = await agent.run("two", middleware=agent_hooks_middleware([guard])) + first = await agent.run("one", middleware=[create_agent_hooks_middleware([guard])]) + second = await agent.run("two", middleware=[create_agent_hooks_middleware([guard])]) assert first.text == "test response - one" assert second.text == "test response - two" - # Identity semantics: fresh trios never compare equal, and the objects stay hashable. - trio_a = agent_hooks_middleware([guard]) - trio_b = agent_hooks_middleware([guard]) - assert all(a != b for a, b in zip(trio_a, trio_b)) - assert all(isinstance(hash(middleware), int) for middleware in trio_a) + # Identity semantics: fresh bundles' members never compare equal, and stay hashable. + members_a = categorize_middleware([create_agent_hooks_middleware([guard])]) + members_b = categorize_middleware([create_agent_hooks_middleware([guard])]) + for category in ("agent", "chat", "function"): + assert members_a[category][0] != members_b[category][0] # type: ignore[literal-required] + assert isinstance(hash(members_a[category][0]), int) # type: ignore[literal-required] @requires_sdk @@ -1173,8 +1203,8 @@ async def test_stacked_trios_fail_loudly(chat_client_base: MockBaseChatClient) - agent = Agent( client=chat_client_base, middleware=[ - *agent_hooks_middleware([AllowGuard()], record_sink=records_a.append), - *agent_hooks_middleware([AllowGuard()], record_sink=records_b.append), + create_agent_hooks_middleware([AllowGuard()], record_sink=records_a.append), + create_agent_hooks_middleware([AllowGuard()], record_sink=records_b.append), ], ) @@ -1199,7 +1229,7 @@ async def test_nested_agents_with_their_own_trios_stay_isolated(chat_client_base inner_agent = Agent( client=inner_client, name="inner", - middleware=agent_hooks_middleware([AllowGuard()], record_sink=records_inner.append), + middleware=[create_agent_hooks_middleware([AllowGuard()], record_sink=records_inner.append)], ) @tool(approval_mode="never_require") @@ -1225,7 +1255,7 @@ async def ask_inner(question: str) -> str: client=chat_client_base, name="outer", tools=[ask_inner], - middleware=agent_hooks_middleware([AllowGuard()], record_sink=records_outer.append), + middleware=[create_agent_hooks_middleware([AllowGuard()], record_sink=records_outer.append)], ) response = await outer_agent.run("go ask the inner agent") @@ -1259,7 +1289,7 @@ async def process(self, context: AgentContext, call_next: Callable[[], Awaitable records: list[InterceptionRecord] = [] agent = Agent( client=chat_client_base, - middleware=[*agent_hooks_middleware([AllowGuard()], record_sink=records.append), Boom()], + middleware=[create_agent_hooks_middleware([AllowGuard()], record_sink=records.append), Boom()], ) updates: list[AgentResponseUpdate] = [] @@ -1278,7 +1308,7 @@ async def process(self, context: AgentContext, call_next: Callable[[], Awaitable await call_next() context.result = cast(Any, "plain string") - agent = Agent(client=chat_client_base, middleware=[*agent_hooks_middleware([AllowGuard()]), BadResult()]) + agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware([AllowGuard()]), BadResult()]) with pytest.raises(MiddlewareException, match="cannot guard a run result"): await agent.run("hello") @@ -1288,24 +1318,469 @@ async def process(self, context: AgentContext, call_next: Callable[[], Awaitable async def test_unguardable_chat_result_fails_closed(chat_client_base: MockBaseChatClient) -> None: agent = Agent( client=chat_client_base, - middleware=[*agent_hooks_middleware([AllowGuard()]), ChatShortCircuit(cast(Any, "plain string"))], + middleware=[create_agent_hooks_middleware([AllowGuard()]), ChatShortCircuit(cast(Any, "plain string"))], ) with pytest.raises(MiddlewareException, match="cannot guard a chat result"): await agent.run("hello") +# endregion + +# region Persistence is gated behind verdicts + + +@requires_sdk +@pytest.mark.parametrize("streaming", [False, True], ids=["non_streaming", "streaming"]) +async def test_denied_output_never_becomes_durable_history(streaming: bool) -> None: + from agent_framework import AgentSession, InMemoryHistoryProvider + + client = MockBaseChatClient() + provider = InMemoryHistoryProvider() + session = AgentSession() + guard = PointGuard("output", Verdict.deny(reason="egress_blocked")) + agent = Agent(client=client, context_providers=[provider], middleware=[create_agent_hooks_middleware([guard])]) + + with pytest.raises(InterceptionBlocked): + if streaming: + async for _ in agent.run("hello there", session=session, stream=True): + pass + else: + await agent.run("hello there", session=session) + + # The verdict preceded durability: nothing (input or response) was persisted. + stored = session.state.get(provider.source_id, {}).get("messages", []) + assert stored == [] + + +@requires_sdk +@pytest.mark.parametrize("streaming", [False, True], ids=["non_streaming", "streaming"]) +async def test_transformed_output_is_persisted_post_transform(streaming: bool) -> None: + from agent_framework import AgentSession, InMemoryHistoryProvider + + client = MockBaseChatClient() + provider = InMemoryHistoryProvider() + session = AgentSession() + guard = PointGuard( + "output", + Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target.content", value="[redacted]")), + ) + agent = Agent(client=client, context_providers=[provider], middleware=[create_agent_hooks_middleware([guard])]) + + if streaming: + stream = agent.run("hello there", session=session, stream=True) + async for _ in stream: + pass + else: + await agent.run("hello there", session=session) + + # History stores the redacted response, never the unredacted original. + stored = cast("list[Message]", session.state[provider.source_id]["messages"]) + stored_texts = [message.text for message in stored] + assert "[redacted]" in stored_texts + assert not any("test response" in text for text in stored_texts) + + +@requires_sdk +@pytest.mark.parametrize("streaming", [False, True], ids=["non_streaming", "streaming"]) +async def test_denied_response_never_persists_per_service_call(streaming: bool) -> None: + from agent_framework import AgentSession, InMemoryHistoryProvider + + client = MockBaseChatClient() + provider = InMemoryHistoryProvider() + session = AgentSession() + session.state[provider.source_id] = {"messages": []} + guard = PointGuard("post_model_call", Verdict.deny(reason="response_blocked")) + agent = Agent( + client=client, + context_providers=[provider], + require_per_service_call_history_persistence=True, + middleware=[create_agent_hooks_middleware([guard])], + ) + + with pytest.raises(InterceptionBlocked): + if streaming: + async for _ in agent.run("hi", session=session, stream=True): + pass + else: + await agent.run("hi", session=session) + + # The per-service-call persist was deferred behind the post_model_call verdict + # and dropped on deny: the denied response is not durable and cannot reload. + assert session.state[provider.source_id]["messages"] == [] + + +@requires_sdk +async def test_per_service_call_persistence_still_persists_on_allow() -> None: + from agent_framework import AgentSession, InMemoryHistoryProvider + + client = MockBaseChatClient() + provider = InMemoryHistoryProvider() + session = AgentSession() + session.state[provider.source_id] = {"messages": []} + agent = Agent( + client=client, + context_providers=[provider], + require_per_service_call_history_persistence=True, + middleware=[create_agent_hooks_middleware([AllowGuard()])], + ) + + await agent.run("hi", session=session) + + stored = cast("list[Message]", session.state[provider.source_id]["messages"]) + assert [message.text for message in stored] == ["hi", "test response - hi"] + + +# endregion + +# region Stream hooks cannot escape the gate + + +@requires_sdk +async def test_stream_hooks_cannot_rewrite_egress_after_the_verdict(chat_client_base: MockBaseChatClient) -> None: + seen_at_output: list[Any] = [] + + class RecordingOutputGuard: + def intercept(self, context: dict[str, Any]) -> Any: + if context["interception_point"] == "output": + seen_at_output.append(context["target"]["content"]) + return ALLOW + + class HookInjector(AgentMiddleware): + """Previously: rewrote egressed updates AFTER the output verdict (fail-open).""" + + async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + def sneak(update: AgentResponseUpdate) -> AgentResponseUpdate: + for content in update.contents: + if content.type == "text" and content.text: + content.text = content.text + " INJECTED-AFTER-VERDICT" + return update + + context.stream_transform_hooks.append(sneak) + await call_next() + + agent = Agent( + client=chat_client_base, + middleware=[create_agent_hooks_middleware([RecordingOutputGuard()]), HookInjector()], + ) + + updates: list[str] = [] + stream = agent.run("hi", stream=True) + async for update in stream: + updates.append(update.text) + final = await stream.get_final_response() + + # Streamed egress and the final response match the verdicted content exactly; + # the hook's rewrite could not escape the gate. + assert seen_at_output == ["update - hi"] + assert "".join(updates) == "update - hi" + assert final.text == "update - hi" + + +async def test_gated_response_stream_applies_pending_hooks_before_the_gate_and_seals() -> None: + # Contract test for ResponseStream.buffered_and_gated (no SDK required). + order: list[str] = [] + + async def consume() -> tuple[list[str], str]: + order.append("consume") + return ["a", "b"], "ab" + + async def gate(updates: list[str], final: str, hooks_applied: bool) -> tuple[list[str], str]: + order.append(f"gate(hooks_applied={hooks_applied})") + assert updates == ["a!", "b!"] # pending transform hooks applied pre-gate + assert final == "ab!" # pending result hooks applied pre-gate + return updates, final + + stream = cast("Any", ResponseStream).buffered_and_gated(consume=consume, gate=gate) + + def transform(update: str) -> str: + order.append(f"transform({update})") + return update + "!" + + def result_hook(final: str) -> str: + order.append("result_hook") + return final + "!" + + # Hooks registered before consumption (e.g. by pipelines after unwinding)... + stream.with_transform_hook(transform) + stream.with_result_hook(result_hook) + + released = [update async for update in stream] + assert released == ["a!", "b!"] + assert await stream.get_final_response() == "ab!" + assert order == ["consume", "transform(a)", "transform(b)", "result_hook", "gate(hooks_applied=True)"] + + # ...and once the gate has run, content is sealed: further hooks are rejected. + with pytest.raises(RuntimeError, match="sealed"): + stream.with_transform_hook(transform) + with pytest.raises(RuntimeError, match="sealed"): + stream.with_result_hook(result_hook) + + +# endregion + +# region Approval requests pass through un-bracketed + + +@requires_sdk +async def test_approval_request_on_normal_return_path_passes_through(chat_client_base: MockBaseChatClient) -> None: + class ApprovalGate(FunctionMiddleware): + """Framework pattern: request human approval by substituting a control object.""" + + 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={} + ), + ) + # Normal return (no MiddlewareTermination): the loop still passes the + # approval request through to the caller. + + 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 human-approval pause survived: the control object is + # passed through un-bracketed (no post_tool_call reporting a value for a tool + # that never executed), mirroring the termination-branch handling. + assert weather_tool_calls == [] + assert "post_tool_call" not in points(records) + 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 + + +@requires_sdk +async def test_tool_call_name_transform_is_applied(chat_client_base: MockBaseChatClient) -> None: + executed: list[str] = [] + + @tool(approval_mode="never_require") + def other_tool(location: str) -> str: + """The tool the transform redirects to.""" + executed.append(location) + return "other tool ran" + + guard = PointGuard( + "post_model_call", + lambda ctx: ( + Verdict( + decision=Decision.TRANSFORM, + transform=Transform(path="$target.tool_calls[0].name", value="other_tool"), + ) + if ctx["target"].get("tool_calls") + else ALLOW + ), + ) + chat_client_base.run_responses = [tool_call_response("Seattle"), final_response()] + agent = Agent( + client=chat_client_base, + tools=[weather_tool, other_tool], + middleware=[create_agent_hooks_middleware([guard])], + ) + + await agent.run("get the weather") + + # The rename was applied (not silently dropped): the renamed tool executed. + assert weather_tool_calls == [] + assert executed == ["Seattle"] + + +@requires_sdk +async def test_tool_call_args_transform_must_stay_an_object(chat_client_base: MockBaseChatClient) -> None: + guard = PointGuard( + "post_model_call", + lambda ctx: ( + Verdict( + decision=Decision.TRANSFORM, + transform=Transform(path="$target.tool_calls[0].args", value="oops"), + ) + if ctx["target"].get("tool_calls") + else ALLOW + ), + ) + 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])]) + + with pytest.raises(MiddlewareException, match="args an object"): + await agent.run("get the weather") + + assert weather_tool_calls == [] # the broken transform never reached execution + + +# endregion + +# region Codec unit tests (no Agent required) + + +def _codecs() -> Any: + import agent_framework._agent_hooks as module + + return module + + +def test_input_codec_maps_roles_onto_the_spec_enum() -> None: + codecs = _codecs() + wire = codecs._InputCodec.to_wire([ + Message(role="assistant", contents=["from another agent"]), + Message(role="user", contents=["hi"]), + ]) + assert wire["role"] == "user" + assert [part["role"] for part in wire["content"]] == ["external", "user"] + + +def test_tool_arguments_codec_merges_only_changed_keys() -> None: + codecs = _codecs() + native = {"location": "Seattle", "blob": b"\x00\x01"} + before = codecs._ToolArgumentsCodec.to_wire(native) + after = dict(before) + after["location"] = "Redmond" + + merged, effective = codecs._ToolArgumentsCodec.write_back(native, before, after) + + # Only the transformed key takes the wire value; untouched keys keep their + # original native values (bytes survive, not their base64 projection). + assert merged["location"] == "Redmond" + assert merged["blob"] is native["blob"] + assert effective == after + + # Untouched wire value -> untouched native value. + same, _ = codecs._ToolArgumentsCodec.write_back(native, before, dict(before)) + assert same is native + + # Removed keys are dropped; added keys appear. + shrunk, _ = codecs._ToolArgumentsCodec.write_back(native, before, {"location": "Seattle", "extra": 1}) + assert shrunk == {"location": "Seattle", "extra": 1} + + with pytest.raises(MiddlewareException, match="arguments object"): + codecs._ToolArgumentsCodec.write_back(native, before, "oops") + + +def test_message_list_write_back_matches_by_identity_not_position() -> None: + codecs = _codecs() + originals = [ + Message(role="user", contents=["one"]), + Message(role="user", contents=["two"]), + Message(role="user", contents=["three"]), + ] + before = [codecs._message_to_wire(message) for message in originals] + + # Removing the middle message must not shift "three" onto the "two" original + # (which would duplicate content the interceptor never approved). + removed = codecs._write_back_message_list(originals, before, [before[0], before[2]], point="test") + assert [message.text for message in removed] == ["one", "three"] + assert removed[0] is originals[0] + assert removed[1] is originals[2] + assert originals[1].text == "two" # the removed original was not mutated + + # A changed entry mutates the original it replaces (shared history adoption)... + originals2 = [Message(role="user", contents=["one"]), Message(role="user", contents=["two"])] + before2 = [codecs._message_to_wire(message) for message in originals2] + changed = codecs._write_back_message_list( + originals2, before2, [before2[0], {"role": "user", "content": "TWO"}], point="test" + ) + assert changed[1] is originals2[1] + assert originals2[1].text == "TWO" + + # ...while an insertion before a preserved entry becomes a new message. + originals3 = [Message(role="user", contents=["one"])] + before3 = [codecs._message_to_wire(message) for message in originals3] + inserted = codecs._write_back_message_list( + originals3, before3, [{"role": "user", "content": "new"}, before3[0]], point="test" + ) + assert [message.text for message in inserted] == ["new", "one"] + assert inserted[1] is originals3[0] + assert originals3[0].text == "one" # the preserved original was not mutated + + +def test_model_response_codec_surfaces_hosted_tool_calls_in_content() -> None: + codecs = _codecs() + response = ChatResponse( + messages=[ + Message( + role="assistant", + contents=[ + Content.from_text("checking..."), + Content.from_function_call( + "h1", "hosted_web_search", arguments={"q": "x"}, informational_only=True + ), + Content.from_function_call("c1", "weather_tool", arguments={"location": "Seattle"}), + ], + ) + ] + ) + wire = codecs._ModelResponseCodec.to_wire(response) + # Host-executed calls ride tool_calls; the hosted (service-executed) call is part + # of the response content, so it is still interceptable at post_model_call. + assert [call["name"] for call in wire["tool_calls"]] == ["weather_tool"] + content_names = [part.get("name") for part in wire["content"][0]["content"] if isinstance(part, dict)] + assert "hosted_web_search" in content_names + + +def test_tool_call_name_and_args_write_back_rules() -> None: + codecs = _codecs() + response = ChatResponse( + messages=[Message(role="assistant", contents=[Content.from_function_call("c1", "a_tool", arguments={"x": 1})])] + ) + before = codecs._ModelResponseCodec.to_wire(response) + + renamed = {**before, "tool_calls": [{"id": "c1", "name": "b_tool", "args": {"x": 1}}]} + assert codecs._ModelResponseCodec.write_back(response, before, renamed) is True + assert response.messages[0].contents[0].name == "b_tool" + + broken_args = {**before, "tool_calls": [{"id": "c1", "name": "b_tool", "args": None}]} + with pytest.raises(MiddlewareException, match="args an object"): + codecs._ModelResponseCodec.write_back(response, before, broken_args) + + missing_args = {**before, "tool_calls": [{"id": "c1", "name": "b_tool"}]} + with pytest.raises(MiddlewareException, match="args an object"): + codecs._ModelResponseCodec.write_back(response, before, missing_args) + + +def test_tool_result_codec_round_trip() -> None: + codecs = _codecs() + original = [Content.from_text("weather in Seattle")] + wire = codecs._ToolResultCodec.to_wire(original) + assert wire == "weather in Seattle" + # Untouched wire value -> shape-preserving native value. + written = codecs._ToolResultCodec.write_back(original, "scrubbed") + assert isinstance(written[0], Content) + assert written[0].text == "scrubbed" + + +def test_output_codec_untouched_target_is_a_no_op() -> None: + codecs = _codecs() + response = AgentResponse(messages=[Message(role="assistant", contents=["hello"])]) + before = codecs._OutputCodec.to_wire(response) + assert codecs._OutputCodec.write_back(response, before, {"content": before}) is False + assert response.messages[0].text == "hello" + + # endregion # region Optional dependency -def test_agent_hooks_middleware_importable_without_sdk(monkeypatch: pytest.MonkeyPatch) -> None: +def _hide_agent_hooks(monkeypatch: pytest.MonkeyPatch, *, error: BaseException | None = None) -> None: + """Make imports of ``agent_hooks`` fail (with a custom error to simulate breakage).""" import builtins import sys - import agent_framework._agent_hooks as agent_hooks_module - real_import = builtins.__import__ def _import_without_agent_hooks( @@ -1316,6 +1791,8 @@ def _import_without_agent_hooks( level: int = 0, ) -> object: if name == "agent_hooks" or name.startswith("agent_hooks."): + if error is not None: + raise error raise ModuleNotFoundError(f"No module named '{name}'", name="agent_hooks") return real_import(name, globals_, locals_, fromlist, level) @@ -1324,11 +1801,35 @@ def _import_without_agent_hooks( monkeypatch.delitem(sys.modules, module_name) monkeypatch.setattr(builtins, "__import__", _import_without_agent_hooks) - # The lazy root export and the module itself stay importable without the SDK. - assert agent_framework.agent_hooks_middleware is agent_hooks_module.agent_hooks_middleware + +def test_agent_hooks_middleware_importable_without_sdk(monkeypatch: pytest.MonkeyPatch) -> None: + import agent_framework._agent_hooks as agent_hooks_module + + _hide_agent_hooks(monkeypatch) + + # The lazy root exports and the module itself stay importable without the SDK. + assert agent_framework.create_agent_hooks_middleware is agent_hooks_module.create_agent_hooks_middleware + assert ( + agent_framework.create_agent_hooks_middleware_from_emitter + is agent_hooks_module.create_agent_hooks_middleware_from_emitter + ) with pytest.raises(ModuleNotFoundError, match=r"agent-framework-core\[agent-hooks\]"): - agent_framework.agent_hooks_middleware([cast("Any", object())]) + agent_framework.create_agent_hooks_middleware([cast("Any", object())]) + with pytest.raises(ModuleNotFoundError, match=r"agent-framework-core\[agent-hooks\]"): + agent_framework.create_agent_hooks_middleware_from_emitter(cast("Any", object()), cast("Any", object())) + + +def test_broken_sdk_installation_is_not_masked_as_missing_extra(monkeypatch: pytest.MonkeyPatch) -> None: + # A transitively missing dependency (or any other breakage inside the SDK) must + # propagate unchanged — only a genuinely absent `agent_hooks` package gets the + # install-the-extra hint. + _hide_agent_hooks( + monkeypatch, error=ModuleNotFoundError("No module named 'some_native_dep'", name="some_native_dep") + ) + + with pytest.raises(ModuleNotFoundError, match="some_native_dep"): + agent_framework.create_agent_hooks_middleware([cast("Any", object())]) # endregion diff --git a/python/packages/core/tests/core/test_middleware_with_agent.py b/python/packages/core/tests/core/test_middleware_with_agent.py index 723ef0c81b8..636b1ce1d59 100644 --- a/python/packages/core/tests/core/test_middleware_with_agent.py +++ b/python/packages/core/tests/core/test_middleware_with_agent.py @@ -69,6 +69,28 @@ async def process(self, context: AgentContext, call_next: Callable[[], Awaitable # Verify middleware execution order assert execution_order == ["agent_middleware_before", "agent_middleware_after"] + async def test_bare_middleware_at_construction_is_installed(self, client: SupportsChatGetResponse) -> None: + """A single middleware object passed bare (not in a list) at construction is installed. + + Construction-time middleware mirrors categorize_middleware's single-source + handling, matching the run-level ``middleware=`` behavior instead of silently + dropping the middleware. + """ + execution_order: list[str] = [] + + class TrackingAgentMiddleware(AgentMiddleware): + async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + execution_order.append("before") + await call_next() + execution_order.append("after") + + agent = Agent(client=client, middleware=TrackingAgentMiddleware()) # type: ignore[arg-type] + + response = await agent.run([Message(role="user", contents=["test message"])]) + + assert response is not None + assert execution_order == ["before", "after"] + async def test_class_based_function_middleware_with_chat_agent(self, client: "MockChatClient") -> None: """Test class-based function middleware with Agent.""" diff --git a/python/pyproject.toml b/python/pyproject.toml index 17731b17f69..56abcb8874a 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -52,6 +52,9 @@ dev = [ test = [ "azure-monitor-opentelemetry", "mcp[ws]", + # Optional SDK behind core's `agent-hooks` extra; declared here (like mcp[ws]) so + # isolated source checks (dependency-pyright) can resolve its API. + "agent-hooks-sdk>=0.1.0a4,<0.2", ] [tool.uv] diff --git a/python/uv.lock b/python/uv.lock index 3999ca3ee74..9e63138ef81 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -145,6 +145,7 @@ dev = [ { name = "zuban" }, ] test = [ + { name = "agent-hooks-sdk" }, { name = "azure-monitor-opentelemetry" }, { name = "mcp", extra = ["ws"] }, ] @@ -175,6 +176,7 @@ dev = [ { name = "zuban", specifier = "==0.9.0" }, ] test = [ + { name = "agent-hooks-sdk", specifier = ">=0.1.0a4,<0.2" }, { name = "azure-monitor-opentelemetry" }, { name = "mcp", extras = ["ws"] }, ] @@ -1014,7 +1016,12 @@ version = "0.1.0a4" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/c4/9b/07e9e54941ce60991ce8ba00fe9dc38ea96d885b38a86cd944a409cd7772/agent_hooks_sdk-0.1.0a4.tar.gz", hash = "sha256:e3affd44d7779b8254fd5142f7d41648e48d3d4646acc94b07a28d13c5c07b1d", size = 173745, upload-time = "2026-07-30T02:11:22.235Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/b2/f8/5d27dd5a784c7037cbdd23f8ae7982354440a333e9433b19ea02bf8a0ca3/agent_hooks_sdk-0.1.0a4-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:793cf357413546ddc74c62dac10bf62c05a0ffe69009871fd3484105995a13b7", size = 599140, upload-time = "2026-08-05T05:06:39.637Z" }, + { url = "https://files.pythonhosted.org/packages/0b/82/2f1d273fc4abd882b8a82c0c65574a0bf56c7181bc1a6be52b2bd58d880b/agent_hooks_sdk-0.1.0a4-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:7299cfd93ed455f0bcbfafb759040500e2920c6547fb6d8829f792307d58a6d4", size = 570698, upload-time = "2026-08-05T05:06:40.906Z" }, + { url = "https://files.pythonhosted.org/packages/6f/4c/26c71782f12e691958dcb59d530c5a17e8efe4651af3085c52db578e56f9/agent_hooks_sdk-0.1.0a4-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:fa51c408f29fa907be33e8a542f1025b2a92ccce0ccad0a362c299b6cb6dfdd4", size = 626367, upload-time = "2026-08-05T05:06:42.221Z" }, + { url = "https://files.pythonhosted.org/packages/a8/93/bce858a4e83480373eb55e2923657126e0e736b152f814fcf035b293e3e6/agent_hooks_sdk-0.1.0a4-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:ed556c2d020fffe371eda8055950f7b42c22eee16f3c3ef1dcf1f4ba5704320e", size = 649044, upload-time = "2026-08-05T05:06:43.492Z" }, { url = "https://files.pythonhosted.org/packages/d6/07/dc7f4e7a53c208ba35fcf0bd0208cf77a1e84ffd7a42d2a61c6e9faf274b/agent_hooks_sdk-0.1.0a4-cp310-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:86f5ff8e0c1201a61b113026c688fc98ea590b10f8bbc5ccfa6ac289fdec037d", size = 634247, upload-time = "2026-07-30T02:11:20.497Z" }, + { url = "https://files.pythonhosted.org/packages/34/2f/35caaffec73e0e60829a2ec255da1bfc7fe9564876b97eaa4a063327fd93/agent_hooks_sdk-0.1.0a4-cp310-abi3-win_amd64.whl", hash = "sha256:5e1cd64e3ef9c7447cb56a615f79747b7aa42a0065c8a2c5f0b2ae7f0447c908", size = 479900, upload-time = "2026-08-05T05:06:44.898Z" }, ] [[package]] From c1cfa830f247d131cca08ce729016c8a0f15ce9b Mon Sep 17 00:00:00 2001 From: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> Date: Wed, 5 Aug 2026 09:00:53 +0000 Subject: [PATCH 4/8] refactor(python): second review round for agent-hooks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses the second review round on the agent-hooks feature: - Nested-run persistence ownership: RawAgent.run stamps a run identity over the run's dynamic extent (including streaming pulls and result hooks); the persistence gate binds to its owning run via an offer/adopt handshake keyed to the agent instance and accepts only its owner's persists — nested runs persist inline regardless of how they were started (tool calls, middleware, custom run loops). The tool-seam suspension remains for custom-loop sub-agents invoked as tools; the one residual case (custom loop nested in a custom loop off the tool path) is fail-closed and documented. Fixes a latent pre-existing re-deferral: flush() now drains with the gate context suspended, so a nested hooked run's permitted after-run persistence no longer re-defers into an enclosing gate. - as_tool stream_callback consumes the released (verdicted) stream; observers cannot see denied or pre-transform content. Both directions are regression-tested. - categorize_middleware gained supported_categories: a bundle member landing in a category a call site cannot install raises; bare middleware warns like _add_middleware. Wired at the chat-client sites and the provider seam. - ResponseStream.buffered_and_gated owns the re-derivation rule via a rederive callable (gates cannot choose released updates) and is marked experimental. - Wire codecs compare with bool-aware equality (Python == equates 1 == True, which made bool/number transforms look untouched and get dropped) and _ToolResultCodec.write_back owns the untouched-wire rule via the before value. - middleware parameters accept a bare middleware or bundle everywhere the runtime does (constructors, run overloads, as_agent, telemetry and harness layers, foundry); the bare-source rule has a single owner in categorize_middleware; bare middleware assigned to the attribute now executes (documented behavior change). - MiddlewareBundle is experimental and validates members; approval passthrough, typing-check fixes (ty ignores mypy-coded ignore comments), logging, and documentation updates per review. Test count: 85 feature tests plus 12 new this round across sessions, middleware, agents; full core suite green; typing checked under mypy, pyrefly, ty, zuban, and pyright. Signed-off-by: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> --- python/PACKAGE_STATUS.md | 4 +- .../core/agent_framework/_agent_hooks.py | 214 +++++---- .../packages/core/agent_framework/_agents.py | 88 +++- .../packages/core/agent_framework/_clients.py | 2 +- .../core/agent_framework/_harness/_agent.py | 13 +- .../core/agent_framework/_harness/_agent.pyi | 2 +- .../core/agent_framework/_middleware.py | 124 +++-- .../core/agent_framework/_sessions.py | 218 ++++++++- .../packages/core/agent_framework/_tools.py | 38 +- .../packages/core/agent_framework/_types.py | 48 +- .../core/agent_framework/observability.py | 8 +- .../core/tests/core/test_agent_hooks.py | 429 +++++++++++++++++- .../core/tests/core/test_middleware.py | 27 ++ .../tests/core/test_middleware_with_agent.py | 27 +- .../packages/core/tests/core/test_sessions.py | 185 ++++++++ .../foundry/agent_framework_foundry/_agent.py | 6 +- 16 files changed, 1227 insertions(+), 206 deletions(-) diff --git a/python/PACKAGE_STATUS.md b/python/PACKAGE_STATUS.md index ebb16c62ff0..89b3af8cdd4 100644 --- a/python/PACKAGE_STATUS.md +++ b/python/PACKAGE_STATUS.md @@ -68,7 +68,9 @@ listed below. - `agent-framework-core`: `create_agent_hooks_middleware` and `create_agent_hooks_middleware_from_emitter` from `agent_framework/_agent_hooks.py`, - the AGENT-HOOKS-0.1 enforcement middleware bundle. Requires the opt-in + the AGENT-HOOKS-0.1 enforcement middleware bundle, and the `MiddlewareBundle` + container from `agent_framework/_middleware.py` that both factories produce + (`MiddlewareBundle` itself needs no extra). Requires the opt-in `agent-framework-core[agent-hooks]` extra (`agent-hooks-sdk`), which is deliberately not part of `agent-framework-core[all]`. Known limitation: service-side (hosted) tool execution never passes through the framework's function-invocation seam, so the diff --git a/python/packages/core/agent_framework/_agent_hooks.py b/python/packages/core/agent_framework/_agent_hooks.py index 2a61031c36e..580392efe3c 100644 --- a/python/packages/core/agent_framework/_agent_hooks.py +++ b/python/packages/core/agent_framework/_agent_hooks.py @@ -47,7 +47,21 @@ emission permits the content, so denied content never becomes durable and transformed content is persisted post-transform. Each persist is gated by its own covering verdict: per-service-call history persisted under a permitted ``post_model_call`` - verdict remains durable even if the run's ``output`` is later denied. + verdict remains durable even if the run's ``output`` is later denied. A mid-run deny + is conservative in the other direction: the dropped per-service-call persist carries + that service call's request messages too, so the denied turn's input is not + persisted either. Only the gated run's own persistence defers: nested and + middleware-initiated agent runs (sub-agents invoked as tools, agents run by other + middleware or context providers, even on a shared session) persist inline at their + own run boundaries — the gate is bound to its run's identity (see the run + persistence gate in ``_sessions.py``), and the function-invocation layer + additionally suspends the gate around tool invocations to cover nested agents with + fully custom run loops. An outer deny therefore never discards fully-permitted + inner history, and a second sub-agent call within one outer run reads fresh + history. Residual limitation: an agent whose run loop never stamps a run identity + (a fully custom ``run()`` implementation) that itself nests another such agent + outside the tool seam falls back to deferring both — fail-closed, matching the + pre-ownership behavior. Streaming is supported **fail-closed by buffering** via :meth:`ResponseStream.buffered_and_gated`: the model/agent stream is fully consumed @@ -102,7 +116,10 @@ MiddlewareTermination, ) from ._serialization import make_json_safe -from ._sessions import _DEFERRED_RUN_PERSISTENCE # pyright: ignore[reportPrivateUsage] +from ._sessions import ( + _current_run_identity, # pyright: ignore[reportPrivateUsage] + _RunPersistenceGate, # pyright: ignore[reportPrivateUsage] +) from ._types import ( AgentResponse, AgentResponseUpdate, @@ -218,22 +235,37 @@ class _RunState: _RUN_STATE: ContextVar[_RunState | None] = ContextVar("agent_framework_agent_hooks_run_state", default=None) -_PersistCallback = Callable[[], Awaitable[None]] - -async def _run_pending_persistence(pending: list[_PersistCallback]) -> None: - """Execute persistence work that was deferred behind a permitting verdict. +# endregion - The caller must have reset the persistence gate first, so re-entrant persistence - calls made by these callables run inline. - """ - for persist in pending: - await persist() +# region Wire building blocks (framework values <-> AGENT-HOOKS wire JSON) -# endregion +def _wire_equal(left: Any, right: Any) -> bool: + """Type-aware equality for wire JSON values (the codecs' untouched-target checks). -# region Wire building blocks (framework values <-> AGENT-HOOKS wire JSON) + Python's ``==`` equates ``1 == True`` and ``0 == False``, so a transform that + swaps a value between bool and number would look untouched and be silently + dropped (the untransformed native value would proceed — fail-open for the + transform). Bool-ness is therefore compared explicitly at every nesting level. + """ + if isinstance(left, bool) != isinstance(right, bool): + return False + if isinstance(left, Mapping) and isinstance(right, Mapping): + left_map = cast("Mapping[Any, Any]", left) + right_map = cast("Mapping[Any, Any]", right) + if left_map.keys() != right_map.keys(): + return False + return all(_wire_equal(item, right_map[key]) for key, item in left_map.items()) + left_is_sequence = isinstance(left, Sequence) and not isinstance(left, (str, bytes, bytearray)) + right_is_sequence = isinstance(right, Sequence) and not isinstance(right, (str, bytes, bytearray)) + if left_is_sequence and right_is_sequence: + left_items = cast("Sequence[Any]", left) + right_items = cast("Sequence[Any]", right) + if len(left_items) != len(right_items): + return False + return all(_wire_equal(left_item, right_item) for left_item, right_item in zip(left_items, right_items)) + return bool(left == right) def _role_str(role: Any) -> str: @@ -338,7 +370,7 @@ def _write_back_message_list( for index, item in enumerate(after_items): item_dict = dict(item) match_index = next( - (position for position in range(cursor, len(originals)) if before_dicts[position] == item_dict), + (position for position in range(cursor, len(originals)) if _wire_equal(before_dicts[position], item_dict)), None, ) if match_index is not None: @@ -347,7 +379,7 @@ def _write_back_message_list( continue if cursor < len(originals): candidate_projection = before_dicts[cursor] - preserved_later = any(dict(later) == candidate_projection for later in after_items[index + 1 :]) + preserved_later = any(_wire_equal(dict(later), candidate_projection) for later in after_items[index + 1 :]) role = str(item.get("role") or "user") if not preserved_later and role == str(candidate_projection.get("role")): message = originals[cursor] @@ -424,7 +456,7 @@ def to_wire(cls, messages: Sequence[Message]) -> dict[str, Any]: @classmethod def write_back(cls, messages: list[Message], before: Mapping[str, Any], after: Any) -> None: """Write a transformed ``input`` target back into the run's message list.""" - if after is None or after == before: + if after is None or _wire_equal(after, before): return if not isinstance(after, Mapping): raise _AgentHooksWriteBackError("agent-hooks input transform must produce an input object target.") @@ -440,7 +472,7 @@ def write_back(cls, messages: list[Message], before: Mapping[str, Any], after: A ) messages[0].role = after_role after_content = after_map.get("content") - if after_content == before.get("content"): + if _wire_equal(after_content, before.get("content")): return if len(messages) == 1 and not _looks_like_message_dicts(after_content): messages[0].contents = _wire_to_contents(after_content, point="input") @@ -461,7 +493,7 @@ def write_back( messages: Sequence[Message], before: Sequence[Mapping[str, Any]], after: Any ) -> list[Message] | None: """Return the transformed message list, or None when the target is untouched.""" - if after == list(before): + if _wire_equal(after, list(before)): return None return _write_back_message_list(list(messages), before, after, point="pre_model_call") @@ -517,7 +549,7 @@ def to_wire(cls, response: ChatResponse[Any]) -> dict[str, Any]: @classmethod def write_back(cls, response: ChatResponse[Any], before: Mapping[str, Any], after: Any) -> bool: """Write a transformed ``post_model_call`` target back into the chat response.""" - if after is None or after == before: + if after is None or _wire_equal(after, before): return False if not isinstance(after, Mapping): raise _AgentHooksWriteBackError("agent-hooks post_model_call transform must produce a response object.") @@ -532,10 +564,10 @@ def write_back(cls, response: ChatResponse[Any], before: Mapping[str, Any], afte response.finish_reason = cast(Any, after_finish) changed = True after_calls = after_map.get("tool_calls") - if after_calls != before.get("tool_calls"): + if not _wire_equal(after_calls, before.get("tool_calls")): changed = cls._write_back_tool_calls(response, after_calls) or changed after_content = after_map.get("content") - if after_content != before.get("content"): + if not _wire_equal(after_content, before.get("content")): cls._write_back_content(response, after_content) changed = True return changed @@ -579,7 +611,7 @@ def _write_back_tool_calls(response: ChatResponse[Any], after_calls: Any) -> boo raise _AgentHooksWriteBackError( "agent-hooks post_model_call transform must keep each tool call's args an object." ) - if _arguments_to_wire(content.arguments) != dict(cast("Mapping[str, Any]", wire_args)): + if not _wire_equal(_arguments_to_wire(content.arguments), dict(cast("Mapping[str, Any]", wire_args))): content.arguments = {str(key): item for key, item in cast("Mapping[Any, Any]", wire_args).items()} changed = True kept.append(content) @@ -662,7 +694,7 @@ def write_back(arguments: Any, before: Mapping[str, Any], after: Any) -> tuple[A if not isinstance(after, Mapping): raise _AgentHooksWriteBackError("agent-hooks pre_tool_call transform must produce an arguments object.") effective = {str(key): item for key, item in cast("Mapping[Any, Any]", after).items()} - if effective == dict(before): + if _wire_equal(effective, dict(before)): return arguments, effective if isinstance(arguments, BaseModel): native: dict[str, Any] = {name: getattr(arguments, name) for name in type(arguments).model_fields} @@ -672,7 +704,7 @@ def write_back(arguments: Any, before: Mapping[str, Any], after: Any) -> tuple[A native = {} merged = {key: item for key, item in native.items() if key in effective} for key, item in effective.items(): - if key not in before or before[key] != item: + if key not in before or not _wire_equal(before[key], item): merged[key] = item return merged, effective @@ -711,14 +743,19 @@ def to_wire(value: Any) -> Any: return make_json_safe(value) @staticmethod - def write_back(original: Any, after: Any) -> Any: + def write_back(original: Any, before: Any, after: Any) -> Any: """Convert a transformed ``post_tool_call`` value back into the native result shape. - When the original result is the framework's canonical ``list[Content]`` and the - transformed value is shape-compatible, the Content wrappers are preserved; - otherwise the transformed wire value becomes the result as-is (the function - invocation layer serializes arbitrary JSON-native results faithfully). + A wire value the interceptors left untouched maps back to the untouched native + result (the codec region's shared rule, applied here like in the other five + codecs). When the original result is the framework's canonical + ``list[Content]`` and the transformed value is shape-compatible, the Content + wrappers are preserved; otherwise the transformed wire value becomes the + result as-is (the function invocation layer serializes arbitrary JSON-native + results faithfully). """ + if _wire_equal(after, before): + return original if ( isinstance(original, list) and original @@ -731,7 +768,7 @@ def write_back(original: Any, after: Any) -> Any: if after_items is not None and len(after_items) == len(original_contents): rebuilt: list[Content] = [] for content, item in zip(original_contents, after_items): - if _ToolResultCodec.to_wire(content) == item: + if _wire_equal(_ToolResultCodec.to_wire(content), item): rebuilt.append(content) elif content.type == "text" and isinstance(item, str): rebuilt.append(Content.from_text(item)) @@ -760,7 +797,7 @@ def write_back(response: AgentResponse[Any], before_content: Any, after: Any) -> if not isinstance(after, Mapping): raise _AgentHooksWriteBackError("agent-hooks output transform must produce an output object target.") after_content = cast("Mapping[str, Any]", after).get("content") - if after_content == before_content: + if _wire_equal(after_content, before_content): return False originals = list(response.messages) if isinstance(after_content, str): @@ -972,20 +1009,21 @@ async def process(self, context: AgentContext, call_next: Callable[[], Awaitable token = _RUN_STATE.set(state) shutdown_reason = "completed" try: + gate = _RunPersistenceGate() + # Hand the gate to the pipeline's final handler via the context: the run + # it starts adopts the gate and binds it to its own identity, so only that + # run's persistence defers (middleware-initiated runs persist inline). + context._run_persistence_gate = gate # pyright: ignore[reportPrivateUsage] try: await self._emit_run_start(context, state) termination: MiddlewareTermination | None = None - pending: list[_PersistCallback] = [] - persistence_token = _DEFERRED_RUN_PERSISTENCE.set(pending) - try: + with gate: try: await call_next() except MiddlewareTermination as exc: # A middleware short-circuited the run; any substituted result # still egresses to the caller, so it passes the output point. termination = exc - finally: - _DEFERRED_RUN_PERSISTENCE.reset(persistence_token) if state.halted is not None: raise state.halted result = context.result @@ -996,13 +1034,14 @@ async def process(self, context: AgentContext, call_next: Callable[[], Awaitable f"agent-hooks cannot guard a run result of type {type(result).__name__}; " "the output interception point was not emitted." ) - # The verdict permitted the content (or nothing egresses): run the - # persistence the run deferred behind the gate. A deny above skips - # this entirely, so denied content never becomes durable. - await _run_pending_persistence(pending) + # The verdict permitted the content (or nothing egresses): release the + # persistence the run deferred behind the gate. A deny above drops it + # instead, so denied content never becomes durable. + await gate.flush() if termination is not None: raise termination except InterceptionBlocked: + gate.drop() shutdown_reason = "error" context.result = None raise @@ -1033,6 +1072,10 @@ async def _process_streaming( # the trail as an error; ``None`` means ownership was handed off. shutdown_reason: str | None = "error" try: + gate_handle = _RunPersistenceGate() + # Hand the gate to the pipeline's final handler via the context (see the + # non-streaming branch): the run it starts binds the gate to its identity. + context._run_persistence_gate = gate_handle # pyright: ignore[reportPrivateUsage] await self._emit_run_start(context, state) termination: MiddlewareTermination | None = None try: @@ -1050,7 +1093,7 @@ async def _process_streaming( f"got {type(inner).__name__}." ) context.result = self._gated_agent_stream( - state, cast("ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", inner) + state, cast("ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", inner), gate_handle ) # The gated stream now owns the shutdown emission. shutdown_reason = None @@ -1070,23 +1113,23 @@ def _gated_agent_stream( self, state: _RunState, inner: ResponseStream[AgentResponseUpdate, AgentResponse[Any]], + gate_handle: _RunPersistenceGate, ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: """Guard a streaming run with a fail-closed buffered gate. The run is fully consumed (with the run state and the persistence gate active), middleware stream hooks are applied by the combinator before the verdict, the - ``output`` verdict is applied to the finalized response, deferred persistence - runs only after the verdict permits, and only then are updates released. A - transformed output (or any applied stream hooks) re-derives the released - updates from the verdicted response so streamed egress cannot diverge from it. + ``output`` verdict is applied to the finalized response, and deferred + persistence is released only after the verdict permits. The combinator owns the + no-divergence rule: a transformed output (or any applied stream hooks) + re-derives the released updates from the verdicted response. """ - pending: list[_PersistCallback] = [] async def _consume() -> tuple[Sequence[AgentResponseUpdate], AgentResponse[Any]]: run_token = _RUN_STATE.set(state) - persistence_token = _DEFERRED_RUN_PERSISTENCE.set(pending) try: - final = await inner.get_final_response() + with gate_handle: + final = await inner.get_final_response() except asyncio.CancelledError: await self._emit_shutdown(state, "cancelled") raise @@ -1094,13 +1137,12 @@ async def _consume() -> tuple[Sequence[AgentResponseUpdate], AgentResponse[Any]] await self._emit_shutdown(state, "error") raise finally: - _DEFERRED_RUN_PERSISTENCE.reset(persistence_token) _RUN_STATE.reset(run_token) return list(inner.updates), final async def _gate( - updates: list[AgentResponseUpdate], final: AgentResponse[Any], hooks_applied: bool - ) -> tuple[Sequence[AgentResponseUpdate], AgentResponse[Any]]: + _updates: list[AgentResponseUpdate], final: AgentResponse[Any] + ) -> tuple[AgentResponse[Any], bool]: from agent_hooks import InterceptionBlocked try: @@ -1112,11 +1154,11 @@ async def _gate( "the output interception point was not emitted." ) transformed = await self._emit_output(state, final) - await _run_pending_persistence(pending) + await gate_handle.flush() await self._emit_shutdown(state, "completed") - released = _agent_updates_from_response(final) if (transformed or hooks_applied) else updates - return released, final + return final, transformed except InterceptionBlocked: + gate_handle.drop() await self._emit_shutdown(state, "error") raise except asyncio.CancelledError: @@ -1128,7 +1170,9 @@ async def _gate( return cast( "ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", - cast(Any, ResponseStream).buffered_and_gated(consume=_consume, gate=_gate), + cast(Any, ResponseStream).buffered_and_gated( + consume=_consume, gate=_gate, rederive=_agent_updates_from_response + ), ) @@ -1158,22 +1202,22 @@ async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[ context.messages = transformed_messages termination: MiddlewareTermination | None = None - pending: list[_PersistCallback] = [] - persistence_token = _DEFERRED_RUN_PERSISTENCE.set(pending) - try: + gate = _RunPersistenceGate() + # The chat seam runs inside its run's identity scope, so the gate binds to the + # current run immediately: only this run's per-service-call persists defer. + gate.bind_owner(_current_run_identity()) + with gate: try: await call_next() except MiddlewareTermination as exc: # A chat middleware short-circuited with a substituted result; whatever # was substituted still flows into the agent loop, so it is guarded below. termination = exc - finally: - _DEFERRED_RUN_PERSISTENCE.reset(persistence_token) result = context.result if isinstance(result, ResponseStream): context.result = self._gated_chat_stream( - state, model_id, cast("ResponseStream[ChatResponseUpdate, ChatResponse[Any]]", result), pending + state, model_id, cast("ResponseStream[ChatResponseUpdate, ChatResponse[Any]]", result), gate ) elif isinstance(result, ChatResponse): try: @@ -1181,9 +1225,10 @@ async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[ except InterceptionBlocked: # §6.1: the denied response must not be incorporated (and the deferred # per-service-call persistence for it is dropped, never executed). + gate.drop() context.result = None raise - await _run_pending_persistence(pending) + await gate.flush() elif result is not None: raise MiddlewareException( f"agent-hooks cannot guard a chat result of type {type(result).__name__}; " @@ -1214,40 +1259,46 @@ def _gated_chat_stream( state: _RunState, model_id: str, inner: ResponseStream[ChatResponseUpdate, ChatResponse[Any]], - pending: list[_PersistCallback], + gate_handle: _RunPersistenceGate, ) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: """Guard a streaming model call with a fail-closed buffered gate. Spec §12.1: the complete response is assembled before ``post_model_call`` is emitted, and nothing (updates or tool calls) is released beforehand. A deny raises before any update egresses, and per-service-call history persistence - deferred by the run persistence gate executes only after the verdict permits. + deferred by the run persistence gate is released only after the verdict + permits. The combinator owns the no-divergence rule: a transformed response + (or any applied stream hooks) re-derives the released updates from it. """ async def _consume() -> tuple[Sequence[ChatResponseUpdate], ChatResponse[Any]]: - persistence_token = _DEFERRED_RUN_PERSISTENCE.set(pending) - try: + with gate_handle: response = await inner.get_final_response() - finally: - _DEFERRED_RUN_PERSISTENCE.reset(persistence_token) return list(inner.updates), response - async def _gate( - updates: list[ChatResponseUpdate], final: ChatResponse[Any], hooks_applied: bool - ) -> tuple[Sequence[ChatResponseUpdate], ChatResponse[Any]]: + async def _gate(_updates: list[ChatResponseUpdate], final: ChatResponse[Any]) -> tuple[ChatResponse[Any], bool]: + from agent_hooks import InterceptionBlocked + if not isinstance(final, ChatResponse): raise MiddlewareException( f"agent-hooks cannot guard a streamed chat result of type {type(final).__name__}; " "the post_model_call interception point was not emitted." ) - changed = await self._emit_post_model_call(state, model_id, final) - await _run_pending_persistence(pending) - released = _chat_updates_from_response(final) if (changed or hooks_applied) else updates - return released, final + try: + changed = await self._emit_post_model_call(state, model_id, final) + except InterceptionBlocked: + # §6.1: the deferred per-service-call persistence for the denied + # response is dropped, never executed. + gate_handle.drop() + raise + await gate_handle.flush() + return final, changed return cast( "ResponseStream[ChatResponseUpdate, ChatResponse[Any]]", - cast(Any, ResponseStream).buffered_and_gated(consume=_consume, gate=_gate), + cast(Any, ResponseStream).buffered_and_gated( + consume=_consume, gate=_gate, rederive=_chat_updates_from_response + ), ) @@ -1362,8 +1413,7 @@ async def process(self, context: FunctionInvocationContext, call_next: Callable[ outcome = await state.emitter.emit( state.builder.post_tool_call(call_id=call_id, name=name, args=args, value=value) ) - if outcome.target != value: - context.result = _ToolResultCodec.write_back(context.result, outcome.target) + 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") @@ -1413,8 +1463,12 @@ def create_agent_hooks_middleware( (outermost), and install exactly one bundle per agent (stacked bundles are rejected fail-closed). Middleware listed before the bundle runs outside the enforcement boundary: a function middleware placed before it, for example, can - substitute a tool result that the tool seam never brackets. Outer position is - outer trust — the final ``output`` point still guards whatever egresses. + substitute a tool result that the tool seam never brackets, and stream hooks + registered by outer-position middleware are applied to buffered streaming + content ahead of the verdict, granting that middleware pre-verdict *read* + access (its rewrites remain covered by the verdict, and nothing egresses to + the caller before the verdict). Outer position is outer trust — the final + ``output`` point still guards whatever egresses. Args: interceptors: The agent-hooks interceptors to register, either as a sequence or diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index 5f1c9574e6f..7c374b070cf 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -9,6 +9,7 @@ from contextlib import AbstractAsyncContextManager, AsyncExitStack from copy import deepcopy from functools import partial +from inspect import isawaitable from itertools import chain from typing import ( TYPE_CHECKING, @@ -35,7 +36,9 @@ PerServiceCallHistoryPersistingMiddleware, ServiceSessionId, SessionContext, + _adopt_run_persistence_gate_claim, # pyright: ignore[reportPrivateUsage] _defer_run_persistence, # pyright: ignore[reportPrivateUsage] + _run_identity_scope, # pyright: ignore[reportPrivateUsage] is_local_history_conversation_id, ) from ._telemetry import FeatureIndex, mark_feature_used @@ -420,7 +423,7 @@ def __init__( name: str | None = None, description: str | None = None, context_providers: Sequence[ContextProvider] | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, additional_properties: MutableMapping[str, Any] | None = None, ) -> None: """Initialize a BaseAgent instance. @@ -431,7 +434,11 @@ def __init__( name: The name of the agent, can be None. description: The description of the agent. context_providers: Context providers to include during agent invocation. - middleware: List of middleware. + middleware: List of middleware, or a single middleware object (including a + ``MiddlewareBundle``) which is treated as a one-element list. Note: + a bare middleware object — passed here or assigned directly to the + ``middleware`` attribute — used to be silently ignored by ``run()`` + and now executes. additional_properties: Additional properties set on the agent. """ if id is None: @@ -440,10 +447,10 @@ def __init__( self.name = name self.description = description self.context_providers: list[ContextProvider] = list(context_providers or []) - # Normalize a bare middleware source (a single middleware object or a - # MiddlewareBundle) into a single-element list, mirroring how - # categorize_middleware treats non-sequence sources, so construction-time - # and run-time middleware arguments behave identically. + # Canonicalize storage: a bare middleware source (a single middleware object + # or a MiddlewareBundle) is stored as a single-element list so the declared + # attribute type stays `list[MiddlewareTypes] | None`. Interpreting bare + # sources is owned by categorize_middleware; this mirrors it for storage only. if middleware is not None and not isinstance(middleware, Sequence): middleware = [middleware] self.middleware: list[MiddlewareTypes] | None = ( @@ -666,7 +673,15 @@ async def _agent_wrapper(ctx: FunctionInvocationContext, **kwargs: Any) -> str: function_invocation_kwargs=dict(ctx.kwargs), ) if stream_callback is not None: - stream.with_transform_hook(stream_callback) + # The callback is a host-facing observer: feed it the *released* + # updates by consuming the stream, never by registering a transform + # hook on it. Hooks can end up applied to buffered content ahead of an + # egress gate's verdict (see ResponseStream.buffered_and_gated), so a + # hook-registered observer could see denied or unredacted content. + async for update in stream: + callback_result = stream_callback(update) + if isawaitable(callback_result): + await callback_result final_response = await stream.get_final_response() if final_response.user_input_requests: raise UserInputRequiredException(contents=final_response.user_input_requests) @@ -779,7 +794,7 @@ def __init__( tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, default_options: OptionsCoT | None = None, context_providers: Sequence[ContextProvider] | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, require_per_service_call_history_persistence: bool = False, compaction_strategy: CompactionStrategy | None = None, tokenizer: TokenizerProtocol | None = None, @@ -798,6 +813,10 @@ def __init__( description: A brief description of the agent's purpose. context_providers: Context providers to include during agent invocation. middleware: List of middleware to intercept agent and function invocations. + A single middleware object (including a ``MiddlewareBundle``) is + treated as a one-element list. Note: a bare middleware object — + passed here or assigned directly to the ``middleware`` attribute — + used to be silently ignored by ``run()`` and now executes. require_per_service_call_history_persistence: When True (and a HistoryProvider is present), the provider always persists history via per-service-call middleware, regardless of whether the client stores history server-side. If the client does @@ -1067,19 +1086,30 @@ async def _prepare_run_context() -> _RunContext: client_kwargs=client_kwargs, ) + # Stamp a fresh identity for this run and adopt a pending run-persistence gate + # claim targeted at this agent (offered by the middleware layer's final + # handler). The identity marks this run's dynamic extent — including the + # streaming consumption below — so an active gate defers exactly this run's + # own persistence, while nested or middleware-initiated runs (which stamp + # their own identities here) persist inline at their own run boundaries. + run_identity: object = object() + _adopt_run_persistence_gate_claim(self, run_identity) + if not stream: async def _run_non_streaming() -> AgentResponse[Any]: - ctx = await _prepare_run_context() - response = await self._call_chat_client(ctx, stream=False) - return await self._parse_non_streaming_response(ctx, response) + with _run_identity_scope(run_identity): + ctx = await _prepare_run_context() + response = await self._call_chat_client(ctx, stream=False) + return await self._parse_non_streaming_response(ctx, response) return _run_non_streaming() async def _run_streaming() -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: - ctx = await _prepare_run_context() - stream_response = self._call_chat_client(ctx, stream=True) - return self._parse_streaming_response(ctx, stream_response) + with _run_identity_scope(run_identity): + ctx = await _prepare_run_context() + stream_response = self._call_chat_client(ctx, stream=True) + return self._parse_streaming_response(ctx, stream_response, run_identity=run_identity) return cast( ResponseStream[AgentResponseUpdate, AgentResponse[Any]], @@ -1166,6 +1196,8 @@ def _parse_streaming_response( self, context: _RunContext, stream_response: ResponseStream[ChatResponseUpdate, ChatResponse[Any]], + *, + run_identity: object, ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: """Finalize a streaming chat response into an agent response stream.""" @@ -1192,7 +1224,11 @@ async def _post_hook(response: AgentResponse) -> None: if context["suppress_response_id"]: response.response_id = None session_context._response = response # type: ignore[assignment] - await self._run_after_providers(session=session, context=session_context) + # Result hooks run during finalization, outside the per-pull identity + # scope registered below, so re-stamp this run's identity around its + # run-end persistence. + with _run_identity_scope(run_identity): + await self._run_after_providers(session=session, context=session_context) def _propagate_conversation_id(update: AgentResponseUpdate) -> AgentResponseUpdate: """Eagerly propagate conversation_id to session as updates arrive.""" @@ -1231,6 +1267,13 @@ def _finalizer(updates: Sequence[AgentResponseUpdate]) -> AgentResponse[Any]: if context["suppress_response_id"]: stream = stream.with_transform_hook(_suppress_response_id) + # Streaming consumption happens in the consumer's context, outside the + # _run_identity_scope that wrapped this run's setup. Stamp the run identity + # around every underlying pull so persistence issued mid-consumption (e.g. + # per-service-call history persists inside the function-invocation loop) + # carries this run's identity; nested runs re-stamp their own within theirs. + stream = stream.with_pull_context_manager(partial(_run_identity_scope, run_identity)) + return stream.with_transform_hook(_propagate_conversation_id).with_result_hook(_post_hook) def _finalize_response_updates( @@ -1448,7 +1491,10 @@ async def _prepare_run_context( ) provider_middleware = session_context.get_middleware() if provider_middleware: - middleware_list = categorize_middleware(provider_middleware) + # Providers may only contribute chat/function middleware (enforced by + # SessionContext.extend_middleware); declare the same contract here so a + # bundle member outside these categories fails loudly at this seam too. + middleware_list = categorize_middleware(provider_middleware, supported_categories=("chat", "function")) provider_function_chat_middleware = [ *middleware_list["function"], *middleware_list["chat"], @@ -1733,7 +1779,7 @@ def run( *, stream: Literal[False] = ..., session: AgentSession | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, options: ChatOptions[ResponseModelBoundT], compaction_strategy: CompactionStrategy | None = None, @@ -1749,7 +1795,7 @@ def run( *, stream: Literal[False] = ..., session: AgentSession | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, options: OptionsCoT | ChatOptions[None] | None = None, compaction_strategy: CompactionStrategy | None = None, @@ -1765,7 +1811,7 @@ def run( *, stream: Literal[True], session: AgentSession | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, options: OptionsCoT | ChatOptions[Any] | None = None, compaction_strategy: CompactionStrategy | None = None, @@ -1780,7 +1826,7 @@ def run( *, stream: bool = False, session: AgentSession | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, options: OptionsCoT | ChatOptions[Any] | None = None, compaction_strategy: CompactionStrategy | None = None, @@ -1818,7 +1864,7 @@ def __init__( tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, default_options: OptionsCoT | None = None, context_providers: Sequence[ContextProvider] | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, require_per_service_call_history_persistence: bool = False, compaction_strategy: CompactionStrategy | None = None, tokenizer: TokenizerProtocol | None = None, diff --git a/python/packages/core/agent_framework/_clients.py b/python/packages/core/agent_framework/_clients.py index 6119d3dcd3e..d4b758af26f 100644 --- a/python/packages/core/agent_framework/_clients.py +++ b/python/packages/core/agent_framework/_clients.py @@ -578,7 +578,7 @@ def as_agent( tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, default_options: OptionsCoT | Mapping[str, Any] | None = None, context_providers: Sequence[Any] | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, require_per_service_call_history_persistence: bool = False, compaction_strategy: CompactionStrategy | None = None, tokenizer: TokenizerProtocol | None = None, diff --git a/python/packages/core/agent_framework/_harness/_agent.py b/python/packages/core/agent_framework/_harness/_agent.py index 870b8a493bf..e8567092829 100644 --- a/python/packages/core/agent_framework/_harness/_agent.py +++ b/python/packages/core/agent_framework/_harness/_agent.py @@ -14,7 +14,7 @@ import sys from collections.abc import Callable, Sequence from pathlib import Path -from typing import TYPE_CHECKING, Any, TypedDict +from typing import TYPE_CHECKING, Any, TypedDict, cast from .._agents import Agent, SupportsAgentRun from .._clients import SupportsShellTool, SupportsWebSearchTool @@ -339,7 +339,7 @@ def create_harness_agent( loop_max_iterations: int | None = DEFAULT_MAX_ITERATIONS, otel_provider_name: str | None = None, context_providers: Sequence[ContextProvider] | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, default_options: Mapping[str, Any] | None = None, ) -> Agent[OptionsCoT]: """Create a pre-configured agent with batteries included. @@ -655,8 +655,13 @@ def create_harness_agent( # Message injection is always on. It is a no-op when no messages are queued for the session, # so there is no opt-out. assembled_middleware.append(MessageInjectionMiddleware()) - if middleware: - assembled_middleware.extend(middleware) + if middleware is not None: + # A bare middleware object (including a MiddlewareBundle, which is + # deliberately not a sequence) counts as a one-element source. + if isinstance(middleware, Sequence): + assembled_middleware.extend(cast("Sequence[MiddlewareTypes]", middleware)) + else: + assembled_middleware.append(middleware) agent = Agent( client, diff --git a/python/packages/core/agent_framework/_harness/_agent.pyi b/python/packages/core/agent_framework/_harness/_agent.pyi index 26a57288a74..c9f4f318007 100644 --- a/python/packages/core/agent_framework/_harness/_agent.pyi +++ b/python/packages/core/agent_framework/_harness/_agent.pyi @@ -87,6 +87,6 @@ def create_harness_agent( loop_max_iterations: int | None = DEFAULT_MAX_ITERATIONS, otel_provider_name: str | None = None, context_providers: Sequence[ContextProvider] | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, default_options: Mapping[str, Any] | None = None, ) -> Agent[OptionsCoT]: ... diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 4098ca8ee08..46e8d49be16 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -7,7 +7,7 @@ import logging import sys from abc import ABC, abstractmethod -from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequence +from collections.abc import AsyncIterable, Awaitable, Callable, Collection, Mapping, Sequence from enum import Enum from typing import TYPE_CHECKING, Any, Generic, Literal, TypeAlias, cast, overload @@ -41,7 +41,7 @@ from ._agents import SupportsAgentRun from ._compaction import CompactionStrategy, TokenizerProtocol - from ._sessions import AgentSession + from ._sessions import AgentSession, _RunPersistenceGate # pyright: ignore[reportPrivateUsage] from ._tools import FunctionTool, ToolTypes from ._types import ChatOptions @@ -202,6 +202,11 @@ def __init__( self.stream_transform_hooks = list(stream_transform_hooks or []) self.stream_result_hooks = list(stream_result_hooks or []) self.stream_cleanup_hooks = list(stream_cleanup_hooks or []) + # Set by egress-enforcement middleware (agent-hooks): the run-persistence gate + # covering this pipeline's run. The final handler offers it for adoption by + # the run it starts (see _sessions._offer_run_persistence_gate_claim), so the + # gate binds to that run's identity and never to middleware-initiated runs. + self._run_persistence_gate: _RunPersistenceGate | None = None class FunctionInvocationContext: @@ -672,16 +677,22 @@ async def process( ) +@experimental(feature_id=ExperimentalFeature.AGENT_HOOKS) class MiddlewareBundle: """An indivisible group of middleware that forms one coherent feature. Some features (for example the agent-hooks enforcement middleware) consist of several middleware objects that only uphold their contract when installed together. A bundle carries those objects as one opaque unit: pass the bundle - itself anywhere middleware is accepted, and :func:`categorize_middleware` - splits its members into their agent/function/chat categories while the bundle - guarantees the members cannot be installed partially — it is deliberately not - a sequence, so it cannot be unpacked or sliced. + itself (agent-level ``Agent(middleware=[...])`` or per-run + ``agent.run(middleware=[...])``), and :func:`categorize_middleware` splits its + members into their agent/function/chat categories while the bundle guarantees + the members cannot be installed partially — it is deliberately not a sequence, + so it cannot be unpacked or sliced. Middleware seams that install only some + categories (chat-client seams install chat and function middleware only) + enforce the same guarantee by raising ``MiddlewareException`` when a bundle + member falls into a category they cannot install, instead of silently dropping + that member. Examples: .. code-block:: python @@ -1349,7 +1360,7 @@ class AgentMiddlewareLayer: def __init__( self, *args: Any, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, **kwargs: Any, ) -> None: middleware_list = categorize_middleware(middleware) @@ -1380,7 +1391,7 @@ def run( *, stream: Literal[False] = ..., session: AgentSession | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, options: ChatOptions[ResponseModelBoundT], compaction_strategy: CompactionStrategy | None = None, @@ -1396,7 +1407,7 @@ def run( *, stream: Literal[False] = ..., session: AgentSession | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, options: ChatOptions[None] | None = None, compaction_strategy: CompactionStrategy | None = None, @@ -1412,7 +1423,7 @@ def run( *, stream: Literal[True], session: AgentSession | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, options: ChatOptions[Any] | None = None, compaction_strategy: CompactionStrategy | None = None, @@ -1427,7 +1438,7 @@ def run( *, stream: bool = False, session: AgentSession | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, options: ChatOptions[Any] | None = None, compaction_strategy: CompactionStrategy | None = None, @@ -1436,19 +1447,13 @@ def run( client_kwargs: Mapping[str, Any] | None = None, ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: """MiddlewareTypes-enabled unified run method.""" - # Re-categorize self.middleware at runtime to support dynamic changes. A bare - # single source (one middleware object or a MiddlewareBundle assigned directly - # to the attribute) is treated as a one-element list, mirroring how - # categorize_middleware handles non-sequence sources — never silently dropped. - base_middleware_attr = getattr(self, "middleware", None) - base_middleware: Sequence[MiddlewareTypes] - if isinstance(base_middleware_attr, Sequence) and not isinstance(base_middleware_attr, (str, bytes)): - base_middleware = cast(Sequence[MiddlewareTypes], base_middleware_attr) - elif base_middleware_attr is not None: - base_middleware = [cast(MiddlewareTypes, base_middleware_attr)] - else: - base_middleware = [] - base_middleware_list = categorize_middleware(base_middleware) + # Re-categorize self.middleware at runtime to support dynamic changes. The raw + # attribute is passed straight through: categorize_middleware owns the rule + # that a bare single source (one middleware object or a MiddlewareBundle + # assigned directly to the attribute) is one element — never silently dropped. + base_middleware_list = categorize_middleware( + cast("MiddlewareTypes | Sequence[MiddlewareTypes] | None", getattr(self, "middleware", None)) + ) run_middleware_list = categorize_middleware(middleware) pipeline = self._get_agent_middleware_pipeline([*base_middleware_list["agent"], *run_middleware_list["agent"]]) @@ -1521,6 +1526,12 @@ async def _execute_stream() -> ResponseStream[AgentResponseUpdate, AgentResponse def _middleware_handler( self, context: AgentContext ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]: + from ._sessions import _offer_run_persistence_gate_claim # pyright: ignore[reportPrivateUsage] + + # The final handler starts the run this pipeline (and any egress gate on the + # context) covers. Offer the gate for adoption by that run — always, so a + # gate-less pipeline also clears any stale ticket at this boundary. + _offer_run_persistence_gate_claim(context._run_persistence_gate, self) # pyright: ignore[reportPrivateUsage] return super().run( # type: ignore[misc, no-any-return] context.messages, stream=context.stream, @@ -1613,21 +1624,39 @@ class MiddlewareDict(TypedDict): def categorize_middleware( *middleware_sources: MiddlewareTypes | Sequence[MiddlewareTypes] | None, + supported_categories: Collection[str] | None = None, ) -> MiddlewareDict: """Categorize middleware from multiple sources into agent, function, and chat types. Args: *middleware_sources: Variable number of middleware sources to categorize. + A bare (non-sequence) source — a single middleware object or a + :class:`MiddlewareBundle` — is treated as a one-element list; this + function is the single owner of that rule. + + Keyword Args: + supported_categories: The categories the call site actually installs, e.g. + ``("chat", "function")`` at chat-client seams. When provided, middleware + that categorizes outside these is not returned: a bare middleware object + is skipped with a warning (mirroring pipeline registration's leniency), + while a :class:`MiddlewareBundle` member raises ``MiddlewareException`` — + a bundle is indivisible, so dropping one member would silently install a + partial feature. ``None`` (default) supports every category. Returns: Dict with keys "agent", "function", "chat" containing lists of categorized middleware. + + Raises: + MiddlewareException: If a bundle member falls outside ``supported_categories``. """ result: MiddlewareDict = {"agent": [], "function": [], "chat": []} - # Merge all middleware sources into a single list + # Merge all middleware sources into a single list. The None check is deliberate + # (not truthiness): a bare middleware object with a falsy __bool__/__len__ must + # still be treated as a one-element source, never silently dropped. all_middleware: list[Any] = [] for source in middleware_sources: - if source: + if source is not None: if isinstance(source, Sequence) and not isinstance(source, (str, bytes)): all_middleware.extend(source) # type: ignore else: @@ -1635,34 +1664,57 @@ def categorize_middleware( # Expand bundles first: a bundle's members are categorized individually (in # order) but travel as one unit, so a feature spanning several categories can - # never be partially installed. + # never be partially installed. Membership is remembered so an unsupported + # category can fail loudly for bundle members below. expanded_middleware: list[Any] = [] + bundle_member_ids: set[int] = set() for middleware in all_middleware: if isinstance(middleware, MiddlewareBundle): - expanded_middleware.extend(middleware._middleware) # pyright: ignore[reportPrivateUsage] + members = middleware._middleware # pyright: ignore[reportPrivateUsage] + bundle_member_ids.update(id(member) for member in members) + expanded_middleware.extend(members) else: expanded_middleware.append(middleware) all_middleware = expanded_middleware # Categorize each middleware item for middleware in all_middleware: + category: Literal["agent", "function", "chat"] if isinstance(middleware, AgentMiddleware): - result["agent"].append(middleware) + category = "agent" elif isinstance(middleware, FunctionMiddleware): - result["function"].append(middleware) + category = "function" elif isinstance(middleware, ChatMiddleware): - result["chat"].append(middleware) + category = "chat" elif callable(middleware): # Always call _determine_middleware_type to ensure proper validation middleware_type = _determine_middleware_type(middleware) if middleware_type == MiddlewareType.AGENT: - result["agent"].append(middleware) # type: ignore + category = "agent" elif middleware_type == MiddlewareType.FUNCTION: - result["function"].append(middleware) # type: ignore - elif middleware_type == MiddlewareType.CHAT: - result["chat"].append(middleware) # type: ignore + category = "function" + else: + category = "chat" else: # Fallback to agent middleware for unknown types - result["agent"].append(middleware) + category = "agent" + if supported_categories is not None and category not in supported_categories: + supported_text = ", ".join(sorted(supported_categories)) + if id(middleware) in bundle_member_ids: + raise MiddlewareException( + f"MiddlewareBundle member {type(middleware).__name__} is {category} middleware, but this " + f"middleware seam supports only {supported_text} middleware. A bundle is one indivisible " + "feature and cannot be partially installed; pass the bundle to the agent instead " + "(Agent(middleware=[...]) or agent.run(middleware=[...]))." + ) + logger.warning( + "Ignoring %s middleware of type %s: this middleware seam supports only %s middleware " + "and it will not be executed.", + category, + getattr(middleware, "__name__", type(middleware).__name__), + supported_text, + ) + continue + result[category].append(cast("Any", middleware)) return result diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index 1b4a155eebd..73eae2cdf5a 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -16,6 +16,7 @@ from __future__ import annotations import asyncio +import contextlib import copy import hashlib import json @@ -29,8 +30,8 @@ from abc import abstractmethod from base64 import urlsafe_b64encode from collections import deque -from collections.abc import AsyncIterable, Awaitable, Callable, Iterable, Mapping, Sequence -from contextvars import ContextVar +from collections.abc import AsyncIterable, Awaitable, Callable, Generator, Iterable, Mapping, Sequence +from contextvars import ContextVar, Token from dataclasses import dataclass from functools import partial from pathlib import Path @@ -1028,29 +1029,220 @@ async def after_run( # Run-scoped gate for durable persistence side effects. Egress-enforcement middleware -# (agent-hooks) sets this to a collector list around the guarded section of a run; the -# persistence call sites below consult it via ``_defer_run_persistence`` so that history -# only becomes durable after the run's egress verdict permits the content. When no gate -# is active (the default), persistence runs inline exactly as before. -_DEFERRED_RUN_PERSISTENCE: ContextVar[list[Callable[[], Awaitable[None]]] | None] = ContextVar( +# (agent-hooks) activates a ``_RunPersistenceGate`` around the guarded section of a run; +# the persistence call sites below consult it via ``_defer_run_persistence`` so that +# history only becomes durable after the run's egress verdict permits the content. When +# no gate is active (the default), persistence runs inline exactly as before. +_DEFERRED_RUN_PERSISTENCE: ContextVar[_RunPersistenceGate | None] = ContextVar( "agent_framework_deferred_run_persistence", default=None ) +# Identity of the agent run whose dynamic extent the current code executes in. +# ``RawAgent.run`` stamps a fresh identity over each run (nested runs re-stamp within +# their own extent); agents with fully custom run loops may never stamp one, in which +# case the identity is ``None``. The active gate compares this identity against its +# owner so that only the gated run's own persistence defers (see ``_RunPersistenceGate``). +_CURRENT_RUN_IDENTITY: ContextVar[object | None] = ContextVar("agent_framework_current_run_identity", default=None) + +# One-shot claim handshake between the middleware layer's final handler and the run it +# starts: ``(gate_or_None, agent)``. The final handler offers the gate carried by its +# AgentContext (or ``None``, which also clears any stale ticket at every pipeline +# boundary); the first ``RawAgent.run`` on the *same agent instance* adopts it and +# binds the gate's owner to its fresh run identity. Keying the ticket to the agent +# instance means middleware-initiated sibling runs and nested runs can never claim a +# gate that belongs to another run. +_PENDING_GATE_CLAIM: ContextVar[tuple[_RunPersistenceGate | None, object] | None] = ContextVar( + "agent_framework_pending_run_persistence_gate_claim", default=None +) + + +def _current_run_identity() -> object | None: # pyright: ignore[reportUnusedFunction] + """Return the identity of the agent run currently executing, if any.""" + return _CURRENT_RUN_IDENTITY.get() + + +@contextlib.contextmanager +def _run_identity_scope(identity: object) -> Generator[None]: # pyright: ignore[reportUnusedFunction] + """Stamp ``identity`` as the current run identity for the enclosed extent.""" + token = _CURRENT_RUN_IDENTITY.set(identity) + try: + yield + finally: + _CURRENT_RUN_IDENTITY.reset(token) + + +def _offer_run_persistence_gate_claim( # pyright: ignore[reportUnusedFunction] + gate: _RunPersistenceGate | None, agent: object +) -> None: + """Offer ``gate`` for adoption by the run that ``agent`` is about to start. + + Called by the middleware layer's final handler for every pipeline execution (with + ``gate=None`` when the run carries no gate, which clears any stale ticket left by + an agent whose custom run loop never adopts). + """ + _PENDING_GATE_CLAIM.set((gate, agent)) + + +def _adopt_run_persistence_gate_claim(agent: object, identity: object) -> None: # pyright: ignore[reportUnusedFunction] + """Adopt a pending gate claim targeted at ``agent``, binding the gate to ``identity``. + + Called from ``RawAgent.run``. Tickets targeted at a different agent instance are + left untouched: a nested or sibling run must never claim a gate that covers another + run's verdict. + """ + ticket = _PENDING_GATE_CLAIM.get() + if ticket is None: + return + gate, target = ticket + if target is not agent: + return + _PENDING_GATE_CLAIM.set(None) + if gate is not None: + gate.bind_owner(identity) + def _defer_run_persistence(persist: Callable[[], Awaitable[None]]) -> bool: """Queue a persistence side effect behind the active run-persistence gate, if any. - Returns True when ``persist`` was deferred to the gate owner (which MUST reset the - context variable before executing the collected callables, so re-entrant calls run - inline), and False when no gate is active and the caller must persist inline. + Returns True when ``persist`` was deferred to the gate owner (which drains it via + :meth:`_RunPersistenceGate.flush` only after the gate scope has been exited, so + re-entrant calls run inline), and False when the caller must persist inline — + either because no gate is active, or because the persist belongs to a different + run than the one the gate covers (see :meth:`_RunPersistenceGate.accepts`). """ - pending = _DEFERRED_RUN_PERSISTENCE.get() - if pending is None: + gate = _DEFERRED_RUN_PERSISTENCE.get() + if gate is None: + return False + if not gate.accepts(_CURRENT_RUN_IDENTITY.get()): return False - pending.append(persist) + gate.collect(persist) return True +class _RunPersistenceGate: + """Owner handle for one guarded run section's deferred persistence. + + The handle is a context manager: while entered, durable persistence side effects + routed through :func:`_defer_run_persistence` are collected instead of executed. + After the scope has been exited, the owner either releases them with :meth:`flush` + (the covering verdict permitted the content) or discards them with :meth:`drop` + (the content was denied and must never become durable). + + Ownership: the gate covers exactly one run's verdict, so it only collects that + run's own persistence (see :meth:`accepts`). The agent-seam gate is created before + its run starts and is bound lazily through the + :func:`_offer_run_persistence_gate_claim` / :func:`_adopt_run_persistence_gate_claim` + handshake; the chat-seam gate is created inside its run and bound immediately. + While unbound (an agent whose custom run loop never adopts the claim), the gate + falls back fail-closed: persists carrying no run identity defer, while persists + from identity-stamped (nested or sibling) runs execute inline at their own run + boundaries. + + Reset-before-drain is enforced by construction: :meth:`flush` refuses to run while + the gate scope is still entered, and executes the collected callables with the + gate context suspended, so a re-entrant persistence call made by a flushed + callable always runs inline (each deferred persist is covered by exactly one + verdict — it is never re-deferred into an enclosing gate). + + Usage (the gate owner is egress-enforcement middleware such as agent-hooks):: + + gate = _RunPersistenceGate() + with gate: + ... # the guarded section: persists issued here are collected + # verdict is issued here, outside the gate scope + await gate.flush() # permitted: run the collected persists inline + # or gate.drop() # denied: the content never becomes durable + + Gates nest naturally through the context variable: an inner gate (for example a + hooked sub-agent run inside a hooked outer run) collects only its own section's + persists and restores the outer gate on exit. + """ + + __slots__ = ("_owner", "_owner_bound", "_pending", "_token") + + def __init__(self) -> None: + self._pending: list[Callable[[], Awaitable[None]]] = [] + self._token: Token[_RunPersistenceGate | None] | None = None + self._owner: object | None = None + self._owner_bound: bool = False + + def bind_owner(self, owner: object | None) -> None: + """Bind the gate to the run identity whose persistence it defers (first bind wins).""" + if not self._owner_bound: + self._owner = owner + self._owner_bound = True + + def collect(self, persist: Callable[[], Awaitable[None]]) -> None: + """Collect one deferred persistence callable (see :func:`_defer_run_persistence`).""" + self._pending.append(persist) + + def accepts(self, identity: object | None) -> bool: + """Whether a persist issued under ``identity`` belongs to this gate's run. + + A bound gate accepts exactly its owner identity. An unbound gate (its run + never adopted the claim — a custom run loop) accepts only identity-less + persists: the covered run's own persists carry no identity there (fail-closed, + matching the pre-ownership behavior), while identity-stamped persists come + from other runs that own their content's verdicts themselves. + """ + if self._owner_bound: + return identity is self._owner + return identity is None + + def __enter__(self) -> _RunPersistenceGate: + if self._token is not None: + raise RuntimeError("This run-persistence gate is already active; gates are not re-entrant.") + self._token = _DEFERRED_RUN_PERSISTENCE.set(self) + return self + + def __exit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None: + token = self._token + self._token = None + if token is not None: + _DEFERRED_RUN_PERSISTENCE.reset(token) + + async def flush(self) -> None: + """Execute the deferred persistence in order (the covering verdict permitted it).""" + if self._token is not None: + raise RuntimeError( + "Cannot flush a run-persistence gate while its scope is still active: " + "re-entrant persistence would be re-deferred instead of running inline. " + "Exit the gate's context manager first." + ) + # Suspend the gate context while draining: a flushed callable that re-checks + # the gate (e.g. _run_after_providers) must run inline, never re-defer into an + # enclosing gate — its own covering verdict already permitted it. + with _suspend_run_persistence_gate(): + while self._pending: + await self._pending.pop(0)() + + def drop(self) -> None: + """Discard the deferred persistence (the covering verdict denied the content).""" + self._pending.clear() + + +@contextlib.contextmanager +def _suspend_run_persistence_gate() -> Generator[None]: + """Run a nested section with no active run-persistence gate. + + A run-persistence gate defers only the gated run's *own* persistence; nested agent + runs own their content's verdicts themselves (or are unguarded), so their + persistence must run inline: deferring it into the outer run's gate would silently + drop fully-permitted inner history on an outer deny, and a second nested read in + the same outer run would see stale history. Run-identity ownership (see + :meth:`_RunPersistenceGate.accepts`) enforces this for every identity-stamped run; + the function-invocation layer additionally wraps each tool invocation in this + suspension so that nested agents with fully custom run loops (which never stamp an + identity and would otherwise inherit the outer run's) also persist inline on the + tool path. + """ + token = _DEFERRED_RUN_PERSISTENCE.set(None) + try: + yield + finally: + _DEFERRED_RUN_PERSISTENCE.reset(token) + + LOCAL_HISTORY_CONVERSATION_ID = "agent_framework_local_history_persistence" diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index ce9606d3de6..39dd3e20003 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -1656,18 +1656,26 @@ async def _execute_single_function_call( live_tools: list[ToolTypes] | None, ) -> tuple[list[Content], bool]: from ._middleware import MiddlewareTermination + from ._sessions import _suspend_run_persistence_gate # pyright: ignore[reportPrivateUsage] from ._types import Content try: - result = await _auto_invoke_function( - function_call_content=function_call, - custom_args=custom_args, - tool_map=tool_map, - invocation_session=invocation_session, - middleware_pipeline=middleware_pipeline, - config=config, - live_tools=live_tools, - ) + # A run-persistence gate defers only the gated run's own persistence; nested + # agent runs persist inline at their own boundaries. Run-identity ownership + # (see _sessions._RunPersistenceGate.accepts) enforces that for every run that + # stamps an identity; suspending the gate around the tool invocation (the most + # common nesting seam) additionally covers nested agents with fully custom run + # loops, which never stamp one and would otherwise inherit the outer identity. + with _suspend_run_persistence_gate(): + result = await _auto_invoke_function( + function_call_content=function_call, + custom_args=custom_args, + tool_map=tool_map, + invocation_session=invocation_session, + middleware_pipeline=middleware_pipeline, + config=config, + live_tools=live_tools, + ) return [result], False except MiddlewareTermination as exc: if isinstance(exc.result, Content): @@ -2845,7 +2853,10 @@ def __init__( ) -> None: from ._middleware import categorize_middleware - categorized_middleware = categorize_middleware(middleware) + # Chat clients install only chat and function middleware. Agent middleware in + # a bundle raises (a bundle must never be partially installed); bare agent + # middleware is warned about and skipped inside categorize_middleware. + categorized_middleware = categorize_middleware(middleware, supported_categories=("chat", "function")) self.function_middleware: list[FunctionMiddlewareTypes] = list(categorized_middleware["function"]) self._cached_function_middleware_pipeline: FunctionMiddlewarePipeline | None = None self.function_invocation_configuration = normalize_function_invocation_configuration( @@ -3235,7 +3246,12 @@ def get_response( ), *middleware, ] - categorized_runtime_middleware = categorize_middleware(request_kwargs.pop("middleware", [])) + # Same contract as the constructor: this seam installs chat and function + # middleware only; a bundle carrying an agent member fails loudly instead of + # silently losing that member. + categorized_runtime_middleware = categorize_middleware( + request_kwargs.pop("middleware", []), supported_categories=("chat", "function") + ) function_middleware_pipeline = self._get_function_middleware_pipeline( categorized_runtime_middleware["function"] diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 1d3d797d618..6d8c1521b1d 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -28,6 +28,7 @@ from typing_extensions import TypedDict +from ._feature_stage import ExperimentalFeature, experimental from ._serialization import SerializationMixin from .exceptions import AdditionItemMismatch, ContentError @@ -3225,10 +3226,12 @@ def from_awaitable( return stream @classmethod + @experimental(feature_id=ExperimentalFeature.AGENT_HOOKS) def buffered_and_gated( cls, consume: Callable[[], Awaitable[tuple[Sequence[UpdateT], FinalT]]], - gate: Callable[[list[UpdateT], FinalT, bool], Awaitable[tuple[Sequence[UpdateT], FinalT]]], + gate: Callable[[list[UpdateT], FinalT], Awaitable[tuple[FinalT, bool]]], + rederive: Callable[[FinalT], Sequence[UpdateT]], ) -> ResponseStream[UpdateT, FinalT]: """Create a fully buffered stream whose content is finalized by a gate. @@ -3242,29 +3245,38 @@ def buffered_and_gated( far (for example hooks attached by middleware pipelines after they unwind) is applied to the buffered updates and result now — **before** the gate — so the gate's verdict covers their effect. The hooks are consumed: they are - not applied again during replay. - 3. ``gate`` receives the post-hook updates, the post-hook result, and whether - any pending transform/result hooks were applied. It returns the updates and - result to release (it may raise to block egress entirely). A gate that - issues its verdict over the final result should re-derive the released - updates from that result whenever hooks were applied, so streamed egress - cannot diverge from the verdicted content. - 4. The stream is then sealed: the gated updates and result are replayed + not applied again during replay. Because they run ahead of the verdict, + these hooks also *see* pre-verdict content: they are rewriters inside the + enforcement boundary. A host-facing observer must consume the released + stream instead of registering a hook here. + 3. ``gate`` receives the post-hook updates and the post-hook result. It + returns the result to release and whether it transformed that result (it + may raise to block egress entirely). + 4. The combinator owns the no-divergence rule: whenever pending hooks were + applied or the gate reports a transform, the released updates are + re-derived from the gated result via ``rederive``, so streamed egress can + never diverge from the verdicted content. Otherwise the buffered updates + are replayed as-is. + 5. The stream is then sealed: the released updates and result are replayed verbatim, and registering further transform or result hooks raises ``RuntimeError`` — nothing can rewrite content past the gate. Args: consume: Produces the buffered updates and finalized result. Runs inside the caller's context (the caller owns any context-variable scoping). - gate: Receives ``(updates, final, hooks_applied)`` after pending hooks are - applied; returns the ``(updates, final)`` to release. + gate: Receives ``(updates, final)`` after pending hooks are applied; + returns ``(final, transformed)`` — the result to release and whether + the gate changed it. + rederive: Rebuilds the released updates from the gated result; applied by + the combinator whenever hooks ran or the gate transformed, so no + caller can accidentally egress un-verdicted updates. Returns: A sealed, fully buffered ResponseStream. """ return cast( "ResponseStream[UpdateT, FinalT]", - cast(Any, _GatedResponseStream).create_buffered_and_gated(consume, gate), + cast(Any, _GatedResponseStream).create_buffered_and_gated(consume, gate, rederive), ) async def _get_stream(self) -> AsyncIterable[UpdateT]: @@ -3565,7 +3577,8 @@ def with_result_hook( def create_buffered_and_gated( cls, consume: Callable[[], Awaitable[tuple[Sequence[UpdateT], FinalT]]], - gate: Callable[[list[UpdateT], FinalT, bool], Awaitable[tuple[Sequence[UpdateT], FinalT]]], + gate: Callable[[list[UpdateT], FinalT], Awaitable[tuple[FinalT, bool]]], + rederive: Callable[[FinalT], Sequence[UpdateT]], ) -> _GatedResponseStream[UpdateT, FinalT]: """Build the gated stream for :meth:`ResponseStream.buffered_and_gated`.""" holder: dict[str, Any] = {} @@ -3602,10 +3615,15 @@ async def _materialize() -> AsyncGenerator[UpdateT]: cleanup_result = cleanup_hook() if isawaitable(cleanup_result): await cleanup_result - gated_updates, gated_final = await gate(hooked_updates, final, bool(transform_hooks or result_hooks)) + gated_final, gate_transformed = await gate(hooked_updates, final) holder["final"] = gated_final stream._gate_sealed = True - for update in gated_updates: + # No-divergence rule (contract step 4), owned here: hooks or a gate + # transform mean the buffered updates may no longer match the verdicted + # result, so the released updates are re-derived from it. + hooks_applied = bool(transform_hooks or result_hooks) + released = rederive(gated_final) if (hooks_applied or gate_transformed) else hooked_updates + for update in released: yield update def _finalizer(_: Sequence[UpdateT]) -> FinalT: diff --git a/python/packages/core/agent_framework/observability.py b/python/packages/core/agent_framework/observability.py index 20629b76b81..ad9a6f4e574 100644 --- a/python/packages/core/agent_framework/observability.py +++ b/python/packages/core/agent_framework/observability.py @@ -2050,7 +2050,7 @@ def run( *, stream: Literal[False] = ..., session: AgentSession | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, options: ChatOptions[ResponseModelBoundT], compaction_strategy: CompactionStrategy | None = None, @@ -2066,7 +2066,7 @@ def run( *, stream: Literal[False] = ..., session: AgentSession | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, options: ChatOptions[None] | None = None, compaction_strategy: CompactionStrategy | None = None, @@ -2082,7 +2082,7 @@ def run( *, stream: Literal[True], session: AgentSession | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, options: ChatOptions[Any] | None = None, compaction_strategy: CompactionStrategy | None = None, @@ -2097,7 +2097,7 @@ def run( *, stream: bool = False, session: AgentSession | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, options: ChatOptions[Any] | None = None, compaction_strategy: CompactionStrategy | None = None, diff --git a/python/packages/core/tests/core/test_agent_hooks.py b/python/packages/core/tests/core/test_agent_hooks.py index 31811a38272..c4d4a085532 100644 --- a/python/packages/core/tests/core/test_agent_hooks.py +++ b/python/packages/core/tests/core/test_agent_hooks.py @@ -180,7 +180,7 @@ async def test_bare_bundle_at_construction_is_fully_enforced(chat_client_base: M guard = PointGuard("output", Verdict.deny(reason="egress_blocked")) agent = Agent( client=chat_client_base, - middleware=cast("Any", create_agent_hooks_middleware([guard], record_sink=records.append)), + middleware=create_agent_hooks_middleware([guard], record_sink=records.append), ) with pytest.raises(InterceptionBlocked) as exc_info: @@ -223,11 +223,23 @@ async def test_factory_returns_an_indivisible_bundle() -> None: assert len(categorized["agent"]) == 1 assert len(categorized["chat"]) == 1 assert len(categorized["function"]) == 1 - # ...but cannot be partially installed: it is opaque (not a sequence). + # ...but cannot be partially installed: it is opaque (not a sequence). These are + # deliberate runtime probes of operations the static types also reject, so they go + # through an Any-typed alias instead of ignore comments (which the three test + # typing checkers spell differently). + opaque = cast("Any", bundle) with pytest.raises(TypeError): - iter(bundle) # type: ignore[call-overload] + iter(opaque) with pytest.raises(TypeError): - bundle[0] # type: ignore[index] + opaque[0] + + +def test_middleware_bundle_is_experimental() -> None: + # Both bundle producers carry @experimental; the bundle type itself does too. + # Asserting on the stage metadata is deterministic (the runtime warning dedups + # per feature id per process, so warn-order would make a warns() assertion flaky). + assert getattr(MiddlewareBundle, "__feature_stage__", None) == "experimental" + assert getattr(MiddlewareBundle, "__feature_id__", None) == "AGENT_HOOKS" # endregion @@ -840,7 +852,7 @@ async def test_chat_seam_without_run_state_fails_closed(chat_client_base: MockBa # without an active agent-hooks run (its agent sibling never ran) must fail # closed, not silently skip enforcement. bundle = create_agent_hooks_middleware([AllowGuard()]) - chat_client_base.chat_middleware = cast("list[Any]", [_bundle_member(bundle, "ChatMiddleware")]) + chat_client_base.chat_middleware = [_bundle_member(bundle, "ChatMiddleware")] with pytest.raises(MiddlewareException, match="without an active agent-hooks run"): await chat_client_base.get_response([Message(role="user", contents=["hi"])]) @@ -1161,7 +1173,7 @@ async def test_function_seam_without_run_state_blocks_tool(chat_client_base: Moc agent = Agent( client=chat_client_base, tools=[weather_tool], - middleware=[cast("Any", _bundle_member(bundle, "FunctionMiddleware"))], + middleware=[_bundle_member(bundle, "FunctionMiddleware")], ) response = await agent.run("get the weather") @@ -1173,6 +1185,27 @@ async def test_function_seam_without_run_state_blocks_tool(chat_client_base: Moc assert "without an active agent-hooks run" in transcript +@requires_sdk +async def test_bundle_passed_to_chat_client_call_raises_instead_of_dropping_the_gate( + 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 + # 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"): + chat_client_base.get_response([Message(role="user", contents=["hi"])], middleware=cast("Any", [bundle])) + + +@requires_sdk +async def test_bundle_passed_to_chat_client_constructor_raises_instead_of_dropping_the_gate() -> None: + from agent_framework._tools import FunctionInvocationLayer + + bundle = create_agent_hooks_middleware([AllowGuard()]) + with pytest.raises(MiddlewareException, match="cannot be partially installed"): + FunctionInvocationLayer(middleware=cast("Any", [bundle])) + + @requires_sdk async def test_fresh_bundle_per_run_is_not_conflated_by_pipeline_caching(chat_client_base: MockBaseChatClient) -> None: # The agent middleware pipeline is cached and compared with ==; the bundle members @@ -1431,6 +1464,231 @@ async def test_per_service_call_persistence_still_persists_on_allow() -> None: assert [message.text for message in stored] == ["hi", "test response - hi"] +def _sub_agent_call_response(task: str, call_id: str) -> ChatResponse: + """An outer-model response that calls the ``sub`` sub-agent tool.""" + return ChatResponse( + messages=[ + Message( + role="assistant", + contents=[Content.from_function_call(call_id=call_id, name="sub", arguments=f'{{"task": "{task}"}}')], + ) + ] + ) + + +@requires_sdk +@pytest.mark.parametrize("streaming", [False, True], ids=["non_streaming", "streaming"]) +async def test_outer_deny_never_drops_permitted_nested_run_history(streaming: bool) -> None: + # An unhooked sub-agent run (as_tool, session shared with the parent) owns its own + # persistence: it must land inline at the inner run's boundary, not defer into the + # outer gate's pending list where an outer output deny would silently drop it. + from agent_framework import AgentSession, InMemoryHistoryProvider + + inner_provider = InMemoryHistoryProvider(source_id="inner_history") + session = AgentSession() + sub_agent = Agent(client=MockBaseChatClient(), name="sub", context_providers=[inner_provider]) + + outer_client = MockBaseChatClient() + if streaming: + outer_client.streaming_responses = [ + [ + ChatResponseUpdate( + contents=[ + Content.from_function_call(call_id="c1", name="sub", arguments='{"task": "look this up"}') + ], + role="assistant", + finish_reason="tool_calls", + ) + ], + [ChatResponseUpdate(contents=[Content.from_text("outer summary")], role="assistant", finish_reason="stop")], + ] + else: + outer_client.run_responses = [_sub_agent_call_response("look this up", "c1"), final_response("outer summary")] + outer_agent = Agent( + client=outer_client, + name="outer", + tools=[sub_agent.as_tool(propagate_session=True)], + middleware=[create_agent_hooks_middleware([PointGuard("output", Verdict.deny(reason="egress_blocked"))])], + ) + + with pytest.raises(InterceptionBlocked): + if streaming: + async for _ in outer_agent.run("go delegate", session=session, stream=True): + pass + else: + await outer_agent.run("go delegate", session=session) + + # The fully-permitted inner history is durable despite the outer deny... + # (as_tool always runs the sub-agent streaming; the mock streams "update - ...") + inner_stored = cast("list[Message]", session.state["inner_history"]["messages"]) + assert [message.text for message in inner_stored] == ["look this up", "update - look this up"] + # ...while the denied outer run's own history was dropped with the outer gate. + outer_stored = cast("list[Message]", session.state.get("in_memory", {}).get("messages", [])) + assert outer_stored == [] + + +@requires_sdk +async def test_second_sub_agent_call_in_one_outer_run_reads_fresh_history() -> None: + # Two sequential calls to the same sub-agent within one hooked outer run: the + # second call must load the history the first call persisted (inline at the inner + # run boundary), not a stale pre-first-call snapshot. + from agent_framework import AgentSession, InMemoryHistoryProvider + + inner_requests: list[list[str]] = [] + + class CaptureInnerRequests(ChatMiddleware): + async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None: + inner_requests.append([message.text for message in context.messages]) + await call_next() + + inner_provider = InMemoryHistoryProvider(source_id="inner_history") + session = AgentSession() + sub_agent = Agent( + client=MockBaseChatClient(), + name="sub", + context_providers=[inner_provider], + middleware=[CaptureInnerRequests()], + ) + + outer_client = MockBaseChatClient() + outer_client.run_responses = [ + _sub_agent_call_response("task one", "c1"), + _sub_agent_call_response("task two", "c2"), + final_response("outer summary"), + ] + outer_agent = Agent( + client=outer_client, + name="outer", + tools=[sub_agent.as_tool(propagate_session=True)], + middleware=[create_agent_hooks_middleware([AllowGuard()])], + ) + + response = await outer_agent.run("go delegate twice", session=session) + + assert response.text == "outer summary" + assert inner_requests[0] == ["task one"] + # The second inner model call sees the first call's persisted exchange. + # (as_tool always runs the sub-agent streaming; the mock streams "update - ...") + assert inner_requests[1] == ["task one", "update - task one", "task two"] + + +@requires_sdk +@pytest.mark.parametrize("streaming", [False, True], ids=["non_streaming", "streaming"]) +@pytest.mark.parametrize("when", ["before_call_next", "after_call_next"]) +async def test_middleware_initiated_sub_agent_history_survives_outer_deny(streaming: bool, when: str) -> None: + # Nested runs are not limited to the tool seam: user agent middleware below the + # bundle can run a sub-agent directly inside the gated section. That run owns its + # own persistence, so it must land inline at the inner run boundary (run-identity + # ownership), never defer into the outer gate where an outer deny would drop it. + from agent_framework import AgentSession, InMemoryHistoryProvider + + inner_provider = InMemoryHistoryProvider(source_id="inner_history") + session = AgentSession() + sub_agent = Agent(client=MockBaseChatClient(), name="sub", context_providers=[inner_provider]) + + class DelegatingMiddleware(AgentMiddleware): + async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + if when == "before_call_next": + await sub_agent.run("delegated task", session=session) + await call_next() + if when == "after_call_next": + await sub_agent.run("delegated task", session=session) + + outer_agent = Agent( + client=MockBaseChatClient(), + name="outer", + middleware=[ + create_agent_hooks_middleware([PointGuard("output", Verdict.deny(reason="egress_blocked"))]), + DelegatingMiddleware(), + ], + ) + + with pytest.raises(InterceptionBlocked): + if streaming: + async for _ in outer_agent.run("go delegate", session=session, stream=True): + pass + else: + await outer_agent.run("go delegate", session=session) + + # The middleware-initiated run's fully-permitted history is durable... + inner_stored = cast("list[Message]", session.state["inner_history"]["messages"]) + assert [message.text for message in inner_stored] == ["delegated task", "test response - delegated task"] + # ...while the denied outer run's own history was dropped with the outer gate. + outer_stored = cast("list[Message]", session.state.get("in_memory", {}).get("messages", [])) + assert outer_stored == [] + + +@requires_sdk +async def test_custom_run_loop_agent_nested_sub_agent_persists_inline() -> None: + # GitHubCopilotAgent-shaped composition: a hookable agent (AgentMiddlewareLayer in + # its MRO) whose fully custom run loop invokes tools directly via + # FunctionTool.invoke — never passing the framework's function-invocation seam. + # The nested core-agent run stamps its own identity, so its permitted history + # persists inline even though the outer gate never binds (the custom loop never + # adopts the claim); the custom outer run's own persistence stays gated + # fail-closed and is dropped on the outer deny. + from agent_framework import AgentSession, BaseAgent, InMemoryHistoryProvider, SessionContext + from agent_framework._middleware import AgentMiddlewareLayer + + inner_provider = InMemoryHistoryProvider(source_id="inner_history") + outer_provider = InMemoryHistoryProvider(source_id="outer_history") + session = AgentSession() + sub_agent = Agent(client=MockBaseChatClient(), name="sub", context_providers=[inner_provider]) + sub_tool = sub_agent.as_tool(propagate_session=True) + + class CustomLoopRawAgent(BaseAgent): + def run( # type: ignore[override] + self, + messages: Any = None, + *, + stream: bool = False, + session: Any = None, + **kwargs: Any, + ) -> Any: + assert not stream + + async def _run() -> AgentResponse[Any]: + # The custom loop invokes the sub-agent tool directly, bypassing + # _execute_single_function_call (like GitHubCopilotAgent's loop). + tool_context = FunctionInvocationContext( + function=sub_tool, arguments={"task": "delegated task"}, session=session + ) + await sub_tool.invoke(arguments={"task": "delegated task"}, context=tool_context) + response: AgentResponse[Any] = AgentResponse( + messages=[Message(role="assistant", contents=["custom outer answer"])] + ) + session_context = SessionContext( + session_id=session.session_id if session is not None else None, + input_messages=[Message(role="user", contents=["go delegate"])], + ) + session_context._response = response # pyright: ignore[reportPrivateUsage] + await self._run_after_providers(session=session, context=session_context) + return response + + return _run() + + class CustomLoopAgent( # type: ignore[misc] # ty: ignore[invalid-method-override] # same shape as GitHubCopilotAgent + AgentMiddlewareLayer, CustomLoopRawAgent + ): + pass + + outer_agent = CustomLoopAgent( + name="outer", + context_providers=[outer_provider], + middleware=[create_agent_hooks_middleware([PointGuard("output", Verdict.deny(reason="egress_blocked"))])], + ) + + with pytest.raises(InterceptionBlocked): + await cast("Any", outer_agent).run("go delegate", session=session) + + # The nested sub-agent's permitted history persisted inline (as_tool streams the + # sub-agent; the mock streams "update - ...")... + inner_stored = cast("list[Message]", session.state["inner_history"]["messages"]) + assert [message.text for message in inner_stored] == ["delegated task", "update - delegated task"] + # ...while the denied custom outer run's own history stayed gated and was dropped. + assert session.state.get("outer_history", {}).get("messages", []) == [] + + # endregion # region Stream hooks cannot escape the gate @@ -1477,6 +1735,74 @@ def sneak(update: AgentResponseUpdate) -> AgentResponseUpdate: assert final.text == "update - hi" +@requires_sdk +async def test_as_tool_stream_callback_sees_nothing_on_deny() -> None: + # The observe direction of the gate contract: `as_tool(stream_callback=...)` is a + # host-facing observer, so on a denied sub-agent run it must never receive the + # (complete, buffered) denied content — not even transiently. + seen: list[str] = [] + + def observe(update: AgentResponseUpdate) -> None: + seen.append(update.text) + + sub_agent = Agent( + client=MockBaseChatClient(), + name="sub", + middleware=[create_agent_hooks_middleware([PointGuard("output", Verdict.deny(reason="egress_blocked"))])], + ) + sub_tool = sub_agent.as_tool(stream_callback=observe) + + with pytest.raises(InterceptionBlocked) as exc_info: + await sub_tool.invoke(arguments={"task": "hello"}) + + assert exc_info.value.result.verdict.reason == "egress_blocked" + assert seen == [] + + +@requires_sdk +async def test_as_tool_stream_callback_sees_only_transformed_egress() -> None: + # Observe direction, transform case: the callback receives the redacted updates + # only, never the unredacted original. + seen: list[str] = [] + + async def observe(update: AgentResponseUpdate) -> None: + seen.append(update.text) + + guard = PointGuard( + "output", + Verdict(decision=Decision.TRANSFORM, transform=Transform(path="$target.content", value="[redacted]")), + ) + sub_agent = Agent( + client=MockBaseChatClient(), + name="sub", + middleware=[create_agent_hooks_middleware([guard])], + ) + sub_tool = sub_agent.as_tool(stream_callback=observe) + + result = await sub_tool.invoke(arguments={"task": "hello"}) + + assert [content.text for content in result] == ["[redacted]"] + assert "".join(seen) == "[redacted]" + assert not any("test response" in text for text in seen) + + +async def test_as_tool_stream_callback_still_observes_unhooked_streams() -> None: + # Behavior preservation for the common (unhooked) case: the callback observes + # every released update of the sub-agent's stream. + seen: list[str] = [] + + def observe(update: AgentResponseUpdate) -> None: + seen.append(update.text) + + sub_agent = Agent(client=MockBaseChatClient(), name="sub") + sub_tool = sub_agent.as_tool(stream_callback=observe) + + result = await sub_tool.invoke(arguments={"task": "hello"}) + + assert "".join(content.text or "" for content in result) == "".join(seen) + assert seen # the stream produced updates and the observer saw them + + async def test_gated_response_stream_applies_pending_hooks_before_the_gate_and_seals() -> None: # Contract test for ResponseStream.buffered_and_gated (no SDK required). order: list[str] = [] @@ -1485,13 +1811,17 @@ async def consume() -> tuple[list[str], str]: order.append("consume") return ["a", "b"], "ab" - async def gate(updates: list[str], final: str, hooks_applied: bool) -> tuple[list[str], str]: - order.append(f"gate(hooks_applied={hooks_applied})") + async def gate(updates: list[str], final: str) -> tuple[str, bool]: + order.append("gate") assert updates == ["a!", "b!"] # pending transform hooks applied pre-gate assert final == "ab!" # pending result hooks applied pre-gate - return updates, final + return final, False + + def rederive(final: str) -> list[str]: + order.append(f"rederive({final})") + return list(final) - stream = cast("Any", ResponseStream).buffered_and_gated(consume=consume, gate=gate) + stream = cast("Any", ResponseStream).buffered_and_gated(consume=consume, gate=gate, rederive=rederive) def transform(update: str) -> str: order.append(f"transform({update})") @@ -1506,9 +1836,11 @@ def result_hook(final: str) -> str: stream.with_result_hook(result_hook) released = [update async for update in stream] - assert released == ["a!", "b!"] + # Hooks ran, so the combinator itself re-derived the released updates from the + # gated result — the hooked buffer cannot egress un-verdicted. + assert released == ["a", "b", "!"] assert await stream.get_final_response() == "ab!" - assert order == ["consume", "transform(a)", "transform(b)", "result_hook", "gate(hooks_applied=True)"] + assert order == ["consume", "transform(a)", "transform(b)", "result_hook", "gate", "rederive(ab!)"] # ...and once the gate has run, content is sealed: further hooks are rejected. with pytest.raises(RuntimeError, match="sealed"): @@ -1517,6 +1849,59 @@ def result_hook(final: str) -> str: stream.with_result_hook(result_hook) +async def test_gated_response_stream_combinator_owns_the_rederive_rule() -> None: + # The no-divergence rule lives in the combinator, not in each gate: whenever the + # gate reports a transform, the released updates are re-derived from the gated + # result even though the gate never touches the update list; when nothing changed, + # the buffered updates replay verbatim and rederive is never consulted. + rederived: list[str] = [] + + def rederive(final: str) -> list[str]: + rederived.append(final) + return list(final) + + async def consume() -> tuple[list[str], str]: + return ["a", "b"], "ab" + + async def transforming_gate(updates: list[str], final: str) -> tuple[str, bool]: + return "XY", True + + stream = cast("Any", ResponseStream).buffered_and_gated(consume=consume, gate=transforming_gate, rederive=rederive) + assert [update async for update in stream] == ["X", "Y"] + assert await stream.get_final_response() == "XY" + assert rederived == ["XY"] + + async def passthrough_gate(updates: list[str], final: str) -> tuple[str, bool]: + return final, False + + rederived.clear() + stream = cast("Any", ResponseStream).buffered_and_gated(consume=consume, gate=passthrough_gate, rederive=rederive) + assert [update async for update in stream] == ["a", "b"] + assert rederived == [] + + +async def test_gated_response_stream_raising_rederive_releases_nothing() -> None: + # A failing rederive is fail-closed for streamed egress: the iteration raises + # before anything is released. The gate already sealed its verdicted result, so + # non-streaming consumption still returns it. + async def consume() -> tuple[list[str], str]: + return ["a", "b"], "ab" + + async def gate(updates: list[str], final: str) -> tuple[str, bool]: + return "XY", True + + def rederive(final: str) -> list[str]: + raise RuntimeError("rederive failed") + + stream = cast("Any", ResponseStream).buffered_and_gated(consume=consume, gate=gate, rederive=rederive) + released: list[str] = [] + with pytest.raises(RuntimeError, match="rederive failed"): + async for update in stream: + released.append(update) + assert released == [] # zero egress + assert await stream.get_final_response() == "XY" # the verdicted final survives + + # endregion # region Approval requests pass through un-bracketed @@ -1757,12 +2142,28 @@ def test_tool_result_codec_round_trip() -> None: original = [Content.from_text("weather in Seattle")] wire = codecs._ToolResultCodec.to_wire(original) assert wire == "weather in Seattle" - # Untouched wire value -> shape-preserving native value. - written = codecs._ToolResultCodec.write_back(original, "scrubbed") + # The codec owns the untouched-wire rule: an untouched wire value maps back to + # the identical native object. + assert codecs._ToolResultCodec.write_back(original, wire, wire) is original + # A transformed wire value maps back shape-preservingly onto the native value. + written = codecs._ToolResultCodec.write_back(original, wire, "scrubbed") assert isinstance(written[0], Content) assert written[0].text == "scrubbed" +def test_wire_equality_distinguishes_bool_from_number() -> None: + # Python's == equates 1 == True; the untouched-wire checks must not, or a + # bool<->number transform would be silently dropped (fail-open for the transform). + codecs = _codecs() + assert codecs._wire_equal(1, True) is False + assert codecs._wire_equal(True, 1) is False + assert codecs._wire_equal({"flag": [0]}, {"flag": [False]}) is False + assert codecs._wire_equal({"flag": [1, "x"]}, {"flag": [1, "x"]}) is True + # The tool-result codec treats 1 -> True as a genuine transform, not untouched. + assert codecs._ToolResultCodec.write_back(1, 1, True) is True + assert codecs._ToolResultCodec.write_back(1, 1, 1) == 1 + + def test_output_codec_untouched_target_is_a_no_op() -> None: codecs = _codecs() response = AgentResponse(messages=[Message(role="assistant", contents=["hello"])]) diff --git a/python/packages/core/tests/core/test_middleware.py b/python/packages/core/tests/core/test_middleware.py index 8522adde923..0f519d234bf 100644 --- a/python/packages/core/tests/core/test_middleware.py +++ b/python/packages/core/tests/core/test_middleware.py @@ -1728,3 +1728,30 @@ def test_categorize_middleware_with_string_does_not_decompose(self) -> None: total_items = len(result["chat"]) + len(result["function"]) + len(result["agent"]) assert total_items == 1 assert result["agent"] == ["not_a_middleware"] + + def test_categorize_middleware_supported_categories_skips_bare_with_warning( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """A bare middleware outside the supported categories is warned about, not installed.""" + chat_mw = TestChatMiddleware() + agent_mw = TestAgentMiddleware() + with caplog.at_level("WARNING", logger="agent_framework._middleware"): + result = categorize_middleware([chat_mw, agent_mw], supported_categories=("chat", "function")) + assert result["chat"] == [chat_mw] + assert result["agent"] == [] + assert "will not be executed" in caplog.text + + @pytest.mark.filterwarnings("ignore::agent_framework._feature_stage.ExperimentalWarning") + def test_categorize_middleware_supported_categories_raises_for_bundle_members(self) -> None: + """A bundle member outside the supported categories raises: bundles are indivisible.""" + from agent_framework import MiddlewareBundle + from agent_framework.exceptions import MiddlewareException + + bundle = MiddlewareBundle([TestAgentMiddleware(), TestChatMiddleware(), TestFunctionMiddleware()]) + with pytest.raises(MiddlewareException, match="cannot be partially installed"): + categorize_middleware([bundle], supported_categories=("chat", "function")) + # A bundle whose members all fit the supported categories expands normally. + chat_only = MiddlewareBundle([TestChatMiddleware(), TestFunctionMiddleware()]) + result = categorize_middleware([chat_only], supported_categories=("chat", "function")) + assert len(result["chat"]) == 1 + assert len(result["function"]) == 1 diff --git a/python/packages/core/tests/core/test_middleware_with_agent.py b/python/packages/core/tests/core/test_middleware_with_agent.py index 636b1ce1d59..708951a88f6 100644 --- a/python/packages/core/tests/core/test_middleware_with_agent.py +++ b/python/packages/core/tests/core/test_middleware_with_agent.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. from collections.abc import Awaitable, Callable -from typing import Any +from typing import Any, cast import pytest @@ -84,7 +84,30 @@ async def process(self, context: AgentContext, call_next: Callable[[], Awaitable await call_next() execution_order.append("after") - agent = Agent(client=client, middleware=TrackingAgentMiddleware()) # type: ignore[arg-type] + agent = Agent(client=client, middleware=TrackingAgentMiddleware()) + + response = await agent.run([Message(role="user", contents=["test message"])]) + + assert response is not None + assert execution_order == ["before", "after"] + + async def test_bare_middleware_assigned_to_attribute_is_installed(self, client: SupportsChatGetResponse) -> None: + """A single middleware object assigned bare to ``agent.middleware`` executes. + + categorize_middleware owns the bare-source rule (a non-sequence source is a + one-element list) and ``run()`` passes the raw attribute straight to it, so a + bare attribute assignment — which used to be silently ignored — now executes. + """ + execution_order: list[str] = [] + + class TrackingAgentMiddleware(AgentMiddleware): + async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + execution_order.append("before") + await call_next() + execution_order.append("after") + + agent = Agent(client=client) + agent.middleware = cast("Any", TrackingAgentMiddleware()) response = await agent.run([Message(role="user", contents=["test message"])]) diff --git a/python/packages/core/tests/core/test_sessions.py b/python/packages/core/tests/core/test_sessions.py index 5a229b2a413..cf22041d4b2 100644 --- a/python/packages/core/tests/core/test_sessions.py +++ b/python/packages/core/tests/core/test_sessions.py @@ -35,7 +35,13 @@ ) from agent_framework._sessions import ( LOCAL_HISTORY_CONVERSATION_ID, + _adopt_run_persistence_gate_claim, + _defer_run_persistence, _filter_approval_control_messages, + _offer_run_persistence_gate_claim, + _run_identity_scope, + _RunPersistenceGate, + _suspend_run_persistence_gate, is_local_history_conversation_id, ) from agent_framework._telemetry import FeatureIndex @@ -1658,3 +1664,182 @@ def tracked_open(path: Path, *args: Any, **kwargs: Any) -> Any: assert not overlap_detected loaded = await provider.get_messages(session_id) assert [message.text for message in loaded] == ["first", "second"] + + +# --------------------------------------------------------------------------- +# Run-persistence gate tests +# --------------------------------------------------------------------------- + + +class TestRunPersistenceGate: + """The deferral gate owned by _sessions (used by egress-enforcement middleware).""" + + @staticmethod + def _recording_persist(executed: list[str], label: str) -> Callable[[], Awaitable[None]]: + async def persist() -> None: + executed.append(label) + + return persist + + async def test_gate_defers_and_flush_runs_in_order(self) -> None: + executed: list[str] = [] + gate = _RunPersistenceGate() + with gate: + assert _defer_run_persistence(self._recording_persist(executed, "a")) is True + assert _defer_run_persistence(self._recording_persist(executed, "b")) is True + assert executed == [] + await gate.flush() + assert executed == ["a", "b"] + # Outside any gate scope, callers persist inline. + assert _defer_run_persistence(self._recording_persist(executed, "c")) is False + + async def test_flush_inside_the_scope_raises(self) -> None: + # Reset-before-drain by construction: draining while the gate is active would + # re-defer re-entrant persistence into a list nobody drains. + gate = _RunPersistenceGate() + with gate, pytest.raises(RuntimeError, match="still active"): + await gate.flush() + await gate.flush() # after exit the same flush is legal + + def test_reentering_an_active_gate_raises(self) -> None: + gate = _RunPersistenceGate() + with gate, pytest.raises(RuntimeError, match="not re-entrant"), gate: + pass + + async def test_gate_can_be_reused_sequentially(self) -> None: + # The streaming chat gate enters once around call_next and once around stream + # consumption; both sections collect into the same handle. + executed: list[str] = [] + gate = _RunPersistenceGate() + with gate: + _defer_run_persistence(self._recording_persist(executed, "first")) + with gate: + _defer_run_persistence(self._recording_persist(executed, "second")) + await gate.flush() + assert executed == ["first", "second"] + + async def test_drop_discards_deferred_persistence(self) -> None: + executed: list[str] = [] + gate = _RunPersistenceGate() + with gate: + _defer_run_persistence(self._recording_persist(executed, "denied")) + gate.drop() + await gate.flush() + assert executed == [] + + async def test_nested_gates_collect_independently(self) -> None: + executed: list[str] = [] + outer = _RunPersistenceGate() + inner = _RunPersistenceGate() + with outer: + _defer_run_persistence(self._recording_persist(executed, "outer-1")) + with inner: + _defer_run_persistence(self._recording_persist(executed, "inner")) + _defer_run_persistence(self._recording_persist(executed, "outer-2")) + await inner.flush() + assert executed == ["inner"] + await outer.flush() + assert executed == ["inner", "outer-1", "outer-2"] + + async def test_suspension_makes_nested_persistence_inline(self) -> None: + # The function-invocation layer suspends the gate around tool invocations so + # nested agent runs persist inline instead of deferring into the outer gate. + executed: list[str] = [] + gate = _RunPersistenceGate() + with gate: + with _suspend_run_persistence_gate(): + assert _defer_run_persistence(self._recording_persist(executed, "nested")) is False + assert _defer_run_persistence(self._recording_persist(executed, "own")) is True + await gate.flush() + assert executed == ["own"] + + async def test_reentrant_persist_during_flush_runs_inline(self) -> None: + # Mirrors _run_after_providers: the deferred callable re-checks the gate when + # executed. Because flush only runs after the scope was exited, the re-entrant + # check finds no gate and the persist runs inline instead of re-deferring. + executed: list[str] = [] + gate = _RunPersistenceGate() + + async def reentrant() -> None: + if _defer_run_persistence(reentrant): + return + executed.append("ran") + + with gate: + assert _defer_run_persistence(reentrant) is True + await gate.flush() + assert executed == ["ran"] + + async def test_bound_gate_defers_only_its_owner_run(self) -> None: + executed: list[str] = [] + gate = _RunPersistenceGate() + owner = object() + other = object() + gate.bind_owner(owner) + with gate: + with _run_identity_scope(owner): + assert _defer_run_persistence(self._recording_persist(executed, "own")) is True + with _run_identity_scope(other): + # A different (nested or sibling) run's persist runs inline. + assert _defer_run_persistence(self._recording_persist(executed, "other")) is False + # An identity-less persist under a bound gate comes from a custom-loop + # nested run (the owner always carries its identity): inline. + assert _defer_run_persistence(self._recording_persist(executed, "identity-less")) is False + await gate.flush() + assert executed == ["own"] + + async def test_unbound_gate_defers_only_identity_less_persists(self) -> None: + # An unbound gate (its run never adopted the claim: a fully custom run loop) + # stays fail-closed for identity-less persists — the covered run's own — and + # inline for identity-stamped (nested/sibling) runs. + executed: list[str] = [] + gate = _RunPersistenceGate() + with gate: + assert _defer_run_persistence(self._recording_persist(executed, "custom-own")) is True + with _run_identity_scope(object()): + assert _defer_run_persistence(self._recording_persist(executed, "stamped-nested")) is False + await gate.flush() + assert executed == ["custom-own"] + + def test_bind_owner_first_bind_wins(self) -> None: + gate = _RunPersistenceGate() + first = object() + gate.bind_owner(first) + gate.bind_owner(object()) + assert gate.accepts(first) is True + + def test_claim_handshake_is_keyed_to_the_agent_instance(self) -> None: + gate = _RunPersistenceGate() + target = object() + stranger = object() + identity = object() + _offer_run_persistence_gate_claim(gate, target) + # A nested or sibling run (a different agent instance) must not claim the gate... + _adopt_run_persistence_gate_claim(stranger, object()) + assert gate.accepts(identity) is False + # ...the targeted run does, and adoption consumes the ticket. + _adopt_run_persistence_gate_claim(target, identity) + assert gate.accepts(identity) is True + _adopt_run_persistence_gate_claim(target, object()) # stale ticket is gone + assert gate.accepts(identity) is True + + async def test_flush_runs_callables_with_the_gate_context_suspended(self) -> None: + # A flushed callable that re-checks the gate (like _run_after_providers) must + # run inline even when an enclosing gate is active at flush time: its own + # covering verdict already permitted it, so it is never re-deferred. + executed: list[str] = [] + outer = _RunPersistenceGate() + inner = _RunPersistenceGate() + + async def reentrant() -> None: + if _defer_run_persistence(reentrant): + return + executed.append("ran-inline") + + with inner: + assert _defer_run_persistence(reentrant) is True + with outer: # an enclosing gate is active while the inner gate flushes + await inner.flush() + assert executed == ["ran-inline"] + await outer.flush() + assert executed == ["ran-inline"] diff --git a/python/packages/foundry/agent_framework_foundry/_agent.py b/python/packages/foundry/agent_framework_foundry/_agent.py index 2013d0401af..28a840c9e3e 100644 --- a/python/packages/foundry/agent_framework_foundry/_agent.py +++ b/python/packages/foundry/agent_framework_foundry/_agent.py @@ -297,7 +297,7 @@ def as_agent( tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, default_options: FoundryAgentOptionsT | Mapping[str, Any] | None = None, context_providers: Sequence[ContextProvider] | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, require_per_service_call_history_persistence: bool = False, function_invocation_configuration: FunctionInvocationConfiguration | None = None, compaction_strategy: CompactionStrategy | None = None, @@ -658,7 +658,7 @@ def __init__( default_headers: Mapping[str, str] | None = None, tools: FunctionTool | Callable[..., Any] | Sequence[FunctionTool | Callable[..., Any]] | None = None, context_providers: Sequence[ContextProvider] | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, client_type: type[RawFoundryAgentChatClient] | None = None, env_file_path: str | None = None, env_file_encoding: str | None = None, @@ -973,7 +973,7 @@ def __init__( default_headers: Mapping[str, str] | None = None, tools: FunctionTool | Callable[..., Any] | Sequence[FunctionTool | Callable[..., Any]] | None = None, context_providers: Sequence[ContextProvider] | None = None, - middleware: Sequence[MiddlewareTypes] | None = None, + middleware: MiddlewareTypes | Sequence[MiddlewareTypes] | None = None, client_type: type[RawFoundryAgentChatClient] | None = None, env_file_path: str | None = None, env_file_encoding: str | None = None, From 2c66d6111bf42fa6ad66d7e17cdbdbaa0f118209 Mon Sep 17 00:00:00 2001 From: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:36:41 +0000 Subject: [PATCH 5/8] docs(python): drop previous-behavior notes from middleware docstrings Per review: docstrings describe current behavior only. The bare-middleware behavior change stays recorded in the PR description and commit history. Signed-off-by: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> --- python/packages/core/agent_framework/_agents.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index 7c374b070cf..4b7a1d3f544 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -435,10 +435,7 @@ def __init__( description: The description of the agent. context_providers: Context providers to include during agent invocation. middleware: List of middleware, or a single middleware object (including a - ``MiddlewareBundle``) which is treated as a one-element list. Note: - a bare middleware object — passed here or assigned directly to the - ``middleware`` attribute — used to be silently ignored by ``run()`` - and now executes. + ``MiddlewareBundle``) which is treated as a one-element list. additional_properties: Additional properties set on the agent. """ if id is None: @@ -814,9 +811,7 @@ def __init__( context_providers: Context providers to include during agent invocation. middleware: List of middleware to intercept agent and function invocations. A single middleware object (including a ``MiddlewareBundle``) is - treated as a one-element list. Note: a bare middleware object — - passed here or assigned directly to the ``middleware`` attribute — - used to be silently ignored by ``run()`` and now executes. + treated as a one-element list. require_per_service_call_history_persistence: When True (and a HistoryProvider is present), the provider always persists history via per-service-call middleware, regardless of whether the client stores history server-side. If the client does From 863ff03e4cb8e437a17332492734d98992d0376b Mon Sep 17 00:00:00 2001 From: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> Date: Thu, 6 Aug 2026 02:44:00 +0000 Subject: [PATCH 6/8] fix(python): gate ownership survives retrying middleware MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A retry or fallback middleware issuing a second call_next() gave the new attempt a fresh run identity that the persistence gate's first-bind-wins ownership rejected, so the retried attempt's history persisted inline before the output verdict — a denied response became durable again. The gate now accumulates every identity adopted through its own offer ticket: all attempts' persistence stays behind the one final verdict (deny drops all of it, allow flushes all of it). Accumulation over rebind-replace is deliberate: rebinding would flip an earlier attempt's still-running background work from deferred to inline, which is the fail-open direction. A foreign agent still cannot bind: tickets are minted only by the covered pipeline's final handler and adoption is instance-keyed. Also consolidates the bare-middleware-source rule into a single _as_middleware_list owner used by every interpretation site (the harness merge, BaseAgent.__init__, categorize_middleware, both client-kwargs merges, get_response, SessionContext.extend_middleware), including the str/bytes exclusion the stray copies missed. The constructor now stores a copy of the caller's sequence; assign to the middleware attribute for post-construction changes. Retry regression tests cover denied and allowed retried runs in both stream modes and fail with first-bind-wins restored. Signed-off-by: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> --- .../core/agent_framework/_agent_hooks.py | 5 +- .../packages/core/agent_framework/_agents.py | 60 +++++----- .../core/agent_framework/_harness/_agent.py | 14 +-- .../core/agent_framework/_middleware.py | 33 ++++-- .../core/agent_framework/_sessions.py | 81 +++++++------ .../packages/core/agent_framework/_tools.py | 10 +- .../core/tests/core/test_agent_hooks.py | 109 ++++++++++++++++++ .../core/tests/core/test_middleware.py | 12 ++ .../packages/core/tests/core/test_sessions.py | 13 ++- 9 files changed, 236 insertions(+), 101 deletions(-) diff --git a/python/packages/core/agent_framework/_agent_hooks.py b/python/packages/core/agent_framework/_agent_hooks.py index 580392efe3c..d0867c5319d 100644 --- a/python/packages/core/agent_framework/_agent_hooks.py +++ b/python/packages/core/agent_framework/_agent_hooks.py @@ -1450,7 +1450,10 @@ def create_agent_hooks_middleware( written back into the framework's messages, arguments, and results so execution uses exactly the values the interceptors approved, streaming runs are buffered and only released after the ``output`` verdict permits, and durable history persistence is - deferred until the covering verdict permits the content. + deferred until the covering verdict permits the content. When middleware retries a + run, every attempt's persistence stays behind the one final verdict; note that + ``post_model_call`` is the content-complete audit point — a discarded attempt's + response passes its own ``post_model_call`` verdict but never reaches ``output``. Every agent run is one agent-hooks session: a fresh ``InterceptionEmitter`` and ``AgentContextBuilder`` pair is created per run and ``agent_startup`` / diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index 4b7a1d3f544..05606d0a5fc 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -26,7 +26,13 @@ from ._clients import BaseChatClient, SupportsChatGetResponse from ._docstrings import apply_layered_docstring -from ._middleware import AgentMiddlewareLayer, FunctionInvocationContext, MiddlewareTypes, categorize_middleware +from ._middleware import ( + AgentMiddlewareLayer, + FunctionInvocationContext, + MiddlewareTypes, + _as_middleware_list, # pyright: ignore[reportPrivateUsage] + categorize_middleware, +) from ._serialization import SerializationMixin from ._sessions import ( AgentSession, @@ -435,7 +441,9 @@ def __init__( description: The description of the agent. context_providers: Context providers to include during agent invocation. middleware: List of middleware, or a single middleware object (including a - ``MiddlewareBundle``) which is treated as a one-element list. + ``MiddlewareBundle``) which is treated as a one-element list. The + constructor copies the sequence; assign to or mutate the + ``middleware`` attribute for post-construction changes. additional_properties: Additional properties set on the agent. """ if id is None: @@ -444,14 +452,11 @@ def __init__( self.name = name self.description = description self.context_providers: list[ContextProvider] = list(context_providers or []) - # Canonicalize storage: a bare middleware source (a single middleware object - # or a MiddlewareBundle) is stored as a single-element list so the declared - # attribute type stays `list[MiddlewareTypes] | None`. Interpreting bare - # sources is owned by categorize_middleware; this mirrors it for storage only. - if middleware is not None and not isinstance(middleware, Sequence): - middleware = [middleware] + # Canonicalize storage: the bare-source rule (a single middleware object or a + # MiddlewareBundle is one element) is owned by _as_middleware_list; storing a + # normalized list keeps the declared attribute type honest. self.middleware: list[MiddlewareTypes] | None = ( - cast(list[MiddlewareTypes], middleware) if middleware is not None else None + _as_middleware_list(middleware) if middleware is not None else None ) self.additional_properties: dict[str, Any] = cast(dict[str, Any], additional_properties or {}) @@ -1495,31 +1500,22 @@ async def _prepare_run_context( *middleware_list["chat"], ] if provider_function_chat_middleware: - existing_middleware = effective_client_kwargs.get("middleware") - if isinstance(existing_middleware, Sequence) and not isinstance(existing_middleware, (str, bytes)): - effective_client_kwargs["middleware"] = [ - *existing_middleware, - *provider_function_chat_middleware, - ] - elif existing_middleware is not None: - effective_client_kwargs["middleware"] = [ - cast(MiddlewareTypes, existing_middleware), - *provider_function_chat_middleware, - ] - else: - effective_client_kwargs["middleware"] = provider_function_chat_middleware - - if per_service_call_history_middleware is not None: - existing_middleware = effective_client_kwargs.get("middleware") - if isinstance(existing_middleware, Sequence) and not isinstance(existing_middleware, (str, bytes)): - effective_client_kwargs["middleware"] = [*existing_middleware, per_service_call_history_middleware] - elif existing_middleware is not None: + existing_middleware = cast( + "MiddlewareTypes | Sequence[MiddlewareTypes] | None", effective_client_kwargs.get("middleware") + ) effective_client_kwargs["middleware"] = [ - cast(MiddlewareTypes, existing_middleware), - per_service_call_history_middleware, + *_as_middleware_list(existing_middleware), + *provider_function_chat_middleware, ] - else: - effective_client_kwargs["middleware"] = [per_service_call_history_middleware] + + if per_service_call_history_middleware is not None: + existing_middleware = cast( + "MiddlewareTypes | Sequence[MiddlewareTypes] | None", effective_client_kwargs.get("middleware") + ) + effective_client_kwargs["middleware"] = [ + *_as_middleware_list(existing_middleware), + per_service_call_history_middleware, + ] return { "session": active_session, diff --git a/python/packages/core/agent_framework/_harness/_agent.py b/python/packages/core/agent_framework/_harness/_agent.py index e8567092829..8ad199069f4 100644 --- a/python/packages/core/agent_framework/_harness/_agent.py +++ b/python/packages/core/agent_framework/_harness/_agent.py @@ -14,7 +14,7 @@ import sys from collections.abc import Callable, Sequence from pathlib import Path -from typing import TYPE_CHECKING, Any, TypedDict, cast +from typing import TYPE_CHECKING, Any, TypedDict from .._agents import Agent, SupportsAgentRun from .._clients import SupportsShellTool, SupportsWebSearchTool @@ -655,13 +655,11 @@ def create_harness_agent( # Message injection is always on. It is a no-op when no messages are queued for the session, # so there is no opt-out. assembled_middleware.append(MessageInjectionMiddleware()) - if middleware is not None: - # A bare middleware object (including a MiddlewareBundle, which is - # deliberately not a sequence) counts as a one-element source. - if isinstance(middleware, Sequence): - assembled_middleware.extend(cast("Sequence[MiddlewareTypes]", middleware)) - else: - assembled_middleware.append(middleware) + # Bare-source normalization (a single middleware object or a MiddlewareBundle is + # one element) is owned by _as_middleware_list. + from .._middleware import _as_middleware_list # pyright: ignore[reportPrivateUsage] + + assembled_middleware.extend(_as_middleware_list(middleware)) agent = Agent( client, diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 46e8d49be16..77b452214ea 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -1622,6 +1622,24 @@ class MiddlewareDict(TypedDict): chat: list[ChatMiddleware | ChatMiddlewareCallable] +def _as_middleware_list( + source: MiddlewareTypes | Sequence[MiddlewareTypes] | None, +) -> list[MiddlewareTypes]: + """Normalize one middleware source into a list — the bare-source rule's single owner. + + ``None`` is empty; a sequence (never str/bytes) is taken element-wise; any other + bare source — a single middleware object or a :class:`MiddlewareBundle`, which is + deliberately not a sequence — is one element. The ``None`` check is deliberate + (not truthiness): a bare middleware object with a falsy ``__bool__``/``__len__`` + still counts as one element, never silently dropped. + """ + if source is None: + return [] + if isinstance(source, Sequence) and not isinstance(source, (str, bytes)): + return list(cast("Sequence[MiddlewareTypes]", source)) + return [cast("MiddlewareTypes", source)] + + def categorize_middleware( *middleware_sources: MiddlewareTypes | Sequence[MiddlewareTypes] | None, supported_categories: Collection[str] | None = None, @@ -1631,8 +1649,8 @@ def categorize_middleware( Args: *middleware_sources: Variable number of middleware sources to categorize. A bare (non-sequence) source — a single middleware object or a - :class:`MiddlewareBundle` — is treated as a one-element list; this - function is the single owner of that rule. + :class:`MiddlewareBundle` — is treated as a one-element list + (normalization is owned by :func:`_as_middleware_list`). Keyword Args: supported_categories: The categories the call site actually installs, e.g. @@ -1651,16 +1669,11 @@ def categorize_middleware( """ result: MiddlewareDict = {"agent": [], "function": [], "chat": []} - # Merge all middleware sources into a single list. The None check is deliberate - # (not truthiness): a bare middleware object with a falsy __bool__/__len__ must - # still be treated as a one-element source, never silently dropped. + # Merge all middleware sources into a single list (bare-source normalization is + # owned by _as_middleware_list). all_middleware: list[Any] = [] for source in middleware_sources: - if source is not None: - if isinstance(source, Sequence) and not isinstance(source, (str, bytes)): - all_middleware.extend(source) # type: ignore - else: - all_middleware.append(source) + all_middleware.extend(_as_middleware_list(source)) # Expand bundles first: a bundle's members are categorized individually (in # order) but travel as one unit, so a feature spanning several categories can diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index 73eae2cdf5a..ef075d64dd2 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -35,7 +35,7 @@ from dataclasses import dataclass from functools import partial from pathlib import Path -from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias, TypeGuard, TypeVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias, TypeVar, cast import msgspec @@ -187,18 +187,6 @@ def _deduplicate_origin_session_ids(origin_session_ids: Iterable[str]) -> list[s return unique_origin_session_ids -def _is_middleware_sequence( - middleware: MiddlewareTypes | Sequence[MiddlewareTypes], -) -> TypeGuard[Sequence[MiddlewareTypes]]: - return isinstance(middleware, Sequence) and not isinstance(middleware, (str, bytes)) - - -def _is_single_middleware( - middleware: MiddlewareTypes | Sequence[MiddlewareTypes], -) -> TypeGuard[MiddlewareTypes]: - return not _is_middleware_sequence(middleware) - - @dataclass(frozen=True, slots=True) class _StateTypeRegistration: cls: type[Any] @@ -691,15 +679,10 @@ def extend_middleware( source_id: The provider source_id adding this middleware. middleware: A single chat/function middleware object/callable or sequence of middleware. """ - from ._middleware import categorize_middleware + from ._middleware import _as_middleware_list, categorize_middleware # pyright: ignore[reportPrivateUsage] from .exceptions import MiddlewareException - if _is_middleware_sequence(middleware): - middleware_items = list(middleware) - elif _is_single_middleware(middleware): - middleware_items = [middleware] - else: - raise TypeError("middleware must be a middleware object or a sequence of middleware objects.") + middleware_items = _as_middleware_list(middleware) middleware_list = categorize_middleware(middleware_items) if middleware_list["agent"]: raise MiddlewareException("Context providers may only add chat or function middleware.") @@ -1128,14 +1111,19 @@ class _RunPersistenceGate: (the covering verdict permitted the content) or discards them with :meth:`drop` (the content was denied and must never become durable). - Ownership: the gate covers exactly one run's verdict, so it only collects that - run's own persistence (see :meth:`accepts`). The agent-seam gate is created before - its run starts and is bound lazily through the - :func:`_offer_run_persistence_gate_claim` / :func:`_adopt_run_persistence_gate_claim` - handshake; the chat-seam gate is created inside its run and bound immediately. - While unbound (an agent whose custom run loop never adopts the claim), the gate - falls back fail-closed: persists carrying no run identity defer, while persists - from identity-stamped (nested or sibling) runs execute inline at their own run + Ownership: the gate covers one pipeline execution's verdict, so it only collects + persistence issued by the run(s) that pipeline's final handler started (see + :meth:`accepts`). The agent-seam gate is created before its run starts and is + bound lazily through the :func:`_offer_run_persistence_gate_claim` / + :func:`_adopt_run_persistence_gate_claim` handshake; the chat-seam gate is created + inside its run and bound immediately. A retrying or fallback middleware may invoke + ``call_next()`` several times, re-offering the same gate: **every** identity + adopted through the gate's own ticket becomes an accepted owner, so each attempt's + persistence stays deferred behind the final verdict (a first-bind-wins rule would + let later attempts persist inline ahead of the verdict — fail-open). While unbound + (an agent whose custom run loop never adopts the claim), the gate falls back + fail-closed: persists carrying no run identity defer, while persists from + identity-stamped (nested or sibling) runs execute inline at their own run boundaries. Reset-before-drain is enforced by construction: :meth:`flush` refuses to run while @@ -1158,19 +1146,26 @@ class _RunPersistenceGate: persists and restores the outer gate on exit. """ - __slots__ = ("_owner", "_owner_bound", "_pending", "_token") + __slots__ = ("_owner_identities", "_pending", "_token") def __init__(self) -> None: self._pending: list[Callable[[], Awaitable[None]]] = [] self._token: Token[_RunPersistenceGate | None] | None = None - self._owner: object | None = None - self._owner_bound: bool = False + self._owner_identities: list[object | None] = [] def bind_owner(self, owner: object | None) -> None: - """Bind the gate to the run identity whose persistence it defers (first bind wins).""" - if not self._owner_bound: - self._owner = owner - self._owner_bound = True + """Add ``owner`` to the run identities whose persistence this gate defers. + + Every bind accumulates (it never replaces): binding is only reachable through + the gate's own instance-keyed claim ticket, so each added identity is a run + started by the covered pipeline's final handler — e.g. successive attempts + made by a retrying middleware. All of those attempts' persists must stay + deferred behind the one final verdict; dropping earlier identities (rebind) or + ignoring later ones (first-bind-wins) would let some attempt's persistence run + inline ahead of the verdict, which is fail-open. + """ + if not any(existing is owner for existing in self._owner_identities): + self._owner_identities.append(owner) def collect(self, persist: Callable[[], Awaitable[None]]) -> None: """Collect one deferred persistence callable (see :func:`_defer_run_persistence`).""" @@ -1179,14 +1174,16 @@ def collect(self, persist: Callable[[], Awaitable[None]]) -> None: def accepts(self, identity: object | None) -> bool: """Whether a persist issued under ``identity`` belongs to this gate's run. - A bound gate accepts exactly its owner identity. An unbound gate (its run - never adopted the claim — a custom run loop) accepts only identity-less - persists: the covered run's own persists carry no identity there (fail-closed, - matching the pre-ownership behavior), while identity-stamped persists come - from other runs that own their content's verdicts themselves. + A bound gate accepts every identity adopted through its claim ticket (each + ``call_next()`` attempt of the covered pipeline — see :meth:`bind_owner`). An + unbound gate (its run never adopted the claim — a custom run loop) accepts + only identity-less persists: the covered run's own persists carry no identity + there (fail-closed, matching the pre-ownership behavior), while + identity-stamped persists come from other runs that own their content's + verdicts themselves. """ - if self._owner_bound: - return identity is self._owner + if self._owner_identities: + return any(owner is identity for owner in self._owner_identities) return identity is None def __enter__(self) -> _RunPersistenceGate: diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 39dd3e20003..e7760c9cbdf 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -72,6 +72,7 @@ FunctionInvocationContext, FunctionMiddlewarePipeline, FunctionMiddlewareTypes, + MiddlewareTypes, ) from ._sessions import AgentSession from ._types import ( @@ -3223,7 +3224,7 @@ def get_response( function_invocation_kwargs: Mapping[str, Any] | None = None, client_kwargs: Mapping[str, Any] | None = None, ) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: - from ._middleware import categorize_middleware + from ._middleware import _as_middleware_list, categorize_middleware # pyright: ignore[reportPrivateUsage] from ._types import ( ChatResponse, ResponseStream, @@ -3237,12 +3238,9 @@ def get_response( # Build the run-local middleware pipeline and recover shared budget/session state for approval re-entry. request_kwargs = dict(client_kwargs) if client_kwargs is not None else {} if middleware is not None: - existing_middleware = request_kwargs.get("middleware", []) request_kwargs["middleware"] = [ - *( - existing_middleware - if isinstance(existing_middleware, Sequence) and not isinstance(existing_middleware, (str, bytes)) - else [existing_middleware] + *_as_middleware_list( + cast("MiddlewareTypes | Sequence[MiddlewareTypes] | None", request_kwargs.get("middleware")) ), *middleware, ] diff --git a/python/packages/core/tests/core/test_agent_hooks.py b/python/packages/core/tests/core/test_agent_hooks.py index c4d4a085532..512f820c91f 100644 --- a/python/packages/core/tests/core/test_agent_hooks.py +++ b/python/packages/core/tests/core/test_agent_hooks.py @@ -1689,6 +1689,115 @@ class CustomLoopAgent( # type: ignore[misc] # ty: ignore[invalid-method-overri assert session.state.get("outer_history", {}).get("messages", []) == [] +class _FlakyOnceClient(MockBaseChatClient): + """Fails the first model call (both shapes), then behaves like the mock.""" + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._failed_once = False + + def _inner_get_response(self, **kwargs: Any) -> Any: + if not self._failed_once: + self._failed_once = True + if kwargs.get("stream"): + + async def _boom() -> AsyncIterable[ChatResponseUpdate]: + raise RuntimeError("attempt 1 failed") + yield # pragma: no cover # pyright: ignore[reportUnreachable] + + return ResponseStream(_boom()) + + async def _fail() -> ChatResponse: + raise RuntimeError("attempt 1 failed") + + return _fail() + return super()._inner_get_response(**kwargs) # pyright: ignore[reportCallIssue] + + +class _RetryOnceMiddleware(AgentMiddleware): + """The documented retry pattern: catch a failing attempt and re-invoke call_next().""" + + async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + if not context.stream: + try: + await call_next() + except RuntimeError: + context.result = None + await call_next() + return + await call_next() + stream = cast("ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", context.result) + try: + await stream.get_final_response() + except RuntimeError: + context.result = None + await call_next() + + +@requires_sdk +@pytest.mark.parametrize("streaming", [False, True], ids=["non_streaming", "streaming"]) +async def test_retried_run_denied_output_never_becomes_durable(streaming: bool) -> None: + # A retrying middleware re-invokes call_next(): attempt 2 runs with a fresh + # identity adopted through the same gate ticket. The gate must accept every + # adopted identity — pinning the first attempt's identity would let attempt 2's + # persistence run inline BEFORE the output verdict, making the denied response + # durable (fail-open). + from agent_framework import AgentSession, InMemoryHistoryProvider + + provider = InMemoryHistoryProvider() + session = AgentSession() + agent = Agent( + client=_FlakyOnceClient(), + name="retried", + context_providers=[provider], + middleware=[ + create_agent_hooks_middleware([PointGuard("output", Verdict.deny(reason="egress_blocked"))]), + _RetryOnceMiddleware(), + ], + ) + + with pytest.raises(InterceptionBlocked): + if streaming: + async for _ in agent.run("hello there", session=session, stream=True): + pass + else: + await agent.run("hello there", session=session) + + # Nothing from either attempt is durable: the failed attempt produced no + # persistence and the retried attempt's persistence stayed behind the denied gate. + assert session.state.get(provider.source_id, {}).get("messages", []) == [] + + +@requires_sdk +@pytest.mark.parametrize("streaming", [False, True], ids=["non_streaming", "streaming"]) +async def test_retried_run_allowed_output_persists_the_retry_attempt(streaming: bool) -> None: + from agent_framework import AgentSession, InMemoryHistoryProvider + + provider = InMemoryHistoryProvider() + session = AgentSession() + agent = Agent( + client=_FlakyOnceClient(), + name="retried", + context_providers=[provider], + middleware=[create_agent_hooks_middleware([AllowGuard()]), _RetryOnceMiddleware()], + ) + + if streaming: + stream = agent.run("hello there", session=session, stream=True) + async for _ in stream: + pass + final = await stream.get_final_response() + else: + final = await agent.run("hello there", session=session) + + expected_response = "update - hello there" if streaming else "test response - hello there" + assert final.text == expected_response + # The permitted retry attempt's exchange is durable (the failed attempt produced + # nothing to persist). + stored = cast("list[Message]", session.state[provider.source_id]["messages"]) + assert [message.text for message in stored] == ["hello there", expected_response] + + # endregion # region Stream hooks cannot escape the gate diff --git a/python/packages/core/tests/core/test_middleware.py b/python/packages/core/tests/core/test_middleware.py index 0f519d234bf..264d13b87a7 100644 --- a/python/packages/core/tests/core/test_middleware.py +++ b/python/packages/core/tests/core/test_middleware.py @@ -1755,3 +1755,15 @@ def test_categorize_middleware_supported_categories_raises_for_bundle_members(se result = categorize_middleware([chat_only], supported_categories=("chat", "function")) assert len(result["chat"]) == 1 assert len(result["function"]) == 1 + + def test_as_middleware_list_owns_the_bare_source_rule(self) -> None: + """One owner for bare-source normalization, including the str/bytes exclusion.""" + from agent_framework._middleware import _as_middleware_list + + agent_mw = TestAgentMiddleware() + assert _as_middleware_list(None) == [] + assert _as_middleware_list(agent_mw) == [agent_mw] + assert _as_middleware_list([agent_mw]) == [agent_mw] + assert _as_middleware_list((agent_mw,)) == [agent_mw] + # Strings are sequences but never element-ized into characters. + assert _as_middleware_list("bare-string") == ["bare-string"] # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] diff --git a/python/packages/core/tests/core/test_sessions.py b/python/packages/core/tests/core/test_sessions.py index cf22041d4b2..d4606f895ec 100644 --- a/python/packages/core/tests/core/test_sessions.py +++ b/python/packages/core/tests/core/test_sessions.py @@ -1801,12 +1801,21 @@ async def test_unbound_gate_defers_only_identity_less_persists(self) -> None: await gate.flush() assert executed == ["custom-own"] - def test_bind_owner_first_bind_wins(self) -> None: + def test_bind_owner_accepts_every_adopted_identity(self) -> None: + # A retrying/fallback middleware re-invokes call_next(): the final handler + # re-offers the same gate and each attempt's run adopts it. Every adopted + # identity must stay accepted — first-bind-wins would let attempt 2's + # persistence run inline ahead of the final verdict (fail-open), and + # rebind-replace would flip attempt 1's still-running work to inline. gate = _RunPersistenceGate() first = object() + second = object() gate.bind_owner(first) - gate.bind_owner(object()) + gate.bind_owner(second) assert gate.accepts(first) is True + assert gate.accepts(second) is True + assert gate.accepts(object()) is False # other runs still persist inline + assert gate.accepts(None) is False # bound gates reject identity-less persists def test_claim_handshake_is_keyed_to_the_agent_instance(self) -> None: gate = _RunPersistenceGate() From 84671f1fe73a844dda81879dbfd08c127b2055be Mon Sep 17 00:00:00 2001 From: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> Date: Thu, 6 Aug 2026 03:38:40 +0000 Subject: [PATCH 7/8] fix(python): streaming seam runs pipeline descent inside the gate The streaming agent seam ran call_next() outside the persistence gate (only _consume entered it later), so a retry middleware that drained a successful attempt with get_final_response() and discarded it persisted that attempt's exchange before any verdict existed; a later deny dropped only the retry attempt's deferred work. The descent is now wrapped in the gate exactly like the non-streaming seam: attempt identities adopted during descent are accepted owners, so in-pipeline draining defers, deny drops every attempt, and a middleware that raises after draining strands the pending persists unexecuted. The bind_owner docstring now states the actual soundness invariant covering both bind sites: every bind comes from a run inside the covered pipeline. New tests cover drained-and-discarded attempts (deny and allow, both stream modes) and a sub-agent tool inside a drained attempt; the streaming deny variant fails with the gate wrap reverted. Signed-off-by: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> --- .../core/agent_framework/_agent_hooks.py | 16 ++- .../core/agent_framework/_sessions.py | 17 ++- .../core/tests/core/test_agent_hooks.py | 133 ++++++++++++++++++ 3 files changed, 155 insertions(+), 11 deletions(-) diff --git a/python/packages/core/agent_framework/_agent_hooks.py b/python/packages/core/agent_framework/_agent_hooks.py index d0867c5319d..86c103e562a 100644 --- a/python/packages/core/agent_framework/_agent_hooks.py +++ b/python/packages/core/agent_framework/_agent_hooks.py @@ -1078,10 +1078,18 @@ async def _process_streaming( context._run_persistence_gate = gate_handle # pyright: ignore[reportPrivateUsage] await self._emit_run_start(context, state) termination: MiddlewareTermination | None = None - try: - await call_next() - except MiddlewareTermination as exc: - termination = exc + # The gate covers the pipeline descent too, exactly like the non-streaming + # branch: a middleware that drains an attempt's stream in-pipeline (e.g. a + # retry that discards a successful attempt) issues that attempt's run-end + # persistence here, and it must defer behind the final verdict — the + # attempt's identity is an accepted owner via the claim ticket. The gate + # is then re-entered (sequential reuse) around the released stream's + # consumption in _consume, and flushed/dropped exactly once by its gate. + with gate_handle: + try: + await call_next() + except MiddlewareTermination as exc: + termination = exc inner = context.result if inner is None and termination is not None: # Terminated without a result: nothing will egress. diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index ef075d64dd2..893bb96e526 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -1156,13 +1156,16 @@ def __init__(self) -> None: def bind_owner(self, owner: object | None) -> None: """Add ``owner`` to the run identities whose persistence this gate defers. - Every bind accumulates (it never replaces): binding is only reachable through - the gate's own instance-keyed claim ticket, so each added identity is a run - started by the covered pipeline's final handler — e.g. successive attempts - made by a retrying middleware. All of those attempts' persists must stay - deferred behind the one final verdict; dropping earlier identities (rebind) or - ignoring later ones (first-bind-wins) would let some attempt's persistence run - inline ahead of the verdict, which is fail-open. + Every bind accumulates (it never replaces). Both bind sites sit inside the + covered pipeline, so every added identity is a run whose persistence the + gate's one covering verdict must gate: the agent seam binds through the + gate's instance-keyed claim ticket (each run the pipeline's final handler + starts — e.g. successive attempts made by a retrying middleware), and the + chat seam binds directly at gate creation to the identity of the run it is + executing in. All of those runs' persists must stay deferred behind the + final verdict; dropping earlier identities (rebind) or ignoring later ones + (first-bind-wins) would let some attempt's persistence run inline ahead of + the verdict, which is fail-open. """ if not any(existing is owner for existing in self._owner_identities): self._owner_identities.append(owner) diff --git a/python/packages/core/tests/core/test_agent_hooks.py b/python/packages/core/tests/core/test_agent_hooks.py index 512f820c91f..f6b4b7b4df1 100644 --- a/python/packages/core/tests/core/test_agent_hooks.py +++ b/python/packages/core/tests/core/test_agent_hooks.py @@ -1798,6 +1798,139 @@ async def test_retried_run_allowed_output_persists_the_retry_attempt(streaming: assert [message.text for message in stored] == ["hello there", expected_response] +class _DrainAndRetryOnceMiddleware(AgentMiddleware): + """Retry pattern that fully drains a SUCCESSFUL first attempt, discards it, retries. + + Unlike the failure-retry pattern, the discarded attempt completes and issues its + run-end persistence during the pipeline descent — on the streaming seam that + draining happens inside the middleware, so the gate must already be active there. + """ + + async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + await call_next() + if context.stream and context.result is not None: + stream = cast("ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", context.result) + await stream.get_final_response() # attempt 1 fully drained (successful) + context.result = None # ...and discarded + await call_next() + + +@requires_sdk +@pytest.mark.parametrize("streaming", [False, True], ids=["non_streaming", "streaming"]) +async def test_drained_and_discarded_attempt_stays_gated_on_deny(streaming: bool) -> None: + # The drained attempt SUCCEEDS: its run-end persistence is issued during the + # pipeline descent. On the streaming seam call_next must run under the gate just + # like the non-streaming seam, or that exchange persists on the spot, before any + # verdict — a deny would then only drop the retry attempt's deferred work. + from agent_framework import AgentSession, InMemoryHistoryProvider + + provider = InMemoryHistoryProvider() + session = AgentSession() + agent = Agent( + client=MockBaseChatClient(), + name="retried", + context_providers=[provider], + middleware=[ + create_agent_hooks_middleware([PointGuard("output", Verdict.deny(reason="egress_blocked"))]), + _DrainAndRetryOnceMiddleware(), + ], + ) + + with pytest.raises(InterceptionBlocked): + if streaming: + async for _ in agent.run("hello there", session=session, stream=True): + pass + else: + await agent.run("hello there", session=session) + + # NOTHING is durable — including the drained-and-discarded attempt's exchange. + assert session.state.get(provider.source_id, {}).get("messages", []) == [] + + +@requires_sdk +@pytest.mark.parametrize("streaming", [False, True], ids=["non_streaming", "streaming"]) +async def test_drained_and_discarded_attempt_flushes_on_allow(streaming: bool) -> None: + # Accumulation semantics: both attempts' identities are accepted owners, so on a + # permitted verdict the gate flushes both attempts' deferred run-end persistence + # (matching unhooked semantics, where each attempt would have persisted inline). + from agent_framework import AgentSession, InMemoryHistoryProvider + + provider = InMemoryHistoryProvider() + session = AgentSession() + agent = Agent( + client=MockBaseChatClient(), + name="retried", + context_providers=[provider], + middleware=[create_agent_hooks_middleware([AllowGuard()]), _DrainAndRetryOnceMiddleware()], + ) + + if streaming: + stream = agent.run("hello there", session=session, stream=True) + async for _ in stream: + pass + final = await stream.get_final_response() + else: + final = await agent.run("hello there", session=session) + + expected_response = "update - hello there" if streaming else "test response - hello there" + assert final.text == expected_response + stored = cast("list[Message]", session.state[provider.source_id]["messages"]) + assert [message.text for message in stored] == ["hello there", expected_response] * 2 + + +@requires_sdk +async def test_tool_nested_run_inside_drained_attempt_persists_inline() -> None: + # With the streaming gate now covering the pipeline descent, tool invocations + # inside a middleware-drained attempt execute under an ACTIVE gate. The tool-seam + # suspension must still make the nested sub-agent run persist inline there, while + # the drained attempt's own run-end persistence defers and drops on the deny. + from agent_framework import AgentSession, InMemoryHistoryProvider + + inner_provider = InMemoryHistoryProvider(source_id="inner_history") + outer_provider = InMemoryHistoryProvider(source_id="outer_history", load_messages=False) + session = AgentSession() + sub_agent = Agent(client=MockBaseChatClient(), name="sub", context_providers=[inner_provider]) + + outer_client = MockBaseChatClient() + outer_client.streaming_responses = [ + # Attempt 1 (drained by the middleware): calls the sub-agent tool, then answers. + [ + ChatResponseUpdate( + contents=[Content.from_function_call(call_id="c1", name="sub", arguments='{"task": "look this up"}')], + role="assistant", + finish_reason="tool_calls", + ) + ], + [ + ChatResponseUpdate( + contents=[Content.from_text("attempt one answer")], role="assistant", finish_reason="stop" + ) + ], + # Attempt 2 falls through to the mock default ("update - ..."). + ] + outer_agent = Agent( + client=outer_client, + name="outer", + tools=[sub_agent.as_tool(propagate_session=True)], + context_providers=[outer_provider], + middleware=[ + create_agent_hooks_middleware([PointGuard("output", Verdict.deny(reason="egress_blocked"))]), + _DrainAndRetryOnceMiddleware(), + ], + ) + + with pytest.raises(InterceptionBlocked): + async for _ in outer_agent.run("go delegate", session=session, stream=True): + pass + + # The nested run inside the drained attempt persisted inline (suspension under an + # active gate)... + inner_stored = cast("list[Message]", session.state["inner_history"]["messages"]) + assert [message.text for message in inner_stored] == ["look this up", "update - look this up"] + # ...while both outer attempts' own persistence stayed gated and dropped on deny. + assert session.state.get("outer_history", {}).get("messages", []) == [] + + # endregion # region Stream hooks cannot escape the gate From be5101546c63ee2c6dcb0fd01a76cd9614aa71b2 Mon Sep 17 00:00:00 2001 From: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> Date: Thu, 6 Aug 2026 04:10:52 +0000 Subject: [PATCH 8/8] fix(python): flush deferred persistence on streaming no-result termination With the pipeline descent now running inside the persistence gate, a middleware that drains a successful attempt and then terminates without a result left that attempt's deferred persistence stranded: the streaming no-result termination path raised before any flush, so history of exchanges that really happened and passed their own verdicts quietly vanished (streaming only; non-streaming already flushes before its re-raise). The path now flushes before re-raising the termination, with a state.halted guard first so an enforcement failure during the drained attempt still strands pending fail-closed and surfaces the halt, mirroring the non-streaming ordering exactly. The regression test covers both seams; the streaming variant fails without the fix. Signed-off-by: MohammadHaroonAbuomar <40180927+MohammadHaroonAbuomar@users.noreply.github.com> --- .../core/agent_framework/_agent_hooks.py | 11 ++++- .../core/tests/core/test_agent_hooks.py | 43 +++++++++++++++++++ 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_agent_hooks.py b/python/packages/core/agent_framework/_agent_hooks.py index 86c103e562a..669cdb262a1 100644 --- a/python/packages/core/agent_framework/_agent_hooks.py +++ b/python/packages/core/agent_framework/_agent_hooks.py @@ -1092,7 +1092,16 @@ async def _process_streaming( termination = exc inner = context.result if inner is None and termination is not None: - # Terminated without a result: nothing will egress. + 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 + # really happened and passed their own verdicts — mirroring the + # non-streaming branch's flush-before-re-raise. + await gate_handle.flush() shutdown_reason = "completed" raise termination if not isinstance(inner, ResponseStream): diff --git a/python/packages/core/tests/core/test_agent_hooks.py b/python/packages/core/tests/core/test_agent_hooks.py index f6b4b7b4df1..347d7bae4fc 100644 --- a/python/packages/core/tests/core/test_agent_hooks.py +++ b/python/packages/core/tests/core/test_agent_hooks.py @@ -1878,6 +1878,49 @@ async def test_drained_and_discarded_attempt_flushes_on_allow(streaming: bool) - assert [message.text for message in stored] == ["hello there", expected_response] * 2 +class _DrainThenTerminateWithoutResultMiddleware(AgentMiddleware): + """Drains a successful attempt, then short-circuits the run with NO result.""" + + async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + await call_next() + if context.stream and context.result is not None: + stream = cast("ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", context.result) + await stream.get_final_response() # the attempt fully ran (successfully) + context.result = None + raise MiddlewareTermination("nothing egresses") + + +@requires_sdk +@pytest.mark.parametrize("streaming", [False, True], ids=["non_streaming", "streaming"]) +async def test_drained_attempt_history_survives_no_result_termination(streaming: bool) -> None: + # A no-egress termination is a permitted outcome (nothing needs an output + # verdict), so persistence the drained work deferred must be released — on both + # seams. The streaming no-result termination path must flush before re-raising, + # mirroring the non-streaming branch; otherwise history of model calls that + # really happened (and passed their own verdicts) quietly vanishes. + from agent_framework import AgentSession, InMemoryHistoryProvider + + provider = InMemoryHistoryProvider() + session = AgentSession() + agent = Agent( + client=MockBaseChatClient(), + name="terminated", + context_providers=[provider], + middleware=[create_agent_hooks_middleware([AllowGuard()]), _DrainThenTerminateWithoutResultMiddleware()], + ) + + if streaming: + stream = agent.run("hello there", session=session, stream=True) + async for _ in stream: # the terminated run egresses nothing + raise AssertionError("no updates should egress") + else: + assert await agent.run("hello there", session=session) is None + + expected_response = "update - hello there" if streaming else "test response - hello there" + stored = cast("list[Message]", session.state.get(provider.source_id, {}).get("messages", [])) + assert [message.text for message in stored] == ["hello there", expected_response] + + @requires_sdk async def test_tool_nested_run_inside_drained_attempt_persists_inline() -> None: # With the streaming gate now covering the pipeline descent, tool invocations