Python: fix(core): warn when a streamed run registers result gates without buffering

python/packages/core/agent_framework/_middleware.py

@@ -1348,6 +1348,26 @@ def _register_middleware_with_wrapper(
             )
 
 
+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 check reads the final stream state, so gates and
+    buffering a middleware configured directly on the stream are reflected too.
+    """
+    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 that releases updates unbuffered; "
+        "updates reach the consumer before the gates run, so a gate that raises cannot hold "
+        "the answer back. Enable buffering (stream_buffer_updates=True or buffer_updates()) "
+        "to make the gates blocking."
+    )
+
+
 class AgentMiddlewarePipeline(BaseMiddlewarePipeline):
     """Executes agent middleware in a chain.
 
@@ -1458,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
 
 
@@ -1668,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
 
 

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

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
@@ -26,6 +27,7 @@
     MiddlewareFailure,
     MiddlewareTermination,
     MiddlewareType,
+    ResponseStream,
     SupportsChatGetResponse,
     agent_middleware,
     chat_middleware,
@@ -385,6 +387,115 @@ 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_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:
+        """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."""

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,69 @@ 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_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}