Python: Preserve termination results in Agent and Chat middleware

python/packages/core/agent_framework/_middleware.py

@@ -1452,8 +1452,11 @@ 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 termination:
+            if context.result is None:
+                context.result = termination.result
 
         if context.result and isinstance(context.result, ResponseStream):
             if context.stream_buffer_updates:
@@ -1663,8 +1666,11 @@ 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 termination:
+            if context.result is None:
+                context.result = termination.result
 
         if context.result and isinstance(context.result, ResponseStream):
             if context.stream_buffer_updates:

python/packages/core/tests/core/test_middleware_context_result.py

@@ -10,6 +10,8 @@
     Agent,
     AgentResponse,
     AgentResponseUpdate,
+    ChatResponse,
+    ChatResponseUpdate,
     Content,
     Message,
     ResponseStream,
@@ -19,9 +21,13 @@
     AgentContext,
     AgentMiddleware,
     AgentMiddlewarePipeline,
+    ChatContext,
+    ChatMiddleware,
+    ChatMiddlewarePipeline,
     FunctionInvocationContext,
     FunctionMiddleware,
     FunctionMiddlewarePipeline,
+    MiddlewareTermination,
 )
 from agent_framework._tools import FunctionTool
 
@@ -37,6 +43,193 @@ class FunctionTestArgs(BaseModel):
 class TestResultOverrideMiddleware:
     """Test cases for middleware result override functionality."""
 
