diff --git a/pyproject.toml b/pyproject.toml index 4364409f6..8a5b29034 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath-langchain" -version = "0.16.19" +version = "0.16.20" 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 39c97e0f7..01c0fab2e 100644 --- a/uv.lock +++ b/uv.lock @@ -4671,7 +4671,7 @@ wheels = [ [[package]] name = "uipath-langchain" -version = "0.16.19" +version = "0.16.20" source = { editable = "." } dependencies = [ { name = "a2a-sdk" }, @@ -4988,8 +4988,8 @@ name = "uvicorn" version = "0.49.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "click" }, - { name = "h11" }, + { name = "click", marker = "python_full_version != '3.12.*' or sys_platform != 'emscripten'" }, + { name = "h11", marker = "python_full_version != '3.12.*' or sys_platform != 'emscripten'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/c4/1f/fa18009dea8469069cca78a4e877a008ab78f08b064bfc9ab891579077ff/uvicorn-0.49.0.tar.gz", hash = "sha256:ebf4271aa580d9de97f93192d4595176df6e91f9aae919ca73e4fc07df1e66a3", size = 91284, upload-time = "2026-06-03T22:01:30.448Z" } wheels = [