From c21fca42bdb0f9cfa95fb4f5cd31164c421191a9 Mon Sep 17 00:00:00 2001 From: Popescu Tudor-Cristian Date: Wed, 16 Sep 2026 16:03:48 +0300 Subject: [PATCH] fix(mcp): keep the MCP connection on its own task (0.16 line) Backport of #1087 onto the mcp 1.26 client. The HTTP client, transport and ClientSession live in one AsyncExitStack whose anyio cancel scopes have to be exited by the task that entered them, LIFO. Callers can't do that: langgraph opens the connection inside a tool task while dispose() runs on the teardown task, and an agent with several servers closes them oldest-first. Result was "Attempted to exit a cancel scope that isn't the current task's current cancel scope", swallowed on the way out and coming back as a CancelledError that failed a job which had already produced the right answer. Now a dedicated task opens the whole stack and holds it until dispose() signals it to close (or cancels it mid-handshake), so enter and exit happen on the same task and order stops mattering. Recovery still re-runs the handshake on the same session. Bumps to 0.16.7.post2 (hotfix on exactly what release/s201 ships; nothing else moves). Co-Authored-By: Claude Fable 5.1 --- pyproject.toml | 2 +- .../agent/tools/mcp/claude.md | 60 +++-- .../agent/tools/mcp/mcp_client.py | 204 +++++++++----- .../test_mcp_client_task_ownership.py | 253 ++++++++++++++++++ uv.lock | 4 +- 5 files changed, 439 insertions(+), 84 deletions(-) create mode 100644 tests/agent/tools/test_mcp/test_mcp_client_task_ownership.py diff --git a/pyproject.toml b/pyproject.toml index 951fc52db..4ae507ff9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath-langchain" -version = "0.16.7.post1" +version = "0.16.7.post2" description = "Python SDK that enables developers to build and deploy LangGraph agents to the UiPath Cloud Platform" readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.11" diff --git a/src/uipath_langchain/agent/tools/mcp/claude.md b/src/uipath_langchain/agent/tools/mcp/claude.md index 107c70cbb..e36a637be 100644 --- a/src/uipath_langchain/agent/tools/mcp/claude.md +++ b/src/uipath_langchain/agent/tools/mcp/claude.md @@ -182,14 +182,17 @@ MCP connections for tool invocations with **two distinct initialization phases** │ ─────────────── │ │ _lock: asyncio.Lock # Protects both init phases │ ├─────────────────────────────────────────────────────────────┤ -│ Client State (created once, reused on session reinit) │ -│ ───────────────────────────────────────────────────── │ +│ Connection State (owned by the connection task) │ +│ ─────────────────────────────────────────────── │ +│ _connection_task: asyncio.Task | None │ +│ _ready: asyncio.Future | None # handshake done │ +│ _close_requested: asyncio.Event | None │ +│ _stack: AsyncExitStack | None # set while it runs │ │ _http_client: httpx.AsyncClient | None │ │ _read_stream: MemoryObjectReceiveStream | None │ │ _write_stream: MemoryObjectSendStream | None │ │ _session_info: SessionInfo | None │ -│ _stack: AsyncExitStack | None │ -│ _client_initialized: bool │ +│ _client_initialized: bool # after ready only │ ├─────────────────────────────────────────────────────────────┤ │ Session State (can be reinitialized without recreating) │ │ ─────────────────────────────────────────────────────── │ @@ -362,8 +365,10 @@ Phase 2: Session Initialization (lightweight, can repeat) └──────┬───────┘ │ 1. UiPath SDK retrieves MCP URL │ 2. Factory creates SessionInfo - │ 3. Creates HTTP client, streams, session - │ 4. Calls _initialize_session() + │ 3. Starts the connection task, which opens + │ HTTP client, streams and session on itself + │ 4. The task runs _initialize_session(), + │ then resolves `ready` ▼ ┌──────────────┐ │ Session │ @@ -463,13 +468,14 @@ tool invocation, we ensure the bindings are properly loaded and applied. ### 2. HTTP Client Configuration -The HTTP client MUST use `get_httpx_client_kwargs()` for proper SSL/proxy configuration: +The HTTP client MUST use `get_httpx_client_kwargs()` for proper SSL/proxy configuration. +It is entered inside `_run_connection`, on the connection task's own stack: ```python from uipath._utils._ssl_context import get_httpx_client_kwargs default_client_kwargs = get_httpx_client_kwargs() -self._http_client = await self._stack.enter_async_context( +self._http_client = await stack.enter_async_context( # inside _run_connection httpx.AsyncClient( **default_client_kwargs, headers=self._headers, @@ -500,22 +506,28 @@ async def _reinitialize_session(self) -> None: await self._initialize_session() # Lightweight! ``` -### 4. No `with` Statement for AsyncExitStack +### 4. The Connection Lives on Its Own Task -Manual lifecycle management: +The HTTP client, the streamable HTTP transport and the `ClientSession` all sit +in one `AsyncExitStack`. Its anyio cancel scopes must be exited by the task that +entered them, in reverse order of entry, and neither is under the caller's +control: langgraph opens the connection inside a tool task, and an agent +disposes several servers oldest-first. So `_run_connection` owns the stack with +a plain `async with` and stays inside it until told to leave: ```python -# Correct - manual management -self._stack = AsyncExitStack() -await self._stack.__aenter__() -# ... use stack ... -await self._stack.__aexit__(None, None, None) - -# Wrong - exits too early async with AsyncExitStack() as stack: - ... # Stack closes here! + ... # enter HTTP client, transport and session + await self._initialize_session() + ready.set_result(None) + await close_requested.wait() # dispose() sets this ``` +`dispose()` never touches the stack. It signals the task, or cancels it if the +handshake has not finished yet, and waits for it to exit. Recovery +(`_reinitialize_session`) still re-runs the handshake on the same session; that +is a request over the streams and can run from any task. + ### 5. Reinitialization Reuses Client On 404, only `_initialize_session()` is called — the HTTP client, streams, @@ -599,11 +611,17 @@ When the upstream MCP SDK changes its transport: ### Modifying Client Initialization -1. Changes go in `_initialize_client()` -2. All resources must be added to `_stack` via `enter_async_context()` -3. Set `_client_initialized = True` before calling `_initialize_session()` +1. Resolving the server (SDK lookup, `SessionInfo`) stays in `_initialize_client()`; + opening resources goes in `_run_connection()` +2. All resources must be entered on the connection task's `AsyncExitStack` via + `enter_async_context()` — never on the caller's task +3. `_client_initialized = True` is set only after `ready` resolves, i.e. after + `_initialize_session()` succeeded; the task's `finally` clears it again, so a + connection that dies on its own is rebuilt by the next call 4. Always use `get_httpx_client_kwargs()` for HTTP client 5. The `SessionInfo` is created via the factory — do not construct it directly +6. `dispose()` never touches the stack: `_close_connection()` signals the task (or + cancels it mid-handshake) and waits for it ### Modifying Session Initialization diff --git a/src/uipath_langchain/agent/tools/mcp/mcp_client.py b/src/uipath_langchain/agent/tools/mcp/mcp_client.py index 600f5b3ab..d2265ddfc 100644 --- a/src/uipath_langchain/agent/tools/mcp/mcp_client.py +++ b/src/uipath_langchain/agent/tools/mcp/mcp_client.py @@ -119,6 +119,11 @@ def __init__( self._write_stream: MemoryObjectSendStream[SessionMessage] | None = None self._session_info: SessionInfo | None = None self._stack: AsyncExitStack | None = None + # The connection (HTTP client, transport, session) lives on its own task, + # so its anyio cancel scopes are entered and exited by the same task. + self._connection_task: asyncio.Task[None] | None = None + self._ready: asyncio.Future[None] | None = None + self._close_requested: asyncio.Event | None = None # Session state (can be reinitialized without recreating client) self._session: ClientSession | None = None @@ -141,15 +146,15 @@ def is_client_initialized(self) -> bool: return self._client_initialized async def _initialize_client(self) -> None: - """Initialize the HTTP client and streamable connection. - - This is called once on first use. Creates: - - UiPath SDK instance to retrieve MCP server URL - - httpx.AsyncClient with authorization headers - - Streamable HTTP connection (read/write streams) - - ClientSession - - Then calls _initialize_session() to complete the MCP handshake. + """Resolve the server, then open the connection on a task that owns it. + + The HTTP client, the streamable HTTP transport and the ``ClientSession`` + all sit in one ``AsyncExitStack`` whose anyio cancel scopes must be exited + by the task that entered them, in reverse order of entry. Callers satisfy + neither: langgraph opens the connection inside a tool task while teardown + runs on the main one, and an agent disposes its servers oldest-first. + Keeping the whole stack on a dedicated task makes the caller's task and + ordering irrelevant -- closing is a signal, not an unwind. """ folder_path = get_execution_folder_path() logger.debug( @@ -175,52 +180,134 @@ async def _initialize_client(self) -> None: logger.debug(f"Retrieved MCP server URL: {self._url}") - # Create exit stack for resource management - self._stack = AsyncExitStack() - await self._stack.__aenter__() - - # Create HTTP client with SSL, proxy, and redirect settings - client_kwargs = get_httpx_client_kwargs(headers=self._headers) - client_kwargs["timeout"] = self._timeout - self._http_client = await self._stack.enter_async_context( - httpx.AsyncClient(**client_kwargs) - ) - # Create session info for tracking session ID - self._session_info = self._session_info_factory.create_session(mcp_server) + session_info = self._session_info_factory.create_session(mcp_server) + self._session_info = session_info - # Load previously stored session ID (no-op for base SessionInfo, - # triggers lazy load from debug state for SessionInfoDebugState) - existing = await self._session_info.get_session_id() + # Load a session ID persisted by the AgentHub debug-state integration. + existing = await session_info.get_session_id() if existing: logger.info(f"Loaded existing session ID from session info: {existing}") - # Create streamable HTTP connection - ( - self._read_stream, - self._write_stream, - ) = await self._stack.enter_async_context( - streamable_http_client( - url=self._url, - http_client=self._http_client, - session_info=self._session_info, - terminate_on_close=self._terminate_on_close, - ) - ) - - # Create ClientSession (but don't initialize yet) - # These are guaranteed to be set by the context manager above - assert self._read_stream is not None - assert self._write_stream is not None - self._session = await self._stack.enter_async_context( - ClientSession(self._read_stream, self._write_stream) + ready: asyncio.Future[None] = asyncio.get_running_loop().create_future() + self._ready = ready + self._close_requested = asyncio.Event() + self._connection_task = asyncio.create_task( + self._run_connection( + ready, self._close_requested, self._url, self._headers, session_info + ), + name=f"mcp-connection-{self._config.slug}", ) + try: + # Shielded so a cancelled caller does not cancel `ready` with it: + # _close_connection reads it to tell a finished handshake (signal) + # from one still in flight (cancel). + await asyncio.shield(ready) + except BaseException: + await self._close_connection() + self._session_info = None + raise self._client_initialized = True logger.info("MCP client initialized") - # Now initialize the MCP session - await self._initialize_session() + async def _run_connection( + self, + ready: asyncio.Future[None], + close_requested: asyncio.Event, + url: str, + headers: dict[str, str], + session_info: SessionInfo, + ) -> None: + """Hold the HTTP client, transport and session open until asked to close. + + Resolves *ready* once the session is negotiated, so ``_initialize_client`` + fails the same way it did when it opened the stack inline. + """ + # Unwinding the transport's task group wraps whatever went wrong in a + # BaseExceptionGroup. Callers used to see the original error, so keep it. + setup_error: BaseException | None = None + try: + async with AsyncExitStack() as stack: + self._stack = stack + try: + client_kwargs = get_httpx_client_kwargs(headers=headers) + client_kwargs["timeout"] = self._timeout + self._http_client = await stack.enter_async_context( + httpx.AsyncClient(**client_kwargs) + ) + streams = await stack.enter_async_context( + streamable_http_client( + url=url, + http_client=self._http_client, + session_info=session_info, + terminate_on_close=self._terminate_on_close, + ) + ) + self._read_stream, self._write_stream = streams + self._session = await stack.enter_async_context( + ClientSession(self._read_stream, self._write_stream) + ) + await self._initialize_session() + except BaseException as error: + setup_error = error + raise + if not ready.done(): + ready.set_result(None) + await close_requested.wait() + except BaseException as error: + failure = setup_error if setup_error is not None else error + if isinstance(failure, asyncio.CancelledError): + # _close_connection cancelled us mid-handshake, or the loop is + # going down. End as cancelled, not as failed. + if not ready.done(): + ready.cancel() + if failure is error: + raise + raise failure from error + if not ready.done(): + ready.set_exception(failure) + else: + logger.debug("MCP connection ended with an error: %s", failure) + finally: + # Whatever ended the task, the client no longer has a connection. + # Clearing the flag makes the next call rebuild it instead of + # handing a missing session to the operation. + self._client_initialized = False + self._stack = None + self._session = None + self._read_stream = None + self._write_stream = None + self._http_client = None + + async def _close_connection(self) -> None: + """Ask the connection task to unwind, and wait for it to finish. + + A task still in the handshake would not see the signal until the + handshake returned -- up to the transport timeout -- so it is cancelled + instead. + """ + task = self._connection_task + ready = self._ready + close_requested = self._close_requested + self._connection_task = None + self._ready = None + self._close_requested = None + self._session = None + if task is None: + return + + if ready is not None and ready.done() and close_requested is not None: + close_requested.set() + else: + task.cancel() + + # asyncio.wait rather than `await task`: the task's own cancellation + # must not read as ours, while a real cancellation of this task still + # propagates. + await asyncio.wait({task}) + if not task.cancelled() and (error := task.exception()) is not None: + logger.debug("Error closing MCP connection: %s", error) async def _initialize_session(self) -> None: """Initialize or reinitialize the MCP session. @@ -267,7 +354,10 @@ async def _ensure_session(self) -> ClientSession: if not self._client_initialized: await self._initialize_client() - return self._session # type: ignore[return-value] + session = self._session + if session is None: + raise RuntimeError("MCP client initialized without a session") + return session async def _reinitialize_session(self) -> None: """Reinitialize only the MCP session after a disconnect error. @@ -408,18 +498,12 @@ async def dispose(self) -> None: async with self._tools_lock: self._tools_cache = None async with self._lock: - if self._stack is not None: - try: - await self._stack.__aexit__(None, None, None) - except Exception as e: - logger.debug(f"Error during cleanup: {e}") - finally: - self._stack = None - self._session = None - self._http_client = None - self._read_stream = None - self._write_stream = None - self._session_info = None - self._client_initialized = False - + try: + await self._close_connection() + finally: + # The connection fields are cleared by the task itself; + # this is the client-level state, and it must not survive + # a dispose() that gets cancelled halfway. + self._session_info = None + self._client_initialized = False logger.info("MCP client disposed") diff --git a/tests/agent/tools/test_mcp/test_mcp_client_task_ownership.py b/tests/agent/tools/test_mcp/test_mcp_client_task_ownership.py new file mode 100644 index 000000000..192bbe588 --- /dev/null +++ b/tests/agent/tools/test_mcp/test_mcp_client_task_ownership.py @@ -0,0 +1,253 @@ +"""``McpClient`` teardown must not depend on which task opened the session. + +The transport and ``ClientSession`` live in an anyio task group inside the +client's exit stack, so whoever exits it has to be the task that entered it, +and scopes on one task have to unwind LIFO. Neither holds for a client: +langgraph runs each tool call in its own task (``asyncio.gather`` in +``ToolNode``) while ``dispose()`` runs on the teardown task, and an agent with +several servers disposes them in creation order. + +Before the connection moved onto its own task, the scope violation was caught +under ``except Exception`` and logged at DEBUG, so a test that only checked +state after ``dispose()`` passed anyway. These assert nothing was swallowed, +against a real server that records what reached it. +""" + +import asyncio +import logging +import socket +import threading +import time +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import Any +from unittest.mock import patch + +import pytest +from uipath.agent.models.agent import AgentMcpResourceConfig, AgentMcpTool + +from uipath_langchain.agent.tools.mcp import McpClient + +MCP_CLIENT_LOGGER = "uipath_langchain.agent.tools.mcp.mcp_client" + + +@dataclass +class MathServer: + """A running MCP server plus the HTTP methods it has seen, in order.""" + + url: str + methods: list[str] = field(default_factory=list) + + +@pytest.fixture +def math_server() -> Iterator[MathServer]: + """Host a real ``FastMCP`` server on an ephemeral port, on its own thread.""" + import uvicorn + from mcp.server.fastmcp import FastMCP + + server = FastMCP("Math") + + @server.tool() + def add(a: int, b: int) -> int: + """Add two numbers""" + return a + b + + sock = socket.socket() + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + sock.close() + + running = MathServer(url=f"http://127.0.0.1:{port}/mcp") + app = server.streamable_http_app() + + async def recording_app(scope: dict[str, Any], receive: Any, send: Any) -> None: + if scope["type"] == "http": + running.methods.append(scope["method"]) + await app(scope, receive, send) + + uv = uvicorn.Server( + uvicorn.Config(recording_app, host="127.0.0.1", port=port, log_level="error") + ) + thread = threading.Thread(target=uv.run, daemon=True) + thread.start() + deadline = time.monotonic() + 10 + while not uv.started and time.monotonic() < deadline: + time.sleep(0.05) + assert uv.started, "MCP test server did not start" + try: + yield running + finally: + uv.should_exit = True + thread.join(5) + + +@contextmanager +def patched_sdk(url: str) -> Iterator[None]: + """Point the client's lazy ``UiPath`` lookup at the local server.""" + + class _Server: + mcp_url = url + slug = "math" + folder_key = "folder-key" + name = "Math" + + class _Mcp: + async def retrieve_async(self, name: str, folder_path: str | None) -> _Server: + return _Server() + + class _Config: + secret = "test-token" + + class FakeUiPath: + def __init__(self, *args: object, **kwargs: object) -> None: + self.mcp = _Mcp() + self._config = _Config() + + with patch("uipath.platform.UiPath", FakeUiPath): + yield + + +def make_client() -> McpClient: + return McpClient( + config=AgentMcpResourceConfig( + name="Math", + description="Math MCP server", + folder_path="Shared", + slug="math", + available_tools=[ + AgentMcpTool( + name="add", + description="Add two numbers", + input_schema={ + "type": "object", + "properties": { + "a": {"type": "integer"}, + "b": {"type": "integer"}, + }, + "required": ["a", "b"], + }, + ) + ], + ) + ) + + +def swallowed_errors(caplog: pytest.LogCaptureFixture) -> list[str]: + """Errors the client logged instead of raising.""" + return [ + r.getMessage() + for r in caplog.records + if r.name == MCP_CLIENT_LOGGER and "error" in r.getMessage().lower() + ] + + +@pytest.mark.asyncio +async def test_session_opened_on_a_tool_task_disposes_from_the_teardown_task( + math_server: MathServer, caplog: pytest.LogCaptureFixture +) -> None: + """Cached discovery connects on first tool call, i.e. inside a tool task.""" + caplog.set_level(logging.DEBUG, logger=MCP_CLIENT_LOGGER) + with patched_sdk(math_server.url): + client = make_client() + + # asyncio.gather wraps the coroutine in its own Task, the way + # langgraph's ToolNode dispatches tool calls. + (result,) = await asyncio.gather(client.call_tool("add", {"a": 2, "b": 3})) + assert "5" in str(result) + + await client.dispose() + + assert not client.is_client_initialized + + assert swallowed_errors(caplog) == [] + # terminate_on_close still reaches the server through the task-owned teardown: + # the HTTP client is entered before the transport, so it is open when the + # transport exits and sends the session DELETE. + assert "DELETE" in math_server.methods + + +@pytest.mark.asyncio +async def test_clients_dispose_cleanly_in_creation_order( + math_server: MathServer, caplog: pytest.LogCaptureFixture +) -> None: + """Dynamic discovery connects every server up front, on one task.""" + caplog.set_level(logging.DEBUG, logger=MCP_CLIENT_LOGGER) + with patched_sdk(math_server.url): + clients = [make_client() for _ in range(3)] + for client in clients: + await client.list_tools() + + # Oldest-first: the order a plain `for disposable in ...` loop uses. + for client in clients: + await client.dispose() + + assert all(not c.is_client_initialized for c in clients) + + assert swallowed_errors(caplog) == [] + + +@pytest.mark.asyncio +async def test_cancelling_the_caller_mid_handshake_unwinds_promptly( + math_server: MathServer, caplog: pytest.LogCaptureFixture +) -> None: + """A caller cancelled during the handshake must not wait for it to finish. + + Signalling a connection task does nothing until the handshake returns, and + the transport timeout is ten minutes. The task has to be cancelled. + """ + caplog.set_level(logging.DEBUG, logger=MCP_CLIENT_LOGGER) + with patched_sdk(math_server.url): + client = make_client() + handshake_started = asyncio.Event() + + async def slow_handshake() -> None: + handshake_started.set() + await asyncio.sleep(5) + + with patch.object(client, "_initialize_session", slow_handshake): + caller = asyncio.create_task(client.call_tool("add", {"a": 2, "b": 3})) + await asyncio.wait_for(handshake_started.wait(), timeout=5) + + loop = asyncio.get_running_loop() + started = loop.time() + caller.cancel() + with pytest.raises(asyncio.CancelledError): + await caller + elapsed = loop.time() - started + + assert elapsed < 1.0, f"cancellation waited {elapsed:.2f}s" + assert client._connection_task is None + assert not client.is_client_initialized + + assert swallowed_errors(caplog) == [] + + +@pytest.mark.asyncio +async def test_a_connection_that_dies_on_its_own_is_rebuilt_on_the_next_call( + math_server: MathServer, caplog: pytest.LogCaptureFixture +) -> None: + """The connection task ending by itself must not poison the client. + + A transport child task failing after the handshake cancels the task + group's scope, which ends the connection task. If that only cleared the + session while leaving the client flagged as initialized, every later call + would fail on a missing session with no retry. The next call has to + rebuild the connection instead. + """ + caplog.set_level(logging.DEBUG, logger=MCP_CLIENT_LOGGER) + with patched_sdk(math_server.url): + client = make_client() + await client.call_tool("add", {"a": 1, "b": 1}) + + task = client._connection_task + assert task is not None + task.cancel() # what the transport's task group does when a child fails + await asyncio.wait({task}) + + result = await client.call_tool("add", {"a": 2, "b": 3}) + assert "5" in str(result) + + await client.dispose() + + assert swallowed_errors(caplog) == [] diff --git a/uv.lock b/uv.lock index d0c25996a..30ab489ab 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-09-01T19:52:33.763506Z" +exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. exclude-newer-span = "P2D" [options.exclude-newer-package] @@ -4546,7 +4546,7 @@ wheels = [ [[package]] name = "uipath-langchain" -version = "0.16.7.post1" +version = "0.16.7.post2" source = { editable = "." } dependencies = [ { name = "a2a-sdk" },