+    async def test_agent_middleware_termination_returns_exception_result(self, mock_agent: SupportsAgentRun) -> None:
+        expected = AgentResponse(messages=[Message(role="assistant", contents=["terminated response"])])
+
+        class TerminatingMiddleware(AgentMiddleware):
+            async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+                raise MiddlewareTermination(result=expected)
+
+        context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["test"])])
+
+        async def final_handler(_: AgentContext) -> AgentResponse:
+            raise AssertionError("final handler should not run after termination")
+
+        result = await AgentMiddlewarePipeline(TerminatingMiddleware()).execute(context, final_handler)
+
+        assert result is expected
+        assert isinstance(result, AgentResponse)
+        assert result.text == "terminated response"
+
+    async def test_agent_middleware_termination_applies_stream_hooks(self, mock_agent: SupportsAgentRun) -> None:
+        async def updates() -> AsyncIterable[AgentResponseUpdate]:
+            yield AgentResponseUpdate(contents=[Content.from_text("agent update")])
+
+        response_stream = ResponseStream(updates(), finalizer=AgentResponse.from_updates)
+
+        class TerminatingMiddleware(AgentMiddleware):
+            async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+                context.stream_update_transforms.append(
+                    lambda update: AgentResponseUpdate(contents=[Content.from_text(f"{update.text} transformed")])
+                )
+                context.stream_result_transforms.append(
+                    lambda response: AgentResponse(
+                        messages=[Message(role="assistant", contents=[f"{response.text} finalized"])]
+                    )
+                )
+                raise MiddlewareTermination(result=response_stream)
+
+        context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["test"])], stream=True)
+
+        async def final_handler(_: AgentContext) -> ResponseStream[AgentResponseUpdate, AgentResponse]:
+            raise AssertionError("final handler should not run after termination")
+
+        result = await AgentMiddlewarePipeline(TerminatingMiddleware()).execute(context, final_handler)
+
+        assert isinstance(result, ResponseStream)
+        streamed_updates = [update async for update in result]
+        final_response = await result.get_final_response()
+
+        assert [update.text for update in streamed_updates] == ["agent update transformed"]
+        assert all(isinstance(update, AgentResponseUpdate) for update in streamed_updates)
+        assert isinstance(final_response, AgentResponse)
+        # Finalizers aggregate the original updates; result hooks transform that response separately.
+        assert final_response.text == "agent update finalized"
+
+    async def test_chat_middleware_termination_returns_exception_result(self) -> None:
+        expected = ChatResponse(messages=[Message(role="assistant", contents=["terminated response"])])
+
+        class TerminatingMiddleware(ChatMiddleware):
+            async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
+                raise MiddlewareTermination(result=expected)
+
+        context = ChatContext(client=MagicMock(), messages=[Message(role="user", contents=["test"])], options={})
+
+        async def final_handler(_: ChatContext) -> ChatResponse:
+            raise AssertionError("final handler should not run after termination")
+
+        result = await ChatMiddlewarePipeline(TerminatingMiddleware()).execute(context, final_handler)
+
+        assert result is expected
+        assert isinstance(result, ChatResponse)
+        assert result.text == "terminated response"
+
+    async def test_chat_middleware_termination_applies_stream_hooks(self) -> None:
+        async def updates() -> AsyncIterable[ChatResponseUpdate]:
+            yield ChatResponseUpdate(contents=[Content.from_text("chat update")])
+
+        response_stream = ResponseStream(updates(), finalizer=ChatResponse.from_updates)
+
+        class TerminatingMiddleware(ChatMiddleware):
+            async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
+                context.stream_update_transforms.append(
+                    lambda update: ChatResponseUpdate(contents=[Content.from_text(f"{update.text} transformed")])
+                )
+                context.stream_result_transforms.append(
+                    lambda response: ChatResponse(
+                        messages=[Message(role="assistant", contents=[f"{response.text} finalized"])]
+                    )
+                )
+                raise MiddlewareTermination(result=response_stream)
+
+        context = ChatContext(
+            client=MagicMock(),
+            messages=[Message(role="user", contents=["test"])],
+            options={},
+            stream=True,
+        )
+
+        async def final_handler(_: ChatContext) -> ResponseStream[ChatResponseUpdate, ChatResponse]:
+            raise AssertionError("final handler should not run after termination")
+
+        result = await ChatMiddlewarePipeline(TerminatingMiddleware()).execute(context, final_handler)
+
+        assert isinstance(result, ResponseStream)
+        streamed_updates = [update async for update in result]
+        final_response = await result.get_final_response()
+
+        assert [update.text for update in streamed_updates] == ["chat update transformed"]
+        assert all(isinstance(update, ChatResponseUpdate) for update in streamed_updates)
+        assert isinstance(final_response, ChatResponse)
+        # Finalizers aggregate the original updates; result hooks transform that response separately.
+        assert final_response.text == "chat update finalized"
+
+    async def test_agent_context_result_precedes_termination_result(self, mock_agent: SupportsAgentRun) -> None:
+        context_response = AgentResponse(messages=[Message(role="assistant", contents=["context response"])])
+        termination_response = AgentResponse(messages=[Message(role="assistant", contents=["exception response"])])
+
+        class TerminatingMiddleware(AgentMiddleware):
+            async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+                await call_next()
+                raise MiddlewareTermination(result=termination_response)
+
+        context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["test"])])
+
+        async def final_handler(_: AgentContext) -> AgentResponse:
+            return context_response
+
+        result = await AgentMiddlewarePipeline(TerminatingMiddleware()).execute(context, final_handler)
+
+        assert result is context_response
+
+    async def test_chat_context_result_precedes_termination_result(self) -> None:
+        context_response = ChatResponse(messages=[Message(role="assistant", contents=["context response"])])
+        termination_response = ChatResponse(messages=[Message(role="assistant", contents=["exception response"])])
+
+        class TerminatingMiddleware(ChatMiddleware):
+            async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
+                await call_next()
+                raise MiddlewareTermination(result=termination_response)
+
+        context = ChatContext(client=MagicMock(), messages=[Message(role="user", contents=["test"])], options={})
+
+        async def final_handler(_: ChatContext) -> ChatResponse:
+            return context_response
+
+        result = await ChatMiddlewarePipeline(TerminatingMiddleware()).execute(context, final_handler)
+
+        assert result is context_response
+
+    @pytest.mark.parametrize("termination_result", [None, "exception response"])
+    async def test_agent_context_result_precedes_optional_termination_result(
+        self, mock_agent: SupportsAgentRun, termination_result: object | None
+    ) -> None:
+        expected = AgentResponse(messages=[Message(role="assistant", contents=["context response"])])
+
+        class TerminatingMiddleware(AgentMiddleware):
+            async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+                context.result = expected
+                raise MiddlewareTermination(result=termination_result)
+
+        context = AgentContext(agent=mock_agent, messages=[Message(role="user", contents=["test"])])
+
+        async def final_handler(_: AgentContext) -> AgentResponse:
+            raise AssertionError("final handler should not run after termination")
+
+        result = await AgentMiddlewarePipeline(TerminatingMiddleware()).execute(context, final_handler)
+
+        assert result is expected
+
+    @pytest.mark.parametrize("termination_result", [None, "exception response"])
+    async def test_chat_context_result_precedes_optional_termination_result(
+        self, termination_result: object | None
+    ) -> None:
+        expected = ChatResponse(messages=[Message(role="assistant", contents=["context response"])])
+
+        class TerminatingMiddleware(ChatMiddleware):
+            async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
+                context.result = expected
+                raise MiddlewareTermination(result=termination_result)
+
+        context = ChatContext(client=MagicMock(), messages=[Message(role="user", contents=["test"])], options={})
+
+        async def final_handler(_: ChatContext) -> ChatResponse:
+            raise AssertionError("final handler should not run after termination")
+
+        result = await ChatMiddlewarePipeline(TerminatingMiddleware()).execute(context, final_handler)
+
+        assert result is expected
+
     async def test_agent_middleware_response_override_non_streaming(self, mock_agent: SupportsAgentRun) -> None:
         """Test that agent middleware can override response for non-streaming execution."""
         override_response = AgentResponse(messages=[Message(role="assistant", contents=["overridden response"])])

python/packages/core/tests/core/test_middleware_with_agent.py

