diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 8af5c885c4c..073644cbbcd 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/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 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..d04767c904f 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 @@ -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.""" 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..0271ed7c11d 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,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}