feat: Reuse MCP session between tool calls

[edis-uipath]

Summary

MCP tools created from metadata now reuse the same session across multiple tool invocations, significantly improving performance by avoiding repeated initialization overhead.

Changes

Session Reuse Implementation

Previously, each tool invocation created a new MCP session, requiring:

• New HTTP connection establishment
• Session initialization handshake
• Redundant authentication

Now, all tools created from the same AgentMcpResourceConfig share a single session using a closure variable:

session_id: str | None = None

def build_mcp_tool(mcp_tool: AgentMcpTool) -> Any:
    async def tool_fn(**kwargs: Any) -> Any:
        nonlocal session_id  # Share session across all tool invocations
        
        # Reuse existing session if available
        read, write, getSessionId = await streamable_http_client(
            url=mcpServer.mcp_url,
            http_client=http_client,
            session_id=session_id,  # Pass existing session ID
        )
        
        # Only initialize on first call
        if not session_id:
            await session.initialize()
            session_id = getSessionId()
            logger.info(f"session {session_id} created")

Custom StreamableHTTP Client

Added a custom streamable_http_client() implementation that:

• Accepts an optional session_id parameter for session reuse
• Properly manages HTTP client lifecycle
• Extracts and tracks session IDs from HTTP response headers
• Provides get_session_id() callback for retrieving the current session

Benefits

• Performance: Single initialization per resource instead of per-tool-invocation
• Efficiency: Reduced network overhead and connection setup time
• Statefulness: Maintains conversation context across tool calls
• Resource Usage: Lower memory and connection pool consumption

Future Work

Session Termination: The current implementation does not terminate sessions when tools are done. Session cleanup and proper termination will be added in a future PR to ensure resources are released appropriately.

Tests

Added comprehensive test suite with 14 tests covering all aspects of MCP tool functionality:

TestMcpToolMetadata (5 tests)

Validates that tools have correct metadata for observability and tracing:

• test_mcp_tool_has_metadata - Verifies metadata dict exists
• test_mcp_tool_metadata_has_tool_type - Validates tool_type: "mcp" for span detection
• test_mcp_tool_metadata_has_display_name - Checks display name from tool name
• test_mcp_tool_metadata_has_folder_path - Verifies folder path for span attributes
• test_mcp_tool_metadata_has_slug - Validates server slug for identification

TestMcpToolCreation (3 tests)

Tests tool creation and configuration:

• test_creates_multiple_tools - Verifies multiple tools created from config
• test_tool_has_correct_description - Validates tool descriptions
• test_disabled_config_returns_empty_list - Tests disabled config handling

TestMcpToolFunctionality (2 tests)

Tests tool invocation with mocked HTTP layer:

• test_tool_invocation_calls_mcp_server - Validates complete tool invocation flow with GUID session ID
• test_tool_returns_result_without_content_attribute - Tests handling of results without content attribute

TestMcpToolSessionReuse (2 tests)

Key tests validating the session reuse feature:

• test_tools_share_same_session_id_variable - Verifies session reuse at HTTP level:
  • Intercepts HTTP stream calls and validates only 1 initialize call for 2 tool invocations
  • Checks both tools use the same mcp-session-id header
  • Confirms session ID matches the one returned from initialization
  • Uses real ClientSession code with mocked HTTP client for realistic testing
• test_session_not_terminated_between_calls - Validates session not terminated between tool invocations

TestMcpToolNameSanitization (2 tests)

Tests tool name sanitization:

• test_tool_name_with_spaces - Verifies spaces removed from tool names
• test_tool_name_with_special_chars - Tests special character handling

Testing Approach

Tests use strategic mocking:

• Mock only the HTTP layer (httpx.AsyncClient) to allow real MCP protocol code execution
• Simulate MCP server responses for initialize, tools/call, tools/list, and notification methods
• Track all HTTP stream calls to verify session reuse behavior
• Return proper JSON-RPC responses with correct headers and content-types

Technical Details

Closure Scoping with nonlocal

The session reuse mechanism relies on Python closure scoping:

• session_id variable declared at the outer scope (per resource config)
• Each tool function uses nonlocal session_id to access and modify the shared variable
• First tool invocation initializes the session and stores the ID
• Subsequent invocations reuse the existing session ID

MCP Protocol Flow

1. First Tool Invocation:

   • session_id is None
   • Create streamable HTTP connection without session ID
   • Initialize session with session.initialize()
   • Extract session ID from HTTP response headers via getSessionId()
   • Store session ID in closure variable

2. Subsequent Tool Invocations:

   • session_id is set from previous call
   • Create streamable HTTP connection with existing session ID
   • Skip initialization step
   • Session already established and ready for tool calls

HTTP Headers

The session ID is transmitted via the mcp-session-id HTTP header:

• Server returns it in the initialization response
• Client includes it in subsequent requests
• Transport layer manages header injection automatically

Migration

No breaking changes. Existing code continues to work, but benefits from improved performance through automatic session reuse.

[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]

The logging statement includes user-controlled data (session_id) which could potentially contain sensitive information or be used for log injection. While session IDs are typically GUIDs, it's better to validate the format or use structured logging to prevent any potential log injection issues.

