feat: mcp client with 404 recovery

[edis-uipath]

Summary

Introduces McpClient class for managing MCP session lifecycle with two-phase initialization and automatic session recovery on 404 errors. Also separates MCP tool creation from the general tool factory, providing a dedicated create_mcp_tools_from_agent function that returns both tools and clients for proper lifecycle management.

Changes

Commit 1: McpClient with Two-Phase Initialization

Introduces the core McpClient class with:

• Two-phase initialization: Separates expensive client setup from lightweight session handshake
• Automatic 404 recovery: Reuses HTTP client, only re-runs MCP handshake on session disconnect
• Thread safety: Single asyncio.Lock protects both initialization phases

Commit 2: Separate MCP Tools from Tool Factory

• Made McpClient public: Renamed _mcp_client.py → mcp_client.py, exported from tools/mcp
• Added create_mcp_tools_from_agent: Factory function that creates tools from LowCodeAgentDefinition and returns (tools, clients) tuple for proper cleanup
• Lazy SDK initialization: UiPath() is initialized inside create_mcp_tools_from_agent, not passed as parameter
• Removed MCP from general tool factory: _build_tool_for_resource no longer handles MCP resources

Architecture

McpClient Class

McpClient manages the lifecycle of MCP connections with two distinct initialization phases:

1. Client Initialization (first call): Full stack creation - HTTP client, streams, ClientSession
2. Session Reinitialization (on 404): Lightweight - reuses existing client, only re-runs MCP handshake

┌─────────────────────────────────────────────────────────────┐
│                     McpClient                               │
├─────────────────────────────────────────────────────────────┤
│  Configuration (immutable)                                  │
│  _url, _headers, _timeout, _max_retries                     │
├─────────────────────────────────────────────────────────────┤
│  Synchronization: _lock (asyncio.Lock)                      │
├─────────────────────────────────────────────────────────────┤
│  Client State (created once, reused on session reinit)      │
│  _http_client, _read_stream, _write_stream, _stack          │
├─────────────────────────────────────────────────────────────┤
│  Session State (can be reinitialized)                       │
│  _session, _session_id                                      │
├─────────────────────────────────────────────────────────────┤
│  Public: call_tool(), close(), session_id, is_client_initialized │
└─────────────────────────────────────────────────────────────┘

Factory Functions

# Primary factory - creates tools from agent definition
async def create_mcp_tools_from_agent(
    agent: LowCodeAgentDefinition,
) -> tuple[list[BaseTool], list[McpClient]]:
    """Creates tools and returns clients for lifecycle management.
    
    Caller is responsible for closing clients when done:
        tools, clients = await create_mcp_tools_from_agent(agent)
        try:
            # Use tools...
        finally:
            for client in clients:
                await client.close()
    """

# Lower-level factory - creates tools for a single MCP server
async def create_mcp_tools_from_metadata_for_mcp_server(
    config: AgentMcpResourceConfig,
    mcpClient: McpClient,
) -> list[BaseTool]:
    """Creates tools using an existing McpClient."""

Session Recovery on 404

When an MCP server returns 404 (session terminated), McpClient automatically:

1. Keeps the existing HTTP client and streamable connection
2. Re-runs only session.initialize() to get a new session ID
3. Retries the failed tool call

MCP Protocol Flow

First tool call:

Client                              Server
   │──── initialize ──────────────────►│
   │◄─── result + session-id-1 ────────│
   │──── notifications/initialized ───►│
   │◄─── 204 No Content ───────────────│
   │──── tools/call ──────────────────►│
   │◄─── result ───────────────────────│

On 404 error (automatic recovery):

Client                              Server
   │──── tools/call ──────────────────►│
   │◄─── 404 (session terminated) ─────│
   │  [Reuses existing HTTP client]    │
   │──── initialize ──────────────────►│  ← new session
   │◄─── result + session-id-2 ────────│    (same client)
   │──── notifications/initialized ───►│
   │──── tools/call ──────────────────►│  ← retry
   │◄─── result ───────────────────────│

File Structure

src/uipath_langchain/agent/tools/mcp/
├── __init__.py          # Exports McpClient, create_mcp_tools, create_mcp_tools_from_agent
├── mcp_client.py        # McpClient class (public)
├── mcp_tool.py          # Tool factory functions
└── claude.md            # Implementation documentation

tests/agent/tools/test_mcp/
├── test_mcp_client.py   # 7 tests for McpClient
├── test_mcp_tool.py     # 17 tests for tool creation
└── claude.md            # Testing documentation

