From 5eb4a33e812e8ae8b5a781bdef3bbdcf0140288f Mon Sep 17 00:00:00 2001 From: liweiguang Date: Tue, 10 Feb 2026 13:12:25 +0800 Subject: [PATCH 1/5] fix(result): support pydantic model_rebuild for run results --- src/agents/result.py | 13 +++++++++++++ tests/test_result_cast.py | 12 +++++++++++- 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/src/agents/result.py b/src/agents/result.py index 5e27634f70..365d000793 100644 --- a/src/agents/result.py +++ b/src/agents/result.py @@ -8,6 +8,9 @@ from dataclasses import InitVar, dataclass, field from typing import Any, Literal, TypeVar, cast +from pydantic import GetCoreSchemaHandler +from pydantic_core import core_schema + from .agent import Agent from .agent_output import AgentOutputSchemaBase from .exceptions import ( @@ -124,6 +127,16 @@ class RunResultBase(abc.ABC): _trace_state: TraceState | None = field(default=None, init=False, repr=False) """Serialized trace metadata captured during the run.""" + @classmethod + def __get_pydantic_core_schema__( + cls, + _source_type: Any, + _handler: GetCoreSchemaHandler, + ) -> core_schema.CoreSchema: + # RunResult objects are runtime values; schema generation should treat them as instances + # instead of recursively traversing internal dataclass annotations. + return core_schema.is_instance_schema(cls) + @property @abc.abstractmethod def last_agent(self) -> Agent[Any]: diff --git a/tests/test_result_cast.py b/tests/test_result_cast.py index 63e4d2e8fc..8cbbe038bf 100644 --- a/tests/test_result_cast.py +++ b/tests/test_result_cast.py @@ -7,7 +7,7 @@ import pytest from openai.types.responses import ResponseOutputMessage, ResponseOutputText -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict from agents import ( Agent, @@ -45,6 +45,16 @@ class Foo(BaseModel): bar: int +def test_run_result_streaming_supports_pydantic_model_rebuild() -> None: + class StreamingRunContainer(BaseModel): + query_id: str + run_stream: RunResultStreaming | None + + model_config = ConfigDict(arbitrary_types_allowed=True) + + StreamingRunContainer.model_rebuild() + + def _create_message(text: str) -> ResponseOutputMessage: return ResponseOutputMessage( id="msg", From 47e48f8df39f6ce0e0bba72e934b394489dbd73f Mon Sep 17 00:00:00 2001 From: liweiguang Date: Tue, 10 Feb 2026 17:59:48 +0800 Subject: [PATCH 2/5] fix(mcp): reset client session stack state after cleanup --- src/agents/mcp/server.py | 5 ++++ tests/mcp/test_connect_disconnect.py | 42 ++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/src/agents/mcp/server.py b/src/agents/mcp/server.py index d7c758be7f..3f52ae5e52 100644 --- a/src/agents/mcp/server.py +++ b/src/agents/mcp/server.py @@ -709,7 +709,12 @@ async def cleanup(self): else: logger.error(f"Error cleaning up server: {e}") finally: + # Always reset the exit stack so we don't retain callbacks/references from the + # previous connection. This keeps teardown deterministic and allows reconnecting + # with a fresh stack even if cleanup encountered recoverable errors. + self.exit_stack = AsyncExitStack() self.session = None + self.server_initialize_result = None class MCPServerStdioParams(TypedDict): diff --git a/tests/mcp/test_connect_disconnect.py b/tests/mcp/test_connect_disconnect.py index b001303974..f55c9fee7a 100644 --- a/tests/mcp/test_connect_disconnect.py +++ b/tests/mcp/test_connect_disconnect.py @@ -8,6 +8,18 @@ from .helpers import DummyStreamsContextManager, tee +class CountingStreamsContextManager: + def __init__(self, counter: dict[str, int]): + self.counter = counter + + async def __aenter__(self): + self.counter["enter"] += 1 + return (object(), object()) + + async def __aexit__(self, exc_type, exc_val, exc_tb): + self.counter["exit"] += 1 + + @pytest.mark.asyncio @patch("mcp.client.stdio.stdio_client", return_value=DummyStreamsContextManager()) @patch("mcp.client.session.ClientSession.initialize", new_callable=AsyncMock, return_value=None) @@ -67,3 +79,33 @@ async def test_manual_connect_disconnect_works( await server.cleanup() assert server.session is None, "Server should be disconnected" + + +@pytest.mark.asyncio +@patch("agents.mcp.server.ClientSession.initialize", new_callable=AsyncMock, return_value=None) +@patch("agents.mcp.server.stdio_client") +async def test_cleanup_resets_exit_stack_and_reconnects( + mock_stdio_client: AsyncMock, mock_initialize: AsyncMock +): + counter = {"enter": 0, "exit": 0} + mock_stdio_client.side_effect = lambda params: CountingStreamsContextManager(counter) + + server = MCPServerStdio( + params={ + "command": tee, + }, + cache_tools_list=True, + ) + + await server.connect() + original_exit_stack = server.exit_stack + + await server.cleanup() + assert server.session is None + assert server.exit_stack is not original_exit_stack + assert server.server_initialize_result is None + assert counter == {"enter": 1, "exit": 1} + + await server.connect() + await server.cleanup() + assert counter == {"enter": 2, "exit": 2} From 8cbf14ab5111d079fc5676c046a98cb5feda5511 Mon Sep 17 00:00:00 2001 From: liweiguang Date: Wed, 11 Feb 2026 14:49:33 +0800 Subject: [PATCH 3/5] fix(mcp): keep exit stack on cancelled cleanup --- src/agents/mcp/server.py | 14 +++++----- tests/mcp/test_connect_disconnect.py | 40 ++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 6 deletions(-) diff --git a/src/agents/mcp/server.py b/src/agents/mcp/server.py index 3f52ae5e52..6d3f3faca6 100644 --- a/src/agents/mcp/server.py +++ b/src/agents/mcp/server.py @@ -638,6 +638,7 @@ async def get_prompt( async def cleanup(self): """Cleanup the server.""" async with self._cleanup_lock: + cleanup_cancelled = False # Only raise HTTP errors if we're cleaning up after a failed connection. # During normal teardown (via __aexit__), log but don't raise to avoid # masking the original exception. @@ -646,6 +647,7 @@ async def cleanup(self): try: await self.exit_stack.aclose() except asyncio.CancelledError as e: + cleanup_cancelled = True logger.debug(f"Cleanup cancelled for MCP server '{self.name}': {e}") raise except BaseExceptionGroup as eg: @@ -709,12 +711,12 @@ async def cleanup(self): else: logger.error(f"Error cleaning up server: {e}") finally: - # Always reset the exit stack so we don't retain callbacks/references from the - # previous connection. This keeps teardown deterministic and allows reconnecting - # with a fresh stack even if cleanup encountered recoverable errors. - self.exit_stack = AsyncExitStack() - self.session = None - self.server_initialize_result = None + if not cleanup_cancelled: + # Reset stack state only after a completed cleanup. If cleanup is cancelled, + # keep the existing stack so a follow-up cleanup can finish unwinding it. + self.exit_stack = AsyncExitStack() + self.session = None + self.server_initialize_result = None class MCPServerStdioParams(TypedDict): diff --git a/tests/mcp/test_connect_disconnect.py b/tests/mcp/test_connect_disconnect.py index f55c9fee7a..d8aa09a00e 100644 --- a/tests/mcp/test_connect_disconnect.py +++ b/tests/mcp/test_connect_disconnect.py @@ -1,3 +1,4 @@ +import asyncio from unittest.mock import AsyncMock, patch import pytest @@ -20,6 +21,16 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): self.counter["exit"] += 1 +class CancelThenCloseExitStack: + def __init__(self): + self.close_calls = 0 + + async def aclose(self): + self.close_calls += 1 + if self.close_calls == 1: + raise asyncio.CancelledError("first cleanup interrupted") + + @pytest.mark.asyncio @patch("mcp.client.stdio.stdio_client", return_value=DummyStreamsContextManager()) @patch("mcp.client.session.ClientSession.initialize", new_callable=AsyncMock, return_value=None) @@ -109,3 +120,32 @@ async def test_cleanup_resets_exit_stack_and_reconnects( await server.connect() await server.cleanup() assert counter == {"enter": 2, "exit": 2} + + +@pytest.mark.asyncio +async def test_cleanup_cancellation_preserves_exit_stack_for_retry(): + server = MCPServerStdio( + params={ + "command": tee, + }, + cache_tools_list=True, + ) + cancelled_exit_stack = CancelThenCloseExitStack() + + server.exit_stack = cancelled_exit_stack # type: ignore[assignment] + server.session = object() # type: ignore[assignment] + server.server_initialize_result = object() # type: ignore[assignment] + + with pytest.raises(asyncio.CancelledError): + await server.cleanup() + + assert id(server.exit_stack) == id(cancelled_exit_stack) + assert server.session is not None + assert server.server_initialize_result is not None + + await server.cleanup() + + assert cancelled_exit_stack.close_calls == 2 + assert id(server.exit_stack) != id(cancelled_exit_stack) + assert server.session is None + assert server.server_initialize_result is None From d7f3e096f1b30e63c67cf29da74676e126147d13 Mon Sep 17 00:00:00 2001 From: liweiguang Date: Wed, 11 Feb 2026 20:05:37 +0800 Subject: [PATCH 4/5] fix(mcp): preserve grouped cancellation in cleanup --- src/agents/mcp/server.py | 15 +++++++++ tests/mcp/test_connect_disconnect.py | 48 ++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+) diff --git a/src/agents/mcp/server.py b/src/agents/mcp/server.py index 6d3f3faca6..0327f03b1f 100644 --- a/src/agents/mcp/server.py +++ b/src/agents/mcp/server.py @@ -651,6 +651,21 @@ async def cleanup(self): logger.debug(f"Cleanup cancelled for MCP server '{self.name}': {e}") raise except BaseExceptionGroup as eg: + def contains_cancelled_error(exc: BaseException) -> bool: + if isinstance(exc, asyncio.CancelledError): + return True + if isinstance(exc, BaseExceptionGroup): + return any(contains_cancelled_error(inner) for inner in exc.exceptions) + return False + + if contains_cancelled_error(eg): + cleanup_cancelled = True + logger.debug( + "Cleanup cancelled for MCP server " + f"'{self.name}' with grouped exception: {eg}" + ) + raise + # Extract HTTP errors from ExceptionGroup raised during cleanup # This happens when background tasks fail (e.g., HTTP errors) http_error = None diff --git a/tests/mcp/test_connect_disconnect.py b/tests/mcp/test_connect_disconnect.py index d8aa09a00e..0c5dcbf337 100644 --- a/tests/mcp/test_connect_disconnect.py +++ b/tests/mcp/test_connect_disconnect.py @@ -1,4 +1,5 @@ import asyncio +import sys from unittest.mock import AsyncMock, patch import pytest @@ -8,6 +9,11 @@ from .helpers import DummyStreamsContextManager, tee +if sys.version_info < (3, 11): + from exceptiongroup import BaseExceptionGroup # pyright: ignore[reportMissingImports] +else: + from builtins import BaseExceptionGroup + class CountingStreamsContextManager: def __init__(self, counter: dict[str, int]): @@ -31,6 +37,19 @@ async def aclose(self): raise asyncio.CancelledError("first cleanup interrupted") +class CancelGroupThenCloseExitStack: + def __init__(self): + self.close_calls = 0 + + async def aclose(self): + self.close_calls += 1 + if self.close_calls == 1: + raise BaseExceptionGroup( + "grouped cancellation during cleanup", + [asyncio.CancelledError("grouped cleanup interruption")], + ) + + @pytest.mark.asyncio @patch("mcp.client.stdio.stdio_client", return_value=DummyStreamsContextManager()) @patch("mcp.client.session.ClientSession.initialize", new_callable=AsyncMock, return_value=None) @@ -149,3 +168,32 @@ async def test_cleanup_cancellation_preserves_exit_stack_for_retry(): assert id(server.exit_stack) != id(cancelled_exit_stack) assert server.session is None assert server.server_initialize_result is None + + +@pytest.mark.asyncio +async def test_cleanup_grouped_cancellation_preserves_exit_stack_for_retry(): + server = MCPServerStdio( + params={ + "command": tee, + }, + cache_tools_list=True, + ) + cancelled_exit_stack = CancelGroupThenCloseExitStack() + + server.exit_stack = cancelled_exit_stack # type: ignore[assignment] + server.session = object() # type: ignore[assignment] + server.server_initialize_result = object() # type: ignore[assignment] + + with pytest.raises(BaseExceptionGroup): + await server.cleanup() + + assert id(server.exit_stack) == id(cancelled_exit_stack) + assert server.session is not None + assert server.server_initialize_result is not None + + await server.cleanup() + + assert cancelled_exit_stack.close_calls == 2 + assert id(server.exit_stack) != id(cancelled_exit_stack) + assert server.session is None + assert server.server_initialize_result is None From 928f113fb13ac19489a06a8aeeb1d8f38ad83082 Mon Sep 17 00:00:00 2001 From: liweiguang Date: Wed, 11 Feb 2026 20:11:39 +0800 Subject: [PATCH 5/5] style(mcp): format grouped cancellation cleanup block --- src/agents/mcp/server.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/agents/mcp/server.py b/src/agents/mcp/server.py index 0327f03b1f..3ebcb57949 100644 --- a/src/agents/mcp/server.py +++ b/src/agents/mcp/server.py @@ -651,6 +651,7 @@ async def cleanup(self): logger.debug(f"Cleanup cancelled for MCP server '{self.name}': {e}") raise except BaseExceptionGroup as eg: + def contains_cancelled_error(exc: BaseException) -> bool: if isinstance(exc, asyncio.CancelledError): return True