diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 0372967ccf1..0c483d03fbf 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/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: diff --git a/python/packages/core/tests/core/test_middleware_context_result.py b/python/packages/core/tests/core/test_middleware_context_result.py index 75ebc9e2204..410ccba7afe 100644 --- a/python/packages/core/tests/core/test_middleware_context_result.py +++ b/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"])]) 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..49356539377 100644 --- a/python/packages/core/tests/core/test_middleware_with_agent.py +++ b/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] = [] 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 0271ed7c11d..a44eed78e9e 100644 --- a/python/packages/core/tests/core/test_middleware_with_chat.py +++ b/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] = []