From 1ed4c83de3a9cff2f6f5ba4a597c9931f4b3eb97 Mon Sep 17 00:00:00 2001 From: Somay-kousis Date: Sat, 17 Oct 2026 17:52:32 +0530 Subject: [PATCH 1/2] Python: fix(core): return MiddlewareTermination(result=...) from agent and chat middleware Fixes #9471. --- .../core/agent_framework/_middleware.py | 15 +++- .../core/tests/core/test_middleware.py | 83 +++++++++++++++++++ .../tests/core/test_middleware_with_agent.py | 32 +++++++ 3 files changed, 127 insertions(+), 3 deletions(-) diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 0372967ccf1..45050dfb1fa 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -2,7 +2,6 @@ from __future__ import annotations -import contextlib import inspect import logging import sys @@ -1452,8 +1451,13 @@ async def current_handler() -> None: return current_handler first_handler = create_next_handler(0) - with contextlib.suppress(MiddlewareTermination): + try: await first_handler() + except MiddlewareTermination as exc: + if context.result is None: + context.result = cast( + "AgentResponse | ResponseStream[AgentResponseUpdate, AgentResponse] | None", exc.result + ) if context.result and isinstance(context.result, ResponseStream): if context.stream_buffer_updates: @@ -1663,8 +1667,13 @@ async def current_handler() -> None: return current_handler first_handler = create_next_handler(0) - with contextlib.suppress(MiddlewareTermination): + try: await first_handler() + except MiddlewareTermination as exc: + if context.result is None: + context.result = cast( + "ChatResponse | ResponseStream[ChatResponseUpdate, ChatResponse] | None", exc.result + ) if context.result and isinstance(context.result, ResponseStream): if context.stream_buffer_updates: diff --git a/python/packages/core/tests/core/test_middleware.py b/python/packages/core/tests/core/test_middleware.py index ab473306b5f..ea271a97a97 100644 --- a/python/packages/core/tests/core/test_middleware.py +++ b/python/packages/core/tests/core/test_middleware.py @@ -32,6 +32,27 @@ ) from agent_framework._tools import FunctionTool, ToolTypes +_TERMINATION_TEXT = "Sorry, blocked." + + +def _termination_result(kind: str, stream: bool) -> Any: + """The response a middleware hands to ``MiddlewareTermination(result=...)``.""" + if kind == "agent": + if not stream: + return AgentResponse(messages=[Message(role="assistant", contents=[_TERMINATION_TEXT])]) + + async def _agent_updates() -> AsyncIterable[AgentResponseUpdate]: + yield AgentResponseUpdate(contents=[Content.from_text(_TERMINATION_TEXT)], role="assistant") + + return ResponseStream(_agent_updates(), finalizer=AgentResponse.from_updates) + if not stream: + return ChatResponse(messages=[Message(role="assistant", contents=[_TERMINATION_TEXT])]) + + async def _chat_updates() -> AsyncIterable[ChatResponseUpdate]: + yield ChatResponseUpdate(contents=[Content.from_text(_TERMINATION_TEXT)], role="assistant") + + return ResponseStream(_chat_updates(), finalizer=ChatResponse.from_updates) + class TestAgentContext: """Test cases for AgentContext.""" @@ -1929,3 +1950,65 @@ def test_as_middleware_list_owns_the_bare_source_rule(self) -> None: 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] + + +class TestMiddlewareTerminationResult: + """``MiddlewareTermination(result=...)`` is the run's result in the agent and chat pipelines too.""" + + @pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) + @pytest.mark.parametrize("kind", ["agent", "chat"]) + async def test_termination_result_is_returned( + self, kind: str, stream: bool, mock_agent: SupportsAgentRun, mock_chat_client: Any + ) -> None: + expected = _termination_result(kind, stream) + handler_calls: list[str] = [] + + async def terminate(context: Any, call_next: Callable[[], Awaitable[None]]) -> None: + raise MiddlewareTermination("blocked", result=expected) + + async def final_handler(context: Any) -> Any: + handler_calls.append("final") + return None + + messages = [Message(role="user", contents=["test"])] + result: object + if kind == "agent": + result = await AgentMiddlewarePipeline(terminate).execute( + AgentContext(agent=mock_agent, messages=messages, stream=stream), final_handler + ) + else: + result = await ChatMiddlewarePipeline(terminate).execute( + ChatContext(client=mock_chat_client, messages=messages, options={}, stream=stream), final_handler + ) + + assert result is expected + assert handler_calls == [] + + @pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) + @pytest.mark.parametrize("kind", ["agent", "chat"]) + async def test_context_result_wins_over_termination_result( + self, kind: str, stream: bool, mock_agent: SupportsAgentRun, mock_chat_client: Any + ) -> None: + """A result the middleware already set on the context wins, as in the function pipeline.""" + from_context = _termination_result(kind, stream) + from_exception = _termination_result(kind, stream) + + async def set_then_terminate(context: Any, call_next: Callable[[], Awaitable[None]]) -> None: + context.result = from_context + raise MiddlewareTermination("blocked", result=from_exception) + + async def final_handler(context: Any) -> Any: + return None + + messages = [Message(role="user", contents=["test"])] + result: object + if kind == "agent": + result = await AgentMiddlewarePipeline(set_then_terminate).execute( + AgentContext(agent=mock_agent, messages=messages, stream=stream), final_handler + ) + else: + result = await ChatMiddlewarePipeline(set_then_terminate).execute( + ChatContext(client=mock_chat_client, messages=messages, options={}, stream=stream), final_handler + ) + + assert result is from_context 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 d04767c904f..9ad4503c1b4 100644 --- a/python/packages/core/tests/core/test_middleware_with_agent.py +++ b/python/packages/core/tests/core/test_middleware_with_agent.py @@ -37,6 +37,7 @@ from agent_framework._sessions import InMemoryHistoryProvider from .conftest import MockBaseChatClient, MockChatClient +from .test_middleware import _TERMINATION_TEXT, _termination_result # region Agent Tests @@ -3275,3 +3276,34 @@ async def process(self, context: AgentContext, call_next: Callable[[], Awaitable # endregion + + +class TestMiddlewareTerminationResultWithAgent: + """A middleware that terminates with ``result=`` answers ``Agent.run`` without calling the model.""" + + @pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) + @pytest.mark.parametrize("kind", ["agent", "chat"]) + async def test_agent_run_returns_termination_result( + self, kind: str, stream: bool, chat_client_base: MockBaseChatClient + ) -> None: + @agent_middleware + async def block_agent(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + raise MiddlewareTermination("blocked", result=_termination_result("agent", context.stream)) + + @chat_middleware + async def block_chat(context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None: + raise MiddlewareTermination("blocked", result=_termination_result("chat", context.stream)) + + agent = Agent(client=chat_client_base, middleware=[block_agent if kind == "agent" else block_chat]) + + if stream: + response_stream = agent.run("hello", stream=True) + updates = [update async for update in response_stream] + response = await response_stream.get_final_response() + assert "".join(update.text for update in updates) == _TERMINATION_TEXT + else: + response = await agent.run("hello") + + assert response is not None + assert response.text == _TERMINATION_TEXT + assert chat_client_base.call_count == 0 From ca095655b7be627e3be3ece99164a7e4d1d6d0fe Mon Sep 17 00:00:00 2001 From: Somay-kousis Date: Tue, 27 Oct 2026 15:26:51 +0530 Subject: [PATCH 2/2] Python: test(core): move the termination test helper into tests/core/utils.py CODING_STANDARD.md asks for non-fixture helpers in a helper module, not imported from another test module. --- .../core/tests/core/test_middleware.py | 27 +++------------ .../tests/core/test_middleware_with_agent.py | 10 +++--- python/packages/core/tests/core/utils.py | 33 +++++++++++++++++++ 3 files changed, 42 insertions(+), 28 deletions(-) diff --git a/python/packages/core/tests/core/test_middleware.py b/python/packages/core/tests/core/test_middleware.py index ea271a97a97..b74012c7f7e 100644 --- a/python/packages/core/tests/core/test_middleware.py +++ b/python/packages/core/tests/core/test_middleware.py @@ -32,26 +32,7 @@ ) from agent_framework._tools import FunctionTool, ToolTypes -_TERMINATION_TEXT = "Sorry, blocked." - - -def _termination_result(kind: str, stream: bool) -> Any: - """The response a middleware hands to ``MiddlewareTermination(result=...)``.""" - if kind == "agent": - if not stream: - return AgentResponse(messages=[Message(role="assistant", contents=[_TERMINATION_TEXT])]) - - async def _agent_updates() -> AsyncIterable[AgentResponseUpdate]: - yield AgentResponseUpdate(contents=[Content.from_text(_TERMINATION_TEXT)], role="assistant") - - return ResponseStream(_agent_updates(), finalizer=AgentResponse.from_updates) - if not stream: - return ChatResponse(messages=[Message(role="assistant", contents=[_TERMINATION_TEXT])]) - - async def _chat_updates() -> AsyncIterable[ChatResponseUpdate]: - yield ChatResponseUpdate(contents=[Content.from_text(_TERMINATION_TEXT)], role="assistant") - - return ResponseStream(_chat_updates(), finalizer=ChatResponse.from_updates) +from .utils import termination_result class TestAgentContext: @@ -1960,7 +1941,7 @@ class TestMiddlewareTerminationResult: async def test_termination_result_is_returned( self, kind: str, stream: bool, mock_agent: SupportsAgentRun, mock_chat_client: Any ) -> None: - expected = _termination_result(kind, stream) + expected = termination_result(kind, stream) handler_calls: list[str] = [] async def terminate(context: Any, call_next: Callable[[], Awaitable[None]]) -> None: @@ -1990,8 +1971,8 @@ async def test_context_result_wins_over_termination_result( self, kind: str, stream: bool, mock_agent: SupportsAgentRun, mock_chat_client: Any ) -> None: """A result the middleware already set on the context wins, as in the function pipeline.""" - from_context = _termination_result(kind, stream) - from_exception = _termination_result(kind, stream) + from_context = termination_result(kind, stream) + from_exception = termination_result(kind, stream) async def set_then_terminate(context: Any, call_next: Callable[[], Awaitable[None]]) -> None: context.result = from_context 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 9ad4503c1b4..9b06e69833a 100644 --- a/python/packages/core/tests/core/test_middleware_with_agent.py +++ b/python/packages/core/tests/core/test_middleware_with_agent.py @@ -37,7 +37,7 @@ from agent_framework._sessions import InMemoryHistoryProvider from .conftest import MockBaseChatClient, MockChatClient -from .test_middleware import _TERMINATION_TEXT, _termination_result +from .utils import TERMINATION_TEXT, termination_result # region Agent Tests @@ -3288,11 +3288,11 @@ async def test_agent_run_returns_termination_result( ) -> None: @agent_middleware async def block_agent(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: - raise MiddlewareTermination("blocked", result=_termination_result("agent", context.stream)) + raise MiddlewareTermination("blocked", result=termination_result("agent", context.stream)) @chat_middleware async def block_chat(context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None: - raise MiddlewareTermination("blocked", result=_termination_result("chat", context.stream)) + raise MiddlewareTermination("blocked", result=termination_result("chat", context.stream)) agent = Agent(client=chat_client_base, middleware=[block_agent if kind == "agent" else block_chat]) @@ -3300,10 +3300,10 @@ async def block_chat(context: ChatContext, call_next: Callable[[], Awaitable[Non response_stream = agent.run("hello", stream=True) updates = [update async for update in response_stream] response = await response_stream.get_final_response() - assert "".join(update.text for update in updates) == _TERMINATION_TEXT + assert "".join(update.text for update in updates) == TERMINATION_TEXT else: response = await agent.run("hello") assert response is not None - assert response.text == _TERMINATION_TEXT + assert response.text == TERMINATION_TEXT assert chat_client_base.call_count == 0 diff --git a/python/packages/core/tests/core/utils.py b/python/packages/core/tests/core/utils.py index 61bf277d949..ad32fbfeb63 100644 --- a/python/packages/core/tests/core/utils.py +++ b/python/packages/core/tests/core/utils.py @@ -1,9 +1,42 @@ # Copyright (c) Microsoft. All rights reserved. +from collections.abc import AsyncIterable from copy import deepcopy +from typing import Any from unittest.mock import MagicMock +from agent_framework import ( + AgentResponse, + AgentResponseUpdate, + ChatResponse, + ChatResponseUpdate, + Content, + Message, + ResponseStream, +) + +TERMINATION_TEXT = "Sorry, blocked." + + +def termination_result(kind: str, stream: bool) -> Any: + """The response a middleware hands to ``MiddlewareTermination(result=...)``.""" + if kind == "agent": + if not stream: + return AgentResponse(messages=[Message(role="assistant", contents=[TERMINATION_TEXT])]) + + async def _agent_updates() -> AsyncIterable[AgentResponseUpdate]: + yield AgentResponseUpdate(contents=[Content.from_text(TERMINATION_TEXT)], role="assistant") + + return ResponseStream(_agent_updates(), finalizer=AgentResponse.from_updates) + if not stream: + return ChatResponse(messages=[Message(role="assistant", contents=[TERMINATION_TEXT])]) + + async def _chat_updates() -> AsyncIterable[ChatResponseUpdate]: + yield ChatResponseUpdate(contents=[Content.from_text(TERMINATION_TEXT)], role="assistant") + + return ResponseStream(_chat_updates(), finalizer=ChatResponse.from_updates) + class CopyingMock(MagicMock): def __call__(self, *args, **kwargs):