agent-framework-core 1.18.0__tar.gz → 1.20.0__tar.gz

{agent_framework_core-1.18.0 → agent_framework_core-1.20.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agent-framework-core
-Version: 1.18.0
+Version: 1.20.0
 Summary: Microsoft Agent Framework for building AI Agents with Python. This is the core package that has all the core abstractions and implementations.
 Author-email: Microsoft <af-support@microsoft.com>
 Requires-Python: >=3.10
@@ -16,41 +16,43 @@ Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
 Classifier: Typing :: Typed
 License-File: LICENSE
-Requires-Dist: msgspec>=0.20.0,<0.22
+Requires-Dist: msgspec>=0.20.0,<0.23
 Requires-Dist: typing-extensions>=4.15.0,<5
 Requires-Dist: pydantic>=2,<3
 Requires-Dist: python-dotenv>=1,<2
 Requires-Dist: opentelemetry-api>=1.39.0,<2
+Requires-Dist: pyyaml>=6.0,<7.0
+Requires-Dist: regex>=2024.11.6
 Requires-Dist: mcp>=1.24.0,<2 ; extra == "all"
-Requires-Dist: agent-framework-a2a ; extra == "all"
-Requires-Dist: agent-framework-ag-ui ; extra == "all"
-Requires-Dist: agent-framework-anthropic ; extra == "all"
-Requires-Dist: agent-framework-azure-ai-search ; extra == "all"
-Requires-Dist: agent-framework-azure-contentunderstanding ; extra == "all"
-Requires-Dist: agent-framework-azure-cosmos ; extra == "all"
-Requires-Dist: agent-framework-azurefunctions ; extra == "all"
-Requires-Dist: agent-framework-bedrock ; extra == "all"
-Requires-Dist: agent-framework-chatkit ; extra == "all"
-Requires-Dist: agent-framework-claude ; extra == "all"
-Requires-Dist: agent-framework-copilotstudio ; extra == "all"
-Requires-Dist: agent-framework-declarative ; extra == "all"
-Requires-Dist: agent-framework-devui ; extra == "all"
-Requires-Dist: agent-framework-durabletask ; extra == "all"
-Requires-Dist: agent-framework-foundry ; extra == "all"
-Requires-Dist: agent-framework-foundry-hosting ; extra == "all"
-Requires-Dist: agent-framework-foundry-local ; extra == "all"
-Requires-Dist: agent-framework-gemini ; extra == "all"
-Requires-Dist: agent-framework-github-copilot ; extra == "all" and ( python_version >= '3.11')
-Requires-Dist: agent-framework-hyperlight ; extra == "all" and ( (sys_platform == 'linux' and platform_machine == 'x86_64') or (sys_platform == 'win32' and platform_machine == 'AMD64'))
-Requires-Dist: agent-framework-mem0 ; extra == "all"
-Requires-Dist: agent-framework-mistral ; extra == "all"
-Requires-Dist: agent-framework-monty ; extra == "all"
-Requires-Dist: agent-framework-ollama ; extra == "all"
-Requires-Dist: agent-framework-openai ; extra == "all"
-Requires-Dist: agent-framework-orchestrations ; extra == "all"
-Requires-Dist: agent-framework-purview ; extra == "all"
-Requires-Dist: agent-framework-redis ; extra == "all"
-Requires-Dist: agent-framework-tools ; extra == "all"
+Requires-Dist: agent-framework-a2a>=1.0.0b261002,<2 ; extra == "all"
+Requires-Dist: agent-framework-ag-ui>=1.5.0,<2 ; extra == "all"
+Requires-Dist: agent-framework-anthropic>=1.0.0b261002,<2 ; extra == "all"
+Requires-Dist: agent-framework-azure-ai-search>=1.0.0b261002,<2 ; extra == "all"
+Requires-Dist: agent-framework-azure-contentunderstanding>=1.0.0b261002,<2 ; extra == "all"
+Requires-Dist: agent-framework-azure-cosmos>=1.0.0b261002,<2 ; extra == "all"
+Requires-Dist: agent-framework-azurefunctions>=1.0.0b260922,<2 ; extra == "all"
+Requires-Dist: agent-framework-bedrock>=1.0.0b261002,<2 ; extra == "all"
+Requires-Dist: agent-framework-chatkit>=1.0.0b261002,<2 ; extra == "all"
+Requires-Dist: agent-framework-claude>=1.0.0b261002,<2 ; extra == "all"
+Requires-Dist: agent-framework-copilotstudio>=1.0.0b261002,<2 ; extra == "all"
+Requires-Dist: agent-framework-declarative>=1.2.0,<2 ; extra == "all"
+Requires-Dist: agent-framework-devui>=1.0.0b261002,<2 ; extra == "all"
+Requires-Dist: agent-framework-durabletask>=1.0.0b260922,<2 ; extra == "all"
+Requires-Dist: agent-framework-foundry>=1.14.0,<2 ; extra == "all"
+Requires-Dist: agent-framework-foundry-hosting>=1.0.0b261002,<2 ; extra == "all"
+Requires-Dist: agent-framework-foundry-local>=1.0.0b260730,<2 ; extra == "all"
+Requires-Dist: agent-framework-gemini>=1.0.0b261002,<2 ; extra == "all"
+Requires-Dist: agent-framework-github-copilot>=2.0.1,<3 ; extra == "all" and ( python_version >= '3.11')
+Requires-Dist: agent-framework-hyperlight>=1.0.0b261002,<2 ; extra == "all" and ( python_version < '3.15' and ((sys_platform == 'linux' and platform_machine == 'x86_64') or (sys_platform == 'win32' and platform_machine == 'AMD64')))
+Requires-Dist: agent-framework-mem0>=1.0.0b261002,<2 ; extra == "all"
+Requires-Dist: agent-framework-mistral>=1.0.0b260918,<2 ; extra == "all"
+Requires-Dist: agent-framework-monty>=1.0.0b261002,<2 ; extra == "all"
+Requires-Dist: agent-framework-ollama>=1.0.0b261002,<2 ; extra == "all"
+Requires-Dist: agent-framework-openai>=1.15.0,<2 ; extra == "all"
+Requires-Dist: agent-framework-orchestrations>=1.3.0,<2 ; extra == "all"
+Requires-Dist: agent-framework-purview>=1.0.0b261002,<2 ; extra == "all"
+Requires-Dist: agent-framework-redis>=1.0.0b261002,<2 ; extra == "all" and ( python_version < '3.15')
+Requires-Dist: agent-framework-tools>=1.0.0b261002,<2 ; extra == "all"
 Project-URL: homepage, https://aka.ms/agent-framework
 Project-URL: issues, https://github.com/microsoft/agent-framework/issues
 Project-URL: release_notes, https://github.com/microsoft/agent-framework/releases?q=tag%3Apython-1&expanded=true
@@ -227,6 +229,21 @@ asyncio.run(main())
 
 You can explore additional agent samples [here](https://github.com/microsoft/agent-framework/tree/main/python/samples/02-agents).
 
+### MCP HTTP client ownership and cookies
+
+**Breaking change:** HTTP clients created by `MCPStreamableHTTPTool` no longer persist
+response cookies. This applies with or without a `header_provider`, including repeated
+calls within the same MCP session. Connection reuse, timeout and redirect defaults,
+and cleanup on close or reset are preserved. Explicit `Cookie` headers supplied through
+`static_headers` or a `header_provider` are still supported.
+
+Applications relying on cookie-based authentication, sessions, or load-balancer affinity
+must supply an `httpx.AsyncClient` through `http_client=`. Supplied clients retain their
+configuration and response-cookie handling and must be closed by the caller. Scope
+cookie-bearing clients and MCP tool sessions to one authenticated principal; do not
+share a user-specific cookie jar across users. Rejecting cookies does not partition
+MCP protocol sessions or other server-side state between principals.
+
 ## 5. Multi-Agent Orchestration
 
 Coordinate multiple agents to collaborate on complex tasks using orchestration patterns:

{agent_framework_core-1.18.0 → agent_framework_core-1.20.0}/README.md

@@ -168,6 +168,21 @@ asyncio.run(main())
 
 You can explore additional agent samples [here](https://github.com/microsoft/agent-framework/tree/main/python/samples/02-agents).
 
+### MCP HTTP client ownership and cookies
+
+**Breaking change:** HTTP clients created by `MCPStreamableHTTPTool` no longer persist
+response cookies. This applies with or without a `header_provider`, including repeated
+calls within the same MCP session. Connection reuse, timeout and redirect defaults,
+and cleanup on close or reset are preserved. Explicit `Cookie` headers supplied through
+`static_headers` or a `header_provider` are still supported.
+
+Applications relying on cookie-based authentication, sessions, or load-balancer affinity
+must supply an `httpx.AsyncClient` through `http_client=`. Supplied clients retain their
+configuration and response-cookie handling and must be closed by the caller. Scope
+cookie-bearing clients and MCP tool sessions to one authenticated principal; do not
+share a user-specific cookie jar across users. Rejecting cookies does not partition
+MCP protocol sessions or other server-side state between principals.
+
 ## 5. Multi-Agent Orchestration
 
 Coordinate multiple agents to collaborate on complex tasks using orchestration patterns:

{agent_framework_core-1.18.0 → agent_framework_core-1.20.0}/agent_framework/__init__.py

@@ -33,6 +33,7 @@ from ._telemetry import (
 from .exceptions import (
     AgentFrameworkException,
     MiddlewareException,
+    ResponseInvalidatedException,
     UserInputRequiredException,
     WorkflowCheckpointException,
     WorkflowConvergenceException,
@@ -253,6 +254,7 @@ _LAZY_MODULE_EXPORTS: Final[Mapping[str, tuple[str, ...]]] = {
         "ChatOptions",
         "ChatResponse",
         "ChatResponseUpdate",
+        "ComputerSafetyCheck",
         "Content",
         "ContinuationToken",
         "Embedding",
@@ -309,14 +311,24 @@ _LAZY_MODULE_EXPORTS: Final[Mapping[str, tuple[str, ...]]] = {
         "SearchType",
         "SupportsVectorSearch",
         "SupportsVectorUpsert",
+        "VectorCollectionContextProvider",
         "VectorStoreCollectionDefinition",
         "VectorStoreField",
+        "VectorStoreHistoryProvider",
+        "create_delete_tool",
+        "create_get_tool",
+        "create_upsert_tool",
         "create_vector_search_tool",
         "register_vectorstoremodel",
         "vectorstoremodel",
     ),
     "._workflows._agent": ("WorkflowAgent",),
-    "._workflows._agent_executor": ("AgentExecutor", "AgentExecutorRequest", "AgentExecutorResponse"),
+    "._workflows._agent_executor": (
+        "AgentExecutor",
+        "AgentExecutorCheckpointState",
+        "AgentExecutorRequest",
+        "AgentExecutorResponse",
+    ),
     "._workflows._agent_utils": ("resolve_agent_id",),
     "._workflows._checkpoint": (
         "CheckpointID",
@@ -423,6 +435,7 @@ __all__ = [
     "AgentContext",
     "AgentEvalConverter",
     "AgentExecutor",
+    "AgentExecutorCheckpointState",
     "AgentExecutorRequest",
     "AgentExecutorResponse",
     "AgentFileStore",
@@ -464,6 +477,7 @@ __all__ = [
     "ClassSkill",
     "CompactionProvider",
     "CompactionStrategy",
+    "ComputerSafetyCheck",
     "Content",
     "ContextProvider",
     "ContextWindowCompactionStrategy",
@@ -565,6 +579,7 @@ __all__ = [
     "Param",
     "RawAgent",
     "ReleaseCandidateFeature",
+    "ResponseInvalidatedException",
     "ResponseStream",
     "Role",
     "RoleLiteral",
@@ -631,8 +646,10 @@ __all__ = [
     "UsageDetails",
     "UserInputRequiredException",
     "ValidationTypeEnum",
+    "VectorCollectionContextProvider",
     "VectorStoreCollectionDefinition",
     "VectorStoreField",
+    "VectorStoreHistoryProvider",
     "Workflow",
     "WorkflowAgent",
     "WorkflowBuilder",
@@ -665,8 +682,11 @@ __all__ = [
     "create_agent_hooks_middleware_from_emitter",
     "create_always_approve_tool_response",
     "create_always_approve_tool_with_arguments_response",
+    "create_delete_tool",
     "create_edge_runner",
+    "create_get_tool",
     "create_harness_agent",
+    "create_upsert_tool",
     "create_vector_search_tool",
     "detect_media_type_from_base64",
     "enqueue_messages",

{agent_framework_core-1.18.0 → agent_framework_core-1.20.0}/agent_framework/__init__.pyi

@@ -220,6 +220,7 @@ from ._types import (
     ChatOptions,
     ChatResponse,
     ChatResponseUpdate,
+    ComputerSafetyCheck,
     Content,
     ContinuationToken,
     Embedding,
@@ -266,14 +267,24 @@ from ._vectors import (
     SearchType,
     SupportsVectorSearch,
     SupportsVectorUpsert,
+    VectorCollectionContextProvider,
     VectorStoreCollectionDefinition,
     VectorStoreField,
+    VectorStoreHistoryProvider,
+    create_delete_tool,
+    create_get_tool,
+    create_upsert_tool,
     create_vector_search_tool,
     register_vectorstoremodel,
     vectorstoremodel,
 )
 from ._workflows._agent import WorkflowAgent
-from ._workflows._agent_executor import AgentExecutor, AgentExecutorRequest, AgentExecutorResponse
+from ._workflows._agent_executor import (
+    AgentExecutor,
+    AgentExecutorCheckpointState,
+    AgentExecutorRequest,
+    AgentExecutorResponse,
+)
 from ._workflows._agent_utils import resolve_agent_id
 from ._workflows._checkpoint import (
     CheckpointID,
@@ -335,6 +346,7 @@ from ._workflows._workflow_executor import SubWorkflowRequestMessage, SubWorkflo
 from .exceptions import (
     AgentFrameworkException,
     MiddlewareException,
+    ResponseInvalidatedException,
     UserInputRequiredException,
     WorkflowCheckpointException,
     WorkflowConvergenceException,
@@ -378,6 +390,7 @@ __all__ = [
     "AgentContext",
     "AgentEvalConverter",
     "AgentExecutor",
+    "AgentExecutorCheckpointState",
     "AgentExecutorRequest",
     "AgentExecutorResponse",
     "AgentFileStore",
@@ -419,6 +432,7 @@ __all__ = [
     "ClassSkill",
     "CompactionProvider",
     "CompactionStrategy",
+    "ComputerSafetyCheck",
     "Content",
     "ContextProvider",
     "ContextWindowCompactionStrategy",
@@ -520,6 +534,7 @@ __all__ = [
     "Param",
     "RawAgent",
     "ReleaseCandidateFeature",
+    "ResponseInvalidatedException",
     "ResponseStream",
     "Role",
     "RoleLiteral",
@@ -586,8 +601,10 @@ __all__ = [
     "UsageDetails",
     "UserInputRequiredException",
     "ValidationTypeEnum",
+    "VectorCollectionContextProvider",
     "VectorStoreCollectionDefinition",
     "VectorStoreField",
+    "VectorStoreHistoryProvider",
     "Workflow",
     "WorkflowAgent",
     "WorkflowBuilder",
@@ -620,8 +637,11 @@ __all__ = [
     "create_agent_hooks_middleware_from_emitter",
     "create_always_approve_tool_response",
     "create_always_approve_tool_with_arguments_response",
+    "create_delete_tool",
     "create_edge_runner",
+    "create_get_tool",
     "create_harness_agent",
+    "create_upsert_tool",
     "create_vector_search_tool",
     "detect_media_type_from_base64",
     "enqueue_messages",

{agent_framework_core-1.18.0 → agent_framework_core-1.20.0}/agent_framework/_agent_hooks.py

@@ -69,9 +69,8 @@ Enforcement semantics (``mode="enforce"``):
   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
-internally, middleware stream hooks are applied to the buffered content, the
+Streaming is supported **fail-closed by buffering**: 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``
@@ -103,7 +102,7 @@ import contextlib
 import json
 import logging
 import uuid
-from collections.abc import Awaitable, Callable, Mapping, Sequence
+from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence
 from contextvars import ContextVar
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, NoReturn, cast
@@ -1231,7 +1230,10 @@ class _AgentHooksAgentMiddleware(_AgentHooksMiddlewareBase, AgentMiddleware):
                     f"got {type(inner).__name__}."
                 )
             context.result = self._gated_agent_stream(
-                state, cast("ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", inner), gate_handle
+                context,
+                state,
+                cast("ResponseStream[AgentResponseUpdate, AgentResponse[Any]]", inner),
+                gate_handle,
             )
             # The gated stream now owns the shutdown emission.
             shutdown_reason = None
@@ -1255,6 +1257,7 @@ class _AgentHooksAgentMiddleware(_AgentHooksMiddlewareBase, AgentMiddleware):
 
     def _gated_agent_stream(
         self,
+        context: AgentContext,
         state: _RunState,
         inner: ResponseStream[AgentResponseUpdate, AgentResponse[Any]],
         gate_handle: _RunPersistenceGate,
@@ -1262,18 +1265,19 @@ class _AgentHooksAgentMiddleware(_AgentHooksMiddlewareBase, AgentMiddleware):
         """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
+        middleware stream hooks are applied by the buffered enforcement helper before the verdict, the
         ``output`` verdict is applied to the finalized response, and deferred
-        persistence is released only after the verdict permits. The combinator owns the
+        persistence is released only after the verdict permits. The helper owns the
         no-divergence rule: a transformed output (or any applied stream hooks)
         re-derives the released updates from the verdicted response.
         """
 
-        async def _consume() -> tuple[Sequence[AgentResponseUpdate], AgentResponse[Any]]:
+        @contextlib.asynccontextmanager
+        async def _consumption_scope() -> AsyncGenerator[None]:
             run_token = _RUN_STATE.set(state)
             try:
                 with gate_handle:
-                    final = await inner.get_final_response()
+                    yield
             except asyncio.CancelledError:
                 await self._emit_shutdown(state, "cancelled")
                 raise
@@ -1290,40 +1294,32 @@ class _AgentHooksAgentMiddleware(_AgentHooksMiddlewareBase, AgentMiddleware):
                 raise
             finally:
                 _RUN_STATE.reset(run_token)
-            return list(inner.updates), final
 
-        async def _gate(
-            _updates: list[AgentResponseUpdate], final: AgentResponse[Any]
-        ) -> tuple[AgentResponse[Any], bool]:
-            from agent_hooks import InterceptionBlocked
+        async def _transform(final: AgentResponse[Any]) -> AgentResponse[Any] | None:
+            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)
+            return final if transformed else None
 
-            try:
-                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)
-                await gate_handle.flush()
-                await self._emit_shutdown(state, "completed")
-                return final, transformed
-            except InterceptionBlocked:
-                gate_handle.drop()
-                await self._emit_shutdown(state, "error")
-                raise
-            except asyncio.CancelledError:
-                await self._emit_shutdown(state, "cancelled")
-                raise
-            except BaseException:
-                await self._emit_shutdown(state, "error")
-                raise
+        async def _release() -> None:
+            await gate_handle.flush()
+            await self._emit_shutdown(state, "completed")
 
-        return cast(
-            "ResponseStream[AgentResponseUpdate, AgentResponse[Any]]",
-            cast(Any, ResponseStream).buffered_and_gated(
-                consume=_consume, gate=_gate, rederive=_agent_updates_from_response
-            ),
-        )
+        async def _release_error(exc: BaseException) -> None:
+            gate_handle.drop()
+            await self._emit_shutdown(state, "cancelled" if isinstance(exc, asyncio.CancelledError) else "error")
+
+        context._stream_terminal_result_transforms.append(_transform)  # pyright: ignore[reportPrivateUsage]
+        context._stream_terminal_result_to_updates = _agent_updates_from_response  # pyright: ignore[reportPrivateUsage]
+        context._stream_terminal_result_is_authoritative = True  # pyright: ignore[reportPrivateUsage]
+        context._stream_release_hooks.append(_release)  # pyright: ignore[reportPrivateUsage]
+        context._stream_release_error_hooks.append(_release_error)  # pyright: ignore[reportPrivateUsage]
+        context.stream_buffer_updates = True
+        context.stream_consumption_context_manager_factories.append(_consumption_scope)
+        return inner
 
 
 class _AgentHooksChatMiddleware(_AgentHooksMiddlewareBase, ChatMiddleware):
@@ -1367,7 +1363,11 @@ class _AgentHooksChatMiddleware(_AgentHooksMiddlewareBase, ChatMiddleware):
         result = context.result
         if isinstance(result, ResponseStream):
             context.result = self._gated_chat_stream(
-                state, model_id, cast("ResponseStream[ChatResponseUpdate, ChatResponse[Any]]", result), gate
+                context,
+                state,
+                model_id,
+                cast("ResponseStream[ChatResponseUpdate, ChatResponse[Any]]", result),
+                gate,
             )
         elif isinstance(result, ChatResponse):
             try:
@@ -1406,6 +1406,7 @@ class _AgentHooksChatMiddleware(_AgentHooksMiddlewareBase, ChatMiddleware):
 
     def _gated_chat_stream(
         self,
+        context: ChatContext,
         state: _RunState,
         model_id: str,
         inner: ResponseStream[ChatResponseUpdate, ChatResponse[Any]],
@@ -1417,39 +1418,38 @@ class _AgentHooksChatMiddleware(_AgentHooksMiddlewareBase, ChatMiddleware):
         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 is released only after the verdict
-        permits. The combinator owns the no-divergence rule: a transformed response
+        permits. The buffered enforcement helper 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]]:
+        @contextlib.asynccontextmanager
+        async def _consumption_scope() -> AsyncGenerator[None]:
             with gate_handle:
-                response = await inner.get_final_response()
-            return list(inner.updates), response
-
-        async def _gate(_updates: list[ChatResponseUpdate], final: ChatResponse[Any]) -> tuple[ChatResponse[Any], bool]:
-            from agent_hooks import InterceptionBlocked
+                yield
 
+        async def _transform(final: ChatResponse[Any]) -> ChatResponse[Any] | None:
             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."
                 )
-            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
+            changed = await self._emit_post_model_call(state, model_id, final)
+            return final if changed else None
+
+        async def _release() -> None:
             await gate_handle.flush()
-            return final, changed
 
-        return cast(
-            "ResponseStream[ChatResponseUpdate, ChatResponse[Any]]",
-            cast(Any, ResponseStream).buffered_and_gated(
-                consume=_consume, gate=_gate, rederive=_chat_updates_from_response
-            ),
-        )
+        def _release_error(_: BaseException) -> None:
+            gate_handle.drop()
+
+        context._stream_terminal_result_transforms.append(_transform)  # pyright: ignore[reportPrivateUsage]
+        context._stream_terminal_result_to_updates = _chat_updates_from_response  # pyright: ignore[reportPrivateUsage]
+        context._stream_terminal_result_is_authoritative = True  # pyright: ignore[reportPrivateUsage]
+        context._stream_release_hooks.append(_release)  # pyright: ignore[reportPrivateUsage]
+        context._stream_release_error_hooks.append(_release_error)  # pyright: ignore[reportPrivateUsage]
+        context.stream_buffer_updates = True
+        context.stream_consumption_context_manager_factories.append(_consumption_scope)
+        return inner
 
 
 class _AgentHooksFunctionMiddleware(_AgentHooksMiddlewareBase, FunctionMiddleware):