From c4bc2e0fe4458ce1595c171783c4bedbeb2df34a Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Sun, 4 Oct 2026 01:57:00 +0800 Subject: [PATCH 1/2] Python: fix(core): warn when a streamed run registers result gates without buffering A result gate runs at stream finalization, so with stream_buffer_updates left at its default the consumer has already received every update by the time the gate raises: the guard fails open, silently, and the exception afterwards documents a block that never happened. #8954 has the full analysis and an offline repro. Log a warning from both middleware pipelines when result gates are registered on a streamed run with buffering off, and note on with_result_gate() that buffer_updates() is what makes gates block delivery. Strictly additive: delivery semantics are unchanged. Flipping buffering on when a gate is registered, or adding a per-gate opt-in, changes latency semantics and is a vNext call for maintainers. --- .../core/agent_framework/_middleware.py | 31 +++++++++ .../packages/core/agent_framework/_types.py | 7 ++- .../tests/core/test_middleware_with_agent.py | 63 +++++++++++++++++++ .../tests/core/test_middleware_with_chat.py | 44 +++++++++++++ 4 files changed, 144 insertions(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 8af5c885c4c..3c5bf0df108 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -1348,6 +1348,27 @@ def _register_middleware_with_wrapper( ) +def _warn_unbuffered_result_gates( + gates_before: list[Callable[[Any], object]], + gates_after: list[Callable[[Any], object]], + *, + buffer_updates: bool, +) -> None: + """Warn when result gates are registered on a stream that releases updates unbuffered. + + A result gate runs at finalization, so on an unbuffered stream the consumer has already + received every update by the time the gate raises: the gate fails open. Buffering holds + the updates until the gates pass. + """ + if buffer_updates or not (gates_before or gates_after): + return + logger.warning( + "Result gates are registered on a streamed run with stream_buffer_updates=False; " + "updates reach the consumer before the gates run, so a gate that raises cannot hold " + "the answer back. Set stream_buffer_updates=True to make the gates blocking." + ) + + class AgentMiddlewarePipeline(BaseMiddlewarePipeline): """Executes agent middleware in a chain. @@ -1426,6 +1447,11 @@ async def current_handler() -> None: await first_handler() if context.result and isinstance(context.result, ResponseStream): + _warn_unbuffered_result_gates( + context.stream_result_gates_before, + context.stream_result_gates_after, + buffer_updates=context.stream_buffer_updates, + ) if context.stream_buffer_updates: context.result.buffer_updates(result_to_updates=context.stream_result_to_updates) for factory in context.stream_consumption_context_manager_factories: @@ -1636,6 +1662,11 @@ async def current_handler() -> None: await first_handler() if context.result and isinstance(context.result, ResponseStream): + _warn_unbuffered_result_gates( + context.stream_result_gates_before, + context.stream_result_gates_after, + buffer_updates=context.stream_buffer_updates, + ) if context.stream_buffer_updates: context.result.buffer_updates(result_to_updates=context.stream_result_to_updates) for factory in context.stream_consumption_context_manager_factories: diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index f98743794a4..10f1317f9a1 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -4279,7 +4279,12 @@ def with_result_gate( *, phase: Literal["before_transform", "after_transform"] = "after_transform", ) -> ResponseStream[UpdateT, FinalT]: - """Register a blocking gate for the finalized result.""" + """Register a blocking gate for the finalized result. + + The gate runs at finalization. On a stream that is not buffered the consumer has + already received the updates by then, so a gate that raises cannot hold them back; + call ``buffer_updates()`` to hold updates until the gates pass. + """ self._ensure_content_configuration_mutable() self._validate_gate_phase(phase) gates = self._result_gates_before if phase == "before_transform" else self._result_gates_after 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 1bc77997417..0ed0fc3e09c 100644 --- a/python/packages/core/tests/core/test_middleware_with_agent.py +++ b/python/packages/core/tests/core/test_middleware_with_agent.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. import asyncio +import logging import threading from collections.abc import Awaitable, Callable, Iterator from typing import Any, cast @@ -385,6 +386,68 @@ async def process(self, context: AgentContext, call_next: Callable[[], Awaitable # Verify flags: [non-streaming, streaming] assert streaming_flags == [False, True] + async def test_result_gate_without_buffer_logs_warning( + self, client: "MockChatClient", caplog: pytest.LogCaptureFixture + ) -> None: + """A result gate on an unbuffered streamed run logs a fail-open warning.""" + + class ResultGateMiddleware(AgentMiddleware): + async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + if context.stream: + context.stream_result_gates_after.append(lambda result: None) + await call_next() + + agent = Agent(client=client, middleware=[ResultGateMiddleware()]) # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + client.streaming_responses = [ + [ChatResponseUpdate(contents=[Content.from_text(text="chunk")], role="assistant")] + ] + + with caplog.at_level(logging.WARNING, logger="agent_framework._middleware"): + async for _ in agent.run([Message(role="user", contents=["test message"])], stream=True): + pass + + assert any("stream_buffer_updates" in record.message for record in caplog.records) + + async def test_result_gate_with_buffer_logs_no_warning( + self, client: "MockChatClient", caplog: pytest.LogCaptureFixture + ) -> None: + """A result gate on a buffered streamed run blocks as documented, so no warning.""" + + class BufferedResultGateMiddleware(AgentMiddleware): + async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + if context.stream: + context.stream_buffer_updates = True + context.stream_result_gates_after.append(lambda result: None) + await call_next() + + agent = Agent(client=client, middleware=[BufferedResultGateMiddleware()]) # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + client.streaming_responses = [ + [ChatResponseUpdate(contents=[Content.from_text(text="chunk")], role="assistant")] + ] + + with caplog.at_level(logging.WARNING, logger="agent_framework._middleware"): + updates = [ + update async for update in agent.run([Message(role="user", contents=["test message"])], stream=True) + ] + + assert not any("stream_buffer_updates" in record.message for record in caplog.records) + assert updates + + async def test_no_result_gate_logs_no_warning( + self, client: "MockChatClient", caplog: pytest.LogCaptureFixture + ) -> None: + """A plain streamed run without gates stays quiet.""" + agent = Agent(client=client) # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + client.streaming_responses = [ + [ChatResponseUpdate(contents=[Content.from_text(text="chunk")], role="assistant")] + ] + + with caplog.at_level(logging.WARNING, logger="agent_framework._middleware"): + async for _ in agent.run([Message(role="user", contents=["test message"])], stream=True): + pass + + assert not any("stream_buffer_updates" in record.message for record in caplog.records) + class TestChatAgentMultipleMiddlewareOrdering: """Test cases for multiple middleware execution order with Agent.""" diff --git a/python/packages/core/tests/core/test_middleware_with_chat.py b/python/packages/core/tests/core/test_middleware_with_chat.py index bb5f0c80b35..57df29bae8d 100644 --- a/python/packages/core/tests/core/test_middleware_with_chat.py +++ b/python/packages/core/tests/core/test_middleware_with_chat.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. +import logging from collections.abc import AsyncIterable, Awaitable, Callable, Sequence from typing import Any, cast from unittest.mock import patch @@ -278,6 +279,49 @@ def upper_case_update(update: ChatResponseUpdate) -> ChatResponseUpdate: # Verify middleware executed assert execution_order == ["streaming_before", "streaming_after"] + async def test_chat_result_gate_without_buffer_logs_warning( + self, chat_client_base: "MockBaseChatClient", caplog: pytest.LogCaptureFixture + ) -> None: + """A result gate on an unbuffered streamed run logs a fail-open warning.""" + + @chat_middleware + async def result_gate_middleware(context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None: + if context.stream: + context.stream_result_gates_after.append(lambda result: None) + await call_next() + + chat_client_base.chat_middleware = [cast(ChatMiddlewareTypes, result_gate_middleware)] + messages = [Message(role="user", contents=["test message"])] + + with caplog.at_level(logging.WARNING, logger="agent_framework._middleware"): + async for _ in chat_client_base.get_response(messages, stream=True): + pass + + assert any("stream_buffer_updates" in record.message for record in caplog.records) + + async def test_chat_result_gate_with_buffer_logs_no_warning( + self, chat_client_base: "MockBaseChatClient", caplog: pytest.LogCaptureFixture + ) -> None: + """A result gate on a buffered streamed run blocks as documented, so no warning.""" + + @chat_middleware + async def buffered_result_gate_middleware( + context: ChatContext, call_next: Callable[[], Awaitable[None]] + ) -> None: + if context.stream: + context.stream_buffer_updates = True + context.stream_result_gates_after.append(lambda result: None) + await call_next() + + chat_client_base.chat_middleware = [cast(ChatMiddlewareTypes, buffered_result_gate_middleware)] + messages = [Message(role="user", contents=["test message"])] + + with caplog.at_level(logging.WARNING, logger="agent_framework._middleware"): + updates = [update async for update in chat_client_base.get_response(messages, stream=True)] + + assert not any("stream_buffer_updates" in record.message for record in caplog.records) + assert updates + async def test_run_level_middleware_isolation(self, chat_client_base: "MockBaseChatClient") -> None: """Test that run-level middleware is isolated and doesn't persist across calls.""" execution_count = {"count": 0} From f7b31549a7087716ee98a3266b9733e8f09396c3 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Sun, 4 Oct 2026 12:05:37 +0800 Subject: [PATCH 2/2] Python: fix(core): detect result gates configured directly on the final stream The warning compared the middleware context's gate lists and buffer option before they were mounted, so a middleware that called with_result_gate() or buffer_updates() on the final ResponseStream itself was invisible to it: direct gates went unwarned, direct buffering still warned. Read the mounted stream state instead and run the check after the mount block in both pipelines. --- .../core/agent_framework/_middleware.py | 31 +++++------- .../tests/core/test_middleware_with_agent.py | 48 +++++++++++++++++++ .../tests/core/test_middleware_with_chat.py | 20 ++++++++ 3 files changed, 79 insertions(+), 20 deletions(-) diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 3c5bf0df108..073644cbbcd 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -1348,24 +1348,23 @@ def _register_middleware_with_wrapper( ) -def _warn_unbuffered_result_gates( - gates_before: list[Callable[[Any], object]], - gates_after: list[Callable[[Any], object]], - *, - buffer_updates: bool, -) -> None: +def _warn_unbuffered_result_gates(stream: ResponseStream[Any, Any]) -> None: """Warn when result gates are registered on a stream that releases updates unbuffered. A result gate runs at finalization, so on an unbuffered stream the consumer has already received every update by the time the gate raises: the gate fails open. Buffering holds - the updates until the gates pass. + the updates until the gates pass. The check reads the final stream state, so gates and + buffering a middleware configured directly on the stream are reflected too. """ - if buffer_updates or not (gates_before or gates_after): + if not stream._stream_updates: # pyright: ignore[reportPrivateUsage] + return + if not (stream._result_gates_before or stream._result_gates_after): # pyright: ignore[reportPrivateUsage] return logger.warning( - "Result gates are registered on a streamed run with stream_buffer_updates=False; " + "Result gates are registered on a streamed run that releases updates unbuffered; " "updates reach the consumer before the gates run, so a gate that raises cannot hold " - "the answer back. Set stream_buffer_updates=True to make the gates blocking." + "the answer back. Enable buffering (stream_buffer_updates=True or buffer_updates()) " + "to make the gates blocking." ) @@ -1447,11 +1446,6 @@ async def current_handler() -> None: await first_handler() if context.result and isinstance(context.result, ResponseStream): - _warn_unbuffered_result_gates( - context.stream_result_gates_before, - context.stream_result_gates_after, - buffer_updates=context.stream_buffer_updates, - ) if context.stream_buffer_updates: context.result.buffer_updates(result_to_updates=context.stream_result_to_updates) for factory in context.stream_consumption_context_manager_factories: @@ -1484,6 +1478,7 @@ async def current_handler() -> None: context.result._with_release_error_hook(hook) # pyright: ignore[reportPrivateUsage] for cleanup_hook in context.stream_cleanup_hooks: context.result.with_cleanup_hook(cleanup_hook) + _warn_unbuffered_result_gates(context.result) return context.result @@ -1662,11 +1657,6 @@ async def current_handler() -> None: await first_handler() if context.result and isinstance(context.result, ResponseStream): - _warn_unbuffered_result_gates( - context.stream_result_gates_before, - context.stream_result_gates_after, - buffer_updates=context.stream_buffer_updates, - ) if context.stream_buffer_updates: context.result.buffer_updates(result_to_updates=context.stream_result_to_updates) for factory in context.stream_consumption_context_manager_factories: @@ -1699,6 +1689,7 @@ async def current_handler() -> None: context.result._with_release_error_hook(hook) # pyright: ignore[reportPrivateUsage] for cleanup_hook in context.stream_cleanup_hooks: context.result.with_cleanup_hook(cleanup_hook) + _warn_unbuffered_result_gates(context.result) return context.result 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 0ed0fc3e09c..d04767c904f 100644 --- a/python/packages/core/tests/core/test_middleware_with_agent.py +++ b/python/packages/core/tests/core/test_middleware_with_agent.py @@ -27,6 +27,7 @@ MiddlewareFailure, MiddlewareTermination, MiddlewareType, + ResponseStream, SupportsChatGetResponse, agent_middleware, chat_middleware, @@ -433,6 +434,53 @@ async def process(self, context: AgentContext, call_next: Callable[[], Awaitable assert not any("stream_buffer_updates" in record.message for record in caplog.records) assert updates + async def test_result_gate_attached_directly_to_stream_logs_warning( + self, client: "MockChatClient", caplog: pytest.LogCaptureFixture + ) -> None: + """A gate a middleware attaches straight to the final stream is detected too.""" + + class DirectGateMiddleware(AgentMiddleware): + async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + await call_next() + if isinstance(context.result, ResponseStream): + context.result.with_result_gate(lambda result: None) + + agent = Agent(client=client, middleware=[DirectGateMiddleware()]) # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + client.streaming_responses = [ + [ChatResponseUpdate(contents=[Content.from_text(text="chunk")], role="assistant")] + ] + + with caplog.at_level(logging.WARNING, logger="agent_framework._middleware"): + async for _ in agent.run([Message(role="user", contents=["test message"])], stream=True): + pass + + assert any("stream_buffer_updates" in record.message for record in caplog.records) + + async def test_direct_stream_gate_with_direct_buffer_logs_no_warning( + self, client: "MockChatClient", caplog: pytest.LogCaptureFixture + ) -> None: + """Buffering enabled directly on the stream satisfies the gates, so no warning.""" + + class DirectBufferedGateMiddleware(AgentMiddleware): + async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + await call_next() + if isinstance(context.result, ResponseStream): + context.result.with_result_gate(lambda result: None) + context.result.buffer_updates() + + agent = Agent(client=client, middleware=[DirectBufferedGateMiddleware()]) # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + client.streaming_responses = [ + [ChatResponseUpdate(contents=[Content.from_text(text="chunk")], role="assistant")] + ] + + with caplog.at_level(logging.WARNING, logger="agent_framework._middleware"): + updates = [ + update async for update in agent.run([Message(role="user", contents=["test message"])], stream=True) + ] + + assert not any("stream_buffer_updates" in record.message for record in caplog.records) + assert updates + async def test_no_result_gate_logs_no_warning( self, client: "MockChatClient", caplog: pytest.LogCaptureFixture ) -> None: diff --git a/python/packages/core/tests/core/test_middleware_with_chat.py b/python/packages/core/tests/core/test_middleware_with_chat.py index 57df29bae8d..0271ed7c11d 100644 --- a/python/packages/core/tests/core/test_middleware_with_chat.py +++ b/python/packages/core/tests/core/test_middleware_with_chat.py @@ -322,6 +322,26 @@ async def buffered_result_gate_middleware( assert not any("stream_buffer_updates" in record.message for record in caplog.records) assert updates + async def test_chat_result_gate_attached_directly_to_stream_logs_warning( + self, chat_client_base: "MockBaseChatClient", caplog: pytest.LogCaptureFixture + ) -> None: + """A gate a middleware attaches straight to the final stream is detected too.""" + + @chat_middleware + async def direct_gate_middleware(context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None: + await call_next() + if isinstance(context.result, ResponseStream): + context.result.with_result_gate(lambda result: None) + + chat_client_base.chat_middleware = [cast(ChatMiddlewareTypes, direct_gate_middleware)] + messages = [Message(role="user", contents=["test message"])] + + with caplog.at_level(logging.WARNING, logger="agent_framework._middleware"): + async for _ in chat_client_base.get_response(messages, stream=True): + pass + + assert any("stream_buffer_updates" in record.message for record in caplog.records) + async def test_run_level_middleware_isolation(self, chat_client_base: "MockBaseChatClient") -> None: """Test that run-level middleware is isolated and doesn't persist across calls.""" execution_count = {"count": 0}