@@ -3,7 +3,7 @@
 import asyncio
 import logging
 import threading
-from collections.abc import Awaitable, Callable, Iterator
+from collections.abc import AsyncIterable, Awaitable, Callable, Iterator
 from typing import Any, cast
 
 import pytest
@@ -12,6 +12,7 @@
     Agent,
     AgentContext,
     AgentMiddleware,
+    AgentResponse,
     AgentResponseUpdate,
     ChatContext,
     ChatMiddleware,
@@ -52,6 +53,43 @@ async def process(self, context: AgentContext, call_next: Callable[[], Awaitable
 class TestChatAgentClassBasedMiddleware:
     """Test cases for class-based middleware integration with Agent."""
 
+    async def test_agent_run_returns_termination_result(self) -> None:
+        client = MockBaseChatClient()
+        expected = AgentResponse(messages=[Message(role="assistant", contents=["terminated response"])])
+
+        class TerminatingMiddleware(AgentMiddleware):
+            async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+                raise MiddlewareTermination(result=expected)
+
+        agent = Agent(client=client, middleware=[TerminatingMiddleware()])
+        response = await agent.run("hi")
+
+        assert isinstance(response, AgentResponse)
+        assert response.text == "terminated response"
+        assert client.call_count == 0
+
+    async def test_agent_run_returns_termination_stream(self) -> None:
+        client = MockBaseChatClient()
+        async def updates() -> AsyncIterable[AgentResponseUpdate]:
+            yield AgentResponseUpdate(contents=[Content.from_text("terminated stream")])
+
+        expected = ResponseStream(updates(), finalizer=AgentResponse.from_updates)
+
+        class TerminatingMiddleware(AgentMiddleware):
+            async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+                raise MiddlewareTermination(result=expected)
+
+        agent = Agent(client=client, middleware=[TerminatingMiddleware()])
+        stream = agent.run("hi", stream=True)
+        streamed_updates = [update async for update in stream]
+        final_response = await stream.get_final_response()
+
+        assert [update.text for update in streamed_updates] == ["terminated stream"]
+        assert all(isinstance(update, AgentResponseUpdate) for update in streamed_updates)
+        assert isinstance(final_response, AgentResponse)
+        assert final_response.text == "terminated stream"
+        assert client.call_count == 0
+
     async def test_class_based_agent_middleware_with_chat_agent(self, client: SupportsChatGetResponse) -> None:
         """Test class-based agent middleware with Agent."""
         execution_order: list[str] = []

python/packages/core/tests/core/test_middleware_with_chat.py

@@ -10,6 +10,8 @@
 from agent_framework import (
     MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY,
     Agent,
+    AgentResponse,
+    AgentResponseUpdate,
     AgentSession,
     ChatContext,
     ChatMiddleware,
@@ -21,6 +23,7 @@
     FunctionTool,
     Message,
     MessageInjectionMiddleware,
+    MiddlewareTermination,
     ResponseStream,
     SupportsChatGetResponse,
     chat_middleware,
@@ -37,6 +40,44 @@
 class TestChatMiddleware:
     """Test cases for chat middleware functionality."""
 
+    async def test_agent_run_returns_chat_termination_result(self) -> None:
+        client = MockBaseChatClient()
+        expected = ChatResponse(messages=[Message(role="assistant", contents=["terminated response"])])
+
+        class TerminatingMiddleware(ChatMiddleware):
+            async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
+                raise MiddlewareTermination(result=expected)
+
+        agent = Agent(client=client, middleware=[TerminatingMiddleware()])
+        response = await agent.run("hi")
+
+        assert isinstance(response, AgentResponse)
+        assert response.text == "terminated response"
+        assert client.call_count == 0
+
+    async def test_agent_run_returns_chat_termination_stream(self) -> None:
+        client = MockBaseChatClient()
+
+        async def updates() -> AsyncIterable[ChatResponseUpdate]:
+            yield ChatResponseUpdate(contents=[Content.from_text("terminated chat stream")])
+
+        expected = ResponseStream(updates(), finalizer=ChatResponse.from_updates)
+
+        class TerminatingMiddleware(ChatMiddleware):
+            async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
+                raise MiddlewareTermination(result=expected)
+
+        agent = Agent(client=client, middleware=[TerminatingMiddleware()])
+        stream = agent.run("hi", stream=True)
+        streamed_updates = [update async for update in stream]
+        final_response = await stream.get_final_response()
+
+        assert [update.text for update in streamed_updates] == ["terminated chat stream"]
+        assert all(isinstance(update, AgentResponseUpdate) for update in streamed_updates)
+        assert isinstance(final_response, AgentResponse)
+        assert final_response.text == "terminated chat stream"
+        assert client.call_count == 0
+
     async def test_class_based_chat_middleware(self, chat_client_base: SupportsChatGetResponse) -> None:
         """Test class-based chat middleware with ChatClient."""
         execution_order: list[str] = []