Skip to content

fix(mcp): keep the MCP connection on its own task (0.16 line) #1093

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
PopescuTudor merged 1 commit into from
Sep 16, 2026
Merged

fix(mcp): keep the MCP connection on its own task (0.16 line) #1093

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
60 changes: 39 additions & 21 deletions src/uipath_langchain/agent/tools/mcp/claude.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) │
│ ─────────────────────────────────────────────────────── │
Expand Down Expand Up @@ -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 │
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Comment thread
PopescuTudor marked this conversation as resolved.

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
Comment thread
PopescuTudor marked this conversation as resolved.
```

`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,
Expand Down Expand Up @@ -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

Expand Down
204 changes: 144 additions & 60 deletions src/uipath_langchain/agent/tools/mcp/mcp_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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(
Expand All @@ -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
Comment thread
PopescuTudor marked this conversation as resolved.
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
Comment thread
PopescuTudor marked this conversation as resolved.
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
Comment thread
PopescuTudor marked this conversation as resolved.
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
Comment thread
PopescuTudor marked this conversation as resolved.
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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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")
Loading
Loading