Tests

24 tests total covering session management and tool creation.

test_mcp_client.py (7 tests)

Test	Purpose
test_session_initializes_on_first_call	Lazy initialization on first call_tool()
test_session_reused_across_calls	Session persists across multiple call_tool() invocations
test_session_reinitializes_on_404_error	Key test: HTTP client reused, only session reinit on 404
test_max_retries_exceeded	McpError raised after max retries exhausted
test_close_releases_resources	close() properly releases all resources
test_client_initialized_property	is_client_initialized property reflects actual state
test_session_can_be_reused_after_close	Full reinitialization works after close()

test_mcp_tool.py (17 tests)

Class	Tests
TestMcpToolMetadata	5 tests - tool metadata (tool_type, display_name, folder_path, slug)
TestMcpToolCreation	3 tests - multiple tools, descriptions, disabled config
TestCreateMcpToolsFromAgent	6 tests - agent factory function, client lifecycle
TestMcpToolInvocation	1 test - full MCP protocol flow smoke test
TestMcpToolNameSanitization	2 tests - spaces and special chars in names

Testing Strategy

• Mock only HTTP layer (httpx.AsyncClient)
• Never mock MCP SDK - real ClientSession processes messages
• Patch UiPath class for create_mcp_tools_from_agent tests

Breaking Changes

None. The old create_mcp_tools_from_metadata was internal and has been replaced with create_mcp_tools_from_metadata_for_mcp_server.

🤖 Generated with Claude Code

[Copilot]

Pull request overview

This PR implements session reuse for MCP (Model Context Protocol) tools to improve performance by avoiding repeated initialization overhead when tools from the same resource configuration are invoked multiple times.

Changes:

• Introduced a custom streamable_http_client implementation that accepts an optional session_id parameter for reusing existing server-side sessions
• Added closure variable session_id shared across all tools created from the same AgentMcpResourceConfig to track and reuse sessions
• Modified session initialization logic to skip initialization when connecting with an existing session ID
• Added comprehensive test suite with 14 tests covering metadata, tool creation, functionality, session reuse, and name sanitization

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 9 comments.

File	Description
src/uipath_langchain/agent/tools/mcp_tool.py	Implements session reuse mechanism through custom streamable HTTP client and closure-scoped session_id variable
tests/agent/tools/test_mcp_tool.py	Adds comprehensive test coverage for MCP tool functionality including session reuse validation

Comments suppressed due to low confidence (1)

src/uipath_langchain/agent/tools/mcp_tool.py:275

• The session ID is stored in a closure variable but is never cleared or reset. If a session fails, expires, or becomes invalid, the tool will continue attempting to reuse the stale session_id, causing subsequent tool invocations to fail without recovery.

Consider adding error handling that clears the session_id when session operations fail, allowing the next invocation to create a fresh session. For example, wrap the session operations in a try-except block and set session_id = None on failure to enable recovery.

                    if not session_id:
                        await session.initialize()
                        session_id = getSessionId()
                        logger.info(f"session {session_id} created")

                    # Call the tool
                    result = await session.call_tool(mcp_tool.name, arguments=kwargs)
                    return result.content if hasattr(result, "content") else result

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 7 comments.

[Copilot]

Pull request overview

Copilot reviewed 11 out of 12 changed files in this pull request and generated 3 comments.

Comments suppressed due to low confidence (4)

src/uipath_langchain/agent/tools/mcp/mcp_tool.py:172

• The McpClient instance is never closed, leading to resource leaks. Each call to create_mcp_tools_from_metadata creates a new McpClient with HTTP client, streams, and sessions that are never cleaned up. While the PR description mentions "Session cleanup and proper termination will be added in a future PR", the current implementation creates resource leaks. At minimum, there should be a mechanism to close the client when tools are no longer needed, or the function should document that the caller is responsible for closing the McpClient.
  src/uipath_langchain/agent/tools/mcp/mcp_tool.py:110
• The docstring is misleading. It states "Each tool manages its own session lifecycle - creating, using, and cleaning up the MCP connection within the tool invocation." However, the actual implementation shares a single McpClient instance across all tools from the same resource config, and the session is never cleaned up (no close() is called). The docstring should be updated to reflect that all tools share a single persistent session for performance reasons.
  src/uipath_langchain/agent/tools/mcp/mcp_tool.py:150
