Python: fix(core): return MiddlewareTermination(result=...) from agent and chat middleware

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:

python/packages/core/tests/core/test_middleware.py

@@ -32,6 +32,8 @@
 )
 from agent_framework._tools import FunctionTool, ToolTypes
 
+from .utils import termination_result
+
 
 class TestAgentContext:
     """Test cases for AgentContext."""
@@ -1929,3 +1931,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

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 .utils 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

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):