Consider using structured logging with separate fields, or at minimum validating that the session_id matches expected GUID format before logging.

[Copilot]

The logging statement includes the session_id which could potentially contain sensitive information or be used for log injection. While session IDs are typically GUIDs, it's better to validate the format or use structured logging to prevent any potential log injection issues.

Consider using structured logging with separate fields, or at minimum validating that the session_id matches expected GUID format before logging.

[Copilot]

The docstring parameter session_id is missing from the Args section. The function accepts a session_id parameter but it's not documented in the docstring.

Add documentation for the session_id parameter to explain that it's an optional session ID for reusing an existing MCP session.

Suggested change

	authentication, or other HTTP settings, create an httpx.AsyncClient and pass it here.
	authentication, or other HTTP settings, create an httpx.AsyncClient and pass it here.
	session_id: Optional session ID for reusing an existing MCP session instead of creating a new one.

[Copilot]

Duplicate import: AsyncGenerator is imported from collections.abc here but is already imported from typing at the top of the file (line 6). This creates ambiguity and is unnecessary.

Remove this import line as the typing.AsyncGenerator import at the top of the file is sufficient for the entire module.

Suggested change

from collections.abc import AsyncGenerator

[Copilot]

Duplicate import: asynccontextmanager is imported from contextlib here but is already imported at the top of the file (line 4). This creates ambiguity and is unnecessary.

Remove this import line as the contextlib.asynccontextmanager import at the top of the file is sufficient for the entire module.

Suggested change

from contextlib import asynccontextmanager

[Copilot]

Pull request overview

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

[Copilot]

Empty TYPE_CHECKING block with no imports. This block should either contain type-only imports or be removed entirely.

Suggested change

	if TYPE_CHECKING:
	pass

[Copilot]

The session_lock is created once at module level (per resource config) but each tool function creates its own AsyncExitStack and resources. If two tool invocations from different tools execute concurrently:

1. First tool acquires lock, creates session, releases lock
2. Second tool acquires lock, creates its own separate ClientSession with the same session_id
3. Both tools are now using different ClientSession objects but with the same session_id

This could lead to undefined behavior if the MCP protocol doesn't support multiple concurrent ClientSessions with the same session_id. Consider either:

1. Sharing a single ClientSession across all tool invocations (requiring it to live outside the tool_fn scope)
2. Using a session pool with proper lifecycle management
3. Documenting that the MCP protocol supports this usage pattern

[Copilot]

The lock is held during the entire session initialization process including entering the async context and creating the ClientSession. This prevents concurrent tool invocations from proceeding, even though they could potentially share the same session after initialization completes.

Consider narrowing the lock scope to only protect the session_id check and assignment. The session initialization and ClientSession creation could happen outside the lock, with only the session_id variable access being synchronized. This would allow better concurrency while still ensuring thread-safe session ID management.

[Copilot]

The function-level docstring states "Each tool manages its own session lifecycle - creating, using, and cleaning up the MCP connection within the tool invocation." However, the implementation has changed to reuse sessions across multiple tool invocations via the closure variable session_id. This docstring should be updated to reflect that sessions are now shared and reused across all tool invocations from the same resource config.

Suggested change

	Each tool manages its own session lifecycle - creating, using, and cleaning up
	the MCP connection within the tool invocation.
	MCP sessions are shared and reused across all tool invocations derived from the
	same resource configuration. A shared MCP client/transport (and associated
	session identifier) manages the underlying connection lifecycle rather than
	each tool invocation creating and tearing down its own session.

[Copilot]

Missing error handling for session initialization failure. If session.initialize() throws an exception, the session_id remains None but the lock is released, allowing subsequent tool invocations to attempt initialization again without any coordination. This could lead to multiple concurrent initialization attempts.

Consider adding error handling that either:

1. Sets a flag to indicate initialization failed so subsequent calls don't retry immediately
2. Re-raises the error after cleanup
3. Implements retry logic with exponential backoff

Suggested change

	await session.initialize()
	session_id = getSessionId()
	logger.info(f"session {session_id} created")
	try:
	await session.initialize()
	session_id = getSessionId()
	logger.info(f"session {session_id} created")
	except Exception:
	# Ensure we do not keep a partially initialized session_id
	session_id = None
	logger.exception("Failed to initialize MCP session")
	raise

[Copilot]

The type: ignore[assignment] comment indicates that session_id is being assigned to a private or incompatible attribute on StreamableHTTPTransport. This suggests the implementation is relying on internal implementation details of the transport that may not be part of its public API. If StreamableHTTPTransport doesn't officially support setting session_id, this could break in future versions of the mcp library.

Consider either:

1. Checking if there's an official API for session reuse in the mcp library
2. Adding a comment explaining why this is necessary and what the risk is
3. Implementing a more robust approach that doesn't rely on internal attributes

Suggested change

	transport.session_id = session_id # type: ignore[assignment]
	# Note: StreamableHTTPTransport.session_id is not part of a documented public API
	# in type hints, and may change in future versions of the mcp library. We only
	# set it if present to enable session reuse while avoiding hard runtime failures
	# if the attribute is removed or renamed.
	if session_id is not None and hasattr(transport, "session_id"):
	setattr(transport, "session_id", session_id)