• The docstring is misleading. It mentions "ephemeral session" but the session is actually persistent and shared across all tool invocations from the same resource config. The session is created once and reused, which is the main feature of this PR. The docstring should accurately describe this as a persistent, shared session with automatic reconnection on errors.
  src/uipath_langchain/agent/tools/mcp/mcp_tool.py:127
• The McpClient instantiation does not pass authentication headers. The UiPath SDK is initialized at line 123 which contains authentication credentials in sdk._config.secret, but these are never passed to the McpClient. Without proper authentication headers, the MCP server requests will likely fail with authentication errors. The headers should be passed during McpClient initialization to include the Bearer token.

[Copilot]

Pull request overview

Copilot reviewed 12 out of 13 changed files in this pull request and generated 4 comments.

Comments suppressed due to low confidence (2)

src/uipath_langchain/agent/tools/mcp/mcp_tool.py:129

• The McpClient is created without authentication headers. The old implementation included "headers": {"Authorization": f"Bearer {sdk._config.secret}"} when creating the HTTP client, but line 129 creates McpClient with only the URL. This will cause authentication failures when the MCP server requires authorization.

The McpClient should be initialized with:

mcpClient: McpClient = McpClient(
    mcpServer.mcp_url,
    headers={"Authorization": f"Bearer {sdk._config.secret}"}
)

src/uipath_langchain/agent/tools/mcp/mcp_tool.py:174

• The McpClient instance is never closed, leading to a resource leak. The HTTP client, streams, and session remain open indefinitely since mcpClient.close() is never called. This is especially problematic because the McpClient is created once per tool creation and persists for the lifetime of the tools.

Consider one of these approaches:

1. Return the McpClient from this function and have the caller manage its lifecycle
2. Use an AsyncExitStack or context manager to ensure cleanup
3. Document that the caller is responsible for closing the client

[Copilot]

Non-McpError exceptions are not handled in the retry logic, which could lead to inconsistent state. If session.call_tool() at line 240 raises any exception other than McpError (e.g., httpx.HTTPError, asyncio.TimeoutError, or connection errors), it will propagate immediately without cleanup. This leaves the client in an initialized state but potentially with a broken connection.

Consider wrapping in a broader exception handler to either:

1. Close and reset the client on non-McpError exceptions
2. Document which exceptions are expected and should not be retried
3. Add logging for unexpected exceptions

Suggested change

	raise
	raise
	except asyncio.CancelledError:
	# Preserve task cancellation semantics; do not attempt cleanup here.
	raise
	except Exception as e:
	# Handle unexpected non-McpError exceptions by resetting the client state
	logger.exception(
	"Unexpected error during MCP tool call; closing session to reset state"
	)
	try:
	await self.close()
	except Exception as cleanup_error:
	logger.debug(f"Error during cleanup after unexpected exception: {cleanup_error}")
	raise

[Copilot]

The error checking logic using hasattr on line 207-208 is fragile and could silently fail. If the McpError structure changes or if error.error exists but is None, this will return False instead of raising an appropriate error. Additionally, accessing error.error.code at line 249 without the same checks could raise an AttributeError.

Consider:

1. Using more explicit type checking or isinstance checks
2. Handling the case where error.error is None
3. Being consistent - if you use hasattr in the check, also guard the access at line 249

[Copilot]

If _initialize_session() at line 146 raises an exception, the client will be left in an inconsistent state. The _client_initialized flag is set to True at line 142, but if session initialization fails, future calls to _ensure_session() will return immediately thinking the client is ready (line 176-181), but _session_id will be None and the session may not be properly initialized.

Add exception handling to ensure cleanup on failure:

try:
    # ... client initialization ...
    self._client_initialized = True
    await self._initialize_session()
except Exception:
    # Clean up on failure
    await self.close()
    raise

Suggested change

	self._client_initialized = True
	logger.info("MCP client initialized")
	# Now initialize the MCP session
	await self._initialize_session()
	try:
	self._client_initialized = True
	logger.info("MCP client initialized")
	# Now initialize the MCP session
	await self._initialize_session()
	except Exception:
	# Clean up on failure to avoid leaving the client in an inconsistent state
	await self.close()
	raise

[Copilot]

The test mocks don't validate that authentication headers are passed to the McpClient. The mock SDK at lines 400-404 doesn't expose _config.secret, so the test would fail if authentication were properly implemented. This allowed the authentication bug to slip through.

Consider adding a test assertion that verifies the McpClient is created with proper authentication headers, or mock sdk._config.secret to enable testing of the full authentication flow.

[Copilot]

Pull request overview

