feat: add async wrapper for mcp tools config (#652)

Author: radugheo

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.7.17"
+version = "0.8.0"
 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"

src/uipath_langchain/agent/tools/__init__.py

@@ -5,6 +5,7 @@
 from .extraction_tool import create_ixp_extraction_tool
 from .integration_tool import create_integration_tool
 from .ixp_escalation_tool import create_ixp_escalation_tool
+from .mcp import open_mcp_tools
 from .process_tool import create_process_tool
 from .tool_factory import (
     create_tools_from_resources,
@@ -15,6 +16,7 @@
     "create_tools_from_resources",
     "create_tool_node",
     "create_context_tool",
+    "open_mcp_tools",
     "create_process_tool",
     "create_integration_tool",
     "create_escalation_tool",

src/uipath_langchain/agent/tools/mcp/__init__.py

@@ -3,16 +3,16 @@
 from .mcp_client import McpClient, SessionInfoFactory
 from .mcp_tool import (
     create_mcp_tools,
-    create_mcp_tools_from_agent,
-    create_mcp_tools_from_metadata_for_mcp_server,
+    create_mcp_tools_and_clients,
+    open_mcp_tools,
 )
 from .streamable_http import SessionInfo
 
 __all__ = [
     "McpClient",
     "SessionInfo",
     "SessionInfoFactory",
+    "create_mcp_tools_and_clients",
+    "open_mcp_tools",
     "create_mcp_tools",
-    "create_mcp_tools_from_agent",
-    "create_mcp_tools_from_metadata_for_mcp_server",
 ]

src/uipath_langchain/agent/tools/mcp/claude.md

@@ -32,9 +32,9 @@ src/uipath_langchain/agent/tools/mcp/
 ```python
 from .mcp_client import McpClient, SessionInfoFactory
 from .mcp_tool import (
+    create_mcp_tools_and_clients,
+    open_mcp_tools,
     create_mcp_tools,
-    create_mcp_tools_from_agent,
-    create_mcp_tools_from_metadata_for_mcp_server,
 )
 from .streamable_http import SessionInfo
 ```
@@ -228,12 +228,12 @@ On session reinitialization (404 retry), only steps 5-6 repeat.
 
 ### Tool Factory Functions
 
-#### `create_mcp_tools_from_agent(agent, session_info_factory)` → `tuple[list[BaseTool], list[McpClient]]`
+#### `create_mcp_tools_and_clients(agent, session_info_factory)` → `tuple[list[BaseTool], list[McpClient]]`
 
 **Primary factory function** for creating MCP tools from a LowCodeAgentDefinition.
 
 ```python
-async def create_mcp_tools_from_agent(
+async def create_mcp_tools_and_clients(
     agent: LowCodeAgentDefinition,
     session_info_factory: SessionInfoFactory | None = None,
 ) -> tuple[list[BaseTool], list[McpClient]]:
@@ -245,23 +245,23 @@ defaults to the base `SessionInfoFactory`.  Pass a custom factory (e.g.
 
 **Usage:**
 ```python
-tools, clients = await create_mcp_tools_from_agent(agent, session_info_factory=factory)
+tools, clients = await create_mcp_tools_and_clients(agent, session_info_factory=factory)
 try:
     # Use tools...
 finally:
     for client in clients:
         await client.dispose()
 ```
 
-#### `create_mcp_tools_from_metadata_for_mcp_server(config, mcpClient)` → `list[BaseTool]`
+#### `create_mcp_tools(config, mcpClient)` → `list[BaseTool]`
 
 Creates tools for a single MCP resource config using an existing McpClient.
 
-#### `create_mcp_tools(config)` → Context Manager
+#### `open_mcp_tools(config)` → Context Manager
 
-Async context manager that creates live MCP sessions using the **upstream SDK's**
-`mcp.client.streamable_http.streamable_http_client` (not our local copy).  This
-is a simpler path that does not support `SessionInfo`.
+Async context manager that wraps `create_mcp_tools_and_clients()` with automatic
+client lifecycle management.  Yields a list of `BaseTool` instances and
+disposes all `McpClient` instances on exit.
 
 ## Two-Phase Initialization
 
@@ -487,7 +487,7 @@ and `SessionInfo` instance are all reused.
 uipath-langchain (this package)
 ├── streamable_http.py  → SessionInfo (base class)
 ├── mcp_client.py       → SessionInfoFactory (base factory)
-└── mcp_tool.py         → create_mcp_tools_from_agent(session_info_factory=...)
+└── mcp_tool.py         → create_mcp_tools_and_clients(session_info_factory=...)
 
 uipath-agents (consumer)
 ├── session_info_debug_state.py
@@ -592,7 +592,7 @@ When the upstream MCP SDK changes its transport:
 2. Override `get_session_id` and/or `set_session_id` for custom behavior
 3. Create a corresponding factory that inherits `SessionInfoFactory`
 4. The factory receives `McpServer` — use its `slug`, `folder_key`, etc.
-5. Pass the factory to `create_mcp_tools_from_agent(session_info_factory=...)`
+5. Pass the factory to `create_mcp_tools_and_clients(session_info_factory=...)`
 
 ## Related Files
 

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

@@ -1,24 +1,15 @@
-import asyncio
 import logging
-import os
-from collections import Counter, defaultdict
 from contextlib import AsyncExitStack, asynccontextmanager
-from itertools import chain
 from typing import Any, AsyncGenerator
 
-import httpx
 from langchain_core.tools import BaseTool
-from uipath._utils._ssl_context import get_httpx_client_kwargs
 from uipath.agent.models.agent import (
     AgentMcpResourceConfig,
     AgentMcpTool,
     DynamicToolsMode,
-    LowCodeAgentDefinition,
 )
 from uipath.eval.mocks import mockable
-from uipath.runtime.errors import UiPathErrorCategory
 
-from uipath_langchain.agent.exceptions import AgentStartupError, AgentStartupErrorCode
 from uipath_langchain.agent.tools.base_uipath_structured_tool import (
     BaseUiPathStructuredTool,
 )
@@ -29,97 +20,29 @@
 logger: logging.Logger = logging.getLogger(__name__)
 
 
-def _deduplicate_tools(tools: list[BaseTool]) -> list[BaseTool]:
-    """Deduplicate tools by appending numeric suffix to duplicate names."""
-    counts = Counter(tool.name for tool in tools)
-    seen: defaultdict[str, int] = defaultdict(int)
-
-    for tool in tools:
-        if counts[tool.name] > 1:
-            seen[tool.name] += 1
-            tool.name = f"{tool.name}_{seen[tool.name]}"
-
-    return tools
-
-
-def _filter_tools(tools: list[BaseTool], cfg: AgentMcpResourceConfig) -> list[BaseTool]:
-    """Filter tools to only include those in available_tools."""
-    allowed = {t.name for t in cfg.available_tools}
-    return [t for t in tools if t.name in allowed]
-
-
 @asynccontextmanager
-async def create_mcp_tools(
-    config: AgentMcpResourceConfig | list[AgentMcpResourceConfig],
-    max_concurrency: int = 5,
+async def open_mcp_tools(
+    config: list[AgentMcpResourceConfig],
 ) -> AsyncGenerator[list[BaseTool], None]:
-    """Connect to UiPath MCP server(s) and yield LangChain-compatible tools."""
-    if not (base_url := os.getenv("UIPATH_URL")):
-        raise AgentStartupError(
-            code=AgentStartupErrorCode.INVALID_TOOL_CONFIG,
-            title="Missing UIPATH_URL",
-            detail="UIPATH_URL environment variable is not set.",
-            category=UiPathErrorCategory.SYSTEM,
-        )
-    if not (access_token := os.getenv("UIPATH_ACCESS_TOKEN")):
-        raise AgentStartupError(
-            code=AgentStartupErrorCode.INVALID_TOOL_CONFIG,
-            title="Missing UIPATH_ACCESS_TOKEN",
-            detail="UIPATH_ACCESS_TOKEN environment variable is not set.",
-            category=UiPathErrorCategory.SYSTEM,
-        )
+    """Connect to UiPath MCP server(s) via McpClient and yield LangChain-compatible tools.
 
-    configs = config if isinstance(config, list) else [config]
-    enabled = [c for c in configs if c.is_enabled is not False]
-
-    if not enabled:
-        yield []
-        return
-
-    base_url = base_url.rstrip("/")
-    semaphore = asyncio.Semaphore(max_concurrency)
-
-    default_client_kwargs = get_httpx_client_kwargs()
-    client_kwargs = {
-        **default_client_kwargs,
-        "headers": {"Authorization": f"Bearer {access_token}"},
-        "timeout": httpx.Timeout(60),
-    }
-
-    # Lazy import to improve cold start time
-    from langchain_mcp_adapters.tools import load_mcp_tools
-    from mcp import ClientSession
-    from mcp.client.streamable_http import streamable_http_client
-
-    async def init_session(
-        session: ClientSession, cfg: AgentMcpResourceConfig
-    ) -> list[BaseTool]:
-        async with semaphore:
-            await session.initialize()
-            tools = await load_mcp_tools(session)
-            for tool in tools:
-                tool.metadata = {"tool_type": "mcp", "display_name": tool.name}
-            return _filter_tools(tools, cfg)
-
-    async def create_session(
-        stack: AsyncExitStack, cfg: AgentMcpResourceConfig
-    ) -> ClientSession:
-        url = f"{base_url}/agenthub_/mcp/{cfg.folder_path}/{cfg.slug}"
-        http_client = await stack.enter_async_context(
-            httpx.AsyncClient(**client_kwargs)
-        )
-        read, write, _ = await stack.enter_async_context(
-            streamable_http_client(url=url, http_client=http_client)
-        )
-        return await stack.enter_async_context(ClientSession(read, write))
+    Wraps create_mcp_tools_and_clients() with automatic client lifecycle management.
+    Tools are lazily initialized on first call via the UiPath SDK.
 
+    Args:
+        config: List of MCP resource configurations.
+
+    Yields:
+        List of BaseTool instances for all enabled MCP resources.
+    """
     async with AsyncExitStack() as stack:
-        sessions = [(await create_session(stack, cfg), cfg) for cfg in enabled]
-        results = await asyncio.gather(*[init_session(s, cfg) for s, cfg in sessions])
-        yield _deduplicate_tools(list(chain.from_iterable(results)))
+        tools, clients = await create_mcp_tools_and_clients(config)
+        for client in clients:
+            stack.push_async_callback(client.dispose)
+        yield tools
 
 
-async def create_mcp_tools_from_metadata_for_mcp_server(
+async def create_mcp_tools(
     config: AgentMcpResourceConfig,
     mcpClient: McpClient,
 ) -> list[BaseTool]:
@@ -242,24 +165,25 @@ async def tool_fn(**kwargs: Any) -> Any:
     return tool_fn
 
 
-async def create_mcp_tools_from_agent(
-    agent: LowCodeAgentDefinition,
+async def create_mcp_tools_and_clients(
+    resources: list[AgentMcpResourceConfig],
     session_info_factory: SessionInfoFactory | None = None,
     terminate_on_close: bool = True,
 ) -> tuple[list[BaseTool], list[McpClient]]:
-    """Create MCP tools from a LowCodeAgentDefinition.
+    """Create MCP tools from a list of MCP resource configurations.
 
-    Iterates over all MCP resources in the agent definition and creates tools
-    for each enabled MCP server. Each MCP server gets its own McpClient instance.
+    Iterates over all MCP resources and creates tools for each enabled MCP
+    server. Each MCP server gets its own McpClient instance.
 
     The MCP server URL is loaded lazily on first tool call via the UiPath SDK,
     using environment variables (UIPATH_URL, UIPATH_ACCESS_TOKEN).
 
     Args:
-        agent: The agent definition containing MCP resources.
+        resources: List of MCP resource configurations.
         session_info_factory: Factory for creating SessionInfo instances.
             Defaults to the base ``SessionInfoFactory``.  Pass
             ``SessionInfoDebugStateFactory()`` for playground mode.
+        terminate_on_close: Whether to terminate the MCP session on close.
 
     Returns:
         A tuple of (tools, mcp_clients) where:
@@ -273,10 +197,7 @@ async def create_mcp_tools_from_agent(
     tools: list[BaseTool] = []
     clients: list[McpClient] = []
 
-    for resource in agent.resources:
-        if not isinstance(resource, AgentMcpResourceConfig):
-            continue
-
+    for resource in resources:
         if resource.is_enabled is False:
             logger.info(f"Skipping disabled MCP resource '{resource.name}'")
             continue
@@ -290,9 +211,7 @@ async def create_mcp_tools_from_agent(
         )
         clients.append(mcpClient)
 
-        resource_tools = await create_mcp_tools_from_metadata_for_mcp_server(
-            resource, mcpClient
-        )
+        resource_tools = await create_mcp_tools(resource, mcpClient)
         tools.extend(resource_tools)
         logger.info(
             f"Created {len(resource_tools)} tools for MCP resource '{resource.name}'"

tests/agent/tools/test_mcp/claude.md

@@ -328,7 +328,7 @@ Note: All tests use `patch("uipath_langchain.agent.tools.mcp.mcp_tool.UiPath")`
 **Assertions:**
 ```python
 with patch(..., return_value=mock_uipath_class):
-    tools, clients = await create_mcp_tools_from_agent(agent)
+    tools, clients = await create_mcp_tools_and_clients(agent)
 assert len(tools) == 3  # 2 from server 1 + 1 from server 2
 ```
 
@@ -339,7 +339,7 @@ assert len(tools) == 3  # 2 from server 1 + 1 from server 2
 **Assertions:**
 ```python
 with patch(..., return_value=mock_uipath_class):
-    tools, clients = await create_mcp_tools_from_agent(agent)
+    tools, clients = await create_mcp_tools_and_clients(agent)
 assert len(clients) == 2  # One per MCP server
 ```
 
@@ -350,7 +350,7 @@ assert len(clients) == 2  # One per MCP server
 **Assertions:**
 ```python
 with patch(..., return_value=mock_uipath_class):
-    tools, clients = await create_mcp_tools_from_agent(agent)
+    tools, clients = await create_mcp_tools_and_clients(agent)
 assert len(tools) == 1  # Only enabled server's tool
 assert tools[0].name == "enabled_tool"
 ```
@@ -362,7 +362,7 @@ assert tools[0].name == "enabled_tool"
 **Assertions:**
 ```python
 with patch(..., return_value=mock_uipath_class):
-    tools, clients = await create_mcp_tools_from_agent(agent)
+    tools, clients = await create_mcp_tools_and_clients(agent)
 assert tools == []
 assert clients == []
 ```
@@ -375,7 +375,7 @@ assert clients == []
 ```python
 with patch(..., return_value=mock_sdk_no_url):
     with pytest.raises(ValueError, match="has no URL configured"):
-        await create_mcp_tools_from_agent(agent)
+        await create_mcp_tools_and_clients(agent)
 ```
 
 #### test_tools_have_correct_metadata
@@ -548,7 +548,7 @@ assert tool_call_count[0] == 3
 assert initialize_count[0] == 1  # Session reused
 ```
 
-### Testing create_mcp_tools_from_agent
+### Testing create_mcp_tools_and_clients
 
 The function uses lazy SDK initialization (`sdk = UiPath()`), so we patch the `UiPath` class:
 
@@ -570,7 +570,7 @@ async def test_example(self, agent_fixture, mock_uipath_class):
         "uipath_langchain.agent.tools.mcp.mcp_tool.UiPath",
         return_value=mock_uipath_class,
     ):
-        tools, clients = await create_mcp_tools_from_agent(agent_fixture)
+        tools, clients = await create_mcp_tools_and_clients(agent_fixture)
 ```
 
 ## Debugging Failed Tests