Copilot reviewed 12 out of 13 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (4)

src/uipath_langchain/agent/tools/mcp/mcp_tool.py:129

• The McpClient is being instantiated without authentication headers. The previous implementation (visible in the removed code) included authentication headers with the SDK secret:

"headers": {"Authorization": f"Bearer {sdk._config.secret}"}

The MCP server URL retrieved from sdk.mcp.retrieve_async() likely requires authentication. Without passing the authorization header to McpClient, HTTP requests to the MCP server will fail with 401 Unauthorized errors.

To fix this, pass the SDK credentials to McpClient:

mcpClient: McpClient = McpClient(
    mcpServer.mcp_url,
    headers={"Authorization": f"Bearer {sdk._config.secret}"}
)

src/uipath_langchain/agent/tools/mcp/mcp_tool.py:174

• The McpClient instance is created and shared across all tools but is never explicitly closed. This creates a resource leak where HTTP clients, streams, and connections remain open indefinitely.

The previous implementation used an AsyncExitStack context manager to ensure resources were properly cleaned up. The new design keeps the McpClient alive for the lifetime of the tools, but there's no mechanism to clean up resources when they're no longer needed.

Consider one of these approaches:

1. Make create_mcp_tools_from_metadata return a context manager that yields tools and cleans up McpClient on exit
2. Store the McpClient reference and provide a cleanup function that can be called by the runtime
3. Implement a finalizer or del method in McpClient to ensure cleanup (though this is less reliable)

Without proper cleanup, long-running processes will accumulate TCP connections and other resources.
src/uipath_langchain/agent/tools/mcp/mcp_tool.py:152

• The docstring states "Execute MCP tool call with ephemeral session" but this is no longer accurate. The McpClient instance (and its session) is now shared and persistent across all tool invocations, not ephemeral. The session is only reinitialized on disconnect errors (404), not recreated for each tool call.

Update the docstring to reflect the new behavior, for example:

"""Execute MCP tool call using the shared MCP session.

The session is persistent and automatically recovers from disconnect errors
by reinitializing the MCP handshake while reusing the HTTP client.
"""

src/uipath_langchain/agent/tools/mcp/mcp_tool.py:110

• The docstring states "Each tool manages its own session lifecycle - creating, using, and cleaning up the MCP connection within the tool invocation" but this is no longer accurate.

In the new implementation:

• All tools from the same config share a single McpClient instance
• The session is persistent, not created and destroyed per invocation
• There is no cleanup mechanism for the shared McpClient

Update the docstring to accurately reflect the new shared session model.

[Copilot]

There's a potential race condition when multiple concurrent tool calls encounter 404 errors. The scenario:

1. Tool call A gets 404, enters retry logic
2. Tool call B gets 404 while A is reinitializing, also enters retry logic
3. Both calls try to reinitialize the session

While _reinitialize_session is protected by a lock, the check if self._is_session_error(e) happens outside the lock in call_tool. This means multiple concurrent calls could all increment their retry counters and potentially exhaust retries even though only one session reinitialization is needed.

A more robust approach would be to track the session ID at the start of each call attempt and only reinitialize if the session ID hasn't changed (meaning another concurrent call hasn't already reinitialized). However, given the lock protection in _reinitialize_session, this is a minor efficiency issue rather than a correctness issue - at worst, some calls might retry unnecessarily or fail when they could have succeeded.

[cotovanu-cristian]

should we remove this?

[cotovanu-cristian]

LGTM! 🚀

[cotovanu-cristian]

Future improvement

This PR doesn’t handle session reuse across suspend/resume, which is fine for now, but it’s a gap worth addressing in a future PR.

Given how often suspend is used (Agents + coded SDKs), always losing the session can be expensive.
The potential improvement would be to persist <mcpServer, session> on GraphState and run a lightweight health check before re-initializing.
Since GraphState is checkpointed, it can be restored after resume and only re-init when the session is actually invalid.

Considering the Coded MCP sessions got a 15 mins timeout.
The suspension use cases cover:
• Process invocation (agent processes, RPA, API workflows) — these usually resume <15 mins, so leads to recreation of the session without the need
• ECS context processing (DeepRAG, batch transforms) — mixed behavior here depending on the document size but can also benefit from the session reuse
• Escalations / app tasks — these can easily go past 15 minutes, so session expiration might happen in majority of cases

This gets painful for process tools IMO which are the widely used ones: Agents, RPA, API WFs and ECS DeepRag/BatchTransform.