diff --git a/pyproject.toml b/pyproject.toml index fb1ce3d7a..d0ccdc3a8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath-langchain" -version = "0.16.16" +version = "0.17.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" @@ -18,11 +18,11 @@ dependencies = [ "pydantic-settings>=2.6.0", "python-dotenv>=1.0.1", "httpx>=0.27.0", + "httpx2>=2.5.0, <2.10.0", "openinference-instrumentation-langchain>=0.1.69, <0.2.0", "jsonschema-pydantic-converter>=0.4.0", "jsonpath-ng>=1.7.0", - "mcp==1.26.0", - "langchain-mcp-adapters==0.2.1", + "mcp==2.0.0", "pillow>=12.1.1", "rdflib>=7.0.0, <8.0.0", "a2a-sdk>=1.1.2,<2.0.0", @@ -91,6 +91,11 @@ dev = [ "rust-just>=1.39.0", "types-protobuf<7", "packaging>=24.0", + # tests/agent/tools/test_mcp/real_server.py hosts real MCP servers over real + # HTTP. Both arrive transitively via `mcp`, but the tests import them + # directly, so they are declared here rather than relied on by accident. + "starlette>=0.41.3", + "uvicorn>=0.30.0", ] [tool.hatch.build.targets.wheel] diff --git a/samples/oauth-external-apps-agent/README.md b/samples/oauth-external-apps-agent/README.md index a5f392099..ef384e638 100644 --- a/samples/oauth-external-apps-agent/README.md +++ b/samples/oauth-external-apps-agent/README.md @@ -45,9 +45,9 @@ The workflow follows a ReAct pattern: - Python 3.11+ - `uipath-langchain` -- `langchain-mcp-adapters` +- MCP Python SDK 2.0 - `langgraph` -- `httpx` +- `httpx2` - `python-dotenv` - UiPath OAuth credentials and MCP server URL in environment - UiPath external application configured with `OR.Jobs` scope (or appropriate scope for your MCP server) @@ -81,4 +81,3 @@ For debugging issues: uipath run agent --debug '{"task": "What is 2 + 2?"}' ``` - diff --git a/samples/oauth-external-apps-agent/main.py b/samples/oauth-external-apps-agent/main.py index 97fd32ce9..3210a3f39 100644 --- a/samples/oauth-external-apps-agent/main.py +++ b/samples/oauth-external-apps-agent/main.py @@ -1,21 +1,21 @@ import os -import dotenv -import httpx from contextlib import asynccontextmanager -from typing import Optional, Literal +from typing import Literal, Optional -from pydantic import BaseModel -from langgraph.graph import StateGraph, START, END -from langgraph.types import Command +import dotenv +import httpx2 from langchain.agents import create_agent -from langchain.messages import SystemMessage, HumanMessage - -from uipath_langchain.chat.models import UiPathChat -from langchain_mcp_adapters.tools import load_mcp_tools +from langchain.messages import HumanMessage, SystemMessage +from langgraph.graph import END, START, StateGraph +from langgraph.types import Command from mcp import ClientSession -from mcp.client.streamable_http import streamablehttp_client +from mcp.client.streamable_http import streamable_http_client +from pydantic import BaseModel from uipath.platform import UiPath +from uipath_langchain.agent.tools.mcp import load_mcp_tools +from uipath_langchain.chat.models import UiPathChat + dotenv.load_dotenv() UIPATH_CLIENT_ID = "EXTERNAL_APP_CLIENT_ID_HERE" @@ -24,17 +24,21 @@ UIPATH_URL = "base_url" UIPATH_MCP_SERVER_URL = os.getenv("UIPATH_MCP_SERVER_URL") + class GraphInput(BaseModel): task: str + class GraphOutput(BaseModel): result: str + class State(BaseModel): task: str access_token: Optional[str] = os.getenv("UIPATH_ACCESS_TOKEN") result: Optional[str] = None + async def fetch_new_access_token(state: State) -> Command: try: UiPath( @@ -46,44 +50,58 @@ async def fetch_new_access_token(state: State) -> Command: return Command(update={"access_token": os.getenv("UIPATH_ACCESS_TOKEN")}) except Exception as e: - raise Exception(f"Failed to initialize UiPath SDK: {str(e)}") + raise Exception(f"Failed to initialize UiPath SDK: {str(e)}") from e + @asynccontextmanager async def agent_mcp(access_token: str): - async with streamablehttp_client( - url=UIPATH_MCP_SERVER_URL, + async with httpx2.AsyncClient( headers={"Authorization": f"Bearer {access_token}"}, - timeout=60, - ) as (read, write, _): - async with ClientSession(read, write) as session: - await session.initialize() - tools = await load_mcp_tools(session) - model = UiPathChat(model="anthropic.claude-3-5-sonnet-20240620-v1:0") - agent = create_agent(model, tools=tools) - yield agent + timeout=httpx2.Timeout(60), + follow_redirects=True, + ) as http_client: + async with streamable_http_client( + url=UIPATH_MCP_SERVER_URL, + http_client=http_client, + ) as (read, write): + async with ClientSession(read, write) as session: + await session.initialize() + tools = await load_mcp_tools(session) + model = UiPathChat(model="anthropic.claude-3-5-sonnet-20240620-v1:0") + agent = create_agent(model, tools=tools) + yield agent + async def connect_to_mcp(state: State) -> Command: try: async with agent_mcp(state.access_token) as agent: - agent_response = await agent.ainvoke({ - "messages": [ - SystemMessage(content="You are a helpful assistant."), - HumanMessage(content=state.task), - ], - }) + agent_response = await agent.ainvoke( + { + "messages": [ + SystemMessage(content="You are a helpful assistant."), + HumanMessage(content=state.task), + ], + } + ) return Command(update={"result": agent_response["messages"][-1].content}) except ExceptionGroup as e: for error in e.exceptions: - if isinstance(error, httpx.HTTPStatusError) and error.response.status_code == 401: + if ( + isinstance(error, httpx2.HTTPStatusError) + and error.response.status_code == 401 + ): return Command(update={"access_token": None}) raise + def route_start(state: State) -> Literal["fetch_new_access_token", "connect_to_mcp"]: return "fetch_new_access_token" if state.access_token is None else "connect_to_mcp" + def route_after_connect(state: State): return "fetch_new_access_token" if state.access_token is None else END + builder = StateGraph(State, input=GraphInput, output=GraphOutput) builder.add_node("fetch_new_access_token", fetch_new_access_token) builder.add_node("connect_to_mcp", connect_to_mcp) diff --git a/samples/oauth-external-apps-agent/pyproject.toml b/samples/oauth-external-apps-agent/pyproject.toml index e02acf22f..fe4332739 100644 --- a/samples/oauth-external-apps-agent/pyproject.toml +++ b/samples/oauth-external-apps-agent/pyproject.toml @@ -9,7 +9,8 @@ dependencies = [ "langgraph>=1.0.4", "python-dotenv>=1.0.0", "anthropic>=0.57.1", - "langchain-mcp-adapters>=0.1.14", + "httpx2>=2.5.0,<2.10.0", + "mcp==2.0.0", "mypy>=1.17.1", "uipath", "uipath-langchain", diff --git a/samples/simple-local-mcp/README.md b/samples/simple-local-mcp/README.md index 666eb28f7..454d9cd21 100644 --- a/samples/simple-local-mcp/README.md +++ b/samples/simple-local-mcp/README.md @@ -44,7 +44,7 @@ The workflow follows a ReAct pattern: - Python 3.11+ - `langchain-anthropic` -- `langchain-mcp-adapters` +- MCP Python SDK 2.0 - `langgraph` - Anthropic API key set as an environment variable @@ -91,5 +91,5 @@ For debugging issues: To add a new tool: 1. Create a new MCP-compatible server (similar to math_server.py) -2. Add it to the MultiServerMCPClient configuration dictionary +2. Add its script to the server list in `make_graph` 3. The agent will automatically discover and use the new tool's capabilities diff --git a/samples/simple-local-mcp/pyproject.toml b/samples/simple-local-mcp/pyproject.toml index 99de5aa9d..4f4e32c56 100644 --- a/samples/simple-local-mcp/pyproject.toml +++ b/samples/simple-local-mcp/pyproject.toml @@ -5,8 +5,7 @@ description = "Math and Weather Local MCP Server Agent" authors = [{ name = "John Doe", email = "john.doe@myemail.com" }] dependencies = [ "langchain-anthropic>=1.2.0", - "langchain-mcp-adapters>=0.1.14", - "mcp>=1.15.0", + "mcp==2.0.0", "uipath", "uipath-langchain", ] @@ -16,4 +15,3 @@ requires-python = ">=3.11" dev = [ "uipath-dev", ] - diff --git a/samples/simple-local-mcp/src/simple-local-mcp/graph.py b/samples/simple-local-mcp/src/simple-local-mcp/graph.py index 94cf0ea74..3b10cf03f 100644 --- a/samples/simple-local-mcp/src/simple-local-mcp/graph.py +++ b/samples/simple-local-mcp/src/simple-local-mcp/graph.py @@ -1,26 +1,32 @@ import sys -from contextlib import asynccontextmanager +from contextlib import AsyncExitStack, asynccontextmanager -from langchain_anthropic import ChatAnthropic -from langchain_mcp_adapters.client import MultiServerMCPClient from langchain.agents import create_agent +from langchain_anthropic import ChatAnthropic +from mcp import ClientSession +from mcp.client.stdio import StdioServerParameters, stdio_client + +from uipath_langchain.agent.tools.mcp import load_mcp_tools model = ChatAnthropic(model="claude-3-7-sonnet-latest") @asynccontextmanager async def make_graph(): - client = MultiServerMCPClient({ - "math": { - "command": sys.executable, - "args": ["src/simple-local-mcp/math_server.py"], - "transport": "stdio", - }, - "weather": { - "command": sys.executable, - "args": ["src/simple-local-mcp/weather_server.py"], - "transport": "stdio", - }, - }) - agent = create_agent(model, await client.get_tools()) - yield agent + async with AsyncExitStack() as stack: + tools = [] + for script in ("math_server.py", "weather_server.py"): + read, write = await stack.enter_async_context( + stdio_client( + StdioServerParameters( + command=sys.executable, + args=[f"src/simple-local-mcp/{script}"], + ) + ) + ) + session = await stack.enter_async_context(ClientSession(read, write)) + await session.initialize() + tools.extend(await load_mcp_tools(session)) + + agent = create_agent(model, tools=tools) + yield agent diff --git a/samples/simple-local-mcp/src/simple-local-mcp/math_server.py b/samples/simple-local-mcp/src/simple-local-mcp/math_server.py index 4b781b56b..1d37cca5c 100644 --- a/samples/simple-local-mcp/src/simple-local-mcp/math_server.py +++ b/samples/simple-local-mcp/src/simple-local-mcp/math_server.py @@ -1,11 +1,11 @@ import logging -from mcp.server.fastmcp import FastMCP +from mcp.server.mcpserver import MCPServer logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) -mcp = FastMCP("Math") +mcp = MCPServer("Math") @mcp.tool() def add(a: int, b: int) -> int: diff --git a/samples/simple-local-mcp/src/simple-local-mcp/weather_server.py b/samples/simple-local-mcp/src/simple-local-mcp/weather_server.py index b2045f310..1eb913205 100644 --- a/samples/simple-local-mcp/src/simple-local-mcp/weather_server.py +++ b/samples/simple-local-mcp/src/simple-local-mcp/weather_server.py @@ -1,11 +1,11 @@ import logging -from mcp.server.fastmcp import FastMCP +from mcp.server.mcpserver import MCPServer logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) -mcp = FastMCP("Weather") +mcp = MCPServer("Weather") @mcp.tool() async def get_weather(location: str) -> str: diff --git a/samples/simple-remote-mcp/README.md b/samples/simple-remote-mcp/README.md index 6d2106d27..26fd67fca 100644 --- a/samples/simple-remote-mcp/README.md +++ b/samples/simple-remote-mcp/README.md @@ -42,7 +42,7 @@ The workflow follows a ReAct pattern: - Python 3.11+ - `langchain-anthropic` -- `langchain-mcp-adapters` +- MCP Python SDK 2.0 - `langgraph` - Anthropic API key set as an environment variable @@ -70,4 +70,3 @@ For debugging issues: uipath run agent --debug '{"messages": [{"type": "human", "content": "What is 2+2"}]}' ``` - diff --git a/samples/simple-remote-mcp/main.py b/samples/simple-remote-mcp/main.py index a6f1759f5..759740a27 100644 --- a/samples/simple-remote-mcp/main.py +++ b/samples/simple-remote-mcp/main.py @@ -1,28 +1,35 @@ import os from typing import Any -from langgraph.graph import StateGraph, MessagesState, START, END + +import httpx2 from langchain.agents import create_agent from langchain_anthropic import ChatAnthropic -from langchain_mcp_adapters.tools import load_mcp_tools +from langgraph.graph import END, START, MessagesState, StateGraph from mcp import ClientSession -from mcp.client.streamable_http import streamablehttp_client +from mcp.client.streamable_http import streamable_http_client + +from uipath_langchain.agent.tools.mcp import load_mcp_tools async def mcp_client(state: MessagesState) -> dict[str, Any]: """Agent node that connects to MCP server and processes messages.""" - async with streamablehttp_client( - url=os.getenv("UIPATH_MCP_SERVER_URL"), + async with httpx2.AsyncClient( headers={"Authorization": f"Bearer {os.getenv('UIPATH_ACCESS_TOKEN')}"}, - timeout=60, - ) as (read, write, _): - async with ClientSession(read, write) as session: - await session.initialize() - tools = await load_mcp_tools(session) - print(f"Loaded {len(tools)} tools from MCP server") - model = ChatAnthropic(model="claude-3-7-sonnet-latest") - agent = create_agent(model, tools=tools) - result = await agent.ainvoke(state) - return result + timeout=httpx2.Timeout(60), + follow_redirects=True, + ) as http_client: + async with streamable_http_client( + url=os.environ["UIPATH_MCP_SERVER_URL"], + http_client=http_client, + ) as (read, write): + async with ClientSession(read, write) as session: + await session.initialize() + tools = await load_mcp_tools(session) + print(f"Loaded {len(tools)} tools from MCP server") + model = ChatAnthropic(model="claude-3-7-sonnet-latest") + agent = create_agent(model, tools=tools) + result = await agent.ainvoke(state) + return result builder = StateGraph(MessagesState) diff --git a/samples/simple-remote-mcp/pyproject.toml b/samples/simple-remote-mcp/pyproject.toml index c43f76114..766e78628 100644 --- a/samples/simple-remote-mcp/pyproject.toml +++ b/samples/simple-remote-mcp/pyproject.toml @@ -9,7 +9,8 @@ dependencies = [ "langgraph>=1.0.4", "python-dotenv>=1.0.0", "anthropic>=0.57.1", - "langchain-mcp-adapters>=0.1.14", + "httpx2>=2.5.0,<2.10.0", + "mcp==2.0.0", "uipath", "uipath-langchain", ] diff --git a/src/uipath_langchain/_cli/cli_new.py b/src/uipath_langchain/_cli/cli_new.py index 7d49a8f9a..bd6e6c642 100644 --- a/src/uipath_langchain/_cli/cli_new.py +++ b/src/uipath_langchain/_cli/cli_new.py @@ -11,7 +11,7 @@ # Deliberately a constant: the guard test in tests/cli/test_new.py fails on # every minor bump so the scaffold (pin, template, hints) gets reviewed # alongside the release rather than drifting silently. -UIPATH_LANGCHAIN_SCAFFOLD_MINOR = "0.16" +UIPATH_LANGCHAIN_SCAFFOLD_MINOR = "0.17" def generate_script(target_directory): diff --git a/src/uipath_langchain/agent/tools/mcp/__init__.py b/src/uipath_langchain/agent/tools/mcp/__init__.py index 7b2c33f35..db1b85209 100644 --- a/src/uipath_langchain/agent/tools/mcp/__init__.py +++ b/src/uipath_langchain/agent/tools/mcp/__init__.py @@ -6,6 +6,7 @@ create_mcp_tools_and_clients, open_mcp_tools, ) +from .session_tools import load_mcp_tools from .streamable_http import SessionInfo __all__ = [ @@ -15,4 +16,5 @@ "create_mcp_tools_and_clients", "open_mcp_tools", "create_mcp_tools", + "load_mcp_tools", ] diff --git a/src/uipath_langchain/agent/tools/mcp/claude.md b/src/uipath_langchain/agent/tools/mcp/claude.md index 107c70cbb..62a9585a8 100644 --- a/src/uipath_langchain/agent/tools/mcp/claude.md +++ b/src/uipath_langchain/agent/tools/mcp/claude.md @@ -21,10 +21,12 @@ streamable HTTP transport and provides a factory pattern for session ID tracking ``` src/uipath_langchain/agent/tools/mcp/ -├── __init__.py # Public exports -├── mcp_client.py # SessionInfoFactory, McpClient -├── mcp_tool.py # Tool factory functions -└── streamable_http.py # SessionInfo, StreamableHTTPTransport (copied from MCP SDK) +├── __init__.py # Public exports +├── mcp_client.py # SessionInfoFactory, McpClient +├── mcp_tool.py # Tool factory functions +├── protocol_strategy.py # Per-era negotiation/recovery policy +├── session_tools.py # MCP session -> LangChain tool conversion +└── streamable_http.py # SessionInfo, session identity, SDK transport adapter ``` ### Public Exports (`__init__.py`) @@ -36,51 +38,82 @@ from .mcp_tool import ( open_mcp_tools, create_mcp_tools, ) +from .session_tools import load_mcp_tools from .streamable_http import SessionInfo ``` -`streamable_http_client` is intentionally **not** exported — it is an internal -transport helper used only by `McpClient`. +`streamable_http_client`, `build_protocol_strategy`, and the `SessionIdentity` +types are intentionally **not** exported — they are internal helpers used by +`McpClient` (and by the `simple-http-mcp` testcase, which imports them by module +path to drive the same code the client uses). ## Architecture -### streamable_http.py — Local Copy of MCP SDK Transport +### streamable_http.py — Session-Aware SDK Transport Adapter -This file is a local copy of the **client-side** streamable HTTP transport from -the MCP Python SDK, adapted for session ID tracking via `SessionInfo`. +This file is a thin adapter around the **client-side** streamable HTTP transport +from MCP Python SDK 2.0. The SDK owns protocol parsing, SSE resumption, +cancellation, protocol headers, and session deletion; UiPath adds externally +persistable session ID tracking via `SessionInfo`. **Source**: [`mcp.client.streamable_http`](https://github.com/modelcontextprotocol/python-sdk/blob/main/src/mcp/client/streamable_http.py) -**Why a local copy?** - -The upstream SDK transport has no hook for observing or injecting session IDs. -We need this to support session persistence (e.g. debug state for playground -mode). The local copy adds a `SessionInfo` parameter that receives session ID -updates from the server. - -**Key differences from the upstream SDK:** - -1. **`SessionInfo` class added** — base class for session ID tracking, defined - at the top of the file. The transport delegates all session ID storage to - this object via async methods. -2. **Transport does not own session state** — `StreamableHTTPTransport` has no - `self.session_id`. All reads/writes go through `self._session_info`. -3. **`_prepare_headers` is async** — because it calls - `await self._session_info.get_session_id()`. -4. **`_maybe_extract_session_id_from_response` is async** — calls - `await self._session_info.set_session_id()` so subclasses can persist. -5. **`RequestContext` has no `session_id` field** — it was unused upstream - (headers are built from `_prepare_headers`, not from the context). -6. **`streamable_http_client` accepts `session_info` parameter** — passed - through to the transport constructor. -7. **Returns 2 values, not 3** — yields `(read_stream, write_stream)` instead - of the SDK's `(read_stream, write_stream, get_session_id_callback)`. - -**What was kept identical:** - -The overall request/response flow, SSE handling, reconnection logic, POST/GET -patterns, and error handling are structurally the same as the upstream SDK. -When updating, diff against the upstream source to understand what changed. +**Why an adapter is still needed:** + +The SDK 2.0 transport owns its in-memory session ID but has no asynchronous hook +for loading and saving UiPath debug-state sessions. The adapter installs two +`httpx2` event hooks on the client used by the SDK: + +1. Before every request, call `SessionInfo.get_session_id()` and put the ID on + the headers named by the active `SessionIdentityWire` (removing any it does + not name). +2. After every response, persist an ID returned on the wire's + `capture_response_header` through `SessionInfo.set_session_id()`. + +The hooks are removed when the adapter context exits. This keeps the UiPath +extension small and automatically picks up future SDK transport fixes instead +of maintaining another transport fork. + +#### SessionIdentityWire / SessionIdentity + +The two protocol eras identify a connection differently, so the header behaviour +is data rather than code: + +```python +@dataclass(frozen=True) +class SessionIdentityWire: + request_headers: tuple[str, ...] = (MCP_SESSION_ID,) + capture_response_header: str | None = MCP_SESSION_ID +``` + +The ID travels on the header only. Mirroring it into `params._meta` was tried +and removed: rewriting the request body needs private HTTPX internals, and +gateway routing needs nothing beyond the header. + +| Wire | Sends on | Captures from | Used by | +|------|----------|---------------|---------| +| `LEGACY_IDENTITY` | `mcp-session-id` | `mcp-session-id` | legacy handshake; also `auto` before the era resolves | +| `MODERN_IDENTITY` | `mcp-session-id` | nothing (client mints) | `2026-07-28` | + +**Both eras send the ID on `mcp-session-id`.** In the legacy era it is a real +session the server minted; in the modern era the protocol has no session, so +UiPath mints the value itself and the header carries it purely as a routing key, +which a modern server ignores. Reusing the header rather than inventing one means +the gateway needs no change — it keeps routing on the header it already uses. + +The transport is opened **before** negotiation runs, so the era-specific wire +cannot be fixed at construction time. `streamable_http_client` therefore takes a +mutable `SessionIdentity` holder and reads `identity.wire` on every request; the +strategy narrows it once `connect()` resolves the era. + +`auto` can safely open on the legacy wire: a stored ID is era-ambiguous, but both +eras send it on the same header, and a modern server never sends one back so +there is nothing to capture. + +The eras differ only in `capture_response_header`, and that difference matters. +`MODERN_IDENTITY` sets it to `None` so a proxy or gateway echoing `mcp-session-id` +back cannot overwrite the client-minted routing key mid-connection and scatter +the remaining requests across instances. #### SessionInfo @@ -90,48 +123,200 @@ Base class for MCP session ID tracking. Lives in `streamable_http.py`. class SessionInfo: def __init__(self, session_id: str | None = None) -> None: self.session_id = session_id + self.protocol_version: str | None = None async def get_session_id(self) -> str | None: ... - async def set_session_id(self, session_id: str) -> None: ... + async def set_session_id(self, session_id: str | None) -> None: ... + async def get_protocol_version(self) -> str | None: ... + async def set_protocol_version(self, protocol_version: str | None) -> None: ... ``` The base implementation stores session ID in a plain attribute. Async methods exist so subclasses (e.g. `SessionInfoDebugState` in `uipath-agents`) can add side-effects like HTTP persistence. -**Important:** The transport calls `set_session_id` during `initialize()` when -the server assigns a session ID. `McpClient._initialize_session` then reads -the value via `get_session_id` — it does not call `set_session_id` again. +`SessionInfo` is deliberately era-agnostic: it stores **the ID we persist for +this MCP server**, whoever minted it. In the legacy era that is the server's +session ID; in the modern era it is the client-minted affinity ID. This is why +reaching `2026-07-28` required no change in `uipath-agents-python` — +`SessionInfoDebugState` persists either kind unmodified. -#### StreamableHTTPTransport +`protocol_version` holds **the version the stored session was negotiated at**, +and the legacy strategy both writes and reads it. It cannot be recovered from +the wire -- responses carry only the session ID -- so a store that persists the +ID should persist this too: with it, a resumed session needs no negotiation at +all (see below). A store that does not is not broken, only slower: `None` means +"not known", and the strategy falls back to re-running the handshake. -Handles the MCP streamable HTTP protocol: POST for requests, GET for -server-initiated SSE streams, reconnection with `Last-Event-ID`, and session -termination via DELETE. +A subclass persisting externally should override all four accessors, so the ID +and its version are written and cleared together. -Key methods: +**Important:** The response hook calls `set_session_id` during `initialize()` +when the server assigns an ID. `McpClient._initialize_session` only reads the +stored value afterward. Passing `None` clears a stale session before recovery; +`LegacyHandshakeStrategy.reset` clears the version alongside it, so a +replacement session can never inherit a version it did not negotiate. -| Method | Description | -|--------|-------------| -| `_prepare_headers()` | **async** — builds headers with session ID from `SessionInfo` | -| `_maybe_extract_session_id_from_response()` | **async** — extracts session ID from response, calls `set_session_id` | -| `_handle_post_request()` | POST with JSON or SSE response handling | -| `handle_get_stream()` | GET SSE listener with auto-reconnect | -| `_handle_reconnection()` | Recursive reconnect with `Last-Event-ID` | -| `post_writer()` | Main write loop, dispatches requests to server | -| `terminate_session()` | Sends DELETE to end the session | -| `get_session_id()` | **async** — delegates to `SessionInfo.get_session_id` | +#### Upstream StreamableHTTPTransport + +MCP SDK 2.0's transport handles POST requests, the optional GET SSE channel, +`Last-Event-ID` resumption, 2026 HTTP cancellation, protocol headers, structured +errors from non-2xx responses, and session termination via DELETE. None of those +internals are duplicated locally. #### streamable_http_client (context manager) -Internal async context manager that wires up `StreamableHTTPTransport` with -memory streams and a task group. Used by `McpClient._initialize_client`. +Internal async context manager that attaches the session hooks and delegates to +the SDK context manager. Used by `McpClient._open_connection`. ```python async with streamable_http_client(url, http_client=client, session_info=info) as (read, write): session = ClientSession(read, write) ``` +The adapter yields the SDK's two transport streams unchanged. If no HTTP client +is supplied, it creates and owns an `httpx2.AsyncClient`; `McpClient` normally +supplies its long-lived authenticated client. + +--- + +### protocol_strategy.py — Per-Era Session Lifecycle + +MCP has two negotiation eras, and `ProtocolStrategy` is the seam between them. +Negotiation itself is one call either way; what genuinely differs is the session +*lifecycle*. + +| Concern | Legacy (`2024-11-05`…`2025-11-25`) | Modern (`2026-07-28`) | +|---------|------------------------------------|------------------------| +| Negotiate | `ClientSession.initialize()` | `ClientSession.discover()` | +| Identity | `mcp-session-id`, server-minted | none in-protocol; UiPath-minted affinity ID | +| Resume | adopt the stored version locally; re-handshake only when it is unknown | nothing to resume | +| Terminate on close | `DELETE` | no-op (automatic; no session ID) | +| Recoverable errors | session lost → reopen and re-handshake | only `CONNECTION_CLOSED` | + +```python +class ProtocolStrategy(Protocol): + identity: SessionIdentity + + async def connect(self, session: ClientSession, info: SessionInfo) -> None: ... + def is_recoverable(self, error: MCPError, restored_id: str | None) -> bool: ... + async def reset(self, info: SessionInfo) -> None: ... +``` + +`McpClient` calls `reset` only when the error that triggered recovery is the +server's verdict on the session, which `is_session_rejected(error)` decides. A +dropped transport is not a verdict -- the server never rejected anything -- so +the ID is kept and the reconnect resumes the warm session; anything the server +answered *about the session* clears the ID before the fresh handshake. + +The code alone cannot make that call. `CONNECTION_CLOSED` is JSON-RPC's +implementation-defined server-error code `-32000`, and the TypeScript SDK's +Streamable HTTP transport uses the same code to refuse a session it does not +know (`"Bad Request: No valid session ID provided"`). Treating every `-32000` +as a dropped connection would keep a session the server has just declared dead +and burn the retry resuming it, so a `-32000` whose message names a lost +session counts as a verdict. + +Selected by `McpClient(protocol_mode=...)` via `build_protocol_strategy`: + +- **`"legacy"` (default)** — `LegacyHandshakeStrategy`. Preserves the pre-2026 + wire behaviour exactly. +- **`"modern"`** — `ModernDiscoveryStrategy`. `server/discover` only. +- **`"auto"`** — `AutoStrategy`. Mints the affinity ID first so the probe is + pinned, runs its own `server/discover` probe (`probe_modern_era`, built on the + public `ClientSession.send_discover` / `adopt` seam), then delegates to + whichever era won. Re-resolved on every `connect`, so a server upgraded + mid-run is handled. + +The default stays `"legacy"` on purpose. Defaulting to `"auto"` would silently +move any discovery-capable UiPath MCP server to stateless `2026-07-28` and stop +issuing session IDs, breaking the playground persistence `SessionInfoDebugState` +exists for. + +#### Resuming a legacy session: adopt the version, do not renegotiate + +A restored session ID is useless without the protocol version it was negotiated +at, and that version cannot be recovered from the wire — server responses carry +only the session ID. It is therefore stored *with* the ID, and a resume installs +it locally through `ClientSession.adopt`, which is documented to touch no wire: + +```python +restored = await info.get_session_id() +if restored is None: + await self._handshake(session, info) # cold: negotiate and store the version + return +if await self._adopt_restored_session(session, info, restored): + return # no request sent at all +try: + result = await session.initialize() # version unknown: lands inside the session +except MCPError as error: + if error.code == CONNECTION_CLOSED: + raise # transport died; says nothing about the session + await self.reset(info) # stale, or this server refuses a 2nd handshake + await self._handshake(session, info) # same transport: a refused request does not close it +``` + +**Why adopting beats re-handshaking.** Whether a server accepts a second +`initialize` inside a live session is implementation-defined. The Python SDK +does. The reference TypeScript implementation refuses it outright with +`-32600 "Invalid Request: Server already initialized"`, so a client that resumes +by re-handshaking loses the persisted session on **every** run against such a +server — falling back to a cold one, and with it the gateway affinity the +session ID exists to provide. Adopting asks the server nothing, so it works +either way, and it also restores the pre-SDK-2 wire shape: a resumed run sends +only ordinary requests carrying `mcp-session-id`. + +The handshake path remains for a store written before versions were recorded +(`get_protocol_version()` returns `None`), and for a stored version this client +cannot speak on a legacy wire — a modern version, or one dropped upstream. It is +safe because the server routes purely by the session header and mints a new +session only when the header is **absent** +([`streamable_http_manager.py`](https://github.com/modelcontextprotocol/python-sdk/blob/main/src/mcp/server/streamable_http_manager.py)). +A server that ignores the header instead mints a replacement; the strategy detects +that by comparing the ID before and after and continues with the new session. + +Probing candidate versions with `send_ping` — the original approach — always +matched the *oldest* handshake version, because servers do not validate that +header against what the session negotiated. That silently downgraded every later +request and disabled the server's `2025-11-25` SSE resumability. Do not +reintroduce it: the version is remembered now, not guessed. + +#### Modern-era instance affinity + +`2026-07-28` removes `mcp-session-id` from the protocol, and AgentHub used it to +route to a warm serverless instance. `ModernDiscoveryStrategy` mints its own ID +and keeps sending it on that same header as an opaque routing key — off-spec for +the era, ignored by a modern server, and requiring **no gateway change**. Because +the client mints it *before* negotiating, it is present on the very first request +— `server/discover` included — which a server-assigned session ID never could be. + +In `auto` mode the ID is minted *before* the probe as well, so `server/discover` +reaches the same instance the tool calls will. On a serverless gateway an +unpinned probe warms one instance and the first call lands on another, which is +exactly the scatter the affinity ID exists to prevent. A legacy server never +issued that ID, so when the probe falls back the freshly minted ID is cleared +before the handshake rather than offered as a session to resume -- a routing +server would refuse it. A *restored* ID is not cleared: it may be a live legacy +session, and the handshake resumes it (or replaces it when the server rejects it). + +**Affinity is a hint, not a guarantee.** A fresh client has no ID and every +modern request is self-contained, so any instance must be able to answer any +request. Instance-local state is valid as a warm cache only; a server whose +instances hold state no peer can rebuild should stay on `protocol_mode="legacy"`. + +--- + +### session_tools.py — Session-to-LangChain Tool Conversion + +`load_mcp_tools(session)` paginates `tools/list` and returns `StructuredTool`s +bound to that session. It replaced the `langchain-mcp-adapters` dependency, which +imports `RequestContext` — removed in MCP 2. + +It is **not** a drop-in replacement for that package's function of the same name: +it returns raw MCP content blocks (camelCase, via `model_dump(by_alias=True)`) +under the default `response_format="content"`, where the old one returned +LangChain content blocks plus a `structured_content` artifact. + --- ### SessionInfoFactory @@ -160,7 +345,8 @@ package. They import `SessionInfo` and `SessionInfoFactory` from here. MCP connections for tool invocations with **two distinct initialization phases**: 1. **Client Initialization** (first call): Retrieves MCP server URL via SDK, creates the full stack -2. **Session Reinitialization** (on 404): Lightweight, reuses existing client +2. **Connection Reinitialization** (on session loss): Reuses the HTTP client, + but replaces the transport and `ClientSession` ``` ┌─────────────────────────────────────────────────────────────┐ @@ -169,9 +355,10 @@ MCP connections for tool invocations with **two distinct initialization phases** │ Configuration (immutable after __init__) │ │ ───────────────────────────────────────── │ │ _config: AgentMcpResourceConfig # Contains slug, folder │ -│ _timeout: httpx.Timeout │ +│ _timeout: httpx2.Timeout | float | None │ │ _max_retries: int │ │ _session_info_factory: SessionInfoFactory │ +│ _strategy: ProtocolStrategy # from protocol_mode │ ├─────────────────────────────────────────────────────────────┤ │ Lazy-Resolved State (set during _initialize_client) │ │ ─────────────────────────────────────────────────── │ @@ -182,34 +369,35 @@ 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) │ +│ Client State (created once, reused on connection reinit) │ │ ───────────────────────────────────────────────────── │ -│ _http_client: httpx.AsyncClient | None │ -│ _read_stream: MemoryObjectReceiveStream | None │ -│ _write_stream: MemoryObjectSendStream | None │ +│ _http_client: httpx2.AsyncClient | None │ │ _session_info: SessionInfo | None │ -│ _stack: AsyncExitStack | None │ +│ _stack: AsyncExitStack | None # HTTP client │ │ _client_initialized: bool │ ├─────────────────────────────────────────────────────────────┤ -│ Session State (can be reinitialized without recreating) │ -│ ─────────────────────────────────────────────────────── │ +│ Connection State (replaced after session loss) │ +│ ────────────────────────────────────────────── │ +│ _connection_stack: AsyncExitStack | None │ │ _session: ClientSession | None │ -│ _session_id: str | None │ ├─────────────────────────────────────────────────────────────┤ │ Public Methods │ │ ────────────── │ +│ + list_tools(force_refresh=False) -> ListToolsResult │ │ + call_tool(name, arguments) -> CallToolResult │ │ + dispose() -> None # UiPathDisposableProtocol │ -│ + session_id: str | None (property) │ +│ + get_session_id() -> str | None │ │ + is_client_initialized: bool (property) │ ├─────────────────────────────────────────────────────────────┤ │ Private Methods │ │ ─────────────── │ │ - _initialize_client() -> None # SDK + full init (once) │ -│ - _initialize_session() -> None # MCP handshake only │ +│ - _open_connection() -> None # transport + session │ +│ - _initialize_session() -> None # delegates to strategy │ │ - _ensure_session() -> ClientSession │ -│ - _reinitialize_session() -> None │ -│ - _is_session_error(error) -> bool │ +│ - _reinitialize_session(failed_session, error) -> None │ +│ - _is_recoverable_session_error(error) -> bool │ +│ + is_session_error(error) -> bool # legacy rules, public │ └─────────────────────────────────────────────────────────────┘ ``` @@ -220,11 +408,22 @@ During client initialization, `McpClient`: 1. Retrieves the `McpServer` from the UiPath SDK 2. Calls `self._session_info_factory.create_session(mcp_server)` to get a `SessionInfo` 3. Loads any existing session ID via `await session_info.get_session_id()` -4. Passes the `SessionInfo` to the local `streamable_http_client` -5. Calls `session.initialize()` — the transport calls `set_session_id` internally -6. Reads the new session ID via `await session_info.get_session_id()` - -On session reinitialization (404 retry), only steps 5-6 repeat. +4. Passes the `SessionInfo` to the local adapter, which opens the SDK transport +5. Creates a new `ClientSession` over those streams +6. Calls `strategy.connect(session, session_info)`, which negotiates for its era: + - **legacy** — adopts a restored ID whose version is known, sending nothing; + otherwise `session.initialize()`, whose response hook stores the + server-assigned ID, and whose result is stored as the version + - **modern** — mints an affinity ID if absent, then `session.discover()` + - **auto** — mints an affinity ID if absent, then `probe_modern_era(session)`; + on fallback it clears a freshly minted ID and runs the legacy `connect`; + finally narrows `identity.wire` to the resolved era +7. Reads the current session ID via `await session_info.get_session_id()` + +On recovery, the HTTP client and `SessionInfo` are reused, but the old connection +stack is closed and steps 4-7 run with a fresh transport and `ClientSession`. +This is required because SDK 2.0 makes `ClientSession.initialize()` idempotent +for the lifetime of one `ClientSession`. ### Tool Factory Functions @@ -308,8 +507,8 @@ disposes all `McpClient` instances on exit. The key design principle is separating **client initialization** from **session initialization**: ``` -Phase 1: Client Initialization (expensive, done once) -────────────────────────────────────────────────────── +Phase 1: Base Client Initialization (expensive, done once) +─────────────────────────────────────────────────────────── ┌─────────────────┐ │ UiPath SDK │ ─── Retrieves MCP server URL │ mcp.retrieve() │ and auth token (Bearer) @@ -321,23 +520,21 @@ Phase 1: Client Initialization (expensive, done once) └─────────────────┘ ┌─────────────────┐ -│ httpx.AsyncClient │ ─┐ -└─────────────────┘ │ - │ -┌─────────────────┐ │ Created once via -│ streamable_http │ ├─ AsyncExitStack -│ connection │ │ -└─────────────────┘ │ - │ -┌─────────────────┐ │ -│ ClientSession │ ─┘ +│httpx2.AsyncClient│ ─── Created once via the base AsyncExitStack +└─────────────────┘ + +Phase 2: Connection Initialization (repeated after session loss) +────────────────────────────────────────────────────────────── +┌─────────────────┐ +│ SDK transport + │ ─── Fresh connection AsyncExitStack +│ ClientSession │ └─────────────────┘ -Phase 2: Session Initialization (lightweight, can repeat) -───────────────────────────────────────────────────────── ┌─────────────────┐ │ session. │ ─── Sends initialize request -│ initialize() │ Transport calls set_session_id() +│ initialize() │ Response hook calls set_session_id() +│ │ (skipped when a stored ID *and* version +│ │ are adopted instead) └─────────────────┘ ┌─────────────────┐ │ McpClient reads │ ─── await session_info.get_session_id() @@ -348,42 +545,27 @@ Phase 2: Session Initialization (lightweight, can repeat) ### Session Lifecycle ``` - ┌──────────────┐ - │ Created │ - │ (nothing │ - │ initialized)│ - └──────┬───────┘ - │ call_tool() [first time] - ▼ - ┌──────────────┐ - │ Client │ - │ Initializing │ - │ (Phase 1) │ - └──────┬───────┘ - │ 1. UiPath SDK retrieves MCP URL - │ 2. Factory creates SessionInfo - │ 3. Creates HTTP client, streams, session - │ 4. Calls _initialize_session() - ▼ - ┌──────────────┐ - │ Session │ - │ Initializing │◄────────────────┐ - │ (Phase 2) │ │ - └──────┬───────┘ │ - │ sends initialize, │ - │ transport calls │ 404 error - │ set_session_id() │ (only reinit - ▼ │ session, - ┌──────────────┐ │ not client) - │ Active │─────────────────┘ - │ Session │ - └──────┬───────┘ - │ dispose() - ▼ - ┌──────────────┐ - │ Closed │ - │ (can reuse) │ - └──────────────┘ +┌──────────────┐ first operation ┌────────────────────┐ +│ Created │ ────────────────► │ Base client init │ +└──────────────┘ │ SDK + HTTP client │ + └─────────┬──────────┘ + │ open connection + ▼ + ┌────────────────────┐ + ┌────►│ Session init │ + │ │ transport + session│ + │ └─────────┬──────────┘ + │ │ initialize handshake + │ ▼ + │ ┌────────────────────┐ + │ │ Active session │ + │ └────┬──────────┬────┘ + │ │ │ dispose() + session error │ │ ▼ + close old + └──────────┘ ┌──────────────┐ + clear ID │ Closed │ + │ (can reuse) │ + └──────────────┘ ``` ### MCP Protocol Flow @@ -394,16 +576,16 @@ Phase 2: Session Initialization (lightweight, can repeat) Client Server │ │ │──── initialize ──────────────────►│ - │◄─── result + session-id-1 ────────│ ← transport calls set_session_id() + │◄─── result + session-id-1 ────────│ ← response hook calls set_session_id() │ │ │──── notifications/initialized ───►│ - │◄─── 204 No Content ───────────────│ + │◄─── 202 Accepted / 204 ───────────│ │ │ │──── tools/call ──────────────────►│ │◄─── result ───────────────────────│ ``` -**On 404 error (session reinitialization only):** +**On a terminated session (connection/session replacement):** ``` Client Server @@ -411,14 +593,15 @@ Client Server │──── tools/call ──────────────────►│ │◄─── 404 (session terminated) ─────│ │ │ - │ [Reuses existing HTTP client │ - │ and streamable connection] │ + │ [Closes old transport/session; │ + │ clears stale SessionInfo; │ + │ reuses existing HTTP client] │ │ │ │──── initialize ──────────────────►│ ← new session - │◄─── result + session-id-2 ────────│ (same client) + │◄─── result + session-id-2 ────────│ (same HTTP client) │ │ │──── notifications/initialized ───►│ - │◄─── 204 No Content ───────────────│ + │◄─── 202 Accepted / 204 ───────────│ │ │ │──── tools/call ──────────────────►│ ← retry │◄─── result ───────────────────────│ @@ -426,12 +609,41 @@ Client Server ### Session Error Codes -The following error codes trigger automatic session reinitialization: +Which errors trigger reinitialization is **era-specific**, decided by +`strategy.is_recoverable(error, restored_id)`. + +**Legacy** (`LegacyHandshakeStrategy`): | Code | Meaning | Source | |------|---------|--------| -| `32600` | Session terminated | HTTP 404 converted by transport | -| `-32000` | Server error | Can indicate session not found | +| `CONNECTION_CLOSED` (`-32000`) | Transport connection closed | MCP SDK dispatcher/transport | +| `INVALID_REQUEST` (`-32600`) | Session terminated/expired/invalid | SDK 2 maps a bare session-bound HTTP 404 to this error | +| `32600` | Session terminated | Compatibility with the positive code emitted by the older local transport | +| `METHOD_NOT_FOUND` (`-32601`) + `"Not Found"` | Restored session is invalid | Only while `SessionInfo` still holds the restored ID | + +`CONNECTION_CLOSED` is recovered differently from the other three: the transport +dropped but the server never rejected the session, so `McpClient` skips `reset` +and the reconnect resumes the same session. The other codes are the server's verdict, +and `reset` clears the ID before the fresh handshake. + +`-32000` is both `CONNECTION_CLOSED` and the code the TypeScript SDK's transport +refuses an unknown session with, so `is_session_rejected` reads the *message* to +tell a dropped socket from a verdict. See the `reset` discussion above. + +`INVALID_REQUEST` is retried only when its message explicitly identifies a +terminated, expired, or invalid session. An externally restored session is not +known inside a newly created SDK transport, so its first bare HTTP 404 appears +as `METHOD_NOT_FOUND`/`"Not Found"`; that exact shape is treated as recoverable +only while the restored ID is still in play. Structured JSON-RPC method errors +are not retried. + +**Modern** (`ModernDiscoveryStrategy`): `CONNECTION_CLOSED` **only**. Every +`2026-07-28` request is self-contained, so no server-side session can be lost; +retrying a session-shaped error would spend the retry budget on something a +reconnect cannot fix. + +`McpClient.is_session_error` remains public and keeps the legacy rules — it is +called by `mcp_tool._map_mcp_error`. ## Key Implementation Details @@ -468,13 +680,12 @@ The HTTP client MUST use `get_httpx_client_kwargs()` for proper SSL/proxy config ```python from uipath._utils._ssl_context import get_httpx_client_kwargs -default_client_kwargs = get_httpx_client_kwargs() +self._stack = AsyncExitStack() +await self._stack.__aenter__() +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( - **default_client_kwargs, - headers=self._headers, - timeout=self._timeout, - ) + httpx2.AsyncClient(**client_kwargs) ) ``` @@ -492,12 +703,28 @@ async def _ensure_session(self) -> ClientSession: await self._initialize_client() return self._session -async def _reinitialize_session(self) -> None: +async def _reinitialize_session( + self, + failed_session: ClientSession | None = None, + error: MCPError | None = None, +) -> None: async with self._lock: if not self._client_initialized: await self._initialize_client() else: - await self._initialize_session() # Lightweight! + # Another failing operation may arrive after recovery completed. + if failed_session is not None and self._session is not failed_session: + return + await self._close_connection_for_recovery() + # Discard persisted session state only on the server's verdict; + # a dropped connection keeps it so the reconnect resumes. + if ( + self._session_info is not None + and error is not None + and is_session_rejected(error) + ): + await self._strategy.reset(self._session_info) + await self._open_connection() ``` ### 4. No `with` Statement for AsyncExitStack @@ -509,17 +736,22 @@ Manual lifecycle management: self._stack = AsyncExitStack() await self._stack.__aenter__() # ... use stack ... -await self._stack.__aexit__(None, None, None) +await self._stack.aclose() # Wrong - exits too early async with AsyncExitStack() as stack: ... # Stack closes here! ``` -### 5. Reinitialization Reuses Client +### 5. Reinitialization Reuses the HTTP Client -On 404, only `_initialize_session()` is called — the HTTP client, streams, -and `SessionInfo` instance are all reused. +On a recoverable session error, `_reinitialize_session()` closes the old +connection stack, calls `strategy.reset` to clear the ID unless +`is_session_rejected(error)` is False -- a dropped transport is not the server's +verdict, so the ID is kept and the reconnect resumes the same session -- and +opens a fresh SDK transport and `ClientSession`. The authenticated `httpx2.AsyncClient` +and `SessionInfo` instance are reused. The failed-session identity guard prevents a late failure +from a concurrent operation from tearing down a replacement session. ## Cross-Package Dependencies @@ -553,42 +785,97 @@ For detailed test documentation, mocking strategies, and guidelines for adding n | Test File | Purpose | |-----------|---------| -| `test_mcp_client.py` | McpClient session tests (7 tests) | -| `test_mcp_tool.py` | Tool factory tests (17 tests) | +| `test_mcp_client_real_http.py` | **`McpClient` over real HTTP** against a real `MCPServer`: negotiation per mode, legacy resume, affinity, disposal, per-era retry, every handshake version | +| `real_server.py` | Harness for the above: `serve()`, `build_sdk_app()`, `PinnedVersionServer`, `RecordingGateway`, `patched_sdk()` | +| `test_mcp_client.py` | Pathological legacy servers and concurrency races, over `httpx2.MockTransport` | +| `test_protocol_strategy.py` | Pure per-era policy (`is_recoverable`, `reset`, `build_protocol_strategy`) plus the two server behaviours a cooperative server cannot produce | +| `test_protocol_version_support.py` | Tripwires on the SDK facts the strategies depend on | +| `test_mcp_tool.py` | Tool factory, schema refresh, result serialization, and error mapping | +| `test_session_tools.py` | `load_mcp_tools` discovery, invocation, and error mapping | +| `test_session_info.py` | SessionInfo and SessionInfoFactory contract | + +**Where a behaviour belongs.** Anything a cooperative server can produce is +tested over real HTTP in `test_mcp_client_real_http.py`. `MockTransport` is kept +only for what a real server cannot express: concurrency races (blocking one +`initialize` mid-flight), pathological servers (minting a new session on every +handshake, echoing `mcp-session-id` back in the modern era, repeating the header +on every response), and pure-function matrices. ### Key Test Classes | Class | Tests | |-------|-------| -| `TestMcpClient` | Session lifecycle, 404 retry, client reuse | +| `RecordingGateway` (`real_server.py`) | Pure ASGI middleware recording the JSON-RPC method, `mcp-session-id`, `mcp-protocol-version`, `params._meta` and HTTP method of every request, and optionally injecting a `Session terminated` fault on the Nth `tools/call` | +| Module-level client tests | Pathological session handling, 404 retry, concurrency, client reuse | | `TestMcpToolMetadata` | Tool metadata (tool_type, display_name, etc.) | | `TestMcpToolCreation` | Multiple tools, descriptions, disabled config | | `TestCreateMcpToolsFromAgent` | Agent factory function tests | -| `TestMcpToolInvocation` | Full invocation flow smoke test | | `TestMcpToolNameSanitization` | Tool name sanitization | ### Key Assertion -The most important test verifies client reuse on 404: +The regression guard lives in `test_legacy_resume_keeps_the_originally_negotiated_version`: +a resumed session keeping its ID is only half the contract — every request after +the resume must also carry the version that session was negotiated at. Probing +candidate versions instead settles on the *oldest* handshake version and +silently downgrades every later request, which no session-ID assertion catches. ```python -# HTTP client created only ONCE (not recreated on retry) -assert mock_async_client_class.call_count == 1 +versions = {r.protocol_version for r in after_resume if r.protocol_version} +assert versions == {"2025-11-25"} +``` + +The most important recovery test verifies a fresh session with HTTP client reuse: + +```python +assert endpoint.initialize_count == 2 +assert endpoint.tool_call_count == 2 +assert await client.get_session_id() == "session-2" +assert [h["mcp-session-id"] for h in endpoint.headers_for("tools/call")] == [ + "session-1", + "session-2", +] +``` -# But session initialized TWICE -assert initialize_count[0] == 2 +## Serializing content blocks: always `by_alias=True` + +Any code turning an MCP model into a dict for the model or the wire **must** use: + +```python +block.model_dump(by_alias=True, mode="json", exclude_none=True) ``` +SDK 2.0 renamed the model attributes to snake case and kept the camelCase names +as *serialization aliases*. A plain `model_dump()` therefore silently rewrites +the shape: `mimeType` → `mime_type`, `_meta` → `meta`, for every image, audio, +resource-link and embedded-resource block. + +**Text blocks are byte-identical either way**, which is what makes this hard to +catch — a text-only assertion passes against both. This bug shipped once +(`_normalize_tool_result` in `mcp_tool.py`) and was invisible to the whole suite +until a test used a real `ImageContent`. + +Rules when touching serialization: + +1. Use the call above. Both `session_tools._content_blocks` and + `mcp_tool._dump_block` are correct references. +2. Never assert `model_dump` call arguments on a `MagicMock` — that locks in + whichever call was written. Construct a real `mcp.types` model and assert the + resulting keys. +3. Include at least one non-text block in any serialization test. + ## Guidelines for Changes ### Updating streamable_http.py When the upstream MCP SDK changes its transport: -1. Diff the upstream [`mcp/client/streamable_http.py`](https://github.com/modelcontextprotocol/python-sdk/blob/main/src/mcp/client/streamable_http.py) against our local copy -2. Apply upstream changes while preserving our `SessionInfo` integration -3. Key areas to watch: `_prepare_headers` (must stay async), `_maybe_extract_session_id_from_response` (must use `set_session_id`), `streamable_http_client` (must accept `session_info` param) -4. The transport must never own session state directly — always delegate to `_session_info` +1. Keep delegating to [`mcp.client.streamable_http`](https://github.com/modelcontextprotocol/python-sdk/blob/main/src/mcp/client/streamable_http.py); do not copy the transport again +2. Preserve the async `SessionInfo` request/response hooks and remove both hooks on context exit +3. Confirm the SDK still accepts a supplied `httpx2.AsyncClient` and yields two transport streams +4. Re-run `test_mcp_client_real_http.py` first — it drives a real server, so an + SDK transport change surfaces there before it surfaces in a mock +5. Then re-run the legacy-version, persisted-session, 404 recovery, and DELETE tests ### Adding New Factory Functions @@ -601,17 +888,34 @@ When the upstream MCP SDK changes its transport: 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()` +3. Set `_client_initialized = True` only after `_open_connection()` and the handshake succeed 4. Always use `get_httpx_client_kwargs()` for HTTP client 5. The `SessionInfo` is created via the factory — do not construct it directly ### Modifying Session Initialization -1. Changes go in `_initialize_session()` -2. This should remain lightweight — just the MCP handshake -3. Don't create new HTTP resources here -4. The transport handles `set_session_id` — `_initialize_session` only reads via `get_session_id` -5. Verify tests still show `mock_async_client_class.call_count == 1` on retry +1. Negotiation logic goes in a `ProtocolStrategy`, not in `_initialize_session()`, + which only delegates +2. A strategy's `connect()` runs only on a newly created `ClientSession`; never + call it again on the same SDK 2 session for recovery +3. Don't create HTTP resources there; `_open_connection()` owns the + transport/session stack +4. The response hook handles `set_session_id` — strategies only read via + `get_session_id`, except when minting an affinity ID or clearing a stale one +5. Verify recovery creates two sessions while retaining one HTTP client + +### Adding a Protocol Era + +1. Implement the three `ProtocolStrategy` methods plus a `SessionIdentity` +2. Add the mode to `ProtocolMode` and `build_protocol_strategy` +3. Decide what `is_recoverable` means for it — a stateless era should not retry + session-shaped errors +4. Add real-HTTP `McpClient` tests for the era in + `tests/agent/tools/test_mcp/test_mcp_client_real_http.py`, driving a real + server through `real_server.py` — negotiation, resume, disposal, and retry +5. Add a leg to `testcases/simple-http-mcp` driving a real server on that era +6. Never change the default mode without a major version: it changes the wire + for every existing caller ### Adding New Methods to McpClient @@ -638,7 +942,9 @@ When the upstream MCP SDK changes its transport: | File | Package | Purpose | |------|---------|---------| -| `streamable_http.py` | uipath-langchain | SessionInfo + transport (local SDK copy) | +| `streamable_http.py` | uipath-langchain | SessionInfo, SessionIdentity, thin SDK transport adapter | +| `protocol_strategy.py` | uipath-langchain | Per-era negotiation, recovery, and identity policy | +| `session_tools.py` | uipath-langchain | `load_mcp_tools` session-to-LangChain conversion | | `mcp_client.py` | uipath-langchain | SessionInfoFactory + McpClient | | `mcp_tool.py` | uipath-langchain | Tool factory functions | | `__init__.py` | uipath-langchain | Public exports | @@ -649,30 +955,38 @@ When the upstream MCP SDK changes its transport: The implementation uses these MCP SDK components: -- `mcp.ClientSession` - MCP client session (can call `initialize()` multiple times) -- `mcp.shared.exceptions.McpError` - Error handling +- `mcp.ClientSession` - MCP client session (`initialize()` is idempotent per instance) +- `mcp.shared.exceptions.MCPError` - Error handling - `mcp.types.CallToolResult` - Tool call results -- `mcp.client._transport.TransportStreams` - Type alias used by `streamable_http_client` -- `mcp.shared._httpx_utils.create_mcp_http_client` - Default HTTP client factory -- `mcp.shared.message.SessionMessage` - Message wrapper for JSON-RPC +- `mcp.client.streamable_http.streamable_http_client` - Upstream transport context manager +- `httpx2.AsyncClient` - HTTP and SSE client used by MCP SDK 2 Key SDK behaviors: -- `ClientSession.initialize()` sends initialize request + initialized notification +- `ClientSession.initialize()` sends the latest legacy initialize request and initialized notification - `ClientSession.call_tool()` calls `_validate_tool_result()` on success - `_validate_tool_result()` calls `list_tools()` if output schema not cached -- HTTP 404 is converted to `McpError` with code `32600` by `StreamableHTTPTransport` +- A session-bound bare HTTP 404 is converted to `MCPError(INVALID_REQUEST, "Session terminated")` + +SDK 2 accepts legacy handshake responses for `2024-11-05`, `2025-03-26`, +`2025-06-18`, and `2025-11-25`, and reaches `2026-07-28` through +`ClientSession.discover()`. **Both eras are reachable from this low-level path** — +the high-level `mcp.Client` is not required. `AutoStrategy` owns its era +negotiation in `probe_modern_era`, built on the public `ClientSession.send_discover` +/ `adopt` seam — the same two calls the SDK's private `mode="auto"` helper is +made of. That helper (`mcp.client._probe.negotiate_auto`) is deliberately **not** +imported: private surface can move in a patch release. +`test_protocol_version_support.py` pins the two public seams instead. ## Performance Considerations Session reinitialization is efficient because: 1. **HTTP client reused**: No new TCP connections -2. **Streamable connection reused**: No new task groups or streams +2. **Connection state replaced**: A fresh transport/task group and `ClientSession` 3. **SessionInfo reused**: No new factory calls or debug state loads -4. **Only MCP handshake**: Just 2 HTTP requests (initialize + notification) +4. **Only MCP handshake repeated**: Initialize + initialized notification before retry This is significantly faster than full client reinitialization, which would require: -- Creating new `httpx.AsyncClient` -- Creating new task groups -- Creating new memory streams -- Establishing new connections +- Creating a new `httpx2.AsyncClient` +- Resolving the MCP registration and authorization again +- Re-running the `SessionInfoFactory` and any external debug-state load diff --git a/src/uipath_langchain/agent/tools/mcp/mcp_client.py b/src/uipath_langchain/agent/tools/mcp/mcp_client.py index 600f5b3ab..425b2b393 100644 --- a/src/uipath_langchain/agent/tools/mcp/mcp_client.py +++ b/src/uipath_langchain/agent/tools/mcp/mcp_client.py @@ -11,16 +11,25 @@ from typing import TYPE_CHECKING, Any, TypeVar import httpx -from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream +import httpx2 from mcp import ClientSession -from mcp.shared.exceptions import McpError -from mcp.shared.message import SessionMessage -from mcp.types import CallToolResult, ListToolsResult +from mcp.shared.exceptions import MCPError +from mcp.types import ( + CallToolResult, + ListToolsResult, +) from uipath._utils._ssl_context import get_httpx_client_kwargs from uipath.runtime.base import UiPathDisposableProtocol from uipath_langchain._utils import get_execution_folder_path +from .protocol_strategy import ( + ProtocolMode, + ProtocolStrategy, + build_protocol_strategy, + is_session_rejected, +) +from .protocol_strategy import is_session_error as _is_session_error from .streamable_http import SessionInfo, streamable_http_client if TYPE_CHECKING: @@ -32,6 +41,22 @@ T = TypeVar("T") +def _normalize_timeout( + timeout: httpx.Timeout | httpx2.Timeout | float | None, +) -> httpx2.Timeout | float: + """Convert the pre-upgrade HTTPX timeout type for the MCP 2 transport.""" + if timeout is None: + return httpx2.Timeout(600) + if isinstance(timeout, httpx.Timeout): + return httpx2.Timeout( + connect=timeout.connect, + read=timeout.read, + write=timeout.write, + pool=timeout.pool, + ) + return timeout + + class SessionInfoFactory: """Creates SessionInfo instances for MCP servers. @@ -61,23 +86,23 @@ class McpClient(UiPathDisposableProtocol): - Creates ClientSession - Calls session.initialize() to get session ID - 2. **Session Reinitialization** (on 404 error): - - Reuses existing HTTP client and streamable connection - - Calls session.initialize() again to get new session ID + 2. **Session Reinitialization** (after a terminated session): + - Reuses the existing HTTP client and persisted session store + - Replaces the transport and ``ClientSession`` + - Performs a fresh legacy initialization handshake Thread-safety is ensured via asyncio.Lock for both phases. """ - # Error codes that indicate session disconnect/termination - SESSION_ERROR_CODES = [32600, -32000] - def __init__( self, config: "AgentMcpResourceConfig", - timeout: httpx.Timeout | None = None, + timeout: httpx.Timeout | httpx2.Timeout | float | None = None, max_retries: int = 1, session_info_factory: SessionInfoFactory | None = None, terminate_on_close: bool = True, + *, + protocol_mode: ProtocolMode = "legacy", ) -> None: """Initialize the MCP tool session. @@ -91,12 +116,20 @@ def __init__( max_retries: Maximum number of retries on session disconnect errors. session_info_factory: Factory for creating SessionInfo instances. Defaults to ``SessionInfoFactory`` which returns a plain SessionInfo. + terminate_on_close: Whether to terminate the server session on + disposal. Already a no-op in the modern era, which has no session. + protocol_mode: Which negotiation era to use. ``"legacy"`` (the + default) sends only the ``initialize`` handshake, preserving + pre-existing behavior. ``"modern"`` uses ``server/discover`` + only. ``"auto"`` probes for the modern era and falls back to the + handshake. """ self._config = config - self._timeout = timeout or httpx.Timeout(600) + self._timeout = _normalize_timeout(timeout) self._max_retries = max_retries self._session_info_factory = session_info_factory or SessionInfoFactory() self._terminate_on_close = terminate_on_close + self._strategy: ProtocolStrategy = build_protocol_strategy(protocol_mode) # URL and headers are resolved lazily from SDK self._url: str | None = None @@ -112,15 +145,12 @@ def __init__( self._tools_cache: ListToolsResult | None = None # Client state (created once, reused across session reinitializations) - self._http_client: httpx.AsyncClient | None = None - self._read_stream: ( - MemoryObjectReceiveStream[SessionMessage | Exception] | None - ) = None - self._write_stream: MemoryObjectSendStream[SessionMessage] | None = None + self._http_client: httpx2.AsyncClient | None = None self._session_info: SessionInfo | None = None self._stack: AsyncExitStack | None = None + self._connection_stack: AsyncExitStack | None = None - # Session state (can be reinitialized without recreating client) + # Session state (replaced on recovery while the HTTP client is reused) self._session: ClientSession | None = None self._client_initialized: bool = False @@ -145,7 +175,7 @@ async def _initialize_client(self) -> None: This is called once on first use. Creates: - UiPath SDK instance to retrieve MCP server URL - - httpx.AsyncClient with authorization headers + - httpx2.AsyncClient with authorization headers - Streamable HTTP connection (read/write streams) - ClientSession @@ -175,83 +205,84 @@ 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__() + stack = AsyncExitStack() + await stack.__aenter__() + self._stack = stack + try: + # 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 stack.enter_async_context( + httpx2.AsyncClient(**client_kwargs) + ) - # 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) - # Create session info for tracking session ID - self._session_info = self._session_info_factory.create_session(mcp_server) - - # 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() - 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, - ) - ) + # Load a session ID persisted by the AgentHub debug-state integration. + existing = await self._session_info.get_session_id() + if existing: + logger.info(f"Loaded existing session ID from session info: {existing}") - # 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) - ) + await self._open_connection() + except BaseException: + await stack.aclose() + self._stack = None + self._http_client = None + 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 _open_connection(self) -> None: + """Open a fresh transport and ClientSession over the reusable HTTP client.""" + if self._url is None or self._http_client is None or self._session_info is None: + raise RuntimeError( + "Cannot open MCP connection: client prerequisites missing" + ) + + connection_stack = AsyncExitStack() + await connection_stack.__aenter__() + try: + read_stream, write_stream = await connection_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, + identity=self._strategy.identity, + ) + ) + self._session = await connection_stack.enter_async_context( + ClientSession(read_stream, write_stream) + ) + self._connection_stack = connection_stack + await self._initialize_session() + except BaseException: + await connection_stack.aclose() + self._session = None + self._connection_stack = None + raise async def _initialize_session(self) -> None: - """Initialize or reinitialize the MCP session. + """Negotiate the newly-created MCP session through the protocol strategy. - Calls session.initialize() to perform the MCP handshake and obtain - a session ID from the server. Can be called multiple times on the - same ClientSession to recover from session disconnects. + The strategy owns the era-specific handshake: the legacy ``initialize`` + exchange, or the modern ``server/discover`` probe. MCP 2 makes + ``ClientSession.initialize()`` idempotent per instance, so recovery + creates a new session before calling this again. Requires: Client must be initialized first (_initialize_client). """ - if self._session is None: + if self._session is None or self._session_info is None: raise RuntimeError("Cannot initialize session: client not initialized") - existing_session_id = ( - await self._session_info.get_session_id() if self._session_info else None - ) + existing_session_id = await self._session_info.get_session_id() logger.info( f"Initializing MCP session (session_info id: {existing_session_id})" ) - - if existing_session_id is None: - await self._session.initialize() - - # The transport calls set_session_id during initialize, - # so we just read the current value here. - new_session_id = ( - await self._session_info.get_session_id() - if self._session_info - else None - ) - logger.info(f"MCP session initialized with session ID: {new_session_id}") + await self._strategy.connect(self._session, self._session_info) async def _ensure_session(self) -> ClientSession: """Ensure client and session are initialized, return the session. @@ -262,44 +293,92 @@ async def _ensure_session(self) -> ClientSession: Returns: The initialized ClientSession. """ - if not self._client_initialized: - async with self._lock: - if not self._client_initialized: - await self._initialize_client() - - return self._session # type: ignore[return-value] + # Always cross the lifecycle lock. Recovery creates the replacement + # ClientSession before its initialize handshake completes, so a lock-free + # fast path could expose a half-initialized session to another operation. + async with self._lock: + if not self._client_initialized: + await self._initialize_client() + elif self._session is None: + # A failed replacement leaves the reusable HTTP client intact. + # Reopen on the next operation instead of poisoning the client. + await self._open_connection() + + if self._session is None: + raise RuntimeError("MCP client initialized without a session") + return self._session + + async def _close_connection_for_recovery(self) -> None: + """Detach and best-effort close the current connection stack.""" + connection_stack = self._connection_stack + self._connection_stack = None + self._session = None + if connection_stack is None: + return + try: + await connection_stack.aclose() + except Exception as error: + logger.debug("Error closing failed MCP connection: %s", error) + + async def _reinitialize_session( + self, + failed_session: ClientSession | None = None, + error: MCPError | None = None, + ) -> None: + """Replace the transport/session after a disconnect and initialize again. - async def _reinitialize_session(self) -> None: - """Reinitialize only the MCP session after a disconnect error. + MCP 2 makes ``ClientSession.initialize()`` idempotent, so recovery must + create a fresh ClientSession rather than calling initialize on the old one. + The HTTP client and external ``SessionInfo`` object are reused. - Thread-safe via lock. Reuses existing HTTP client and streamable - connection; only performs a new MCP handshake. - Clears the session info first so initialize() doesn't send a stale session ID. + The strategy is asked to discard persisted session state only when + ``error`` is the server's verdict on the session -- see + ``is_session_rejected``. A dropped transport is not: the ID is kept and + the reconnect resumes the same session. With no ``error`` nothing is + known against the session and it is likewise kept. """ async with self._lock: if not self._client_initialized: # Client not initialized, do full initialization await self._initialize_client() else: - # Clear stale session ID before re-initializing - if self._session_info: - await self._session_info.set_session_id(None) - await self._initialize_session() - - def _is_session_error(self, error: McpError) -> bool: - """Check if an McpError indicates a session disconnect. + if failed_session is not None and self._session is not failed_session: + logger.debug( + "MCP session was already replaced by another operation" + ) + return + await self._close_connection_for_recovery() + if ( + self._session_info is not None + and error is not None + and is_session_rejected(error) + ): + await self._strategy.reset(self._session_info) + await self._open_connection() + + @staticmethod + def is_session_error(error: MCPError) -> bool: + """Check if an MCPError indicates a session disconnect. Args: - error: The McpError to check. + error: The MCPError to check. Returns: True if the error indicates a session disconnect. """ - return ( - hasattr(error, "error") - and hasattr(error.error, "code") - and error.error.code in self.SESSION_ERROR_CODES + return _is_session_error(error) + + async def _is_recoverable_session_error(self, error: MCPError) -> bool: + """Ask the active protocol strategy whether a reconnect could fix ``error``. + + The answer is era-specific: a legacy session can be lost and re-established, + while every modern request is self-contained, so only a dropped connection + is worth retrying there. + """ + restored_id = ( + await self._session_info.get_session_id() if self._session_info else None ) + return self._strategy.is_recoverable(error, restored_id) async def _execute_with_retry( self, @@ -309,7 +388,7 @@ async def _execute_with_retry( """Execute a session operation with automatic retry on session disconnect. On first call, initializes the full client stack. On session - disconnect, reinitializes only the session and retries up to + disconnect, replaces the transport/session and retries up to ``_max_retries`` times. Args: @@ -321,11 +400,12 @@ async def _execute_with_retry( The result of *operation*. Raises: - McpError: If the operation fails after all retries. + MCPError: If the operation fails after all retries. """ retry_count = 0 while retry_count <= self._max_retries: + session: ClientSession | None = None try: session = await self._ensure_session() logger.debug( @@ -333,15 +413,16 @@ async def _execute_with_retry( ) return await operation(session) - except McpError as e: - logger.info(f"McpError during {operation_name}: {e}") + except MCPError as e: + logger.info(f"MCPError during {operation_name}: {e}") - if self._is_session_error(e) and retry_count < self._max_retries: + recoverable = await self._is_recoverable_session_error(e) + if recoverable and retry_count < self._max_retries: logger.warning( - f"Session disconnected (error code: {e.error.code}), " + f"Session disconnected (error code: {e.code}), " f"reinitializing session" ) - await self._reinitialize_session() + await self._reinitialize_session(session, error=e) retry_count += 1 continue else: @@ -408,17 +489,23 @@ async def dispose(self) -> None: async with self._tools_lock: self._tools_cache = None async with self._lock: + if self._connection_stack is not None: + try: + await self._connection_stack.aclose() + except Exception as e: + logger.debug(f"Error during MCP connection cleanup: {e}") + finally: + self._connection_stack = None + self._session = None + if self._stack is not None: try: - await self._stack.__aexit__(None, None, None) + await self._stack.aclose() 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 diff --git a/src/uipath_langchain/agent/tools/mcp/mcp_tool.py b/src/uipath_langchain/agent/tools/mcp/mcp_tool.py index ab9b5f773..5cbf7d57b 100644 --- a/src/uipath_langchain/agent/tools/mcp/mcp_tool.py +++ b/src/uipath_langchain/agent/tools/mcp/mcp_tool.py @@ -3,7 +3,7 @@ from typing import Any, AsyncGenerator from langchain_core.tools import BaseTool -from mcp.shared.exceptions import McpError +from mcp.shared.exceptions import MCPError from uipath.agent.models.agent import ( AgentMcpResourceConfig, AgentMcpTool, @@ -134,7 +134,7 @@ async def _refresh_tool_schema( ) return _tool_removed_message(mcp_tool.name) - if not _breaking_schema_change(mcp_tool.input_schema, fresh.inputSchema): + if not _breaking_schema_change(mcp_tool.input_schema, fresh.input_schema): return None logger.warning( @@ -143,16 +143,16 @@ async def _refresh_tool_schema( ) # Heal: update the cached baseline and the schema the model is bound to, so the # next LLM turn re-binds the live schema and the model can build a valid call. - mcp_tool.input_schema = fresh.inputSchema - mcp_tool.output_schema = fresh.outputSchema + mcp_tool.input_schema = fresh.input_schema + mcp_tool.output_schema = fresh.output_schema if fresh.description: mcp_tool.description = fresh.description tool = tool_holder.get("tool") if tool_holder else None if tool is not None: - tool.args_schema = fresh.inputSchema + tool.args_schema = fresh.input_schema if fresh.description: tool.description = fresh.description - return _schema_change_message(mcp_tool.name, fresh.inputSchema) + return _schema_change_message(mcp_tool.name, fresh.input_schema) @asynccontextmanager @@ -260,8 +260,8 @@ async def create_mcp_tools( AgentMcpTool( name=tool.name, description=tool.description or "", - input_schema=tool.inputSchema, - output_schema=tool.outputSchema, + input_schema=tool.input_schema, + output_schema=tool.output_schema, argument_properties=argument_properties, ) ) @@ -299,15 +299,15 @@ async def create_mcp_tools( def _map_mcp_error( - error: McpError, tool_name: str, server_slug: str + error: MCPError, tool_name: str, server_slug: str ) -> AgentRuntimeError: - """Map a protocol-level McpError to a categorized AgentRuntimeError. + """Map a protocol-level MCPError to a categorized AgentRuntimeError. - MCP tool execution failures come back as ``CallToolResult.isError`` results, - so an McpError raised during a call is a protocol/session/transport failure — + MCP tool execution failures come back as ``CallToolResult.is_error`` results, + so an MCPError raised during a call is a protocol/session/transport failure — hence the SYSTEM category. """ - if error.error.code in McpClient.SESSION_ERROR_CODES: + if McpClient.is_session_error(error): detail = ( f"The connection to MCP server '{server_slug}' was terminated and " f"could not be re-established while calling tool '{tool_name}'. " @@ -316,7 +316,7 @@ def _map_mcp_error( else: detail = ( f"MCP server '{server_slug}' returned an error for tool " - f"'{tool_name}': {error.error.message}" + f"'{tool_name}': {error.message}" ) return AgentRuntimeError( code=AgentRuntimeErrorCode.HTTP_ERROR, @@ -328,16 +328,26 @@ def _map_mcp_error( def _normalize_tool_result(result: Any) -> Any: - """Normalize an MCP ``call_tool`` result into JSON-serializable content.""" + """Normalize an MCP ``call_tool`` result into JSON-serializable content. + + Serialized with ``by_alias=True`` so blocks keep their wire spelling. SDK 2.0 + renamed the model attributes to snake case while keeping the camelCase names + as serialization aliases, so a plain ``model_dump()`` would silently rewrite + every non-text block handed to the model -- ``mimeType`` becoming + ``mime_type``, ``_meta`` becoming ``meta``. Text blocks are unaffected either + way, which is why the difference is easy to miss. + """ content = result.content if hasattr(result, "content") else result if isinstance(content, list): - return [ - item.model_dump(exclude_none=True) if hasattr(item, "model_dump") else item - for item in content - ] - if hasattr(content, "model_dump"): - return content.model_dump(exclude_none=True) - return content + return [_dump_block(item) for item in content] + return _dump_block(content) + + +def _dump_block(item: Any) -> Any: + """Serialize one MCP content block in its wire-compatible representation.""" + if not hasattr(item, "model_dump"): + return item + return item.model_dump(by_alias=True, mode="json", exclude_none=True) def build_mcp_tool( @@ -359,7 +369,7 @@ def build_mcp_tool( output_schema=output_schema, ) async def tool_fn(**kwargs: Any) -> Any: - """Execute MCP tool call with ephemeral session. + """Execute an MCP tool call through the managed client session. When ``refresh_schema_before_call`` is set (cached discovery mode), the live tool schema is checked first against the McpClient's cached tool list (fetched @@ -376,7 +386,7 @@ async def tool_fn(**kwargs: Any) -> Any: return retry_message try: result = await mcpClient.call_tool(mcp_tool.name, arguments=kwargs) - except McpError as e: + except MCPError as e: raise _map_mcp_error(e, mcp_tool.name, mcpClient.server_slug) from e logger.info(f"Tool call successful for {mcp_tool.name}") return _normalize_tool_result(result) diff --git a/src/uipath_langchain/agent/tools/mcp/protocol_strategy.py b/src/uipath_langchain/agent/tools/mcp/protocol_strategy.py new file mode 100644 index 000000000..1ff152606 --- /dev/null +++ b/src/uipath_langchain/agent/tools/mcp/protocol_strategy.py @@ -0,0 +1,497 @@ +"""Era-specific session lifecycle for one MCP connection. + +MCP has two negotiation eras. The legacy handshake (``2024-11-05`` through +``2025-11-25``) agrees a protocol version through ``initialize`` and identifies +the connection with a server-minted ``mcp-session-id``. The modern era +(``2026-07-28``) replaces the handshake with a stateless ``server/discover`` +probe and has no session identity at all: every request re-declares the +protocol version, client info and capabilities. + +Negotiation itself is one call in either era. What genuinely differs is the +*session lifecycle* -- how a connection is negotiated, whether a restored one can +be reused, which errors a reconnect could fix, and how the connection is +identified on the wire. Those four concerns are what :class:`ProtocolStrategy` +abstracts. +""" + +import logging +from typing import Literal, Protocol, runtime_checkable +from uuid import uuid4 + +from mcp import ClientSession +from mcp.shared.exceptions import MCPError +from mcp.types import ( + CONNECTION_CLOSED, + INVALID_REQUEST, + METHOD_NOT_FOUND, + UNSUPPORTED_PROTOCOL_VERSION, + DiscoverResult, + Implementation, + InitializeResult, + ServerCapabilities, + UnsupportedProtocolVersionErrorData, +) +from mcp.types.version import ( + HANDSHAKE_PROTOCOL_VERSIONS, + LATEST_MODERN_VERSION, + MODERN_PROTOCOL_VERSIONS, +) +from pydantic import ValidationError + +from .streamable_http import ( + LEGACY_IDENTITY, + MODERN_IDENTITY, + SessionIdentity, + SessionInfo, +) + +logger = logging.getLogger(__name__) + +ProtocolMode = Literal["legacy", "auto", "modern"] + +#: Session-disconnect markers that make a request-level error recoverable. +#: ``no valid session`` is how the TypeScript SDK's transport phrases it. +_SESSION_LOST_MARKERS = ( + "terminated", + "expired", + "invalid", + "not found", + "no valid session", +) + + +def _names_a_lost_session(error: MCPError) -> bool: + """Report whether ``error``'s message says the session itself is gone.""" + message = error.message.lower() + return "session" in message and any( + marker in message for marker in _SESSION_LOST_MARKERS + ) + + +def is_session_error(error: MCPError) -> bool: + """Check whether an ``MCPError`` reports a lost legacy session. + + Args: + error: The error to classify. + + Returns: + True when the error indicates the session is gone rather than the + request being wrong. + """ + if error.code == CONNECTION_CLOSED: + return True + # ``32600`` is the unsigned spelling some gateways emit for INVALID_REQUEST. + return error.code in (32600, INVALID_REQUEST) and _names_a_lost_session(error) + + +def is_session_rejected(error: MCPError) -> bool: + """Report whether ``error`` is the server's verdict on the stored session. + + The distinction decides whether recovery may keep the persisted session ID. + A transport that simply dropped says nothing about the session, so the ID is + kept and the reconnect resumes it; anything the *server* answered about the + session means the ID is dead and must be discarded. + + The code alone cannot tell those apart, because ``CONNECTION_CLOSED`` is + JSON-RPC's implementation-defined server-error code ``-32000``, which the + TypeScript SDK's Streamable HTTP transport also uses to refuse a session it + does not know (``"Bad Request: No valid session ID provided"``). So a + ``-32000`` naming a lost session is a verdict, not a dropped connection. + + Args: + error: The error that triggered recovery. + + Returns: + True when the persisted session ID must be discarded. + """ + return error.code != CONNECTION_CLOSED or _names_a_lost_session(error) + + +@runtime_checkable +class ProtocolStrategy(Protocol): + """Negotiation and recovery policy for one protocol era.""" + + identity: SessionIdentity + """How this era identifies the connection on the wire.""" + + async def connect(self, session: ClientSession, info: SessionInfo) -> None: + """Bring ``session`` to a negotiated state, reusing persisted state if any.""" + ... + + def is_recoverable(self, error: MCPError, restored_id: str | None) -> bool: + """Report whether opening a fresh connection could plausibly fix ``error``.""" + ... + + async def reset(self, info: SessionInfo) -> None: + """Discard the persisted state that made the last connection fail.""" + ... + + +class LegacyHandshakeStrategy: + """Negotiate through ``initialize`` and reuse server-minted sessions. + + A restored session is resumed without any negotiation traffic when the store + remembers the version it was negotiated at: the version is the only thing a + fresh ``ClientSession`` is missing, and + :meth:`ClientSession.adopt` installs it locally. The session then continues + exactly as it did before this client existed -- the wire sees only the + ``mcp-session-id`` header on ordinary requests, and no server is asked to + re-initialize a session it already initialized. + + When the version is *not* known -- a store written before it was recorded -- + there is nothing to adopt, so the handshake is re-run inside the restored + session. The server routes requests purely by the ``mcp-session-id`` header + and creates a new session only when no header is present, so that + ``initialize`` lands *inside* the session and returns the version it was + negotiated at. Servers differ on whether they allow it (the reference + TypeScript implementation answers "Server already initialized"), which is + why a rejection falls back to a clean session. + """ + + def __init__(self) -> None: + self.identity = SessionIdentity(LEGACY_IDENTITY) + + async def connect(self, session: ClientSession, info: SessionInfo) -> None: + """Run the handshake, resuming a persisted session when one exists.""" + restored_id = await info.get_session_id() + if restored_id is None: + await self._handshake(session, info) + return + + if await self._adopt_restored_session(session, info, restored_id): + return + + try: + result = await session.initialize() + except MCPError as error: + if error.code == CONNECTION_CLOSED: + # The transport died, which says nothing about the session. + # Clearing the ID here would destroy an externally persisted + # session -- permanently, for a store-backed SessionInfo -- over a + # transient failure. Let recovery reopen and resume instead. + raise + # The persisted session is gone, or this server refuses a second + # handshake. Either way a clean session is the correct fallback; the + # transport survives a rejected request, so the same one is reused. + logger.info( + "Persisted MCP session %s was rejected (%s); starting a new session", + restored_id, + error.code, + ) + await self.reset(info) + await self._handshake(session, info) + return + + await info.set_protocol_version(result.protocol_version) + current_id = await info.get_session_id() + if current_id == restored_id: + logger.info( + "Reusing externally persisted MCP session %s at %s", + restored_id, + result.protocol_version, + ) + else: + # A server that ignores the session header mints a replacement. The + # persisted session is lost, but the connection is usable. + logger.info( + "Server replaced persisted MCP session %s with %s at %s", + restored_id, + current_id, + result.protocol_version, + ) + + @staticmethod + async def _handshake(session: ClientSession, info: SessionInfo) -> None: + """Negotiate a brand-new session and remember what it settled on.""" + result = await session.initialize() + await info.set_protocol_version(result.protocol_version) + logger.info( + "MCP session initialized with session ID: %s at %s", + await info.get_session_id(), + result.protocol_version, + ) + + @staticmethod + async def _adopt_restored_session( + session: ClientSession, info: SessionInfo, restored_id: str + ) -> bool: + """Resume a stored session locally, with no request on the wire. + + The stored version is the whole of what a fresh ``ClientSession`` lacks: + ``adopt`` installs it, and every later request is stamped at that + version. The server is never told anything, which is what makes this + safe on a server that refuses a second ``initialize``. + + Args: + session: The un-negotiated session to bring up. + info: The store holding the session ID and its version. + restored_id: The stored session ID, for logging. + + Returns: + True when the session was adopted, False when the version is + unknown and the caller must negotiate instead. + """ + restored_version = await info.get_protocol_version() + if restored_version is None: + return False + if restored_version not in HANDSHAKE_PROTOCOL_VERSIONS: + # A modern version stored against this ID, or a version this client + # no longer speaks. Neither can be adopted onto a legacy wire, and + # the handshake will settle it correctly. + logger.info( + "Stored MCP session %s names version %s, which is not a " + "handshake version; re-negotiating instead", + restored_id, + restored_version, + ) + return False + session.adopt( + InitializeResult( + protocolVersion=restored_version, + # Capabilities and server info are not persisted: nothing in + # this package reads them back, and inventing them here keeps + # the resume free of wire traffic. + capabilities=ServerCapabilities(), + serverInfo=Implementation(name="restored-session", version="0"), + ) + ) + logger.info( + "Resumed MCP session %s at %s without re-initializing", + restored_id, + restored_version, + ) + return True + + def is_recoverable(self, error: MCPError, restored_id: str | None) -> bool: + """Recognize explicit and restored-session disconnect responses. + + The SDK transport only knows session IDs received during its own + lifetime. When UiPath restores an externally persisted ID, the request + hook supplies it but the transport maps a bare HTTP 404 to + ``METHOD_NOT_FOUND``. With a persisted ID on that request, Streamable + HTTP defines the 404 as an invalid session, so a fresh handshake is safe. + + ``restored_id`` is the ID currently stored, not specifically the one + restored when the connection opened. The distinction does not matter in + practice: the SDK only produces this exact shape for a session it does + not know about. + """ + if is_session_error(error): + return True + if error.code != METHOD_NOT_FOUND or error.message != "Not Found": + return False + return restored_id is not None + + async def reset(self, info: SessionInfo) -> None: + """Clear the stale session ID so the next handshake starts clean. + + The version goes with it: it described the session being discarded, and + leaving it behind would let the next connection adopt a version the + replacement session never negotiated. + """ + await info.set_session_id(None) + await info.set_protocol_version(None) + + +class ModernDiscoveryStrategy: + """Negotiate through ``server/discover`` and carry a UiPath affinity ID. + + ``2026-07-28`` has no session identity, so there is nothing to resume and no + session-loss error to recover from. UiPath still needs to reach the same warm + serverless instance across requests and runs, so this strategy mints its own + ID and sends it on ``mcp-session-id`` -- purely as a routing key, since a + modern server has no session to attach it to and ignores it. Reusing that + header rather than inventing one means the gateway needs no change: it keeps + routing on the header it already routes on today. + + Unlike a server-assigned session, the ID is available on the very first + request -- ``server/discover`` included -- so even the cold start is + attributable. + """ + + def __init__(self) -> None: + self.identity = SessionIdentity(MODERN_IDENTITY) + + async def connect(self, session: ClientSession, info: SessionInfo) -> None: + """Mint an affinity ID if needed, then probe ``server/discover``.""" + await mint_affinity_id(info) + result = await session.discover() + logger.info( + "MCP modern discovery negotiated %s (server supports %s)", + session.protocol_version, + list(result.supported_versions), + ) + + def is_recoverable(self, error: MCPError, restored_id: str | None) -> bool: + """Retry only a dropped connection. + + Every modern request is self-contained, so no server-side session can be + lost. Retrying anything but a transport failure spends the retry budget + on an error a reconnect cannot fix. + """ + return error.code == CONNECTION_CLOSED + + async def reset(self, info: SessionInfo) -> None: + """Keep the affinity ID so a reconnect returns to the same instance.""" + return + + +class AutoStrategy: + """Probe for the modern era, falling back to the legacy handshake. + + The affinity ID is minted *before* the probe, so ``server/discover`` reaches + the same instance the tool calls will: on a serverless gateway, an unpinned + probe would warm one instance and the first call would land on another. A + legacy server sees an ID it never issued; it is cleared again before the + handshake so that server is not asked to resume a session that never was. + + The era is re-resolved on every ``connect``, so a server upgraded mid-run is + handled. + """ + + def __init__(self) -> None: + # Both eras send the ID on the same header, so the transport can open on + # the legacy wire before the era is known: a modern server simply never + # sends one back, leaving nothing to capture. + self.identity = SessionIdentity(LEGACY_IDENTITY) + self._legacy = LegacyHandshakeStrategy() + self._modern = ModernDiscoveryStrategy() + self._resolved: ProtocolStrategy = self._legacy + + async def connect(self, session: ClientSession, info: SessionInfo) -> None: + """Negotiate an era, then narrow this strategy to it.""" + self.identity.wire = LEGACY_IDENTITY + # Widen back to the conservative era first: if the probe raises, the + # previous connection's resolution must not decide how this failure is + # recovered from. + self._resolved = self._legacy + restored_id = await info.get_session_id() + if restored_id is None: + await mint_affinity_id(info) + + if await probe_modern_era(session): + self._resolved = self._modern + logger.info("MCP era resolved to modern (%s)", session.protocol_version) + else: + if restored_id is None: + # The ID was minted for this connection and never named a + # session on this server. Sending it into the handshake would + # only earn a rejection; a restored ID, by contrast, may well + # be a live legacy session and is left for the handshake. + await info.set_session_id(None) + await self._legacy.connect(session, info) + self._resolved = self._legacy + logger.info("MCP era resolved to legacy (%s)", session.protocol_version) + self.identity.wire = self._resolved.identity.wire + + def is_recoverable(self, error: MCPError, restored_id: str | None) -> bool: + """Apply the resolved era's recovery policy.""" + return self._resolved.is_recoverable(error, restored_id) + + async def reset(self, info: SessionInfo) -> None: + """Apply the resolved era's reset policy.""" + await self._resolved.reset(info) + + +async def probe_modern_era(session: ClientSession) -> bool: + """Probe ``server/discover`` and adopt the modern era if the peer speaks it. + + Only positive evidence of a modern server counts; anything else reports + ``False`` so the caller can fall back to the ``initialize`` handshake. That + is a denylist, not an allowlist: every JSON-RPC error falls back, including + the HTTP-layer 4xx the transport synthesizes into one, as does a result + that fails to parse or advertises no modern version. A ``-32022`` naming a + mutual modern version earns one re-probe at that version. + + Built on the session's public ``send_discover`` / ``adopt`` seam rather than + the SDK's private ``mode="auto"`` helper, so an SDK patch cannot move this + policy out from under the client. + + Args: + session: The un-negotiated session to probe. + + Returns: + True when ``server/discover`` succeeded and was adopted. + + Raises: + MCPError: The server is modern-only yet shares no version with this + client -- a ``-32022`` whose ``supported`` list has no handshake + version -- so no era can work. + """ + version = LATEST_MODERN_VERSION + for attempt in range(2): + try: + raw = await session.send_discover(version) + except MCPError as error: + supported = _versions_supported_by(error) + if supported is None: + return False + mutual = [v for v in MODERN_PROTOCOL_VERSIONS if v in supported] + if mutual and attempt == 0: + version = mutual[-1] + continue + if not any(v in HANDSHAKE_PROTOCOL_VERSIONS for v in supported): + raise + return False + try: + result = DiscoverResult.model_validate(raw) + except ValidationError: + return False + if not any(v in result.supported_versions for v in MODERN_PROTOCOL_VERSIONS): + # A discover-answering server advertising only handshake versions is + # a legacy advertisement, not an incompatibility. + return False + session.adopt(result) + return True + return False + + +def _versions_supported_by(error: MCPError) -> list[str] | None: + """Read the ``supported`` list off a ``-32022`` error, or ``None``.""" + if error.code != UNSUPPORTED_PROTOCOL_VERSION: + return None + try: + data = UnsupportedProtocolVersionErrorData.model_validate(error.error.data) + except ValidationError: + return None + return data.supported + + +async def mint_affinity_id(info: SessionInfo) -> str: + """Ensure ``info`` holds a client-minted ID, and return it. + + Stored through the same accessors as a server-assigned session ID, so an + external ``SessionInfo`` implementation persists it without modification and + a later run resumes on the same instance. + """ + session_id = await info.get_session_id() + if session_id is None: + session_id = uuid4().hex + await info.set_session_id(session_id) + logger.info("Minted UiPath MCP affinity ID %s", session_id) + return session_id + + +def build_protocol_strategy(mode: ProtocolMode) -> ProtocolStrategy: + """Create the strategy for a protocol mode. + + Args: + mode: ``"legacy"`` for the ``initialize`` handshake only, ``"modern"`` + for ``server/discover`` only, or ``"auto"`` to probe for the modern + era and fall back to the handshake. + + Returns: + The strategy implementing that mode. + + Raises: + ValueError: ``mode`` is not a known protocol mode. + """ + if mode == "legacy": + return LegacyHandshakeStrategy() + if mode == "modern": + return ModernDiscoveryStrategy() + if mode == "auto": + return AutoStrategy() + raise ValueError( + f"Unknown MCP protocol mode {mode!r}; expected 'legacy', 'auto' or 'modern'" + ) diff --git a/src/uipath_langchain/agent/tools/mcp/session_tools.py b/src/uipath_langchain/agent/tools/mcp/session_tools.py new file mode 100644 index 000000000..67dcd3b72 --- /dev/null +++ b/src/uipath_langchain/agent/tools/mcp/session_tools.py @@ -0,0 +1,64 @@ +"""Convert tools from an active MCP SDK session into LangChain tools.""" + +from typing import Any + +from langchain_core.tools import BaseTool, StructuredTool, ToolException +from mcp import ClientSession +from mcp.types import CallToolResult, PaginatedRequestParams, Tool + + +def _content_blocks(result: CallToolResult) -> list[dict[str, Any]]: + """Serialize MCP content blocks in their wire-compatible representation.""" + return [ + block.model_dump(by_alias=True, mode="json", exclude_none=True) + for block in result.content + ] + + +def _error_message(result: CallToolResult) -> str: + """Build a readable LangChain tool error from MCP content blocks.""" + text = [ + block.text + for block in result.content + if getattr(block, "type", None) == "text" and hasattr(block, "text") + ] + return "\n".join(text) if text else str(_content_blocks(result)) + + +def _convert_tool(session: ClientSession, tool: Tool) -> BaseTool: + """Bind one discovered MCP tool to its active client session.""" + + async def call_tool(**arguments: Any) -> list[dict[str, Any]]: + result = await session.call_tool(tool.name, arguments=arguments) + if result.is_error: + raise ToolException(_error_message(result)) + return _content_blocks(result) + + return StructuredTool( + name=tool.name, + description=tool.description or "", + args_schema=tool.input_schema, + coroutine=call_tool, + ) + + +async def load_mcp_tools(session: ClientSession) -> list[BaseTool]: + """Discover all tools from an active MCP session and bind them to LangChain. + + Args: + session: An initialized MCP SDK ``ClientSession`` whose lifetime covers + every invocation of the returned tools. + + Returns: + LangChain tools backed by the supplied MCP session. + """ + tools: list[Tool] = [] + cursor: str | None = None + while True: + params = PaginatedRequestParams(cursor=cursor) if cursor is not None else None + page = await session.list_tools(params=params) + tools.extend(page.tools) + cursor = page.next_cursor + if not cursor: + break + return [_convert_tool(session, tool) for tool in tools] diff --git a/src/uipath_langchain/agent/tools/mcp/streamable_http.py b/src/uipath_langchain/agent/tools/mcp/streamable_http.py index fab856a96..ae15cbfe6 100644 --- a/src/uipath_langchain/agent/tools/mcp/streamable_http.py +++ b/src/uipath_langchain/agent/tools/mcp/streamable_http.py @@ -1,802 +1,243 @@ -"""StreamableHTTP Client Transport Module. - -Adapted from mcp.client.streamable_http (MCP Python SDK 1.26) to support -SessionInfo for external session ID tracking. - -This module implements the StreamableHTTP transport for MCP clients, -providing support for HTTP POST requests with optional SSE streaming responses -and session management. +"""Session-aware adapter for the MCP SDK's Streamable HTTP transport. + +The MCP SDK owns the transport implementation. UiPath only adds asynchronous, +externally-persistable session ID storage through :class:`SessionInfo`, and +describes how that ID travels on the wire through :class:`SessionIdentityWire`. + +The two eras identify a connection differently. Legacy servers mint an +``mcp-session-id`` and return it on the initialize response. ``2026-07-28`` has +no session identity at all, so UiPath mints its own ID and sends it on that same +header purely as a routing key -- a modern server ignores it, and the gateway +keeps routing on the header it already knows. Both are the same operation from +the transport's point of view -- read an ID, put it on the request, maybe read one +back -- so the transport stays ignorant of MCP negotiation and takes the +difference as data. """ -import contextlib +import asyncio +import json import logging -from collections.abc import AsyncGenerator, Awaitable, Callable +from collections.abc import AsyncGenerator from contextlib import asynccontextmanager -from dataclasses import dataclass -from datetime import timedelta -from typing import Any, overload -from warnings import warn - -import anyio -import httpx -from anyio.abc import TaskGroup -from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream -from httpx_sse import EventSource, ServerSentEvent, aconnect_sse -from mcp.shared._httpx_utils import ( - McpHttpClientFactory, - create_mcp_http_client, -) -from mcp.shared.message import ClientMessageMetadata, SessionMessage -from mcp.types import ( - ErrorData, - InitializeResult, - JSONRPCError, - JSONRPCMessage, - JSONRPCNotification, - JSONRPCRequest, - JSONRPCResponse, - RequestId, -) -from typing_extensions import deprecated - -logger = logging.getLogger(__name__) +from dataclasses import dataclass, field +from typing import Any +import httpx2 +from mcp.client.streamable_http import ( + streamable_http_client as sdk_streamable_http_client, +) +from uipath._utils._ssl_context import get_httpx_client_kwargs -SessionMessageOrError = SessionMessage | Exception -StreamWriter = MemoryObjectSendStream[SessionMessageOrError] -StreamReader = MemoryObjectReceiveStream[SessionMessage] MCP_SESSION_ID = "mcp-session-id" MCP_PROTOCOL_VERSION = "mcp-protocol-version" -LAST_EVENT_ID = "last-event-id" - -# Reconnection defaults -DEFAULT_RECONNECTION_DELAY_MS = ( - 1000 # 1 second fallback when server doesn't provide retry -) -MAX_RECONNECTION_ATTEMPTS = 2 # Max retry attempts before giving up -CONTENT_TYPE = "content-type" -ACCEPT = "accept" +#: Every header this adapter manages. Any not selected by the active wire is +#: removed from the request, so a narrowing wire cannot leave a stale header on. +_IDENTITY_HEADERS = (MCP_SESSION_ID,) -JSON = "application/json" -SSE = "text/event-stream" - -# Sentinel value for detecting unset optional parameters -_UNSET = object() +logger = logging.getLogger(__name__) class SessionInfo: - """Base class for MCP session ID tracking. + """Store the MCP session ID and allow subclasses to persist it externally.""" - The transport delegates all session ID storage to this object. - Override ``get_session_id`` / ``set_session_id`` in subclasses to - add side-effects such as HTTP persistence. - """ + #: Class-level default so a subclass that skips ``super().__init__()`` still + #: reads as "version not known" rather than raising. + protocol_version: str | None = None def __init__(self, session_id: str | None = None) -> None: self.session_id = session_id + # The version the stored session was negotiated at. Held alongside the + # ID because it cannot be recovered from the wire: responses carry only + # the session ID. A subclass that persists the ID externally should + # persist this too, so a later run can resume the session without + # re-running the handshake to learn its version. + self.protocol_version: str | None = None async def get_session_id(self) -> str | None: - """Return the current session ID (or None).""" + """Return the current session ID, or ``None`` when no session exists.""" return self.session_id async def set_session_id(self, session_id: str | None) -> None: - """Store a new session ID assigned by the server, or None to clear.""" + """Store a server-assigned session ID, or clear it with ``None``.""" self.session_id = session_id + async def get_protocol_version(self) -> str | None: + """Return the version the stored session was negotiated at, if known. -class StreamableHTTPError(Exception): - """Base exception for StreamableHTTP transport errors.""" - - -class ResumptionError(StreamableHTTPError): - """Raised when resumption request is invalid.""" - - -@dataclass -class RequestContext: - """Context for a request operation.""" - - client: httpx.AsyncClient - session_message: SessionMessage - metadata: ClientMessageMetadata | None - read_stream_writer: StreamWriter - headers: dict[str, str] | None = None # Deprecated - no longer used - sse_read_timeout: float | None = None # Deprecated - no longer used + ``None`` means "not known", not "no session": a store written by an + older revision holds an ID and no version, and a resumed connection + then has to learn the version from the server. + """ + return self.protocol_version + async def set_protocol_version(self, protocol_version: str | None) -> None: + """Record the negotiated version, or clear it with ``None``.""" + self.protocol_version = protocol_version -class StreamableHTTPTransport: - """StreamableHTTP client transport implementation.""" - @overload - def __init__( - self, url: str, *, session_info: SessionInfo | None = None - ) -> None: ... +@dataclass(frozen=True) +class SessionIdentityWire: + """How a stored session ID travels on the wire for one protocol era.""" - @overload - @deprecated( - "Parameters headers, timeout, sse_read_timeout, and auth are deprecated. " - "Configure these on the httpx.AsyncClient instead." - ) - def __init__( - self, - url: str, - headers: dict[str, str] | None = None, - timeout: float | timedelta = 30, - sse_read_timeout: float | timedelta = 60 * 5, - auth: httpx.Auth | None = None, - session_info: SessionInfo | None = None, - ) -> None: ... - - def __init__( - self, - url: str, - headers: Any = _UNSET, - timeout: Any = _UNSET, - sse_read_timeout: Any = _UNSET, - auth: Any = _UNSET, - session_info: SessionInfo | None = None, - ) -> None: - """Initialize the StreamableHTTP transport. - - Args: - url: The endpoint URL. - headers: Optional headers to include in requests. - timeout: HTTP timeout for regular operations. - sse_read_timeout: Timeout for SSE read operations. - auth: Optional HTTPX authentication handler. - session_info: Optional SessionInfo for external session ID tracking. - """ - # Check for deprecated parameters and issue runtime warning - deprecated_params: list[str] = [] - if headers is not _UNSET: - deprecated_params.append("headers") - if timeout is not _UNSET: - deprecated_params.append("timeout") - if sse_read_timeout is not _UNSET: - deprecated_params.append("sse_read_timeout") - if auth is not _UNSET: - deprecated_params.append("auth") - - if deprecated_params: - warn( - f"Parameters {', '.join(deprecated_params)} are deprecated and will be ignored. " - "Configure these on the httpx.AsyncClient instead.", - DeprecationWarning, - stacklevel=2, - ) - - self.url = url - self._session_info = session_info or SessionInfo() - self.protocol_version: str | None = None + request_headers: tuple[str, ...] = (MCP_SESSION_ID,) + """Headers the stored ID is sent on. Empty sends none.""" - async def _prepare_headers(self) -> dict[str, str]: - """Build MCP-specific request headers. + capture_response_header: str | None = MCP_SESSION_ID + """Header a server-assigned ID is read from, or ``None`` when the client mints it. - These headers will be merged with the httpx.AsyncClient's default headers, - with these MCP-specific headers taking precedence. - """ - headers: dict[str, str] = {} - # Add MCP protocol headers - headers[ACCEPT] = f"{JSON}, {SSE}" - headers[CONTENT_TYPE] = JSON - # Add session headers if available - session_id = await self._session_info.get_session_id() - if session_id: - headers[MCP_SESSION_ID] = session_id - if self.protocol_version: - headers[MCP_PROTOCOL_VERSION] = self.protocol_version - return headers - - def _is_initialization_request(self, message: JSONRPCMessage) -> bool: - """Check if the message is an initialization request.""" - return ( - isinstance(message.root, JSONRPCRequest) - and message.root.method == "initialize" - ) - - def _is_initialized_notification(self, message: JSONRPCMessage) -> bool: - """Check if the message is an initialized notification.""" - return ( - isinstance(message.root, JSONRPCNotification) - and message.root.method == "notifications/initialized" - ) - - async def _maybe_extract_session_id_from_response( - self, - response: httpx.Response, - ) -> None: - """Extract and store session ID from response headers.""" - new_session_id = response.headers.get(MCP_SESSION_ID) - if new_session_id: - await self._session_info.set_session_id(new_session_id) - logger.info(f"Received session ID: {new_session_id}") - - def _maybe_extract_protocol_version_from_message( - self, - message: JSONRPCMessage, - ) -> None: - """Extract protocol version from initialization response message.""" - if ( - isinstance(message.root, JSONRPCResponse) and message.root.result - ): # pragma: no branch - try: - # Parse the result as InitializeResult for type safety - init_result = InitializeResult.model_validate(message.root.result) - self.protocol_version = str(init_result.protocolVersion) - logger.info(f"Negotiated protocol version: {self.protocol_version}") - except Exception as exc: # pragma: no cover - logger.warning( - f"Failed to parse initialization response as InitializeResult: {exc}" - ) # pragma: no cover - logger.warning(f"Raw result: {message.root.result}") - - async def _handle_sse_event( - self, - sse: ServerSentEvent, - read_stream_writer: StreamWriter, - original_request_id: RequestId | None = None, - resumption_callback: Callable[[str], Awaitable[None]] | None = None, - is_initialization: bool = False, - ) -> bool: - """Handle an SSE event, returning True if the response is complete.""" - if sse.event == "message": - # Handle priming events (empty data with ID) for resumability - if not sse.data: - # Call resumption callback for priming events that have an ID - if sse.id and resumption_callback: - await resumption_callback(sse.id) - return False - try: - message = JSONRPCMessage.model_validate_json(sse.data) - logger.debug(f"SSE message: {message}") - - # Extract protocol version from initialization response - if is_initialization: - self._maybe_extract_protocol_version_from_message(message) - - # If this is a response and we have original_request_id, replace it - if original_request_id is not None and isinstance( - message.root, JSONRPCResponse | JSONRPCError - ): - message.root.id = original_request_id - - session_message = SessionMessage(message) - await read_stream_writer.send(session_message) - - # Call resumption token callback if we have an ID - if sse.id and resumption_callback: - await resumption_callback(sse.id) - - # If this is a response or error return True indicating completion - # Otherwise, return False to continue listening - return isinstance(message.root, JSONRPCResponse | JSONRPCError) - - except Exception as exc: # pragma: no cover - logger.exception("Error parsing SSE message") - await read_stream_writer.send(exc) - return False - else: # pragma: no cover - logger.warning(f"Unknown SSE event: {sse.event}") - return False - - async def handle_get_stream( - self, - client: httpx.AsyncClient, - read_stream_writer: StreamWriter, - ) -> None: - """Handle GET stream for server-initiated messages with auto-reconnect.""" - last_event_id: str | None = None - retry_interval_ms: int | None = None - attempt: int = 0 - - while attempt < MAX_RECONNECTION_ATTEMPTS: # pragma: no branch - try: - if not await self._session_info.get_session_id(): - return - - headers = await self._prepare_headers() - if last_event_id: - headers[LAST_EVENT_ID] = last_event_id # pragma: no cover - - async with aconnect_sse( - client, - "GET", - self.url, - headers=headers, - ) as event_source: - event_source.response.raise_for_status() - logger.debug("GET SSE connection established") - - async for sse in event_source.aiter_sse(): - # Track last event ID for reconnection - if sse.id: - last_event_id = sse.id # pragma: no cover - # Track retry interval from server - if sse.retry is not None: - retry_interval_ms = sse.retry # pragma: no cover - - await self._handle_sse_event(sse, read_stream_writer) - - # Stream ended normally (server closed) - reset attempt counter - attempt = 0 - - except Exception as exc: # pragma: no cover - logger.debug(f"GET stream error: {exc}") - attempt += 1 - - if attempt >= MAX_RECONNECTION_ATTEMPTS: # pragma: no cover - logger.debug( - f"GET stream max reconnection attempts ({MAX_RECONNECTION_ATTEMPTS}) exceeded" - ) - return - - # Wait before reconnecting - delay_ms = ( - retry_interval_ms - if retry_interval_ms is not None - else DEFAULT_RECONNECTION_DELAY_MS - ) - logger.info(f"GET stream disconnected, reconnecting in {delay_ms}ms...") - await anyio.sleep(delay_ms / 1000.0) - - async def _handle_resumption_request(self, ctx: RequestContext) -> None: - """Handle a resumption request using GET with SSE.""" - headers = await self._prepare_headers() - if ctx.metadata and ctx.metadata.resumption_token: - headers[LAST_EVENT_ID] = ctx.metadata.resumption_token - else: - raise ResumptionError( - "Resumption request requires a resumption token" - ) # pragma: no cover - - # Extract original request ID to map responses - original_request_id = None - if isinstance( - ctx.session_message.message.root, JSONRPCRequest - ): # pragma: no branch - original_request_id = ctx.session_message.message.root.id - - async with aconnect_sse( - ctx.client, - "GET", - self.url, - headers=headers, - ) as event_source: - event_source.response.raise_for_status() - logger.debug("Resumption GET SSE connection established") - - async for sse in event_source.aiter_sse(): # pragma: no branch - is_complete = await self._handle_sse_event( - sse, - ctx.read_stream_writer, - original_request_id, - ctx.metadata.on_resumption_token_update if ctx.metadata else None, - ) - if is_complete: - await event_source.response.aclose() - break - - async def _handle_post_request(self, ctx: RequestContext) -> None: - """Handle a POST request with response processing.""" - headers = await self._prepare_headers() - message = ctx.session_message.message - is_initialization = self._is_initialization_request(message) - - async with ctx.client.stream( - "POST", - self.url, - json=message.model_dump(by_alias=True, mode="json", exclude_none=True), - headers=headers, - ) as response: - if response.status_code == 202: - logger.debug("Received 202 Accepted") - return - - if response.status_code == 404: # pragma: no branch - if isinstance(message.root, JSONRPCRequest): - await self._send_session_terminated_error( # pragma: no cover - ctx.read_stream_writer, # pragma: no cover - message.root.id, # pragma: no cover - ) # pragma: no cover - return # pragma: no cover - - if response.status_code >= 400: - body = await response.aread() - logger.error( - f"HTTP {response.status_code} from POST {self.url}: {body.decode(errors='replace')}" - ) - response.raise_for_status() - if is_initialization: - await self._maybe_extract_session_id_from_response(response) - - # Per https://modelcontextprotocol.io/specification/2025-06-18/basic#notifications: - # The server MUST NOT send a response to notifications. - if isinstance(message.root, JSONRPCRequest): - content_type = response.headers.get(CONTENT_TYPE, "").lower() - if content_type.startswith(JSON): - await self._handle_json_response( - response, ctx.read_stream_writer, is_initialization - ) - elif content_type.startswith(SSE): - await self._handle_sse_response(response, ctx, is_initialization) - else: - await self._handle_unexpected_content_type( # pragma: no cover - content_type, # pragma: no cover - ctx.read_stream_writer, # pragma: no cover - ) # pragma: no cover - - async def _handle_json_response( - self, - response: httpx.Response, - read_stream_writer: StreamWriter, - is_initialization: bool = False, - ) -> None: - """Handle JSON response from the server.""" - try: - content = await response.aread() - message = JSONRPCMessage.model_validate_json(content) - - # Extract protocol version from initialization response - if is_initialization: - self._maybe_extract_protocol_version_from_message(message) - - session_message = SessionMessage(message) - await read_stream_writer.send(session_message) - except Exception as exc: # pragma: no cover - logger.exception("Error parsing JSON response") - await read_stream_writer.send(exc) - - async def _handle_sse_response( - self, - response: httpx.Response, - ctx: RequestContext, - is_initialization: bool = False, - ) -> None: - """Handle SSE response from the server.""" - last_event_id: str | None = None - retry_interval_ms: int | None = None + ``None`` also protects a client-minted ID: nothing on the wire can overwrite + the routing key mid-connection. + """ - try: - event_source = EventSource(response) - async for sse in event_source.aiter_sse(): # pragma: no branch - # Track last event ID for potential reconnection - if sse.id: - last_event_id = sse.id - - # Track retry interval from server - if sse.retry is not None: - retry_interval_ms = sse.retry - - is_complete = await self._handle_sse_event( - sse, - ctx.read_stream_writer, - resumption_callback=( - ctx.metadata.on_resumption_token_update - if ctx.metadata - else None - ), - is_initialization=is_initialization, - ) - # If the SSE event indicates completion, like returning respose/error - # break the loop - if is_complete: - await response.aclose() - return # Normal completion, no reconnect needed - except Exception as e: # pragma: no cover - logger.debug(f"SSE stream ended: {e}") - - # Stream ended without response - reconnect if we received an event with ID - if last_event_id is not None: # pragma: no branch - logger.info("SSE stream disconnected, reconnecting...") - await self._handle_reconnection(ctx, last_event_id, retry_interval_ms) - - async def _handle_reconnection( - self, - ctx: RequestContext, - last_event_id: str, - retry_interval_ms: int | None = None, - attempt: int = 0, - ) -> None: - """Reconnect with Last-Event-ID to resume stream after server disconnect.""" - # Bail if max retries exceeded - if attempt >= MAX_RECONNECTION_ATTEMPTS: # pragma: no cover - logger.debug( - f"Max reconnection attempts ({MAX_RECONNECTION_ATTEMPTS}) exceeded" - ) - return - # Always wait - use server value or default - delay_ms = ( - retry_interval_ms - if retry_interval_ms is not None - else DEFAULT_RECONNECTION_DELAY_MS - ) - await anyio.sleep(delay_ms / 1000.0) +#: Legacy handshake: the server mints the ID and returns it on a response header. +#: Also the starting wire for ``auto``, whose era is unknown when the transport +#: opens -- safe either way, because a modern server simply never sends the header +#: back. +LEGACY_IDENTITY = SessionIdentityWire() - headers = await self._prepare_headers() - headers[LAST_EVENT_ID] = last_event_id +#: Modern era: the same request header, but the client mints the value and no +#: server-assigned ID is ever read back. +MODERN_IDENTITY = SessionIdentityWire( + request_headers=(MCP_SESSION_ID,), + capture_response_header=None, +) - # Extract original request ID to map responses - original_request_id = None - if isinstance( - ctx.session_message.message.root, JSONRPCRequest - ): # pragma: no branch - original_request_id = ctx.session_message.message.root.id - try: - async with aconnect_sse( - ctx.client, - "GET", - self.url, - headers=headers, - ) as event_source: - event_source.response.raise_for_status() - logger.info("Reconnected to SSE stream") - - # Track for potential further reconnection - reconnect_last_event_id: str = last_event_id - reconnect_retry_ms = retry_interval_ms - - async for sse in event_source.aiter_sse(): - if sse.id: # pragma: no branch - reconnect_last_event_id = sse.id - if sse.retry is not None: - reconnect_retry_ms = sse.retry - - is_complete = await self._handle_sse_event( - sse, - ctx.read_stream_writer, - original_request_id, - ctx.metadata.on_resumption_token_update - if ctx.metadata - else None, - ) - if is_complete: - await event_source.response.aclose() - return - - # Stream ended again without response - reconnect again (reset attempt counter) - logger.info("SSE stream disconnected, reconnecting...") - await self._handle_reconnection( - ctx, reconnect_last_event_id, reconnect_retry_ms, 0 - ) - except Exception as e: # pragma: no cover - logger.debug(f"Reconnection failed: {e}") - # Try to reconnect again if we still have an event ID - await self._handle_reconnection( - ctx, last_event_id, retry_interval_ms, attempt + 1 - ) - - async def _handle_unexpected_content_type( - self, - content_type: str, - read_stream_writer: StreamWriter, - ) -> None: # pragma: no cover - """Handle unexpected content type in response.""" - error_msg = f"Unexpected content type: {content_type}" # pragma: no cover - logger.error(error_msg) # pragma: no cover - await read_stream_writer.send(ValueError(error_msg)) # pragma: no cover - - async def _send_session_terminated_error( - self, - read_stream_writer: StreamWriter, - request_id: RequestId, - ) -> None: - """Send a session terminated error response.""" - jsonrpc_error = JSONRPCError( - jsonrpc="2.0", - id=request_id, - error=ErrorData(code=32600, message="Session terminated"), - ) - session_message = SessionMessage(JSONRPCMessage(jsonrpc_error)) - await read_stream_writer.send(session_message) - - async def post_writer( - self, - client: httpx.AsyncClient, - write_stream_reader: StreamReader, - read_stream_writer: StreamWriter, - write_stream: MemoryObjectSendStream[SessionMessage], - start_get_stream: Callable[[], None], - tg: TaskGroup, - ) -> None: - """Handle writing requests to the server.""" - try: - async with write_stream_reader: - async for session_message in write_stream_reader: - message = session_message.message - metadata = ( - session_message.metadata - if isinstance(session_message.metadata, ClientMessageMetadata) - else None - ) - - # Check if this is a resumption request - is_resumption = bool(metadata and metadata.resumption_token) - - logger.debug(f"Sending client message: {message}") - - # Handle initialized notification - if self._is_initialized_notification(message): - start_get_stream() - - ctx = RequestContext( - client=client, - session_message=session_message, - metadata=metadata, - read_stream_writer=read_stream_writer, - ) - - async def handle_request_async( - is_resumption: bool = is_resumption, - ctx: RequestContext = ctx, - ) -> None: - if is_resumption: - await self._handle_resumption_request(ctx) - else: - await self._handle_post_request(ctx) - - # If this is a request, start a new task to handle it - if isinstance(message.root, JSONRPCRequest): - tg.start_soon(handle_request_async) - else: - await handle_request_async() - - except Exception: - logger.exception("Error in post_writer") # pragma: no cover - finally: - await read_stream_writer.aclose() - await write_stream.aclose() - - async def terminate_session( - self, client: httpx.AsyncClient - ) -> None: # pragma: no cover - """Terminate the session by sending a DELETE request.""" - if not await self._session_info.get_session_id(): - return - - try: - headers = await self._prepare_headers() - response = await client.delete(self.url, headers=headers) +@dataclass +class SessionIdentity: + """Mutable holder the transport reads on every request. - if response.status_code == 405: - logger.debug("Server does not allow session termination") - elif response.status_code not in (200, 204): - logger.warning(f"Session termination failed: {response.status_code}") - except Exception as exc: - logger.warning(f"Session termination failed: {exc}") + The transport is opened before negotiation runs, so an era-specific wire + cannot be fixed at construction time. A strategy narrows ``wire`` once it + knows the era, and in-flight requests pick that up on their next call. + """ - async def get_session_id(self) -> str | None: - """Get the current session ID.""" - return await self._session_info.get_session_id() + wire: SessionIdentityWire = field(default_factory=lambda: LEGACY_IDENTITY) @asynccontextmanager async def streamable_http_client( url: str, *, - http_client: httpx.AsyncClient | None = None, + http_client: httpx2.AsyncClient | None = None, terminate_on_close: bool = True, session_info: SessionInfo | None = None, -) -> AsyncGenerator[ - tuple[ - MemoryObjectReceiveStream[SessionMessage | Exception], - MemoryObjectSendStream[SessionMessage], - ], - None, -]: - """ - Client transport for StreamableHTTP. + identity: SessionIdentity | None = None, +) -> AsyncGenerator[tuple[Any, Any], None]: + """Open the SDK transport while synchronizing its session header externally. + + MCP 2 removed the transport's ``get_session_id`` callback. Request and + response hooks preserve UiPath's persisted-session behavior without + maintaining a private copy of the SDK transport. Args: - url: The MCP server endpoint URL. - http_client: Optional pre-configured httpx.AsyncClient. If None, a default - client with recommended MCP timeouts will be created. To configure headers, - authentication, or other HTTP settings, create an httpx.AsyncClient and pass it here. - terminate_on_close: If True, send a DELETE request to terminate the session - when the context exits. - session_info: Optional SessionInfo for external session ID tracking. - - Yields: - Tuple containing: - - read_stream: Stream for reading messages from the server - - write_stream: Stream for sending messages to the server + url: The MCP server endpoint. + http_client: An authenticated client to reuse. One is created and owned + when omitted. + terminate_on_close: Send ``DELETE`` for the session on exit. Already a + no-op in the modern era, which has no session to terminate. + session_info: Store for the session ID. A plain in-memory one is used + when omitted. + identity: How the ID travels on the wire. Defaults to the legacy + handshake behavior. """ - read_stream_writer, read_stream = anyio.create_memory_object_stream[ - SessionMessage | Exception - ](0) - write_stream, write_stream_reader = anyio.create_memory_object_stream[ - SessionMessage - ](0) - - # Determine if we need to create and manage the client - client_provided = http_client is not None - client = http_client - - if client is None: - # Create default client with recommended MCP timeouts - client = create_mcp_http_client() - - transport = StreamableHTTPTransport(url, session_info=session_info) - - async with anyio.create_task_group() as tg: + info = session_info or SessionInfo() + session_identity = identity or SessionIdentity() + owns_client = http_client is None + client = http_client or httpx2.AsyncClient( + # A caller that supplies no client still gets the repo's SSL and proxy + # configuration; only the MCP read timeout is layered on top, since SSE + # streams outlive a default one. + **{ + **get_httpx_client_kwargs(), + "timeout": httpx2.Timeout(30, read=300), + } + ) + restored_session_id = await info.get_session_id() + sdk_session_id: str | None = None + session_persistence_lock = asyncio.Lock() + + async def apply_session_id(request: httpx2.Request) -> None: + wire = session_identity.wire + session_id = await info.get_session_id() + for header in _IDENTITY_HEADERS: + if session_id is not None and header in wire.request_headers: + request.headers[header] = session_id + else: + request.headers.pop(header, None) + + async def capture_session_id(response: httpx2.Response) -> None: + nonlocal sdk_session_id + capture_header = session_identity.wire.capture_response_header + if capture_header is None: + return + session_id = response.headers.get(capture_header) + if session_id is None: + return try: - logger.debug(f"Connecting to StreamableHTTP endpoint: {url}") - - async with contextlib.AsyncExitStack() as stack: - # Only manage client lifecycle if we created it - if not client_provided: - await stack.enter_async_context(client) - - def start_get_stream() -> None: - tg.start_soon( - transport.handle_get_stream, client, read_stream_writer - ) - - tg.start_soon( - transport.post_writer, - client, - write_stream_reader, - read_stream_writer, - write_stream, - start_get_stream, - tg, - ) - + request_body = json.loads(response.request.content) + except (json.JSONDecodeError, UnicodeDecodeError): + request_body = None + if not ( + isinstance(request_body, dict) + and request_body.get("method") == "initialize" + ): + # Only the handshake assigns a session, which is how the SDK's own + # transport reads it too. Persisting from any response would let a + # proxy echoing the header replace a client-minted routing key -- + # reachable in ``auto`` mode, whose probe runs on the legacy wire + # before the era is known. + return + sdk_session_id = session_id + async with session_persistence_lock: + if await info.get_session_id() != session_id: + await info.set_session_id(session_id) + + async def terminate_restored_session() -> None: + if session_identity.wire.capture_response_header is None: + # This era's ID is minted by the client for routing, not assigned by + # the server, so there is no server-side session to terminate. + # Deleting it would tear down a live connection on a restored run. + return + current_session_id = await info.get_session_id() + if ( + not terminate_on_close + or restored_session_id is None + or current_session_id != restored_session_id + or current_session_id == sdk_session_id + ): + return + try: + await client.delete(url) + except Exception as error: # pragma: no cover - best-effort cleanup + logger.warning("Persisted MCP session termination failed: %s", error) + + client.event_hooks["request"].append(apply_session_id) + client.event_hooks["response"].append(capture_session_id) + try: + if owns_client: + async with client: + async with sdk_streamable_http_client( + url, + http_client=client, + terminate_on_close=terminate_on_close, + ) as streams: + try: + yield streams + finally: + await terminate_restored_session() + else: + async with sdk_streamable_http_client( + url, + http_client=client, + terminate_on_close=terminate_on_close, + ) as streams: try: - yield ( - read_stream, - write_stream, - ) + yield streams finally: - if await transport.get_session_id() and terminate_on_close: - await transport.terminate_session(client) - tg.cancel_scope.cancel() - finally: - await read_stream_writer.aclose() - await write_stream.aclose() - - -@asynccontextmanager -@deprecated("Use `streamable_http_client` instead.") -async def streamablehttp_client( - url: str, - headers: dict[str, str] | None = None, - timeout: float | timedelta = 30, - sse_read_timeout: float | timedelta = 60 * 5, - terminate_on_close: bool = True, - httpx_client_factory: McpHttpClientFactory = create_mcp_http_client, - auth: httpx.Auth | None = None, -) -> AsyncGenerator[ - tuple[ - MemoryObjectReceiveStream[SessionMessage | Exception], - MemoryObjectSendStream[SessionMessage], - ], - None, -]: - # Convert timeout parameters - timeout_seconds = ( - timeout.total_seconds() if isinstance(timeout, timedelta) else timeout - ) - sse_read_timeout_seconds = ( - sse_read_timeout.total_seconds() - if isinstance(sse_read_timeout, timedelta) - else sse_read_timeout - ) - - # Create httpx client using the factory with old-style parameters - client = httpx_client_factory( - headers=headers, - timeout=httpx.Timeout(timeout_seconds, read=sse_read_timeout_seconds), - auth=auth, - ) - - # Manage client lifecycle since we created it - async with client: - async with streamable_http_client( - url, - http_client=client, - terminate_on_close=terminate_on_close, - ) as streams: - yield streams + await terminate_restored_session() + finally: + client.event_hooks["request"].remove(apply_session_id) + client.event_hooks["response"].remove(capture_session_id) diff --git a/testcases/simple-http-mcp/input.json b/testcases/simple-http-mcp/input.json new file mode 100644 index 000000000..412f915b4 --- /dev/null +++ b/testcases/simple-http-mcp/input.json @@ -0,0 +1 @@ +{"a": 2, "b": 3} diff --git a/testcases/simple-http-mcp/langgraph.json b/testcases/simple-http-mcp/langgraph.json new file mode 100644 index 000000000..0d728beac --- /dev/null +++ b/testcases/simple-http-mcp/langgraph.json @@ -0,0 +1,6 @@ +{ + "dependencies": ["."], + "graphs": { + "agent": "./src/simple-http-mcp/graph.py:graph" + } +} diff --git a/testcases/simple-http-mcp/pyproject.toml b/testcases/simple-http-mcp/pyproject.toml new file mode 100644 index 000000000..d2285463f --- /dev/null +++ b/testcases/simple-http-mcp/pyproject.toml @@ -0,0 +1,34 @@ +[project] +name = "simple-http-mcp" +version = "0.0.1" +description = "MCP over Streamable HTTP protocol-version matrix testcase" +authors = [{ name = "John Doe", email = "john.doe@myemail.com" }] +dependencies = [ + "langgraph>=0.2.70", + "langchain-core>=0.3.34", + "langgraph-checkpoint-sqlite>=2.0.3", + "python-dotenv>=1.0.1", + "uipath-langchain", + "pydantic>=2.10.6", + "typing-extensions>=4.12.2", + "mcp==2.0.0", + "starlette>=0.40.0", + "uvicorn>=0.30.0", +] +requires-python = ">=3.11" + +[tool.uv] +exclude-newer = "2 days" + +[tool.uv.exclude-newer-package] +uipath = false +uipath-core = false +uipath-platform = false +uipath-runtime = false +uipath-dev = false +uipath-langchain-client = false +uipath-llm-client = false +jsonschema-pydantic-converter = false + +[tool.uv.sources] +uipath-langchain = { path = "../../", editable = true } diff --git a/testcases/simple-http-mcp/run.sh b/testcases/simple-http-mcp/run.sh new file mode 100644 index 000000000..2fd10311a --- /dev/null +++ b/testcases/simple-http-mcp/run.sh @@ -0,0 +1,17 @@ +#!/bin/bash +set -e + +echo "Syncing dependencies..." +uv sync + +echo "Authenticating with UiPath..." +uv run uipath auth --client-id="$CLIENT_ID" --client-secret="$CLIENT_SECRET" --base-url="$BASE_URL" + +echo "Initializing the project..." +uv run uipath init + +echo "Packing agent..." +uv run uipath pack + +echo "Running agent..." +uv run uipath run agent --file input.json diff --git a/testcases/simple-http-mcp/src/assert.py b/testcases/simple-http-mcp/src/assert.py new file mode 100644 index 000000000..92bee8807 --- /dev/null +++ b/testcases/simple-http-mcp/src/assert.py @@ -0,0 +1,291 @@ +import json +import os + +print("Checking MCP-over-HTTP protocol version matrix...") + +# Operands from input.json; the remote add tool must return their sum. +EXPECTED_SUM = "5" +MODERN_VERSION = "2026-07-28" + +# Check NuGet package +uipath_dir = ".uipath" +assert os.path.exists(uipath_dir), "NuGet package directory (.uipath) not found" + +nupkg_files = [f for f in os.listdir(uipath_dir) if f.endswith(".nupkg")] +assert nupkg_files, "NuGet package file (.nupkg) not found in .uipath directory" + +print(f"NuGet package found: {nupkg_files[0]}") + +# Check agent output file +output_file = "__uipath/output.json" +assert os.path.isfile(output_file), "Agent output file not found" + +print("Agent output file found") + +with open(output_file, "r", encoding="utf-8") as f: + output_data = json.load(f) + +status = output_data.get("status") +assert status == "successful", f"Agent execution failed with status: {status}" + +print("Agent execution status: successful") + +assert "output" in output_data, "Missing 'output' field in agent response" +output_content = output_data["output"] + +for field in ("results", "supported_versions", "unsupported_versions"): + assert field in output_content, f"Missing '{field}' field in output" + +results = output_content["results"] +assert isinstance(results, list) and results, "Results field is empty or not a list" + +by_label = {r["label"]: r for r in results} +EXPECTED_LEGS = ( + "legacy-sdk-server", + "legacy-pinned-2025-06-18", + "modern-sdk-server", + "modern-only-endpoint", + "auto-sdk-server", + "auto-pinned-2025-06-18", +) +for label in EXPECTED_LEGS: + assert label in by_label, f"No result recorded for leg {label}" + + +def check_leg_negotiated( + label: str, + version: str, + era: str, + expected_tool: str, + *, + server_session: bool, +) -> None: + """Every leg must negotiate its era, discover tools, and run one.""" + leg = by_label[label] + assert leg["supported"], ( + f"{label} should be supported but failed: {leg.get('error_message')}" + ) + assert leg["era"] == era, f"{label} resolved era {leg['era']}, expected {era}" + assert leg["negotiated_version"] == version, ( + f"{label} negotiated {leg['negotiated_version']} instead of {version}" + ) + # A legacy session ID comes from the server; a modern one is minted by the + # client for routing, so the wire must show no server-assigned session. + assert leg["server_session_id_seen"] is server_session, ( + f"{label} server_session_id_seen={leg['server_session_id_seen']}, " + f"expected {server_session}" + ) + assert expected_tool in leg["tools"], ( + f"{label} did not discover the '{expected_tool}' tool: {leg['tools']}" + ) + assert leg["tool_result"] == EXPECTED_SUM, ( + f"{label} tool returned {leg['tool_result']!r}, expected {EXPECTED_SUM!r}" + ) + print(f"{label}: {era} {version}, tool call returned {EXPECTED_SUM}") + + +# Legacy era: the server mints and returns a session ID. +check_leg_negotiated( + "legacy-sdk-server", "2025-11-25", "legacy", "multiply", server_session=True +) +check_leg_negotiated( + "legacy-pinned-2025-06-18", "2025-06-18", "legacy", "add", server_session=True +) + +# Modern era: session IDs are gone from the protocol entirely. The modern-only +# endpoint refuses the handshake, so this cannot be passing via legacy fallback. +check_leg_negotiated( + "modern-sdk-server", MODERN_VERSION, "modern", "multiply", server_session=False +) +check_leg_negotiated( + "modern-only-endpoint", MODERN_VERSION, "modern", "add", server_session=False +) + +# auto resolves a different era per server, preferring modern and falling back. +check_leg_negotiated( + "auto-sdk-server", MODERN_VERSION, "modern", "multiply", server_session=False +) +check_leg_negotiated( + "auto-pinned-2025-06-18", "2025-06-18", "legacy", "add", server_session=True +) + +assert output_content["unsupported_versions"] == [], ( + f"Every leg should negotiate now: {output_content['unsupported_versions']}" +) + +# --- gateway affinity ------------------------------------------------------- +# 2026-07-28 removes mcp-session-id, so the UiPath affinity ID is what lets a +# gateway keep routing to one warm serverless instance. +print("Checking modern-era instance affinity...") + +affinity = output_content.get("affinity") +assert affinity, "Missing 'affinity' field in output" +assert affinity["tool_results"] == [EXPECTED_SUM, EXPECTED_SUM], ( + f"Affinity leg tool results were {affinity['tool_results']}" +) +ids = affinity["affinity_ids"] +assert len(ids) == 2 and ids[0] and ids[0] == ids[1], ( + f"The affinity ID must persist across clients, got {ids}" +) +assert affinity["instances"] and len(affinity["instances"]) == 1, ( + f"Requests spread across instances {affinity['instances']}; affinity failed" +) +# The client mints the ID before negotiating, so even discovery is routable -- +# unlike a server-assigned session, which cannot reach the first request. +assert affinity["first_request_pinned"], ( + "The first request was not pinned; the affinity ID reached the gateway late" +) +assert affinity["requests"] > 0 and affinity["unpinned_requests"] == 0, ( + f"{affinity['unpinned_requests']} of {affinity['requests']} requests reached " + "the gateway with no affinity header, so it had to route them blind" +) +print( + f"Affinity ID {ids[0]} pinned {affinity['instances'][0]} for all " + f"{affinity['requests']} requests across both clients" +) + +print("Protocol version matrix validation passed") + +# --- uipath-agents-python compatibility ------------------------------------- +# Everything below pins the MCP API surface consumed by uipath-agents-python. +# A failure here means that repository breaks on its next dependency bump. + +print("Checking uipath-agents-python API compatibility...") + +assert "agents_api" in output_content, "Missing 'agents_api' field in output" +agents_api = output_content["agents_api"] +assert agents_api, "Downstream-compatibility leg produced no result" + +assert agents_api["imports"] == [ + "McpClient", + "SessionInfo", + "SessionInfoFactory", + "create_mcp_tools_and_clients", +], f"Unexpected downstream import list: {agents_api['imports']}" + +# SessionInfo.protocol_version is retained for compatibility even though nothing +# reads it now, so a subclass calling super().__init__() must still inherit it. +assert agents_api["session_info_super_init"], ( + "SessionInfo subclass did not inherit protocol_version from super().__init__()" +) + + +def check_leg( + name: str, + expected_tools: list[str], + *, + expected_version: str | None = None, + server_session: bool = True, +) -> dict: + """A downstream-shaped call must build tools, run one, and dispose cleanly. + + Args: + name: Field on the agents_api result holding this leg. + expected_tools: Tool names the leg must have built. + expected_version: Protocol version the live session must have settled + on, or None to skip the check. + server_session: Whether the server is allowed to assign an + ``mcp-session-id``. False for modern legs, which have no session. + """ + leg = agents_api[name] + assert leg["error_type"] is None, ( + f"{name} leg failed: {leg['error_type']}: {leg['error_message']}" + ) + assert leg["tools"] == expected_tools, ( + f"{name} leg built tools {leg['tools']}, expected {expected_tools}" + ) + assert leg["tool_result"] == EXPECTED_SUM, ( + f"{name} leg tool returned {leg['tool_result']!r}, expected {EXPECTED_SUM!r}" + ) + assert leg["session_id"], f"{name} leg did not establish a session id" + assert leg["disposed"], f"{name} leg did not dispose its McpClient" + if expected_version is not None: + assert leg["negotiated_version"] == expected_version, ( + f"{name} leg negotiated {leg['negotiated_version']!r}, expected " + f"{expected_version!r}" + ) + assert leg["server_session_issued"] is server_session, ( + f"{name} leg server_session_issued={leg['server_session_issued']}, " + f"expected {server_session}" + ) + print(f"{name}: tools {leg['tools']}, tool call returned {EXPECTED_SUM}, disposed") + return leg + + +# Production shape: no session_info_factory, terminate_on_close=True, dynamic +# discovery -- which reads the SDK's snake_case Tool.input_schema/output_schema. +check_leg("production", ["add", "multiply"], expected_version="2025-11-25") + +# Playground shape: SessionInfoFactory subclass, terminate_on_close=False, +# cached discovery with refresh_schema_before_call left at its default. +playground = check_leg("playground", ["add"], expected_version="2025-11-25") +resumed = check_leg("playground_resumed", ["add"], expected_version="2025-11-25") + +# The resumed session must speak the version it was originally negotiated at. +# The session ID surviving is only half the contract: probing candidate versions +# instead settles on the oldest handshake version and silently downgrades every +# later request, which no session-ID assertion can catch. +assert resumed["negotiated_version"] == playground["negotiated_version"], ( + f"Resumed session negotiated {resumed['negotiated_version']!r} but the " + f"session it resumed was negotiated at {playground['negotiated_version']!r}; " + "the resumed connection was silently downgraded" +) + +assert agents_api["debug_state_writes"] >= 1, ( + "The SessionInfo subclass never persisted a session id" +) +persisted = agents_api["persisted_session_id"] +assert persisted == playground["session_id"], ( + f"Persisted session id {persisted!r} does not match the playground leg's " + f"{playground['session_id']!r}" +) +# The second client must adopt the persisted session rather than start a new +# one; this is what playground mode relies on across runs. +assert agents_api["session_resumed"], ( + f"Second client did not resume the persisted session: " + f"{resumed['session_id']!r} != {persisted!r}" +) +print(f"Persisted session {persisted} resumed by a second McpClient") + +# --- McpClient on the modern era -------------------------------------------- +# 2026-07-28 has no session identity at all, so no response on these legs may +# assign one; the ID in play is the client-minted affinity ID. +print("Checking McpClient on the 2026-07-28 era...") + +check_leg( + "modern", + ["add", "multiply"], + expected_version=MODERN_VERSION, + server_session=False, +) +affinity_first = check_leg( + "modern_affinity", + ["add"], + expected_version=MODERN_VERSION, + server_session=False, +) +affinity_resumed = check_leg( + "modern_affinity_resumed", + ["add"], + expected_version=MODERN_VERSION, + server_session=False, +) + +# Two clients sharing one SessionInfo must keep the same affinity ID, so a +# gateway routes both runs to the same warm instance. +assert agents_api["modern_affinity_id_reused"], ( + f"The second modern client sent affinity ID " + f"{affinity_resumed['session_id']!r} instead of reusing " + f"{affinity_first['session_id']!r}" +) +assert affinity_resumed["negotiated_version"] == affinity_first["negotiated_version"], ( + f"Resumed modern client negotiated " + f"{affinity_resumed['negotiated_version']!r} instead of " + f"{affinity_first['negotiated_version']!r}" +) +print( + f"Affinity ID {affinity_first['session_id']} reused by a second modern " + f"McpClient at {MODERN_VERSION}" +) + +print("uipath-agents-python API compatibility validation passed") diff --git a/testcases/simple-http-mcp/src/simple-http-mcp/agents_api.py b/testcases/simple-http-mcp/src/simple-http-mcp/agents_api.py new file mode 100644 index 000000000..a87b93952 --- /dev/null +++ b/testcases/simple-http-mcp/src/simple-http-mcp/agents_api.py @@ -0,0 +1,557 @@ +"""Pin the MCP public API that ``uipath-agents-python`` consumes. + +`uipath-agents-python` is the only known downstream consumer of the MCP tool +layer, and it pins ``uipath-langchain`` exactly, so a break there surfaces on +its next bump rather than in this repository's own tests. This leg exercises +the API it actually calls, over real sockets, against a real SDK ``MCPServer``: + +* ``create_mcp_tools_and_clients(resources, session_info_factory=..., terminate_on_close=...)`` + -- both call shapes it uses: production (no factory, terminate on close) and + playground (debug-state factory, session kept alive). +* ``SessionInfo`` subclassed with its own ``__init__`` and async + ``get_session_id`` / ``set_session_id`` overrides, persisting through HTTP. + This is the shape of ``SessionInfoDebugState``; MCP 2 added + ``SessionInfo.protocol_version``, which the transport now reads directly, so + such a subclass fails at the first request unless it calls ``super().__init__()``. +* ``SessionInfoFactory.create_session(mcp_server)`` reading ``McpServer.slug`` + and ``McpServer.folder_key``. +* Session resumption: a second client picking up the session ID the first one + persisted, which is what playground mode relies on across runs. +* Disposal through ``McpClient.dispose()``, the way the caller drains its + ``UiPathDisposableProtocol`` list. +* ``McpClient(protocol_mode="modern")`` against the same real server: the era + has no session identity, so no response may assign one, and every leg records + the version its live session actually negotiated. Note that + ``create_mcp_tools_and_clients`` exposes no ``protocol_mode``, so a downstream + caller reaching ``2026-07-28`` has to construct the client itself. +* A modern affinity pair: two clients sharing one ``SessionInfo``, standing in + for two playground runs that must land on the same warm instance. + +The mirrored source lives in ``uipath_agents/agent_graph_builder/`` -- +``graph.py`` and ``session_info_debug_state.py``. Keep this file in step with it. +""" + +import contextlib +import logging +import os +from collections.abc import Iterator +from enum import Enum, auto +from typing import Any +from urllib.parse import quote + +import httpx +from pydantic import BaseModel, Field +from uipath._utils._ssl_context import get_httpx_client_kwargs +from uipath.agent.models.agent import ( + AgentMcpResourceConfig, + AgentMcpTool, + CachedToolsConfig, + DynamicToolsConfig, + ToolsConfiguration, +) +from uipath.platform.orchestrator.mcp import McpServer + +# Imported exactly as uipath-agents-python imports them. Dropping or renaming +# any of these turns this testcase into an ImportError at graph load. +from uipath_langchain.agent.tools.mcp import ( + McpClient, + SessionInfo, + SessionInfoFactory, + create_mcp_tools_and_clients, +) + +# Not part of the downstream import surface asserted below: the modern legs need +# it because `create_mcp_tools_and_clients` has no `protocol_mode` parameter, so +# reaching 2026-07-28 means constructing the `McpClient` directly. +from uipath_langchain.agent.tools.mcp import create_mcp_tools + +logger = logging.getLogger(__name__) + +DOWNSTREAM_IMPORTS = [ + "McpClient", + "SessionInfo", + "SessionInfoFactory", + "create_mcp_tools_and_clients", +] + +AGENT_ID = "agent-under-test" +FOLDER_KEY = "folder-key" +FOLDER_PATH = "Shared" +SERVER_SLUG = "math" +SERVER_NAME = "Math" +ACCESS_TOKEN = "test-access-token" + + +# --- SessionInfoDebugState, mirrored from uipath-agents-python --------------- + + +class _SessionState(Enum): + """State machine for the session ID lifecycle, as downstream defines it.""" + + NOT_LOADED = auto() + LOADED = auto() + CLEARED = auto() + + +class SessionInfoDebugState(SessionInfo): + """``SessionInfo`` subclass that persists session IDs over HTTP. + + Deliberately keeps downstream's structure: an ``__init__`` of its own that + calls ``super().__init__()``, direct assignment to the inherited + ``session_id`` attribute, and async overrides of both accessors. + """ + + def __init__(self, slug: str, folder_key: str, agent_id: str | None) -> None: + super().__init__() + self._slug = slug + self._folder_key = folder_key + self._agent_id = agent_id + self._state = _SessionState.NOT_LOADED + + @property + def key(self) -> str: + """Debug-state key for this MCP resource.""" + return f"mcpsession:{self._folder_key}:{self._slug}" + + async def get_session_id(self) -> str | None: + """Return the session ID, loading it from debug state on first call.""" + if self._state == _SessionState.NOT_LOADED: + self._state = _SessionState.LOADED + stored = await self._load_from_debug_state() + if stored is not None: + self.session_id = stored + elif self._state == _SessionState.CLEARED: + return None + return self.session_id + + async def set_session_id(self, session_id: str | None) -> None: + """Store the session ID locally and persist it to debug state.""" + if session_id is None: + self.session_id = None + self._state = _SessionState.CLEARED + else: + self.session_id = session_id + self._state = _SessionState.LOADED + await self._save_to_debug_state(session_id) + + def _debug_state_url(self) -> str | None: + base_url = os.getenv("UIPATH_URL") + if not base_url or not self._agent_id: + return None + encoded_key = quote(self.key, safe="") + return f"{base_url}/agenthub_/design/debugstate/{self._agent_id}/{encoded_key}" + + def _auth_headers(self) -> dict[str, str]: + token = os.getenv("UIPATH_ACCESS_TOKEN", "") + return {"Authorization": f"Bearer {token}"} + + async def _load_from_debug_state(self) -> str | None: + url = self._debug_state_url() + if url is None: + return None + async with httpx.AsyncClient( + headers=self._auth_headers(), **get_httpx_client_kwargs() + ) as client: + response = await client.get(url) + if response.status_code == 200: + return response.text + return None + + async def _save_to_debug_state(self, session_id: str) -> None: + url = self._debug_state_url() + if url is None: + return + async with httpx.AsyncClient( + headers=self._auth_headers(), **get_httpx_client_kwargs() + ) as client: + await client.put( + url, content=session_id, headers={"Content-Type": "text/plain"} + ) + + +class SessionInfoDebugStateFactory(SessionInfoFactory): + """Factory returning ``SessionInfoDebugState``, as downstream defines it.""" + + def __init__(self, agent_id: str | None) -> None: + self._agent_id = agent_id or os.getenv("UIPATH_PROJECT_ID") + + def create_session(self, mcp_server: McpServer) -> SessionInfoDebugState: + """Create a SessionInfoDebugState from an McpServer.""" + return SessionInfoDebugState( + slug=mcp_server.slug or "", + folder_key=mcp_server.folder_key or "", + agent_id=self._agent_id, + ) + + +# --- Stand-in for the UiPath SDK lookup McpClient performs lazily ------------ + + +class _FakeMcpService: + def __init__(self, url: str) -> None: + self._url = url + + async def retrieve_async( + self, name: str, folder_path: str | None = None + ) -> McpServer: + return McpServer( + id="mcp-server-id", + name=name, + slug=SERVER_SLUG, + folderKey=FOLDER_KEY, + mcpUrl=self._url, + ) + + +class _FakeConfig: + secret = ACCESS_TOKEN + + +@contextlib.contextmanager +def _patched_sdk(url: str) -> Iterator[None]: + """Point ``McpClient``'s lazy SDK lookup at the local test server. + + ``McpClient._initialize_client`` imports ``UiPath`` from ``uipath.platform`` + at call time, so replacing the module attribute is enough; no tenant or + network access to UiPath Cloud is involved. + """ + import uipath.platform as platform + + class _FakeUiPath: + def __init__(self, *args: Any, **kwargs: Any) -> None: + self._config = _FakeConfig() + self.mcp = _FakeMcpService(url) + + original = platform.UiPath + platform.UiPath = _FakeUiPath # type: ignore[misc] + try: + yield + finally: + platform.UiPath = original # type: ignore[misc] + + +@contextlib.contextmanager +def _agenthub_env(base_url: str) -> Iterator[None]: + """Set the environment ``SessionInfoDebugState`` reads its endpoint from.""" + previous = { + key: os.environ.get(key) for key in ("UIPATH_URL", "UIPATH_ACCESS_TOKEN") + } + os.environ["UIPATH_URL"] = base_url + os.environ["UIPATH_ACCESS_TOKEN"] = ACCESS_TOKEN + try: + yield + finally: + for key, value in previous.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +# --- Resource configs, as an agent definition would carry them --------------- + + +def _cached_resource() -> AgentMcpResourceConfig: + """Design-time tool snapshot; the default cached discovery mode.""" + return AgentMcpResourceConfig( + name=SERVER_NAME, + description="Math MCP server", + folderPath=FOLDER_PATH, + slug=SERVER_SLUG, + availableTools=[ + AgentMcpTool( + name="add", + description="Add two numbers", + inputSchema={ + "type": "object", + "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}}, + "required": ["a", "b"], + }, + ) + ], + toolsConfiguration=ToolsConfiguration(discoveryMode=CachedToolsConfig()), + ) + + +def _dynamic_resource() -> AgentMcpResourceConfig: + """Live tool discovery; reads the SDK's snake_case ``Tool`` attributes.""" + return AgentMcpResourceConfig( + name=SERVER_NAME, + description="Math MCP server", + folderPath=FOLDER_PATH, + slug=SERVER_SLUG, + availableTools=[], + toolsConfiguration=ToolsConfiguration( + discoveryMode=DynamicToolsConfig(allowAll=True) + ), + ) + + +# --- Results ---------------------------------------------------------------- + + +class LegSummary(BaseModel): + """Outcome of one ``create_mcp_tools_and_clients`` call.""" + + label: str + protocol_mode: str = "legacy" + tools: list[str] = Field(default_factory=list) + tool_result: str | None = None + session_id: str | None = None + # What the live MCP session actually settled on. A resumed session that + # silently downgrades still returns the right session ID, so the ID alone + # cannot catch it -- this is the field that can. + negotiated_version: str | None = None + # Whether the server assigned an ``mcp-session-id`` on any response during + # this leg. Modern legs must show none: the ID in play is client-minted. + server_session_issued: bool = False + disposed: bool = False + error_type: str | None = None + error_message: str | None = None + + +class AgentsApiResult(BaseModel): + """Everything the downstream compatibility leg asserts on.""" + + imports: list[str] + session_info_super_init: bool + production: LegSummary + playground: LegSummary + playground_resumed: LegSummary + modern: LegSummary + modern_affinity: LegSummary + modern_affinity_resumed: LegSummary + debug_state_writes: int + debug_state_reads: int + persisted_session_id: str | None = None + session_resumed: bool = False + modern_affinity_id_reused: bool = False + + +def _first_text(blocks: Any) -> str | None: + """Pull the first text block out of a normalized MCP tool result.""" + if isinstance(blocks, list): + for block in blocks: + if isinstance(block, dict) and block.get("type") == "text": + return str(block.get("text")) + return None if blocks is None else str(blocks) + + +def _negotiated_version(client: McpClient) -> str | None: + """Read the protocol version the client's live MCP session settled on. + + ``ClientSession.protocol_version`` is public SDK API, but ``McpClient`` + keeps its session private and exposes no accessor, so this reaches through + ``_session``. Must be read before ``dispose()`` releases the session. + """ + session = getattr(client, "_session", None) + version = getattr(session, "protocol_version", None) + return None if version is None else str(version) + + +async def _run_leg( + label: str, + resource: AgentMcpResourceConfig, + session_info_factory: SessionInfoFactory | None, + terminate_on_close: bool, + operands: tuple[int, int], + protocol_mode: str = "legacy", + recorder: Any = None, +) -> LegSummary: + """Drive one downstream-shaped call from tool creation through disposal. + + Args: + label: Name this leg is reported under. + resource: The MCP resource config the caller would carry. + session_info_factory: Downstream's factory, or None for the default. + terminate_on_close: Whether disposal should terminate the session. + operands: Operands handed to the remote ``add`` tool. + protocol_mode: Negotiation era to drive -- ``legacy`` or ``modern``. + recorder: Optional ``SessionHeaderRecorder`` wrapping the server, used + to tell a server-assigned session ID from a client-minted one. + """ + a, b = operands + seen_before = 0 if recorder is None else len(recorder.response_session_ids) + clients: list[McpClient] = [] + try: + if protocol_mode == "legacy": + # Exactly the call shape uipath-agents-python uses. + tools, clients = await create_mcp_tools_and_clients( + [resource], + session_info_factory=session_info_factory, + terminate_on_close=terminate_on_close, + ) + else: + # `create_mcp_tools_and_clients` exposes no `protocol_mode`, so a + # caller reaching the modern era builds the client itself. This is + # what the same code path looks like from downstream today. + client = McpClient( + config=resource, + session_info_factory=session_info_factory, + terminate_on_close=terminate_on_close, + protocol_mode=protocol_mode, # type: ignore[arg-type] + ) + clients = [client] + tools = await create_mcp_tools(resource, client) + add_tool = next(tool for tool in tools if tool.name == "add") + blocks = await add_tool.ainvoke({"a": a, "b": b}) + summary = LegSummary( + label=label, + protocol_mode=protocol_mode, + tools=sorted(tool.name for tool in tools), + tool_result=_first_text(blocks), + session_id=await clients[0].get_session_id(), + negotiated_version=_negotiated_version(clients[0]), + ) + except Exception as error: # noqa: BLE001 - the failure IS the result + logger.exception("Downstream-compatibility leg %s failed", label) + summary = LegSummary( + label=label, + protocol_mode=protocol_mode, + error_type=type(error).__name__, + error_message=str(error)[:300], + ) + if recorder is not None: + summary.server_session_issued = ( + len(recorder.response_session_ids) > seen_before + ) + + # Mirrors how the caller drains its UiPathDisposableProtocol list. Disposal + # is part of the contract, so a failure here is a leg failure. + try: + for client in clients: + await client.dispose() + summary.disposed = True + except Exception as error: # noqa: BLE001 - the failure IS the result + summary.error_type = summary.error_type or type(error).__name__ + summary.error_message = summary.error_message or str(error)[:300] + return summary + + +class _SharedSessionFactory(SessionInfoFactory): + """Hand every client the same ``SessionInfo``. + + Two clients sharing one store is how a modern affinity ID survives across + runs -- the playground case, where the store outlives the process. + """ + + def __init__(self, session_info: SessionInfo) -> None: + self._session_info = session_info + + def create_session(self, mcp_server: McpServer) -> SessionInfo: + """Return the shared SessionInfo regardless of the server.""" + return self._session_info + + +async def run_agents_api_leg( + serve: Any, + app: Any, + store: Any, + operands: tuple[int, int], + recorder: Any = None, +) -> AgentsApiResult: + """Run every downstream-shaped call against one hosted MCP server. + + Args: + serve: The ``servers.serve`` async context manager. + app: The ASGI app hosting the MCP server and the debug-state route. + store: The ``DebugStateStore`` mounted on the inner app. + operands: Operands handed to the remote ``add`` tool. + recorder: Optional ``servers.SessionHeaderRecorder`` wrapping *app*, + used to tell a server-assigned session ID from a client-minted one. + """ + probe = SessionInfoDebugState( + slug=SERVER_SLUG, folder_key=FOLDER_KEY, agent_id=None + ) + # MCP 2 added this attribute and the transport reads it directly; a subclass + # that skipped super().__init__() would raise AttributeError on first request. + super_init_ok = getattr(probe, "protocol_version", "missing") is None + + async with serve(app) as url: + base_url = url.removesuffix("/mcp") + with _patched_sdk(url), _agenthub_env(base_url): + production = await _run_leg( + "production", + _dynamic_resource(), + session_info_factory=None, + terminate_on_close=True, + operands=operands, + recorder=recorder, + ) + + factory = SessionInfoDebugStateFactory(agent_id=AGENT_ID) + playground = await _run_leg( + "playground", + _cached_resource(), + session_info_factory=factory, + terminate_on_close=False, + operands=operands, + recorder=recorder, + ) + persisted = store.values.get(probe.key) + + # A fresh factory stands in for the next playground run: it must pick + # the persisted session ID back up rather than start a new session. + resumed = await _run_leg( + "playground-resumed", + _cached_resource(), + session_info_factory=SessionInfoDebugStateFactory(agent_id=AGENT_ID), + terminate_on_close=False, + operands=operands, + recorder=recorder, + ) + + # 2026-07-28 through the same public client. The era has no session + # identity, so nothing here may see a server-assigned one. + modern = await _run_leg( + "modern", + _dynamic_resource(), + session_info_factory=None, + terminate_on_close=True, + operands=operands, + protocol_mode="modern", + recorder=recorder, + ) + + # Two modern clients sharing one store: the affinity ID the first + # mints must be the one the second sends, so a gateway keeps routing + # both runs to the same warm instance. + shared = SessionInfo() + shared_factory = _SharedSessionFactory(shared) + modern_affinity = await _run_leg( + "modern-affinity", + _cached_resource(), + session_info_factory=shared_factory, + terminate_on_close=False, + operands=operands, + protocol_mode="modern", + recorder=recorder, + ) + modern_affinity_resumed = await _run_leg( + "modern-affinity-resumed", + _cached_resource(), + session_info_factory=shared_factory, + terminate_on_close=False, + operands=operands, + protocol_mode="modern", + recorder=recorder, + ) + + return AgentsApiResult( + imports=DOWNSTREAM_IMPORTS, + session_info_super_init=super_init_ok, + production=production, + playground=playground, + playground_resumed=resumed, + modern=modern, + modern_affinity=modern_affinity, + modern_affinity_resumed=modern_affinity_resumed, + debug_state_writes=store.writes, + debug_state_reads=store.reads, + persisted_session_id=persisted, + session_resumed=(persisted is not None and resumed.session_id == persisted), + modern_affinity_id_reused=bool( + modern_affinity.session_id + and modern_affinity.session_id == modern_affinity_resumed.session_id + ), + ) diff --git a/testcases/simple-http-mcp/src/simple-http-mcp/graph.py b/testcases/simple-http-mcp/src/simple-http-mcp/graph.py new file mode 100644 index 000000000..4caf538c2 --- /dev/null +++ b/testcases/simple-http-mcp/src/simple-http-mcp/graph.py @@ -0,0 +1,437 @@ +"""Verify the UiPath MCP adapter against HTTP-hosted servers on several protocol versions. + +This testcase deliberately uses no LLM. It drives ``uipath_langchain``'s own +``streamable_http_client`` and ``SessionInfo`` -- not the raw SDK transport -- +over real sockets, so it covers the adapter that the unit tests can only reach +through ``httpx2.MockTransport``. + +Every leg runs through ``build_protocol_strategy``, so the matrix exercises the +same negotiation code ``McpClient`` uses rather than a parallel implementation. + +* ``legacy`` against a genuine SDK ``MCPServer``, which negotiates ``2025-11-25``. +* ``legacy`` against an endpoint pinned to ``2025-06-18``, proving the server's + counter-offer is honored. +* ``modern`` against the same real server, which negotiates ``2026-07-28`` and + issues no session ID. +* ``modern`` against an endpoint that serves only ``server/discover`` and refuses + the handshake, so the modern path cannot be passing by legacy fallback. +* ``auto`` against both a discovery-capable and a handshake-only server, checking + it resolves to a different era for each. + +A further leg drives a gateway stand-in to check that the UiPath affinity ID pins +one warm instance across separate clients -- the routing that ``mcp-session-id`` +used to provide and ``2026-07-28`` removes. + +A final leg pins the public API that ``uipath-agents-python`` consumes -- see +``agents_api.py`` -- so a break in the only known downstream consumer shows up +here rather than on its next dependency bump. That leg also drives ``McpClient`` +itself on both eras, including a modern affinity pair whose two clients share one +``SessionInfo``. +""" + +import importlib.util +import logging +from pathlib import Path +from typing import Any + +import httpx2 +from langgraph.graph import END, START, StateGraph +from mcp import ClientSession +from pydantic import BaseModel, Field + +from uipath_langchain.agent.tools.mcp import load_mcp_tools +from uipath_langchain.agent.tools.mcp.protocol_strategy import build_protocol_strategy +from uipath_langchain.agent.tools.mcp.streamable_http import ( + SessionInfo, + streamable_http_client, +) + +logger = logging.getLogger(__name__) + +MODERN_VERSION = "2026-07-28" + + +def _load_sibling(name: str) -> Any: + """Load a sibling module by path. + + The source directory is not an importable package name, so a path-based + load keeps this working however the graph module itself was loaded. + """ + path = Path(__file__).with_name(f"{name}.py") + spec = importlib.util.spec_from_file_location(f"mcp_testcase_{name}", path) + if spec is None or spec.loader is None: # pragma: no cover - defensive + raise RuntimeError(f"Cannot load MCP testcase module from {path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +servers = _load_sibling("servers") +agents_api = _load_sibling("agents_api") +AgentsApiResult = agents_api.AgentsApiResult + + +class GraphInput(BaseModel): + """Operands passed to the remote ``add`` tool on every supported leg.""" + + a: int = Field(default=2, description="First operand for the add tool") + b: int = Field(default=3, description="Second operand for the add tool") + + +class LegResult(BaseModel): + """Outcome of one protocol-version leg.""" + + label: str + protocol_version: str + server: str + supported: bool + mode: str = "legacy" + era: str | None = None + negotiated_version: str | None = None + session_id_issued: bool = False + server_session_id_seen: bool = False + tools: list[str] = Field(default_factory=list) + tool_result: str | None = None + error_type: str | None = None + error_code: int | None = None + error_message: str | None = None + + +class AffinityResult(BaseModel): + """Outcome of the gateway-routing leg.""" + + affinity_ids: list[str] = Field(default_factory=list) + instances: list[str] = Field(default_factory=list) + requests: int = 0 + unpinned_requests: int = 0 + first_request_pinned: bool = False + tool_results: list[str] = Field(default_factory=list) + + +class GraphOutput(BaseModel): + """Full matrix result.""" + + results: list[LegResult] + supported_versions: list[str] + unsupported_versions: list[str] + agents_api: AgentsApiResult | None = None + affinity: AffinityResult | None = None + + +class GraphState(BaseModel): + """Workflow state.""" + + a: int + b: int + results: list[LegResult] = Field(default_factory=list) + agents_api: AgentsApiResult | None = None + affinity: AffinityResult | None = None + + +def _unwrap_mcp_error(error: BaseException) -> BaseException: + """Return the meaningful error inside nested ``ExceptionGroup`` wrappers. + + A failing handshake surfaces as an ``ExceptionGroup`` from the transport's + task group; the useful ``MCPError`` sits several levels down. + """ + current = error + for _ in range(6): + nested = getattr(current, "exceptions", None) + if not nested: + break + current = nested[0] + return current + + +def _first_text(blocks: Any) -> str | None: + """Pull the first text block out of a LangChain tool result.""" + if isinstance(blocks, list): + for block in blocks: + if isinstance(block, dict) and block.get("type") == "text": + return str(block.get("text")) + return None if blocks is None else str(blocks) + + +async def _run_leg( + label: str, + protocol_version: str, + server_kind: str, + app: Any, + operands: tuple[int, int], + mode: str = "legacy", +) -> LegResult: + """Connect to one HTTP-hosted server and record what the strategy negotiated.""" + session_info = SessionInfo() + strategy = build_protocol_strategy(mode) # type: ignore[arg-type] + a, b = operands + server_session_ids: list[str] = [] + + async def watch_response(response: Any) -> None: + """Note whether the server ever assigned a session ID on this leg.""" + if response.headers.get("mcp-session-id") is not None: + server_session_ids.append(response.headers["mcp-session-id"]) + + async with servers.serve(app) as url: + logger.info("Probing %s at %s (mode=%s)", label, url, mode) + client = httpx2.AsyncClient( + follow_redirects=True, timeout=httpx2.Timeout(30, read=300) + ) + client.event_hooks["response"].append(watch_response) + try: + async with client: + async with streamable_http_client( + url, + http_client=client, + session_info=session_info, + identity=strategy.identity, + ) as (read, write): + async with ClientSession(read, write) as session: + await strategy.connect(session, session_info) + tools = await load_mcp_tools(session) + add_tool = next(t for t in tools if t.name == "add") + blocks = await add_tool.ainvoke({"a": a, "b": b}) + session_id = await session_info.get_session_id() + return LegResult( + label=label, + protocol_version=protocol_version, + server=server_kind, + supported=True, + mode=mode, + era=( + "modern" + if session.discover_result is not None + else "legacy" + ), + negotiated_version=str(session.protocol_version), + session_id_issued=session_id is not None, + server_session_id_seen=bool(server_session_ids), + tools=sorted(t.name for t in tools), + tool_result=_first_text(blocks), + ) + except BaseException as error: # noqa: BLE001 - the failure IS the result + inner = _unwrap_mcp_error(error) + logger.info("%s is unsupported: %s", label, inner) + session_id = await session_info.get_session_id() + return LegResult( + label=label, + protocol_version=protocol_version, + server=server_kind, + supported=False, + mode=mode, + session_id_issued=session_id is not None, + server_session_id_seen=bool(server_session_ids), + error_type=type(inner).__name__, + error_code=getattr(inner, "code", None), + error_message=str(inner)[:300], + ) + + +async def run_matrix(state: GraphState) -> GraphState: + """Run every protocol-version leg in turn.""" + operands = (state.a, state.b) + results = [ + await _run_leg( + "legacy-sdk-server", + "2025-11-25", + "real MCPServer over Streamable HTTP", + servers.build_sdk_app(), + operands, + mode="legacy", + ), + await _run_leg( + "legacy-pinned-2025-06-18", + "2025-06-18", + "endpoint pinned to 2025-06-18", + servers.PinnedVersionServer("2025-06-18").build_app(), + operands, + mode="legacy", + ), + await _run_leg( + "modern-sdk-server", + MODERN_VERSION, + "real MCPServer over Streamable HTTP", + servers.build_sdk_app(), + operands, + mode="modern", + ), + await _run_leg( + "modern-only-endpoint", + MODERN_VERSION, + "endpoint serving only server/discover", + servers.PinnedVersionServer(MODERN_VERSION, modern_only=True).build_app(), + operands, + mode="modern", + ), + await _run_leg( + "auto-sdk-server", + MODERN_VERSION, + "real MCPServer over Streamable HTTP", + servers.build_sdk_app(), + operands, + mode="auto", + ), + await _run_leg( + "auto-pinned-2025-06-18", + "2025-06-18", + "endpoint pinned to 2025-06-18", + servers.PinnedVersionServer("2025-06-18").build_app(), + operands, + mode="auto", + ), + ] + return GraphState( + a=state.a, + b=state.b, + results=results, + agents_api=state.agents_api, + affinity=state.affinity, + ) + + +class _RoutingGateway: + """Stand-in for AgentHub: pins a warm instance by ``mcp-session-id``. + + ``2026-07-28`` drops that header from the protocol, but UiPath keeps sending + it with a client-minted value purely as a routing key -- so the gateway routes + on exactly the header it already uses today, with no change. A request + arriving without one cannot be pinned. + + Written as a pure ASGI middleware rather than a + ``starlette.middleware.base.BaseHTTPMiddleware`` subclass: the latter + buffers through an inner task and breaks the server's SSE stream with + ``ASGI callable returned without completing response``. + """ + + def __init__(self, app: Any) -> None: + self.app = app + self.routed: list[tuple[str, bool]] = [] + self.instances: dict[str, str] = {} + + async def __call__(self, scope: Any, receive: Any, send: Any) -> None: + """Record the instance this request would be routed to, then forward it.""" + if scope["type"] != "http": + await self.app(scope, receive, send) + return + token: str | None = None + for key, value in scope.get("headers", ()): + if key.decode("latin-1").lower() == "mcp-session-id": + token = value.decode("latin-1") + pinned = token is not None + if token is None: + # Unroutable: a real gateway would have to pick an instance blind, so + # give every such request its own to make the spread visible. + token = f"unpinned-{len(self.routed)}" + if token not in self.instances: + self.instances[token] = f"instance-{len(self.instances) + 1}" + self.routed.append((self.instances[token], pinned)) + await self.app(scope, receive, send) + + +async def run_affinity(state: GraphState) -> GraphState: + """Check the affinity ID pins one instance across separate modern clients. + + Two clients share one ``SessionInfo``, standing in for two runs of a + playground agent whose session store survives the process. + """ + gateway: dict[str, _RoutingGateway] = {} + + class _Capturing(_RoutingGateway): + def __init__(self, app: Any) -> None: + super().__init__(app) + gateway["value"] = self + + app = servers.build_sdk_app() + app.add_middleware(_Capturing) + + shared_info = SessionInfo() + affinity_ids: list[str] = [] + tool_results: list[str] = [] + async with servers.serve(app) as url: + for _ in range(2): + strategy = build_protocol_strategy("modern") + async with streamable_http_client( + url, session_info=shared_info, identity=strategy.identity + ) as (read, write): + async with ClientSession(read, write) as session: + await strategy.connect(session, shared_info) + tools = await load_mcp_tools(session) + add_tool = next(t for t in tools if t.name == "add") + blocks = await add_tool.ainvoke({"a": state.a, "b": state.b}) + tool_results.append(str(_first_text(blocks))) + affinity_ids.append(str(await shared_info.get_session_id())) + + observed = gateway["value"].routed + return GraphState( + a=state.a, + b=state.b, + results=state.results, + agents_api=state.agents_api, + affinity=AffinityResult( + affinity_ids=affinity_ids, + instances=sorted({instance for instance, _ in observed}), + requests=len(observed), + unpinned_requests=sum(1 for _, pinned in observed if not pinned), + # The client mints the ID before negotiating, so even the discovery + # probe carries it -- a server-assigned session never could. + first_request_pinned=bool(observed) and observed[0][1], + tool_results=tool_results, + ), + ) + + +async def run_agents_api(state: GraphState) -> GraphState: + """Exercise the MCP API surface that ``uipath-agents-python`` consumes. + + Kept as its own node so a downstream-compatibility break is reported + separately from a protocol-version regression. + """ + app = servers.build_sdk_app() + store = servers.DebugStateStore() + store.attach(app) + # Wrapping the app is the only way to tell a server-assigned session ID from + # the client-minted affinity ID: both are opaque hex on the wire. + recorder = servers.SessionHeaderRecorder(app) + result = await agents_api.run_agents_api_leg( + servers.serve, recorder, store, (state.a, state.b), recorder=recorder + ) + return GraphState( + a=state.a, + b=state.b, + results=state.results, + agents_api=result, + affinity=state.affinity, + ) + + +def build_output(state: GraphState) -> GraphOutput: + """Summarize the matrix for assertions.""" + return GraphOutput( + results=state.results, + supported_versions=[r.protocol_version for r in state.results if r.supported], + unsupported_versions=[ + r.protocol_version for r in state.results if not r.supported + ], + agents_api=state.agents_api, + affinity=state.affinity, + ) + + +def _prepare(graph_input: GraphInput) -> GraphState: + """Seed the workflow state from the graph input.""" + return GraphState(a=graph_input.a, b=graph_input.b) + + +builder = StateGraph(GraphState, input_schema=GraphInput, output_schema=GraphOutput) +builder.add_node("prepare", _prepare) +builder.add_node("run_matrix", run_matrix) +builder.add_node("run_affinity", run_affinity) +builder.add_node("run_agents_api", run_agents_api) +builder.add_node("summarize", build_output) +builder.add_edge(START, "prepare") +builder.add_edge("prepare", "run_matrix") +builder.add_edge("run_matrix", "run_affinity") +builder.add_edge("run_affinity", "run_agents_api") +builder.add_edge("run_agents_api", "summarize") +builder.add_edge("summarize", END) + +graph = builder.compile() diff --git a/testcases/simple-http-mcp/src/simple-http-mcp/servers.py b/testcases/simple-http-mcp/src/simple-http-mcp/servers.py new file mode 100644 index 000000000..740673e39 --- /dev/null +++ b/testcases/simple-http-mcp/src/simple-http-mcp/servers.py @@ -0,0 +1,328 @@ +"""MCP servers hosted over Streamable HTTP for the protocol-version matrix. + +Two kinds of server are needed: + +* :func:`build_sdk_app` hosts a genuine ``MCPServer`` over Streamable HTTP. It + is the realistic leg, but it negotiates whatever the client asks for, so it + always settles on the latest legacy handshake version. +* :class:`PinnedVersionServer` declares one specific ``protocolVersion``. The + SDK server exposes no version knob, so pinning an older version -- or + refusing the legacy handshake the way a modern-only server would -- requires + a small endpoint that speaks the wire protocol directly. +""" + +import asyncio +import contextlib +import json +import socket +from collections.abc import AsyncGenerator +from typing import Any + +import uvicorn +from mcp.server.mcpserver import MCPServer +from starlette.applications import Starlette +from starlette.requests import Request +from starlette.responses import JSONResponse, PlainTextResponse, Response +from starlette.routing import Route + +# JSON-RPC error codes used by the pinned endpoint. +INVALID_REQUEST = -32600 +METHOD_NOT_FOUND = -32601 + +TOOL_NAME = "add" + + +def build_sdk_app() -> Starlette: + """Host a real SDK ``MCPServer`` over Streamable HTTP.""" + server = MCPServer("Math") + + @server.tool() + def add(a: int, b: int) -> int: + """Add two numbers""" + return a + b + + @server.tool() + def multiply(a: int, b: int) -> int: + """Multiply two numbers""" + return a * b + + return server.streamable_http_app(json_response=True) + + +class PinnedVersionServer: + """Serve Streamable HTTP while declaring one fixed protocol version. + + Args: + protocol_version: The version echoed from ``initialize``. + modern_only: When True, serve ``server/discover`` and reject the legacy + ``initialize`` handshake, the way a server that speaks only the + modern protocol would. Requires a modern ``protocol_version``. + """ + + def __init__(self, protocol_version: str, *, modern_only: bool = False) -> None: + self.protocol_version = protocol_version + self.modern_only = modern_only + self.session_ids: list[str] = [] + self.delete_count = 0 + self.initialize_count = 0 + self.discover_count = 0 + + def build_app(self) -> Starlette: + """Return an ASGI app exposing this endpoint at ``/mcp``.""" + return Starlette(routes=[Route("/mcp", self._handle, methods=["GET", "POST", "DELETE"])]) + + async def _handle(self, request: Request) -> Response: + if request.method == "GET": + # No server-initiated stream is needed for this testcase. + return Response(status_code=405) + if request.method == "DELETE": + self.delete_count += 1 + return Response(status_code=204) + + body = json.loads(await request.body()) + method = body.get("method") + + if method == "server/discover": + return self._discover(body) + if method == "initialize": + return self._initialize(body) + if method is not None and method.startswith("notifications/"): + return Response(status_code=202) + if method == "tools/list": + return self._result( + body["id"], + { + "tools": [_tool_schema()], + # Required on the 2026-07-28 wire, ignored by older peers. + "resultType": "complete", + "ttlMs": 0, + "cacheScope": "private", + }, + ) + if method == "tools/call": + return self._call_tool(body) + return self._error( + body.get("id"), METHOD_NOT_FOUND, f"Method not found: {method}", status=404 + ) + + def _discover(self, body: dict[str, Any]) -> Response: + """Answer ``server/discover``, but only for a modern endpoint.""" + if not self.modern_only: + # A handshake-era server has no discovery method to land on. + return self._error( + body.get("id"), + METHOD_NOT_FOUND, + "server/discover is not supported; use the initialize handshake", + status=404, + ) + self.discover_count += 1 + return self._result( + body["id"], + { + "supportedVersions": [self.protocol_version], + "capabilities": {"tools": {"listChanged": True}}, + "resultType": "complete", + "ttlMs": 0, + "cacheScope": "private", + }, + ) + + def _initialize(self, body: dict[str, Any]) -> Response: + if self.modern_only: + # A modern-only server offers no legacy handshake at all. + return self._error( + body.get("id"), + METHOD_NOT_FOUND, + "Legacy initialize is not supported; use modern discovery", + status=404, + ) + self.initialize_count += 1 + session_id = f"session-{self.initialize_count}" + self.session_ids.append(session_id) + return self._result( + body["id"], + { + "protocolVersion": self.protocol_version, + "capabilities": {"tools": {}}, + "serverInfo": {"name": "pinned-server", "version": "1.0.0"}, + }, + headers={"mcp-session-id": session_id}, + ) + + def _call_tool(self, body: dict[str, Any]) -> Response: + params = body.get("params") or {} + if params.get("name") != TOOL_NAME: + return self._error( + body.get("id"), INVALID_REQUEST, f"Unknown tool: {params.get('name')}" + ) + arguments = params.get("arguments") or {} + total = int(arguments.get("a", 0)) + int(arguments.get("b", 0)) + return self._result( + body["id"], + { + "content": [{"type": "text", "text": str(total)}], + "structuredContent": {"result": total}, + "isError": False, + "resultType": "complete", + }, + ) + + @staticmethod + def _result( + request_id: Any, result: dict[str, Any], *, headers: dict[str, str] | None = None + ) -> Response: + return JSONResponse( + {"jsonrpc": "2.0", "id": request_id, "result": result}, headers=headers + ) + + @staticmethod + def _error( + request_id: Any, code: int, message: str, *, status: int = 400 + ) -> Response: + return JSONResponse( + {"jsonrpc": "2.0", "id": request_id, "error": {"code": code, "message": message}}, + status_code=status, + ) + + +def _tool_schema() -> dict[str, Any]: + """Describe the single tool the pinned endpoint exposes.""" + return { + "name": TOOL_NAME, + "description": "Add two numbers", + "inputSchema": { + "type": "object", + "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}}, + "required": ["a", "b"], + }, + } + + +def _reset_sse_shutdown_latch() -> None: + """Clear ``sse-starlette``'s process-global shutdown latch. + + ``sse_starlette`` runs a watcher that polls ``uvicorn.Server.should_exit`` + and latches a *module-global* ``AppStatus.should_exit`` when any server + stops. This testcase hosts one server per leg, so from the second leg + onward every SSE stream would see the latch already set and end the instant + it opened -- logging ``ASGI callable returned without completing response`` + and sending the client into a reconnect loop. Clearing the latch (and the + per-thread watcher bookkeeping) gives every leg the same first-leg + behaviour. + + Best-effort: these are private names, so a changed internal degrades to the + noisy-but-working behaviour rather than breaking the matrix. + """ + try: + from sse_starlette import sse as sse_module + except ImportError: + return + try: + sse_module.AppStatus.should_exit = False + state = getattr(sse_module._thread_state, "shutdown_state", None) + if state is not None: + state.watcher_started = False + state.events.clear() + except AttributeError: + return + + +class SessionHeaderRecorder: + """Pure ASGI middleware noting every ``mcp-session-id`` a response carries. + + ``2026-07-28`` has no session identity, so a modern leg must show the header + on *requests* -- carrying the client-minted affinity ID -- and never on a + response. Recording it at the ASGI layer is the only way to tell a + server-assigned session from a client-minted one, since both are opaque hex. + + ``starlette.middleware.base.BaseHTTPMiddleware`` is deliberately not used: + it buffers through an inner task and breaks the server's SSE stream with + ``ASGI callable returned without completing response``. + """ + + def __init__(self, app: Any) -> None: + self.app = app + self.response_session_ids: list[str] = [] + + async def __call__(self, scope: Any, receive: Any, send: Any) -> None: + """Forward one ASGI request, noting any session ID on its response.""" + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + async def recording_send(message: Any) -> None: + if message["type"] == "http.response.start": + for key, value in message.get("headers", ()): + if key.decode("latin-1").lower() == "mcp-session-id": + self.response_session_ids.append(value.decode("latin-1")) + await send(message) + + await self.app(scope, receive, recording_send) + + +@contextlib.asynccontextmanager +async def serve(app: Any) -> AsyncGenerator[str, None]: + """Run *app* on an ephemeral port and yield its ``/mcp`` URL. + + Binding port 0 keeps parallel CI jobs from colliding, and hosting in-process + means there is no child process left behind if a leg fails. + """ + _reset_sse_shutdown_latch() + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + + server = uvicorn.Server(uvicorn.Config(app, log_level="warning")) + task = None + try: + task = asyncio.create_task(server.serve(sockets=[sock])) + while not server.started: + if task.done(): + task.result() + raise RuntimeError("MCP test server exited before startup") + await asyncio.sleep(0.02) + yield f"http://127.0.0.1:{port}/mcp" + finally: + server.should_exit = True + if task is not None: + with contextlib.suppress(BaseException): + await task + sock.close() + + +class DebugStateStore: + """In-memory stand-in for the AgentHub debug-state endpoint. + + ``uipath-agents-python`` persists MCP session IDs by PUT/GET against + ``{UIPATH_URL}/agenthub_/design/debugstate/{agent_id}/{key}``. Serving that + route here keeps the downstream ``SessionInfo`` subclass on its real code + path -- a legacy ``httpx.AsyncClient`` round trip -- instead of a stub. + """ + + def __init__(self) -> None: + self.values: dict[str, str] = {} + self.writes = 0 + self.reads = 0 + + def attach(self, app: Starlette) -> None: + """Mount the debug-state route onto an existing app.""" + app.router.routes.append( + Route( + "/agenthub_/design/debugstate/{agent_id}/{key:path}", + self._handle, + methods=["GET", "PUT"], + ) + ) + + async def _handle(self, request: Request) -> Response: + key = request.path_params["key"] + if request.method == "PUT": + self.values[key] = (await request.body()).decode() + self.writes += 1 + return Response(status_code=204) + self.reads += 1 + value = self.values.get(key) + if value is None: + return Response(status_code=404) + return PlainTextResponse(value) diff --git a/testcases/simple-local-mcp/pyproject.toml b/testcases/simple-local-mcp/pyproject.toml index 71a10aacd..52388943b 100644 --- a/testcases/simple-local-mcp/pyproject.toml +++ b/testcases/simple-local-mcp/pyproject.toml @@ -13,9 +13,8 @@ dependencies = [ "pydantic>=2.10.6", "aiohttp>=3.11.12", "typing-extensions>=4.12.2", - "langchain-mcp-adapters>=0.1.0", "ipython>=8.32.0", - "mcp>=1.4.1", + "mcp==2.0.0", ] requires-python = ">=3.11" diff --git a/testcases/simple-local-mcp/src/simple-local-mcp/graph.py b/testcases/simple-local-mcp/src/simple-local-mcp/graph.py index 093aa4f88..e4d12de93 100644 --- a/testcases/simple-local-mcp/src/simple-local-mcp/graph.py +++ b/testcases/simple-local-mcp/src/simple-local-mcp/graph.py @@ -1,26 +1,31 @@ import sys -import os -from contextlib import asynccontextmanager +from contextlib import AsyncExitStack, asynccontextmanager -from langchain_mcp_adapters.client import MultiServerMCPClient from langchain.agents import create_agent +from mcp import ClientSession +from mcp.client.stdio import StdioServerParameters, stdio_client + +from uipath_langchain.agent.tools.mcp import load_mcp_tools from uipath_langchain.chat import UiPathChat model = UiPathChat(model="gpt-4o-2024-11-20", streaming=False) @asynccontextmanager async def make_graph(): - client = MultiServerMCPClient({ - "math": { - "command": sys.executable, - "args": ["src/simple-local-mcp/math_server.py"], - "transport": "stdio", - }, - "weather": { - "command": sys.executable, - "args": ["src/simple-local-mcp/weather_server.py"], - "transport": "stdio", - }, - }) - agent = create_agent(model, tools=await client.get_tools()) - yield agent + async with AsyncExitStack() as stack: + tools = [] + for script in ("math_server.py", "weather_server.py"): + read, write = await stack.enter_async_context( + stdio_client( + StdioServerParameters( + command=sys.executable, + args=[f"src/simple-local-mcp/{script}"], + ) + ) + ) + session = await stack.enter_async_context(ClientSession(read, write)) + await session.initialize() + tools.extend(await load_mcp_tools(session)) + + agent = create_agent(model, tools=tools) + yield agent diff --git a/testcases/simple-local-mcp/src/simple-local-mcp/math_server.py b/testcases/simple-local-mcp/src/simple-local-mcp/math_server.py index 9cd78d791..64f464d7d 100644 --- a/testcases/simple-local-mcp/src/simple-local-mcp/math_server.py +++ b/testcases/simple-local-mcp/src/simple-local-mcp/math_server.py @@ -1,6 +1,6 @@ -from mcp.server.fastmcp import FastMCP +from mcp.server.mcpserver import MCPServer -mcp = FastMCP("Math") +mcp = MCPServer("Math") @mcp.tool() def add(a: int, b: int) -> int: diff --git a/testcases/simple-local-mcp/src/simple-local-mcp/weather_server.py b/testcases/simple-local-mcp/src/simple-local-mcp/weather_server.py index b7318114e..250a236f5 100644 --- a/testcases/simple-local-mcp/src/simple-local-mcp/weather_server.py +++ b/testcases/simple-local-mcp/src/simple-local-mcp/weather_server.py @@ -1,6 +1,6 @@ -from mcp.server.fastmcp import FastMCP +from mcp.server.mcpserver import MCPServer -mcp = FastMCP("Weather") +mcp = MCPServer("Weather") @mcp.tool() async def get_weather(location: str) -> str: diff --git a/tests/agent/tools/test_mcp/claude.md b/tests/agent/tools/test_mcp/claude.md index 4822c7110..acc41001f 100644 --- a/tests/agent/tools/test_mcp/claude.md +++ b/tests/agent/tools/test_mcp/claude.md @@ -2,9 +2,11 @@ > **CLAUDE: UPDATE THIS DOCUMENT** > -> When you modify `test_mcp_client.py` or `test_mcp_tool.py`, you MUST update this document to reflect: +> When you modify `real_server.py`, `test_mcp_client_real_http.py`, +> `test_mcp_client.py`, `test_protocol_strategy.py`, or `test_mcp_tool.py`, you +> MUST update this document to reflect: > - New test cases (add to Test File Structure and create explanation section) -> - Changes to MockStreamResponse (update Handled MCP Methods table and examples) +> - Changes to LegacyMcpEndpoint (update Handled MCP Methods table and examples) > - New mocking patterns (add to Common Patterns section) > - New assertion patterns (add to Guidelines for Adding New Tests) > - Changes to test tracking variables (update Tracking Test State section) @@ -17,31 +19,121 @@ This document explains the testing strategy for MCP-related code. Use this as a ## Testing Philosophy -The tests mock **only the HTTP layer** (`httpx.AsyncClient`), allowing the real MCP SDK to process messages. This approach: +There are **two tiers**, and which tier a behaviour belongs in is not a matter of +taste. -- Tests the actual MCP protocol flow -- Validates error handling with real `McpError` exceptions -- Ensures `ClientSession.initialize()` behaves correctly when called multiple times -- Catches integration issues between our code and the SDK +### Tier 1 - real HTTP (default) + +`test_mcp_client_real_http.py` drives the public `McpClient` API against a +genuine `mcp.server.mcpserver.MCPServer` hosted in-process on an ephemeral port, +through a recording gateway. Nothing is simulated except the UiPath SDK lookup +that resolves the server URL. This tier: + +- Runs a real ASGI server, a real socket, and the real SDK transport on both ends +- Exercises `McpClient` itself, not just the strategies underneath it +- Observes the wire (`mcp-session-id`, `mcp-protocol-version`, `params._meta`, + HTTP method) exactly as a gateway would, so assertions survive refactors +- Catches server behaviour a hand-written mock silently gets wrong - session + routing, `DELETE` on teardown, the optional GET SSE channel + +**Anything a cooperative server can produce belongs here.** + +### Tier 2 - `httpx2.MockTransport` (exceptions only) + +`test_mcp_client.py` and `test_protocol_strategy.py` keep mocks only for +conditions a cooperative real server cannot express: + +| Condition | Why a real server cannot do it | +|-----------|--------------------------------| +| Concurrency races | Requires blocking one `initialize` mid-flight and releasing it on cue | +| `mints_new_session_on_initialize` | A conforming server routes by the session header instead | +| `echo_session_id` in the modern era | A modern server has no session to echo | +| `repeat_session_header` | A conforming server sends the header once | +| `fail_initialize_on` | A deterministic handshake failure on the *second* attempt only | +| A bare, body-less HTTP 404 | The SDK server always returns a JSON-RPC body | +| Pure-function matrices | `is_recoverable`, `reset`, `build_protocol_strategy` need no server at all | + +If you are adding a test and it does not fall into that table, write it in +`test_mcp_client_real_http.py`. ## Test File Structure ``` tests/agent/tools/test_mcp/ -├── test_mcp_client.py # McpClient session + tool-list caching tests -│ └── TestMcpClient (class) -│ ├── create_mock_stream_response() -│ ├── create_mock_http_client() -│ ├── test_session_initializes_on_first_call -│ ├── test_session_reused_across_calls -│ ├── test_session_reinitializes_on_404_error ← Key test -│ ├── test_max_retries_exceeded -│ ├── test_dispose_releases_resources -│ ├── test_client_initialized_property -│ ├── test_session_can_be_reused_after_dispose -│ ├── test_list_tools_caches_result_across_calls ← list_tools fetched once per lifetime -│ ├── test_list_tools_force_refresh_bypasses_cache -│ └── test_dispose_clears_tools_cache +├── real_server.py # Real-HTTP harness (no tests of its own) +│ ├── serve(app) # ephemeral port + in-process uvicorn +│ ├── build_sdk_app() # real MCPServer with add/multiply +│ ├── PinnedVersionServer # one fixed protocolVersion, or modern-only +│ ├── RecordingGateway # pure ASGI recorder + fault injector +│ ├── patched_sdk(url) # redirects McpClient's lazy UiPath lookup +│ └── connected_client / make_client / pinned_session_factory +│ +├── test_mcp_client_real_http.py # McpClient over real HTTP ← default tier +│ ├── negotiation per mode +│ │ ├── test_legacy_mode_negotiates_the_newest_handshake_version +│ │ ├── test_modern_mode_negotiates_without_any_server_session +│ │ ├── test_auto_mode_resolves_to_modern_against_a_real_server +│ │ ├── test_auto_mode_falls_back_to_legacy_against_a_handshake_only_server +│ │ └── test_modern_mode_works_against_a_server_that_refuses_the_handshake +│ ├── resume across clients +│ │ ├── test_legacy_resume_keeps_the_originally_negotiated_version ← Key test +│ │ └── test_unknown_persisted_session_falls_back_to_a_fresh_session +│ ├── affinity and disposal +│ │ ├── test_modern_affinity_pins_one_instance_across_clients +│ │ ├── test_auto_mode_pins_the_first_request +│ │ ├── test_modern_disposal_does_not_delete_a_restored_affinity_id +│ │ └── test_legacy_disposal_deletes_a_restored_session +│ ├── retry semantics per era +│ │ ├── test_legacy_recovers_from_an_injected_session_termination +│ │ └── test_modern_does_not_retry_an_injected_session_termination +│ ├── version breadth +│ │ └── test_legacy_negotiates_every_supported_handshake_version[4 versions] +│ └── lifecycle +│ ├── test_list_tools_is_cached_until_force_refresh +│ └── test_dispose_then_reuse_reinitializes_the_client +│ +├── test_mcp_client.py # MockTransport: pathological servers + races +│ ├── LegacyMcpEndpoint # httpx2.MockTransport request handler +│ ├── test_legacy_httpx_timeout_is_normalized_for_final_client +│ ├── test_replaces_transport_and_session_after_404 (bare, body-less 404) +│ ├── test_replaces_session_after_official_session_not_found_error +│ ├── test_dropped_connection_resumes_the_persisted_session (CONNECTION_CLOSED keeps the ID) +│ ├── test_auto_mode_does_not_offer_a_minted_id_to_a_legacy_handshake +│ ├── test_persisted_session_replaced_when_server_ignores_the_header +│ ├── test_rejected_persisted_session_is_initialized_and_deleted_once +│ ├── test_repeated_session_headers_do_not_repeat_external_persistence +│ ├── test_max_retries_exceeded_raises_mcp_error +│ ├── test_concurrent_recovery_does_not_replace_a_new_session +│ ├── test_recovery_continues_when_failed_connection_cleanup_raises +│ ├── test_concurrent_call_waits_for_recovery_initialization +│ ├── test_later_call_recovers_after_replacement_initialization_failure +│ ├── test_raises_on_missing_mcp_url +│ ├── test_initialization_failure_cleans_state_and_allows_retry +│ └── test_only_session_specific_invalid_request_is_retryable +│ +├── test_protocol_strategy.py # Pure per-era policy + two odd server behaviours +│ ├── EraMcpEndpoint # Serves either era, or both +│ ├── test_auto_mode_sends_a_restored_id_before_the_era_resolves +│ ├── test_modern_mode_ignores_a_server_assigned_session_id +│ ├── test_legacy_recovers_from_session_loss_but_not_from_bad_requests +│ ├── test_modern_recovers_only_from_a_dropped_connection +│ ├── test_modern_reset_keeps_the_affinity_id +│ ├── test_legacy_reset_clears_the_stale_session_id +│ ├── test_auto_applies_the_legacy_policy_before_an_era_is_resolved +│ ├── test_legacy_keeps_a_persisted_session_when_the_connection_drops +│ ├── test_auto_does_not_carry_a_stale_era_through_a_failed_probe +│ ├── test_build_protocol_strategy_maps_every_mode +│ └── test_legacy_is_the_default_mode +│ +├── test_session_info.py # SessionInfo + SessionInfoFactory contract +│ +├── test_session_tools.py # load_mcp_tools discovery/invocation/errors +│ +├── test_protocol_version_support.py # SDK facts the strategies depend on (tripwires) +│ ├── test_the_auto_probe_builds_on_public_session_methods +│ ├── test_the_low_level_session_reaches_the_modern_era +│ ├── test_initialize_cannot_choose_a_protocol_version +│ └── test_the_two_eras_share_no_protocol_version │ └── test_mcp_tool.py # Tool factory tests (17 tests) ├── TestMcpToolMetadata (class) @@ -60,12 +152,12 @@ tests/agent/tools/test_mcp/ │ ├── test_creates_tools_from_multiple_mcp_servers │ ├── test_returns_mcp_clients_for_each_server │ ├── test_skips_disabled_mcp_resources - │ ├── test_returns_empty_for_agent_without_mcp + │ ├── test_returns_empty_for_empty_resources │ ├── test_raises_on_missing_mcp_url │ └── test_tools_have_correct_metadata │ - ├── TestMcpToolInvocation (class) - │ └── test_tool_invocation_initializes_session_and_returns_result + ├── TestMcpToolResultSerialization (class) + ├── TestMcpToolErrorHandling (class) │ ├── TestMcpToolNameSanitization (class) │ ├── test_tool_name_with_spaces @@ -104,30 +196,109 @@ tool. `tool_fn` tests mock `mcpClient.list_tools` directly, so they exercise the refresh logic per invocation independent of the client's caching. The once-per-run caching -itself lives in `McpClient.list_tools` and is covered in `test_mcp_client.py` -(`test_list_tools_caches_result_across_calls`, `..._force_refresh_bypasses_cache`, -`test_dispose_clears_tools_cache`). +itself lives in `McpClient.list_tools` and is covered over real HTTP in +`test_mcp_client_real_http.py` (`test_list_tools_is_cached_until_force_refresh`; +disposal/reuse is covered by `test_dispose_then_reuse_reinitializes_the_client`). + +## The Real-HTTP Harness (`real_server.py`) + +### serve(app) + +Async context manager. Binds `("127.0.0.1", 0)` for an ephemeral port - parallel +CI jobs must never collide on a fixed one - starts uvicorn **in-process** (so a +failing test leaves no orphan), and yields the `/mcp` URL. + +```python +gateway = RecordingGateway(build_sdk_app()) +async with serve(gateway) as url: + async with connected_client(url, protocol_mode="modern") as client: + await client.call_tool("add", {"a": 2, "b": 3}) +``` + +`serve()` also calls `_reset_sse_shutdown_latch()`. `sse-starlette` polls +`uvicorn.Server.should_exit` and latches a **module-global** +`AppStatus.should_exit` when any server stops; without the reset, every server +after the first would see the latch already set and kill its SSE stream the +instant it opened, logging `ASGI callable returned without completing response` +and sending the client into a reconnect loop. It is process-global state, not a +per-server flag - do not remove the reset. + +### RecordingGateway - pure ASGI, never BaseHTTPMiddleware + +Records per request: JSON-RPC `method`, the `mcp-session-id` and +`mcp-protocol-version` **request** headers, `params._meta`, the HTTP method (so +`DELETE` is observable), the `mcp-session-id` on the **response** (i.e. a +server-assigned session), and the instance a gateway routing on `mcp-session-id` +would have picked. + +```python +gateway.count("tools/call") # JSON-RPC method counts +gateway.http_count("DELETE") # session teardown +gateway.for_rpc("tools/call")[0] # one RecordedRequest +gateway.server_session_ids() # [] in the modern era +gateway.unpinned() # requests with no affinity header +``` + +Fault injection replaces the Nth `tools/call` with a JSON-RPC `-32600 +"Session terminated"` at HTTP 404, which is how the recovery paths are driven +without `MockTransport`: + +```python +gateway = RecordingGateway(build_sdk_app(), fault_on_tool_call=1) +``` + +**Use a pure ASGI middleware (`async def __call__(self, scope, receive, send)`), +never `starlette.middleware.base.BaseHTTPMiddleware`.** The latter buffers +through an inner task and breaks streaming/SSE responses here with `ASGI +callable returned without completing response`. A pure ASGI middleware composes +cleanly and can still read request bodies by wrapping `receive` (see +`_buffer_body`, which replays the buffered body downstream). + +### patched_sdk(url) + +`McpClient._initialize_client` does `from uipath.platform import UiPath` at call +time, so replacing that module attribute is enough to point the client at a local +server. No tenant, credentials, or network access to UiPath Cloud is involved, +and the client still walks its real resolution path. + +### Assert on the wire, not on internals + +There is no public accessor for the negotiated version, so read it from the +`mcp-protocol-version` header the SDK stamps on every post-negotiation request: + +```python +def negotiated_version(gateway, rpc_method): + records = gateway.for_rpc(rpc_method) + return records[0].protocol_version if records else None +``` ## Mocking Strategy ### What We Mock -Only `httpx.AsyncClient` is mocked at the module level: +`LegacyMcpEndpoint` is an async handler installed on a real +`httpx2.AsyncClient` through `httpx2.MockTransport`: ```python -@patch("httpx.AsyncClient") -async def test_something(self, mock_async_client_class): - # mock_async_client_class is the patched class - # We configure it to return our mock client - mock_http_client = self.create_mock_http_client(MockStreamResponse) - mock_async_client_class.return_value = mock_http_client +endpoint = LegacyMcpEndpoint(protocol_version="2025-06-18") +http_kwargs = { + "headers": {"Authorization": "Bearer test-secret-token"}, + "transport": endpoint.transport, + "follow_redirects": True, +} +with patch( + "uipath_langchain.agent.tools.mcp.mcp_client.get_httpx_client_kwargs", + return_value=http_kwargs, +): + result = await client.call_tool("test_tool", {"query": "test"}) ``` ### What We DON'T Mock - `mcp.ClientSession` - Real SDK session handling - `mcp.client.streamable_http.streamable_http_client` - Real transport setup -- `mcp.shared.exceptions.McpError` - Real error types +- `mcp.shared.exceptions.MCPError` - Real error types +- UiPath's `streamable_http_client` event hooks - Real session persistence adapter ### Why This Approach? @@ -137,7 +308,7 @@ async def test_something(self, mock_async_client_class): ├─────────────────────────────────────────────────────────────┤ │ │ │ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ │ -│ │ McpClient │ ──► │ MCP SDK │ ──► │ HTTP Mock │ │ +│ │ McpClient │ ──► │ MCP SDK 2 │ ──► │ MockTransport│ │ │ │ │ │ (real) │ │ (mocked) │ │ │ └─────────────┘ └─────────────┘ └─────────────┘ │ │ ▲ │ │ │ @@ -148,62 +319,93 @@ async def test_something(self, mock_async_client_class): └─────────────────────────────────────────────────────────────┘ ``` -## MockStreamResponse Class +## LegacyMcpEndpoint Class -The core of our mocking - simulates an MCP server's HTTP responses. +The core test endpoint simulates an MCP legacy Streamable HTTP server while +recording methods, headers, initialization count, tool calls, and DELETEs. ### Structure ```python -class MockStreamResponse: - def __init__(self, method: str, url: str, **kwargs): - # method: "GET" or "POST" - # url: The endpoint URL - # kwargs: Contains json (request body), headers, etc. - - def _build_response(self) -> tuple[int, Any, dict[str, str] | None]: - # Returns: (status_code, json_body, headers) - - async def __aenter__(self): ... # Context manager entry - async def __aexit__(self, ...): ... # Context manager exit - async def aread(self) -> bytes: ... # Read response body - def raise_for_status(self): ... # Check HTTP status -``` +class LegacyMcpEndpoint: + def __init__( + self, + protocol_version: str = "2025-11-25", + *, + failed_tool_calls: int = 0, + known_session_ids: set[str] | None = None, + mints_new_session_on_initialize: bool = False, + ) -> None: + self.protocol_version = protocol_version + self.failed_tool_calls = failed_tool_calls + self.known_session_ids = set(known_session_ids or ()) + self.mints_new_session_on_initialize = mints_new_session_on_initialize + self.methods: list[str] = [] + self.request_headers: list[tuple[str, httpx2.Headers]] = [] + self.initialize_count = 0 # initialize requests handled + self.session_mint_count = 0 # sessions actually created + self.tool_call_count = 0 + self.delete_count = 0 + self.transport = httpx2.MockTransport(self.handle) + + async def handle(self, request: httpx2.Request) -> httpx2.Response: ... +``` + +**Session routing.** `_session_for_initialize` mirrors the SDK server: an +`initialize` naming a live session is handled *inside* it, and a new session is +minted only when no session header is present. An unknown or expired ID is +rejected with `"Session not found"` rather than silently replaced. + +This fidelity matters — the legacy resume path re-runs `initialize` inside a +restored session, so an endpoint that minted a fresh ID on every handshake could +not express the behaviour under test. + +- Seed `known_session_ids={"persisted-session"}` to stand in for a session a + previous process established and persisted externally. +- Set `mints_new_session_on_initialize=True` to model a server that ignores the + header instead. +- `initialize_count` counts requests; `session_mint_count` counts sessions + created. A rejected handshake increments only the former. ### Handled MCP Methods | Method | Response | Notes | |--------|----------|-------| -| `initialize` | 200 + session ID | Returns different IDs for each call | -| `notifications/initialized` | 204 No Content | Notification, no body | +| `initialize` | 200 + session ID, or 404 | Routes by `mcp-session-id`; mints only when absent, rejects unknown/expired IDs | +| `notifications/initialized` | 202 Accepted | Notification, no body | | `tools/list` | 200 + tool definitions | For SDK output validation | -| `tools/call` | 200 + result OR 404 | Configurable via `fail_first_tool_call` | +| `tools/call` | 200 + result OR bare 404 | Configurable via `failed_tool_calls` | | GET requests | 405 | Server doesn't support GET streaming | +| DELETE requests | 204 | Records session termination | ### Response Format Examples **Initialize response:** ```python -return ( +return httpx2.Response( 200, - { + headers={ + "content-type": "application/json", + "mcp-session-id": f"session-{self.initialize_count}", + }, + json={ "jsonrpc": "2.0", "id": request_id, "result": { - "protocolVersion": "2025-06-18", + "protocolVersion": self.protocol_version, "capabilities": {"tools": {}}, "serverInfo": {"name": "test-server", "version": "1.0.0"}, }, }, - {"mcp-session-id": session_id}, # Header with session ID ) ``` **Tool call success:** ```python -return ( +return httpx2.Response( 200, - { + headers={"content-type": "application/json"}, + json={ "jsonrpc": "2.0", "id": request_id, "result": { @@ -212,149 +414,117 @@ return ( "isError": False, }, }, - {}, ) ``` **Tool call 404 (session terminated):** ```python -return (404, None, None) +return httpx2.Response(404) ``` ## Tracking Test State -Tests use mutable lists to track state across mock calls: - -```python -method_call_sequence: list[str] = [] # Order of MCP methods called -initialize_count = [0] # How many times initialize was called -tool_call_count = [0] # How many times tools/call was called -``` - -Why lists? Because they're mutable and can be modified inside the mock class closure: +Tests inspect state recorded directly on `LegacyMcpEndpoint`: ```python -def create_mock_stream_response(self, method_call_sequence, initialize_count, ...): - class MockStreamResponse: - def _build_response(self): - if self.method == "initialize": - initialize_count[0] += 1 # Modifies outer list - method_call_sequence.append(self.method) # Tracks call order +assert endpoint.initialize_count == 2 +assert endpoint.tool_call_count == 2 +assert endpoint.delete_count == 1 +assert endpoint.methods.count("tools/list") == 2 +assert endpoint.headers_for("tools/call")[0]["mcp-session-id"] == "session-1" ``` ## Test Cases Explained ### TestMcpClient Tests -#### test_session_initializes_on_first_call +#### test_replaces_transport_and_session_after_404 ⭐ -**Purpose:** Verify lazy initialization on first `call_tool()` +The first tool request returns a bare HTTP 404. The test verifies two +initializations, two tool calls, one DELETE for the old session, a new session +ID, and correct session headers on both attempts. This catches the SDK 2 +idempotent-`initialize()` breaking change: recovery must create a fresh +`ClientSession`, not call `initialize()` again on the old one. -**Assertions:** ```python -assert session.session_id is None # Before call -result = await session.call_tool(...) -assert session.session_id == "test-session-first" # After call -assert session.is_client_initialized -assert mock_async_client_class.call_count == 1 # HTTP client created +assert endpoint.initialize_count == 2 +assert endpoint.tool_call_count == 2 +assert endpoint.delete_count == 1 +assert await client.get_session_id() == "session-2" ``` -#### test_session_reused_across_calls - -**Purpose:** Verify session persists across multiple tool calls - -**Assertions:** -```python -await session.call_tool(...) # First call -assert initialize_count[0] == 1 - -await session.call_tool(...) # Second call -assert initialize_count[0] == 1 # Still 1! No reinit -assert tool_call_count[0] == 2 # But 2 tool calls -``` +#### Persisted-session tests -#### test_session_reinitializes_on_404_error ⭐ +Resuming a persisted session is covered over real HTTP +(`test_legacy_resume_keeps_the_originally_negotiated_version`, +`test_unknown_persisted_session_falls_back_to_a_fresh_session`, +`test_legacy_disposal_deletes_a_restored_session`). What stays here is the +behaviour a conforming server will not produce: -**Purpose:** THE KEY TEST - verify client reuse on session reinit +`test_persisted_session_replaced_when_server_ignores_the_header` covers a server +that mints on every handshake: the persisted session is lost, but the connection +stays usable and the client continues with the replacement ID. -**Setup:** -```python -MockStreamResponse = self.create_mock_stream_response( - ..., - fail_first_tool_call=True, # First tools/call returns 404 -) -``` +`test_rejected_persisted_session_is_initialized_and_deleted_once` covers a server +that explicitly rejects one known-shaped session ID, and asserts only the fresh +SDK session is deleted on close. -**Critical Assertions:** -```python -# Session was reinitialized (initialize called twice) -assert initialize_count[0] == 2 +#### Retry and concurrency tests -# Tool call was retried -assert tool_call_count[0] == 2 +- `test_max_retries_exceeded_raises_mcp_error` expects the real `MCPError` + after the configured retry is consumed. +- `test_concurrent_recovery_does_not_replace_a_new_session` verifies a late + failure from an old `ClientSession` cannot tear down a replacement created by + another operation. +- `test_only_session_specific_invalid_request_is_retryable` verifies an ordinary + `INVALID_REQUEST` is not misclassified as a disconnect. -# Session ID changed -assert session.session_id == "test-session-retry" +#### Cache, disposal, and configuration tests -# KEY: HTTP client created only ONCE (not recreated) -assert mock_async_client_class.call_count == 1 -``` +Caching and disposal-then-reuse moved to real HTTP +(`test_list_tools_is_cached_until_force_refresh`, +`test_dispose_then_reuse_reinitializes_the_client`). What stays here: -#### test_max_retries_exceeded +- `test_raises_on_missing_mcp_url` verifies endpoint validation happens before + HTTP resources are allocated. +- `test_initialization_failure_cleans_state_and_allows_retry` patches + `_initialize_session` to fail, which no server response can cause. +- `test_legacy_httpx_timeout_is_normalized_for_final_client` pins the + pre-upgrade public timeout type; its subject is `_normalize_timeout`, not the + server. -**Purpose:** Verify `McpError` is raised after max retries - -**Setup:** Custom mock that ALWAYS returns 404 for tool calls - -**Assertions:** -```python -with pytest.raises(McpError): - await session.call_tool(...) - -assert initialize_count[0] == 2 # Tried to reinit -assert tool_call_count[0] == 2 # Tried twice -assert mock_async_client_class.call_count == 1 # Still only one client -``` - -#### test_dispose_releases_resources - -**Purpose:** Verify `dispose()` cleans up properly - -**Assertions:** -```python -await session.dispose() -assert session.session_id is None -assert session._session is None -assert session._stack is None -assert not session.is_client_initialized -``` +### Real-HTTP tests explained -#### test_client_initialized_property +#### test_legacy_resume_keeps_the_originally_negotiated_version ⭐ -**Purpose:** Verify `is_client_initialized` property accuracy +Two `McpClient` instances share one `SessionInfo`, standing in for two runs of a +playground agent whose session store outlives the process. The first connects +with `terminate_on_close=False` and disposes; the second restores the ID. -**Assertions:** -```python -assert not session.is_client_initialized # Before -await session.call_tool(...) -assert session.is_client_initialized # After call -await session.dispose() -assert not session.is_client_initialized # After dispose -``` +The session ID surviving is only half the contract. **Every request after the +resume must also carry `mcp-protocol-version` equal to the version the session +was originally negotiated at.** Probing candidate versions instead — the +pre-existing approach — always matched the *oldest* handshake version, silently +downgrading every later request and disabling the server's `2025-11-25` SSE +resumability. No session-ID assertion catches that; this one does, and its +failure message says so. -#### test_session_can_be_reused_after_dispose +#### Retry semantics per era -**Purpose:** Verify session can be fully reinitialized after `dispose()` +`test_legacy_recovers_from_an_injected_session_termination` and +`test_modern_does_not_retry_an_injected_session_termination` send the *identical* +injected `-32600 "Session terminated"` at HTTP 404 through both eras. Legacy +re-handshakes and retries (two `initialize`s, two `tools/call`s, different +session IDs); modern surfaces it immediately (one `tools/call`). Driving both +from one fault injector is what makes the contrast meaningful. -**Assertions:** -```python -await session.call_tool(...) -await session.dispose() -await session.call_tool(...) # Should work! +#### test_modern_disposal_does_not_delete_a_restored_affinity_id -# HTTP client created TWICE (once before dispose, once after) -assert mock_async_client_class.call_count == 2 -``` +A restored affinity ID looks exactly like a restored session to the transport. +Deleting it would reach the gateway as a teardown for a *live* instance on every +run after the first — this was a real bug. The paired +`test_legacy_disposal_deletes_a_restored_session` proves the legacy era still +does send the `DELETE`. ### TestCreateMcpToolsFromAgent Tests @@ -394,7 +564,7 @@ assert len(tools) == 1 # Only enabled server's tool assert tools[0].name == "enabled_tool" ``` -#### test_returns_empty_for_agent_without_mcp +#### test_returns_empty_for_empty_resources **Purpose:** Verify empty lists for agent without MCP resources @@ -432,47 +602,64 @@ for tool in tools: ## Guidelines for Adding New Tests -### 1. Use the Factory Methods +### 1. Start in the real-HTTP tier + +Unless the behaviour is in the Tier 2 table above, write the test in +`test_mcp_client_real_http.py` against a real server: + +```python +gateway = RecordingGateway(build_sdk_app()) +async with serve(gateway) as url: + async with connected_client(url, protocol_mode="legacy") as client: + await client.call_tool("add", {"a": 2, "b": 3}) + +assert gateway.count("initialize") == 1 +assert negotiated_version(gateway, "tools/call") == LEGACY_VERSION +``` -Always use the provided factory methods: +Only if a cooperative server cannot produce the condition, fall back to +`LegacyMcpEndpoint`/`EraMcpEndpoint` with `configured_client`, which keep the +real SDK transport/session path over `httpx2.MockTransport`: ```python -MockStreamResponse = self.create_mock_stream_response( - method_call_sequence, - initialize_count, - tool_call_count, - fail_first_tool_call=False, # Configure behavior +endpoint = LegacyMcpEndpoint( + protocol_version="2025-11-25", + failed_tool_calls=0, ) -mock_http_client = self.create_mock_http_client(MockStreamResponse) -mock_async_client_class.return_value = mock_http_client +async with configured_client(config, mock_uipath_sdk, endpoint) as client: + await client.call_tool("test_tool", {"query": "test"}) ``` -### 2. Add New MCP Methods to MockStreamResponse +### 2. Add New MCP Methods to LegacyMcpEndpoint -If testing a new MCP method, add it to `_build_response()`: +If testing a new MCP method, add it to `handle()` and return an +`httpx2.Response` with wire-format JSON: ```python -elif self.method == "resources/list": - return ( +if method == "resources/list": + return httpx2.Response( 200, - { + headers={"content-type": "application/json"}, + json={ "jsonrpc": "2.0", - "id": request_id, + "id": body["id"], "result": {"resources": [...]}, }, - {}, ) ``` ### 3. Always Verify Client Reuse -For any retry-related test, assert HTTP client count: +For retry tests, assert the base client is reused while connection/session state +is replaced. The endpoint counters and headers are the observable contract: ```python -# After retry logic -assert mock_async_client_class.call_count == 1, ( - "HTTP client should be created only once" -) +assert endpoint.initialize_count == 2 +assert endpoint.delete_count == 1 +assert [h["mcp-session-id"] for h in endpoint.headers_for("tools/call")] == [ + "session-1", + "session-2", +] ``` ### 4. Track Method Sequences @@ -480,7 +667,7 @@ assert mock_async_client_class.call_count == 1, ( For protocol flow tests, verify the sequence: ```python -assert method_call_sequence == [ +assert endpoint.methods == [ "initialize", "notifications/initialized", "tools/call", @@ -493,33 +680,26 @@ assert method_call_sequence == [ When adding error tests: ```python -# Create custom mock for specific error -class CustomErrorMock: - def _build_response(self): - if self.method == "tools/call": - return ( - 200, - { - "jsonrpc": "2.0", - "id": request_id, - "error": {"code": -32000, "message": "Custom error"}, - }, - {}, - ) +# Add a branch to LegacyMcpEndpoint.handle(). +if method == "tools/call": + return httpx2.Response( + 400, + headers={"content-type": "application/json"}, + json={ + "jsonrpc": "2.0", + "id": body["id"], + "error": {"code": -32602, "message": "Invalid parameters"}, + }, + ) ``` ### 6. Clean Up After Tests -Always dispose the session: +Prefer `configured_client`, which disposes in `finally`: ```python -try: - # ... test logic ... -finally: - await session.dispose() - -# Or simply: -await session.dispose() # At end of test +async with configured_client(config, sdk, endpoint) as client: + await client.call_tool("test_tool", {}) ``` ### 7. Use Proper AgentSettings @@ -543,25 +723,19 @@ agent = LowCodeAgentDefinition( ### Testing Different Session IDs -The mock returns different session IDs based on initialize count: - -```python -session_id = ( - session_guid_1 if initialize_count[0] == 1 else session_guid_2 -) -``` - -Use this to verify session ID changes: +The endpoint returns `session-{initialize_count}`. Verify both the external +store and request headers: ```python -assert session.session_id == "test-session-first" # After first init -# ... trigger reinit ... -assert session.session_id == "test-session-retry" # After reinit +assert await client.get_session_id() == "session-2" +assert endpoint.headers_for("tools/call")[1]["mcp-session-id"] == "session-2" ``` ### Testing Structured Content -The SDK validates `structuredContent` against `outputSchema`. Ensure mock returns matching data: +The SDK validates `structuredContent` against `outputSchema`. Ensure mock returns +matching data. Wire JSON stays camelCase; SDK 2 Python attributes are snake_case +(`tool.input_schema`, `tool.output_schema`). ```python # In tools/list response @@ -583,8 +757,8 @@ await session.call_tool("tool1", {...}) await session.call_tool("tool2", {...}) await session.call_tool("tool1", {...}) -assert tool_call_count[0] == 3 -assert initialize_count[0] == 1 # Session reused +assert endpoint.tool_call_count == 3 +assert endpoint.initialize_count == 1 # Session reused ``` ### Testing create_mcp_tools_and_clients @@ -627,7 +801,7 @@ uv run pytest tests/agent/tools/test_mcp/ -v -s --log-cli-level=DEBUG Print the sequence to understand what happened: ```python -logger.info(f"Method sequence: {method_call_sequence}") +logger.info(f"Method sequence: {endpoint.methods}") # Output: ['initialize', 'notifications/initialized', 'tools/call', ...] ``` @@ -636,8 +810,9 @@ logger.info(f"Method sequence: {method_call_sequence}") Add debug logging in mock: ```python -def _build_response(self): - logger.debug(f"Building response for {self.method}, id={request_id}") +async def handle(self, request: httpx2.Request): + body = json.loads(request.content) + logger.debug(f"Building response for {body['method']}, id={body.get('id')}") # ... ``` @@ -645,8 +820,12 @@ def _build_response(self): | File | Purpose | |------|---------| -| `test_mcp_client.py` | McpClient session tests (7 tests) | -| `test_mcp_tool.py` | Tool factory tests (17 tests) | +| `real_server.py` | Real-HTTP harness: `serve`, `build_sdk_app`, `PinnedVersionServer`, `RecordingGateway`, `patched_sdk` | +| `test_mcp_client_real_http.py` | `McpClient` over real HTTP: negotiation per mode, resume, affinity, disposal, retry, every handshake version | +| `test_mcp_client.py` | MockTransport: pathological legacy servers, concurrency races, bare-404 mapping | +| `test_session_info.py` | Async session ID store and factory | +| `test_protocol_version_support.py` | Guards the SDK version constraints that dictate what `McpClient` can negotiate; each failure names the follow-up it unblocks | +| `test_mcp_tool.py` | Tool factories, schemas, result/error mapping, metadata | | `src/.../mcp/mcp_client.py` | McpClient implementation | | `src/.../mcp/mcp_tool.py` | Tool factory implementation | | `src/.../mcp/claude.md` | Implementation documentation | diff --git a/tests/agent/tools/test_mcp/real_server.py b/tests/agent/tools/test_mcp/real_server.py new file mode 100644 index 000000000..832624979 --- /dev/null +++ b/tests/agent/tools/test_mcp/real_server.py @@ -0,0 +1,689 @@ +"""Real-HTTP MCP servers for driving ``McpClient`` over a socket. + +``httpx2.MockTransport`` covers the wire faithfully but never exercises a real +server, a real ASGI stack, or the gateway hop UiPath puts in front of an MCP +endpoint. This module hosts genuine servers on an ephemeral port so +``McpClient`` -- not just the strategies underneath it -- can be driven end to +end. + +Three kinds of endpoint are provided: + +* :func:`build_sdk_app` hosts a genuine ``MCPServer`` over Streamable HTTP. It + answers both eras, so it is the realistic leg for legacy, modern, and ``auto``. +* :class:`PinnedVersionServer` declares one specific ``protocolVersion``. The + SDK server exposes no version knob, so pinning an older handshake version -- + or refusing the handshake the way a modern-only server would -- requires an + endpoint that speaks the wire directly. +* :class:`RecordingGateway` is a pure ASGI middleware standing in for the + AgentHub gateway: it records what every request carried and can inject a + fault, so recovery paths are driven by a real HTTP response rather than a + mocked transport. + +``starlette.middleware.base.BaseHTTPMiddleware`` is deliberately **not** used. +It buffers through an inner task and breaks streaming responses here with +``ASGI callable returned without completing response``; a pure ASGI middleware +composes cleanly and can still read request bodies by wrapping ``receive``. +""" + +import asyncio +import contextlib +import json +import socket +from collections.abc import ( + AsyncGenerator, + Awaitable, + Callable, + Iterator, + MutableMapping, +) +from contextlib import asynccontextmanager, contextmanager +from dataclasses import dataclass, field +from typing import Any + +from mcp.server.mcpserver import MCPServer +from mcp.types.version import HANDSHAKE_PROTOCOL_VERSIONS +from starlette.applications import Starlette +from starlette.requests import Request +from starlette.responses import JSONResponse, Response +from starlette.routing import Route +from uipath.agent.models.agent import AgentMcpResourceConfig, AgentMcpTool + +from uipath_langchain.agent.tools.mcp import McpClient, SessionInfo, SessionInfoFactory + +#: Newest handshake-era version; what a real SDK server settles on for legacy. +LEGACY_VERSION = "2025-11-25" + +#: The only modern-era version; reached through ``server/discover``. +MODERN_VERSION = "2026-07-28" + +#: Every handshake version SDK 2 still accepts, oldest first. +#: Every version reachable through the ``initialize`` handshake, taken from the +#: SDK rather than restated, so a version added upstream widens the matrix +#: instead of silently narrowing the "every handshake version" claim. +HANDSHAKE_VERSIONS: tuple[str, ...] = tuple(HANDSHAKE_PROTOCOL_VERSIONS) + +MCP_SESSION_ID = "mcp-session-id" +MCP_PROTOCOL_VERSION = "mcp-protocol-version" + +# JSON-RPC error codes the hand-written endpoints emit. +INVALID_REQUEST = -32600 +METHOD_NOT_FOUND = -32601 + +PINNED_TOOL_NAME = "add" + +Scope = MutableMapping[str, Any] +Message = MutableMapping[str, Any] +Receive = Callable[[], Awaitable[Message]] +Send = Callable[[Message], Awaitable[None]] +ASGIApp = Callable[[Scope, Receive, Send], Awaitable[None]] + + +def build_sdk_app() -> Starlette: + """Host a real SDK ``MCPServer`` with ``add`` and ``multiply`` tools. + + Returns: + A Streamable HTTP ASGI app serving the server at ``/mcp``. + """ + server = MCPServer("Math") + + @server.tool() + def add(a: int, b: int) -> int: + """Add two numbers""" + return a + b + + @server.tool() + def multiply(a: int, b: int) -> int: + """Multiply two numbers""" + return a * b + + return server.streamable_http_app(json_response=True) + + +class PinnedVersionServer: + """Serve Streamable HTTP while declaring one fixed protocol version. + + Adapted from ``testcases/simple-http-mcp``. Modern-era results carry + ``resultType``/``ttlMs``/``cacheScope``, which the ``2026-07-28`` wire + requires and older peers ignore. + + Args: + protocol_version: The version echoed from ``initialize`` (or advertised + by ``server/discover`` in modern-only mode). + modern_only: When True, serve ``server/discover`` and reject the legacy + ``initialize`` handshake, the way a server that speaks only the + modern protocol would. Requires a modern ``protocol_version``. + refuse_reinitialize: When True, refuse an ``initialize`` that arrives on + an existing session, the way the reference TypeScript + implementation does (``"Invalid Request: Server already + initialized"``). Resuming a session on such a server must therefore + cost no handshake. + """ + + def __init__( + self, + protocol_version: str, + *, + modern_only: bool = False, + refuse_reinitialize: bool = False, + ) -> None: + self.protocol_version = protocol_version + self.modern_only = modern_only + self.refuse_reinitialize = refuse_reinitialize + self.refused_reinitialize_count = 0 + self.session_ids: list[str] = [] + self.delete_count = 0 + self.initialize_count = 0 + self.discover_count = 0 + + def build_app(self) -> Starlette: + """Return an ASGI app exposing this endpoint at ``/mcp``.""" + return Starlette( + routes=[Route("/mcp", self._handle, methods=["GET", "POST", "DELETE"])] + ) + + async def _handle(self, request: Request) -> Response: + if request.method == "GET": + # No server-initiated stream is needed for these tests. + return Response(status_code=405) + if request.method == "DELETE": + self.delete_count += 1 + return Response(status_code=204) + + body = json.loads(await request.body()) + method = body.get("method") + session_header = request.headers.get(MCP_SESSION_ID) + + if method == "server/discover": + return self._discover(body) + if method == "initialize": + if self.refuse_reinitialize and session_header in self.session_ids: + self.refused_reinitialize_count += 1 + return self._error( + body.get("id"), + INVALID_REQUEST, + "Invalid Request: Server already initialized", + ) + return self._initialize(body) + if method is not None and method.startswith("notifications/"): + return Response(status_code=202) + if method == "tools/list": + return self._result( + body["id"], + { + "tools": [_pinned_tool_schema()], + "resultType": "complete", + "ttlMs": 0, + "cacheScope": "private", + }, + ) + if method == "tools/call": + return self._call_tool(body) + return self._error( + body.get("id"), METHOD_NOT_FOUND, f"Method not found: {method}", status=404 + ) + + def _discover(self, body: dict[str, Any]) -> Response: + """Answer ``server/discover``, but only for a modern endpoint.""" + if not self.modern_only: + # A handshake-era server has no discovery method to land on. + return self._error( + body.get("id"), + METHOD_NOT_FOUND, + "server/discover is not supported; use the initialize handshake", + status=404, + ) + self.discover_count += 1 + return self._result( + body["id"], + { + "supportedVersions": [self.protocol_version], + "capabilities": {"tools": {"listChanged": True}}, + "resultType": "complete", + "ttlMs": 0, + "cacheScope": "private", + }, + ) + + def _initialize(self, body: dict[str, Any]) -> Response: + if self.modern_only: + # A modern-only server offers no legacy handshake at all. + return self._error( + body.get("id"), + METHOD_NOT_FOUND, + "Legacy initialize is not supported; use modern discovery", + status=404, + ) + self.initialize_count += 1 + session_id = f"session-{self.initialize_count}" + self.session_ids.append(session_id) + return self._result( + body["id"], + { + "protocolVersion": self.protocol_version, + "capabilities": {"tools": {}}, + "serverInfo": {"name": "pinned-server", "version": "1.0.0"}, + }, + headers={MCP_SESSION_ID: session_id}, + ) + + def _call_tool(self, body: dict[str, Any]) -> Response: + params = body.get("params") or {} + if params.get("name") != PINNED_TOOL_NAME: + return self._error( + body.get("id"), INVALID_REQUEST, f"Unknown tool: {params.get('name')}" + ) + arguments = params.get("arguments") or {} + total = int(arguments.get("a", 0)) + int(arguments.get("b", 0)) + return self._result( + body["id"], + { + "content": [{"type": "text", "text": str(total)}], + "structuredContent": {"result": total}, + "isError": False, + "resultType": "complete", + }, + ) + + @staticmethod + def _result( + request_id: Any, + result: dict[str, Any], + *, + headers: dict[str, str] | None = None, + ) -> Response: + return JSONResponse( + {"jsonrpc": "2.0", "id": request_id, "result": result}, headers=headers + ) + + @staticmethod + def _error( + request_id: Any, code: int, message: str, *, status: int = 400 + ) -> Response: + return JSONResponse( + { + "jsonrpc": "2.0", + "id": request_id, + "error": {"code": code, "message": message}, + }, + status_code=status, + ) + + +def _pinned_tool_schema() -> dict[str, Any]: + """Describe the single tool the pinned endpoint exposes.""" + return { + "name": PINNED_TOOL_NAME, + "description": "Add two numbers", + "inputSchema": { + "type": "object", + "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}}, + "required": ["a", "b"], + }, + } + + +@dataclass +class RecordedRequest: + """One HTTP request as the gateway saw it.""" + + http_method: str + """``POST``, ``DELETE``, ``GET`` -- so a session teardown is observable.""" + + rpc_method: str | None + """JSON-RPC ``method``, or ``None`` for a body-less request.""" + + session_id: str | None + """The ``mcp-session-id`` request header, or ``None`` when unpinned.""" + + protocol_version: str | None + """The ``mcp-protocol-version`` request header, stamped after negotiation.""" + + meta: dict[str, Any] = field(default_factory=dict) + """``params._meta`` as it arrived on the wire.""" + + instance: str = "" + """Backend the gateway would have routed to, derived from ``session_id``.""" + + response_session_id: str | None = None + """``mcp-session-id`` on the response, i.e. a server-issued session.""" + + faulted: bool = False + """True when the gateway answered this request itself with an injected fault.""" + + +class RecordingGateway: + """Pure ASGI middleware that records every request and can inject a fault. + + Stands in for the AgentHub gateway. It routes by ``mcp-session-id`` exactly + as the real one does, so the instance spread it records is the observable + consequence of the affinity identity. + + Args: + app: The ASGI app to wrap. + fault_on_tool_call: 1-based index of the ``tools/call`` to answer with a + fault instead of forwarding, or ``None`` to forward everything. + fault_message: JSON-RPC error message for the injected fault. + fault_code: JSON-RPC error code for the injected fault. + fault_status: HTTP status the injected fault is returned with. + """ + + def __init__( + self, + app: ASGIApp, + *, + fault_on_tool_call: int | None = None, + fault_message: str = "Session terminated", + fault_code: int = INVALID_REQUEST, + fault_status: int = 404, + ) -> None: + self.app = app + self.fault_on_tool_call = fault_on_tool_call + self.fault_message = fault_message + self.fault_code = fault_code + self.fault_status = fault_status + self.records: list[RecordedRequest] = [] + self.instances: dict[str, str] = {} + self._tool_calls = 0 + + # --- observation helpers ------------------------------------------------ + + def rpc_methods(self) -> list[str]: + """Every JSON-RPC method seen, in order.""" + return [r.rpc_method for r in self.records if r.rpc_method is not None] + + def count(self, rpc_method: str) -> int: + """How many times one JSON-RPC method was received.""" + return self.rpc_methods().count(rpc_method) + + def http_count(self, http_method: str) -> int: + """How many requests used one HTTP method (``DELETE`` in particular).""" + return sum(1 for r in self.records if r.http_method == http_method) + + def for_rpc(self, rpc_method: str) -> list[RecordedRequest]: + """Every record for one JSON-RPC method, in order.""" + return [r for r in self.records if r.rpc_method == rpc_method] + + def server_session_ids(self) -> list[str]: + """Session IDs the server assigned on a response header.""" + return [ + r.response_session_id + for r in self.records + if r.response_session_id is not None + ] + + def unpinned(self) -> list[RecordedRequest]: + """Requests that reached the gateway with no affinity/session header.""" + return [r for r in self.records if r.session_id is None] + + # --- ASGI --------------------------------------------------------------- + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + """Record, optionally fault, and otherwise forward one ASGI request.""" + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + headers = { + key.decode("latin-1").lower(): value.decode("latin-1") + for key, value in scope.get("headers", ()) + } + http_method = str(scope.get("method", "")) + + body = b"" + forward_receive = receive + if http_method == "POST": + body, forward_receive = await _buffer_body(receive) + + payload = _parse_json(body) + params = payload.get("params") if isinstance(payload, dict) else None + session_id = headers.get(MCP_SESSION_ID) + record = RecordedRequest( + http_method=http_method, + rpc_method=payload.get("method") if isinstance(payload, dict) else None, + session_id=session_id, + protocol_version=headers.get(MCP_PROTOCOL_VERSION), + meta=dict((params or {}).get("_meta") or {}) + if isinstance(params, dict) + else {}, + instance=self._instance_for(session_id), + ) + self.records.append(record) + + if record.rpc_method == "tools/call": + self._tool_calls += 1 + if self._tool_calls == self.fault_on_tool_call: + record.faulted = True + await self._send_fault(payload, send) + return + + async def recording_send(message: Message) -> None: + if message["type"] == "http.response.start": + for key, value in message.get("headers", ()): + if key.decode("latin-1").lower() == MCP_SESSION_ID: + record.response_session_id = value.decode("latin-1") + await send(message) + + await self.app(scope, forward_receive, recording_send) + + def _instance_for(self, session_id: str | None) -> str: + """Map an affinity/session ID onto the backend a gateway would pick. + + An unpinned request cannot be routed, so it gets an instance of its own + -- which is what makes a missing affinity header visible as a spread. + """ + key = session_id if session_id is not None else f"unpinned-{len(self.records)}" + if key not in self.instances: + self.instances[key] = f"instance-{len(self.instances) + 1}" + return self.instances[key] + + async def _send_fault(self, payload: Any, send: Send) -> None: + """Answer a request with a JSON-RPC error, as a failing gateway would.""" + request_id = payload.get("id") if isinstance(payload, dict) else None + raw = json.dumps( + { + "jsonrpc": "2.0", + "id": request_id, + "error": {"code": self.fault_code, "message": self.fault_message}, + } + ).encode() + await send( + { + "type": "http.response.start", + "status": self.fault_status, + "headers": [ + (b"content-type", b"application/json"), + (b"content-length", str(len(raw)).encode()), + ], + } + ) + await send({"type": "http.response.body", "body": raw}) + + +async def _buffer_body(receive: Receive) -> tuple[bytes, Receive]: + """Read a request body fully and return it with a replaying ``receive``.""" + body = b"" + while True: + message = await receive() + if message["type"] != "http.request": + break + body += bytes(message.get("body", b"")) + if not message.get("more_body", False): + break + + replayed = False + + async def replay() -> Message: + nonlocal replayed + if not replayed: + replayed = True + return {"type": "http.request", "body": body, "more_body": False} + return await receive() + + return body, replay + + +def _parse_json(body: bytes) -> Any: + """Decode a JSON body, returning ``None`` when there is nothing to decode.""" + if not body: + return None + try: + return json.loads(body) + except (json.JSONDecodeError, UnicodeDecodeError): + return None + + +def _reset_sse_shutdown_latch() -> None: + """Clear ``sse-starlette``'s process-global shutdown latch. + + ``sse_starlette`` runs a watcher that polls ``uvicorn.Server.should_exit`` + and latches a *module-global* ``AppStatus.should_exit`` when any server + stops. Tests host one server per case, so from the second server onward + every SSE stream would see the latch already set and end the instant it + opened -- logging ``ASGI callable returned without completing response`` and + sending the client into a reconnect loop. Clearing the latch (and the + per-thread watcher bookkeeping, which is stranded on the previous event + loop) gives every server the same first-server behaviour. + + Best-effort: these are private names, so a changed internal degrades to the + noisy-but-working behaviour rather than breaking the harness. + """ + try: + from sse_starlette import sse as sse_module + except ImportError: # pragma: no cover - sse-starlette ships with mcp + return + try: + sse_module.AppStatus.should_exit = False + state = getattr(sse_module._thread_state, "shutdown_state", None) + if state is not None: + state.watcher_started = False + state.events.clear() + except AttributeError: # pragma: no cover - upstream internals moved + return + + +@asynccontextmanager +async def serve(app: ASGIApp) -> AsyncGenerator[str, None]: + """Run *app* on an ephemeral port and yield its ``/mcp`` URL. + + Binding port 0 keeps parallel CI jobs from colliding, and hosting in-process + means no child process is left behind when a test fails. + + Args: + app: Any ASGI application. + + Yields: + The ``http://127.0.0.1:/mcp`` endpoint URL. + """ + # Imported here so the module can be collected even if the optional dev + # dependency is missing, and to keep import cost off the test session. + import uvicorn + + _reset_sse_shutdown_latch() + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + + server = uvicorn.Server(uvicorn.Config(app, log_level="warning")) + task: asyncio.Task[Any] | None = None + try: + task = asyncio.create_task(server.serve(sockets=[sock])) + while not server.started: + if task.done(): + task.result() + raise RuntimeError("MCP test server exited before startup") + await asyncio.sleep(0.01) + yield f"http://127.0.0.1:{port}/mcp" + finally: + server.should_exit = True + if task is not None: + with contextlib.suppress(BaseException): + await task + sock.close() + + +SERVER_NAME = "Math" +SERVER_SLUG = "math" +FOLDER_KEY = "folder-key" +FOLDER_PATH = "Shared" +ACCESS_TOKEN = "test-access-token" + + +@contextmanager +def patched_sdk(url: str) -> Iterator[None]: + """Point ``McpClient``'s lazy SDK lookup at a locally hosted server. + + ``McpClient._initialize_client`` imports ``UiPath`` from ``uipath.platform`` + at call time, so replacing the module attribute is enough: no tenant, + credentials, or network access to UiPath Cloud are involved, and the client + still walks its real resolution path. + + Args: + url: The MCP endpoint the fake registration should resolve to. + """ + import uipath.platform as platform + from uipath.platform.orchestrator.mcp import McpServer + + class _FakeMcpService: + async def retrieve_async( + self, name: str, folder_path: str | None = None + ) -> McpServer: + # Recorded so a test can assert *how* the server was resolved, not + # just that it was: the client must look up by display name and pass + # the execution folder through. + SDK_LOOKUPS.append({"name": name, "folder_path": folder_path}) + return McpServer( + id="mcp-server-id", + name=name, + slug=SERVER_SLUG, + folderKey=FOLDER_KEY, + mcpUrl=url, + ) + + class _FakeConfig: + secret = ACCESS_TOKEN + + class _FakeUiPath: + def __init__(self, *args: Any, **kwargs: Any) -> None: + self._config = _FakeConfig() + self.mcp = _FakeMcpService() + + SDK_LOOKUPS.clear() + original = platform.UiPath + platform.UiPath = _FakeUiPath # type: ignore[misc,assignment] + try: + yield + finally: + platform.UiPath = original # type: ignore[misc] + + +#: Every ``retrieve_async`` call made through :func:`patched_sdk`, in order. +#: Cleared by that context manager on entry. +SDK_LOOKUPS: list[dict[str, Any]] = [] + + +def make_resource_config() -> AgentMcpResourceConfig: + """Build the MCP resource config every real-HTTP test drives.""" + return AgentMcpResourceConfig( + name=SERVER_NAME, + description="Math MCP server", + folder_path=FOLDER_PATH, + slug=SERVER_SLUG, + available_tools=[ + AgentMcpTool( + name=PINNED_TOOL_NAME, + description="Add two numbers", + input_schema={ + "type": "object", + "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}}, + "required": ["a", "b"], + }, + ) + ], + ) + + +def pinned_session_factory(session_info: SessionInfo) -> SessionInfoFactory: + """Return a factory handing every client the same ``SessionInfo``. + + Two clients sharing one store is how a persisted session (or affinity ID) + survives across runs, so this is what a resume test drives. + + Args: + session_info: The store every created client should use. + """ + + class _PinnedFactory(SessionInfoFactory): + def create_session(self, mcp_server: Any) -> SessionInfo: + return session_info + + return _PinnedFactory() + + +def make_client(**kwargs: Any) -> McpClient: + """Create an ``McpClient`` for the shared resource config. + + The URL is resolved through :func:`patched_sdk`, so that context manager + must be active when the returned client first connects. + + Args: + **kwargs: Forwarded to ``McpClient`` (``protocol_mode``, + ``session_info_factory``, ``terminate_on_close``, ...). + """ + return McpClient(config=make_resource_config(), **kwargs) + + +@asynccontextmanager +async def connected_client(url: str, **kwargs: Any) -> AsyncGenerator[McpClient, None]: + """Yield an ``McpClient`` wired to *url*, disposing it on exit. + + Args: + url: The endpoint returned by :func:`serve`. + **kwargs: Forwarded to ``McpClient``. + """ + with patched_sdk(url): + client = make_client(**kwargs) + try: + yield client + finally: + await client.dispose() diff --git a/tests/agent/tools/test_mcp/test_mcp_client.py b/tests/agent/tools/test_mcp/test_mcp_client.py index dc3d117a3..0917ac1c0 100644 --- a/tests/agent/tools/test_mcp/test_mcp_client.py +++ b/tests/agent/tools/test_mcp/test_mcp_client.py @@ -1,883 +1,769 @@ -"""Tests for McpClient class.""" +"""Tests for the MCP 2 Streamable HTTP client integration.""" +import asyncio import json -import logging -import os +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager from typing import Any from unittest.mock import AsyncMock, MagicMock, patch +import httpx +import httpx2 import pytest +from mcp.shared.exceptions import MCPError +from mcp.types import CONNECTION_CLOSED, INVALID_REQUEST, METHOD_NOT_FOUND from uipath.agent.models.agent import AgentMcpResourceConfig, AgentMcpTool from uipath_langchain.agent.tools.mcp import McpClient, SessionInfo, SessionInfoFactory -logger = logging.getLogger(__name__) - - -class TestMcpClient: - """Test MCP client behavior with mocked HTTP.""" - - @pytest.fixture - def mcp_resource_config(self): - """Create a minimal MCP resource config for testing.""" - return AgentMcpResourceConfig( - name="test_server", - description="Test MCP server", - folder_path="/Shared/TestFolder", - slug="test-server", - available_tools=[ - AgentMcpTool( - name="test_tool", - description="A test tool", - input_schema={ - "type": "object", - "properties": {"query": {"type": "string"}}, - "required": ["query"], - }, - ) - ], - ) - - @pytest.fixture - def mock_uipath_sdk(self): - """Create a mock UiPath SDK for patching.""" - mock_sdk = MagicMock() - mock_server = MagicMock() - mock_server.mcp_url = "https://test.uipath.com/mcp" - mock_sdk.mcp.retrieve_async = AsyncMock(return_value=mock_server) - mock_sdk._config = MagicMock() - mock_sdk._config.secret = "test-secret-token" - return mock_sdk - - def create_mock_stream_response( - self, - method_call_sequence: list[str], - initialize_count: list[int], - tool_call_count: list[int], - session_guid_1: str = "test-session-first", - session_guid_2: str = "test-session-retry", - fail_first_tool_call: bool = False, - ): - """Create a MockStreamResponse class for testing. - - Args: - method_call_sequence: List to track method calls. - initialize_count: Mutable counter for initialize calls. - tool_call_count: Mutable counter for tool calls. - session_guid_1: Session ID for first initialization. - session_guid_2: Session ID for retry initialization. - fail_first_tool_call: If True, first tool call returns 404. - """ - - class MockStreamResponse: - """Mock HTTP stream response for MCP protocol.""" - - def __init__(self, method: str, url: str, **kwargs: Any): - self.request_method = method - self.url = url - self.kwargs = kwargs - if method == "GET": - self.status_code = 405 - self.headers = {} - self._content = b"" - return +class LegacyMcpEndpoint: + """Small Streamable HTTP endpoint used to exercise the real MCP SDK transport.""" - json_body = kwargs.get("json", {}) - request_headers = kwargs.get("headers", {}) - - self.json_body = json_body - self.method = json_body.get("method", "") - self.request_headers = request_headers - self.request_mcp_session_id = request_headers.get("mcp-session-id", "") - - logger.debug( - f"Responding to method {self.method} for session {self.request_mcp_session_id}" + def __init__( + self, + protocol_version: str = "2025-11-25", + *, + failed_tool_calls: int = 0, + failed_tool_message: str | None = None, + failed_tool_code: int = INVALID_REQUEST, + block_initialize_on: int | None = None, + fail_initialize_on: set[int] | None = None, + rejected_session_ids: set[str] | None = None, + repeat_session_header: bool = False, + known_session_ids: set[str] | None = None, + mints_new_session_on_initialize: bool = False, + ) -> None: + self.protocol_version = protocol_version + self.failed_tool_calls = failed_tool_calls + self.failed_tool_message = failed_tool_message + # JSON-RPC code for the injected failure. The SDK synthesizes + # CONNECTION_CLOSED client-side when a transport dies, which a mock + # transport cannot stage cleanly; returning the code in a body drives + # the same recovery decision through the real transport. + self.failed_tool_code = failed_tool_code + self.block_initialize_on = block_initialize_on + self.fail_initialize_on = fail_initialize_on or set() + self.rejected_session_ids = rejected_session_ids or set() + self.repeat_session_header = repeat_session_header + # Sessions this endpoint will route to. Seed it to stand in for a session + # a previous process established and persisted externally. + self.known_session_ids = set(known_session_ids or ()) + # Real servers route by the session header and mint only when it is + # absent. Set this to model a server that ignores the header instead. + self.mints_new_session_on_initialize = mints_new_session_on_initialize + self.initialize_blocked = asyncio.Event() + self.release_initialize = asyncio.Event() + self.methods: list[str] = [] + self.request_headers: list[tuple[str, httpx2.Headers]] = [] + self.initialize_count = 0 + self.session_mint_count = 0 + self.tool_call_count = 0 + self.delete_count = 0 + self.transport = httpx2.MockTransport(self.handle) + + async def handle(self, request: httpx2.Request) -> httpx2.Response: + """Return protocol-correct JSON responses for the MCP methods under test.""" + if request.method == "GET": + return httpx2.Response(405) + if request.method == "DELETE": + self.delete_count += 1 + self.request_headers.append(("DELETE", request.headers)) + return httpx2.Response(204) + + body = json.loads(request.content) + method = body["method"] + self.methods.append(method) + self.request_headers.append((method, request.headers)) + + if method == "initialize": + self.initialize_count += 1 + if self.initialize_count == self.block_initialize_on: + self.initialize_blocked.set() + await self.release_initialize.wait() + if self.initialize_count in self.fail_initialize_on: + return httpx2.Response( + 400, + headers={"content-type": "application/json"}, + json={ + "jsonrpc": "2.0", + "id": body["id"], + "error": { + "code": INVALID_REQUEST, + "message": "Replacement initialization failed", + }, + }, ) - method_call_sequence.append(self.method) - - status_code, response_json, headers = self._build_response() - self.headers = headers or {} - self._response_json = response_json - self.status_code = status_code - - if response_json: - self._content = json.dumps(self._response_json).encode("utf-8") - self.headers["content-type"] = "application/json" - else: - self._content = b"" - - def _build_response(self) -> tuple[int, Any, dict[str, str] | None]: - """Build JSON-RPC response based on method.""" - request_id = self.json_body.get("id") - - if self.method == "initialize": - initialize_count[0] += 1 - session_id = ( - session_guid_1 if initialize_count[0] == 1 else session_guid_2 - ) - logger.debug(f"MCP initializes new session {session_id}") - return ( - 200, - { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "protocolVersion": "2025-06-18", - "capabilities": {"tools": {}}, - "serverInfo": { - "name": "test-server", - "version": "1.0.0", - }, - }, + session_id = self._session_for_initialize(request) + if session_id is None: + return httpx2.Response( + 404, + headers={"content-type": "application/json"}, + json={ + "jsonrpc": "2.0", + "id": body["id"], + "error": { + "code": INVALID_REQUEST, + "message": "Session not found", }, - {"mcp-session-id": session_id}, - ) - - elif self.method == "notifications/initialized": - return (204, None, {}) - - elif self.method == "tools/list": - return ( - 200, - { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "tools": [ - { - "name": "test_tool", - "description": "A test tool", - "inputSchema": { - "type": "object", - "properties": {"query": {"type": "string"}}, - "required": ["query"], - }, - "outputSchema": { - "type": "object", - "properties": { - "result": {"type": "string"} - }, - }, - } - ], - }, + }, + ) + return self._json_response( + body["id"], + { + "protocolVersion": self.protocol_version, + "capabilities": {"tools": {}}, + "serverInfo": {"name": "test-server", "version": "1.0.0"}, + }, + headers={"mcp-session-id": session_id}, + ) + if method == "notifications/initialized": + return httpx2.Response(202) + if method == "ping": + if request.headers.get("mcp-session-id") in self.rejected_session_ids: + return httpx2.Response( + 404, + headers={"content-type": "application/json"}, + json={ + "jsonrpc": "2.0", + "id": body["id"], + "error": { + "code": INVALID_REQUEST, + "message": "Session not found", }, - {}, - ) - - elif self.method == "tools/call": - tool_call_count[0] += 1 - - if fail_first_tool_call and tool_call_count[0] == 1: - # Return HTTP 404 to trigger session re-initialization - return (404, None, None) - - # Success response with structured content - params = self.json_body.get("params", {}) - tool_name = params.get("name", "unknown") - structured_result = {"result": f"Success from {tool_name}"} - - return ( - 200, - { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "content": [ - { - "type": "text", - "text": json.dumps(structured_result), - } - ], - "structuredContent": structured_result, - "isError": False, - }, + }, + ) + if request.headers.get("mcp-protocol-version") != self.protocol_version: + return httpx2.Response( + 400, + headers={"content-type": "application/json"}, + json={ + "jsonrpc": "2.0", + "id": body["id"], + "error": { + "code": INVALID_REQUEST, + "message": "Unsupported protocol version", }, - {}, - ) - - else: - if request_id is None: - return (204, None, {}) - return ( - 500, + }, + ) + return self._json_response(body["id"], {}) + if method == "tools/list": + return self._json_response( + body["id"], + { + "tools": [ { + "name": "test_tool", + "description": "A test tool", + "inputSchema": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + "outputSchema": { + "type": "object", + "properties": {"result": {"type": "string"}}, + }, + } + ] + }, + headers=self._repeated_session_header(), + ) + if method == "tools/call": + self.tool_call_count += 1 + if self.tool_call_count <= self.failed_tool_calls: + if self.failed_tool_message is not None: + return httpx2.Response( + 404, + headers={"content-type": "application/json"}, + json={ "jsonrpc": "2.0", - "id": request_id, - "error": {"code": -32601, "message": "Method not found"}, + "id": body["id"], + "error": { + "code": self.failed_tool_code, + "message": self.failed_tool_message, + }, }, - {}, ) - - async def __aenter__(self): - return self - - async def __aexit__(self, *args: Any, **kwargs: Any): - pass - - async def aread(self) -> bytes: - """Return the response content.""" - return self._content - - def raise_for_status(self) -> None: - """Check response status.""" - if self.status_code >= 400: - raise Exception(f"HTTP {self.status_code}") - - return MockStreamResponse - - def create_mock_http_client(self, mock_stream_response_class: type) -> MagicMock: - """Create a mock HTTP client that uses the given stream response class.""" - mock_client = MagicMock() - mock_client.stream = lambda method, url, **kwargs: mock_stream_response_class( - method, url, **kwargs + return httpx2.Response(404) + result = {"result": f"Success from {body['params']['name']}"} + return self._json_response( + body["id"], + { + "content": [{"type": "text", "text": json.dumps(result)}], + "structuredContent": result, + "isError": False, + }, + headers=self._repeated_session_header(), + ) + return httpx2.Response( + 404, + json={ + "jsonrpc": "2.0", + "id": body.get("id"), + "error": {"code": METHOD_NOT_FOUND, "message": "Method not found"}, + }, ) - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock() - return mock_client - - @pytest.mark.asyncio - @patch("httpx.AsyncClient") - async def test_session_initializes_on_first_call( - self, mock_async_client_class, mcp_resource_config, mock_uipath_sdk - ): - """Test that session is initialized lazily on first tool call.""" - method_call_sequence: list[str] = [] - initialize_count = [0] - tool_call_count = [0] - MockStreamResponse = self.create_mock_stream_response( - method_call_sequence, initialize_count, tool_call_count + @staticmethod + def _json_response( + request_id: int, + result: dict[str, Any], + *, + headers: dict[str, str] | None = None, + ) -> httpx2.Response: + response_headers = {"content-type": "application/json"} + response_headers.update(headers or {}) + return httpx2.Response( + 200, + headers=response_headers, + json={"jsonrpc": "2.0", "id": request_id, "result": result}, ) - mock_http_client = self.create_mock_http_client(MockStreamResponse) - mock_async_client_class.return_value = mock_http_client - - session = McpClient(config=mcp_resource_config) - - # Session should not be initialized yet - assert await session.get_session_id() is None - assert not session.is_client_initialized - - # Call tool - should trigger initialization (with SDK mocked) - with patch( - "uipath.platform.UiPath", - return_value=mock_uipath_sdk, - ): - result = await session.call_tool("test_tool", {"query": "test"}) - - # Verify initialization happened - assert initialize_count[0] == 1 - assert await session.get_session_id() == "test-session-first" - assert session.is_client_initialized - assert tool_call_count[0] == 1 - assert result is not None - - # Verify HTTP client was created once - assert mock_async_client_class.call_count == 1 - - # Verify method sequence - assert "initialize" in method_call_sequence - assert "notifications/initialized" in method_call_sequence - assert "tools/call" in method_call_sequence - - await session.dispose() - - @pytest.mark.asyncio - @patch("httpx.AsyncClient") - async def test_session_reused_across_calls( - self, mock_async_client_class, mcp_resource_config, mock_uipath_sdk - ): - """Test that session is reused for multiple tool calls.""" - method_call_sequence: list[str] = [] - initialize_count = [0] - tool_call_count = [0] - - MockStreamResponse = self.create_mock_stream_response( - method_call_sequence, initialize_count, tool_call_count - ) - - mock_http_client = self.create_mock_http_client(MockStreamResponse) - mock_async_client_class.return_value = mock_http_client - - session = McpClient(config=mcp_resource_config) - - with patch( - "uipath.platform.UiPath", - return_value=mock_uipath_sdk, - ): - # First call - await session.call_tool("test_tool", {"query": "first"}) - assert initialize_count[0] == 1 - - # Second call - should reuse session - await session.call_tool("test_tool", {"query": "second"}) - assert initialize_count[0] == 1 # Still only one initialization - assert tool_call_count[0] == 2 # But two tool calls - - # HTTP client should still be created only once - assert mock_async_client_class.call_count == 1 - - await session.dispose() + def headers_for(self, method: str) -> list[httpx2.Headers]: + """Return captured headers for one protocol or HTTP method.""" + return [headers for name, headers in self.request_headers if name == method] - @pytest.mark.asyncio - @patch("httpx.AsyncClient") - async def test_session_reinitializes_on_404_error( - self, mock_async_client_class, mcp_resource_config, mock_uipath_sdk - ): - """Test that only session (not client) is reinitialized on 404 error. + def _session_for_initialize(self, request: httpx2.Request) -> str | None: + """Resolve the session an ``initialize`` belongs to, or None to reject it. - This verifies the key behavior: when a 404 error occurs, we should: - - Keep the existing HTTP client (not create a new one) - - Keep the existing streamable connection - - Only call session.initialize() again to get a new session ID + Mirrors the SDK server: a request naming a live session is handled inside + it, and a new session is minted only when no session header is present. + An unknown or expired ID is rejected rather than silently replaced. """ - method_call_sequence: list[str] = [] - initialize_count = [0] - tool_call_count = [0] - - MockStreamResponse = self.create_mock_stream_response( - method_call_sequence, - initialize_count, - tool_call_count, - fail_first_tool_call=True, - ) - - mock_http_client = self.create_mock_http_client(MockStreamResponse) - mock_async_client_class.return_value = mock_http_client - - session = McpClient(config=mcp_resource_config) - - # Call tool - first call fails with 404, should retry - with patch( - "uipath.platform.UiPath", - return_value=mock_uipath_sdk, - ): - result = await session.call_tool("test_tool", {"query": "test"}) - - logger.info(f"Result: {result}") - logger.info(f"Method sequence: {method_call_sequence}") - logger.info(f"Initialize count: {initialize_count[0]}") - logger.info(f"Tool call count: {tool_call_count[0]}") - - # Verify session was reinitialized (initialize called twice) - assert initialize_count[0] == 2, ( - f"Expected 2 session initializations, got {initialize_count[0]}" - ) - - # Verify tool call was retried - assert tool_call_count[0] == 2, ( - f"Expected 2 tool calls, got {tool_call_count[0]}" - ) - - # Verify session ID changed to the retry session - assert await session.get_session_id() == "test-session-retry" - assert result is not None - - # KEY ASSERTION: HTTP client should be created only ONCE - # Session reinitialization reuses the existing client - assert mock_async_client_class.call_count == 1, ( - f"Expected HTTP client to be created only once, " - f"but was created {mock_async_client_class.call_count} times" - ) - - # Verify the expected method sequence - expected_init_count = method_call_sequence.count("initialize") - expected_tool_count = method_call_sequence.count("tools/call") - assert expected_init_count == 2, ( - f"Expected 2 initialize calls, got {expected_init_count}" - ) - assert expected_tool_count == 2, ( - f"Expected 2 tools/call, got {expected_tool_count}" - ) - - await session.dispose() - - @pytest.mark.asyncio - @patch("httpx.AsyncClient") - async def test_max_retries_exceeded( - self, mock_async_client_class, mcp_resource_config, mock_uipath_sdk + incoming = request.headers.get("mcp-session-id") + if incoming is not None and not self.mints_new_session_on_initialize: + if ( + incoming in self.rejected_session_ids + or incoming not in self.known_session_ids + ): + return None + return incoming + self.session_mint_count += 1 + minted = f"session-{self.session_mint_count}" + self.known_session_ids.add(minted) + return minted + + def _repeated_session_header(self) -> dict[str, str] | None: + if not self.repeat_session_header or self.session_mint_count == 0: + return None + return {"mcp-session-id": f"session-{self.session_mint_count}"} + + +@pytest.fixture +def mcp_resource_config() -> AgentMcpResourceConfig: + """Create a minimal MCP resource config for testing.""" + return AgentMcpResourceConfig( + name="test_server", + description="Test MCP server", + folder_path="/Shared/TestFolder", + slug="test-server", + available_tools=[ + AgentMcpTool( + name="test_tool", + description="A test tool", + input_schema={ + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + ) + ], + ) + + +@pytest.fixture +def mock_uipath_sdk() -> MagicMock: + """Create a mock UiPath SDK and resolved MCP server.""" + sdk = MagicMock() + server = MagicMock() + server.mcp_url = "https://test.uipath.com/mcp" + server.slug = "test-server" + server.folder_key = "folder-key" + sdk.mcp.retrieve_async = AsyncMock(return_value=server) + sdk._config.secret = "test-secret-token" + return sdk + + +@asynccontextmanager +async def configured_client( + config: AgentMcpResourceConfig, + sdk: MagicMock, + endpoint: LegacyMcpEndpoint, + **kwargs: Any, +) -> AsyncIterator[McpClient]: + """Build an McpClient whose real HTTP client uses the mock transport.""" + client = McpClient(config=config, **kwargs) + http_kwargs = { + "headers": {"Authorization": "Bearer test-secret-token"}, + "transport": endpoint.transport, + "follow_redirects": True, + } + with ( + patch("uipath.platform.UiPath", return_value=sdk), + patch( + "uipath_langchain.agent.tools.mcp.mcp_client.get_httpx_client_kwargs", + return_value=http_kwargs, + ), ): - """Test that exception is raised when max retries are exceeded.""" - from mcp.shared.exceptions import McpError - - method_call_sequence: list[str] = [] - initialize_count = [0] - tool_call_count = [0] - - # Create a response that always fails tool calls - class AlwaysFailMockResponse: - def __init__(self, method: str, url: str, **kwargs: Any): - self.request_method = method - if method == "GET": - self.status_code = 405 - self.headers = {} - self._content = b"" - return - - json_body = kwargs.get("json", {}) - self.method = json_body.get("method", "") - method_call_sequence.append(self.method) - request_id = json_body.get("id") - - if self.method == "initialize": - initialize_count[0] += 1 - self.status_code = 200 - self.headers = {"mcp-session-id": f"session-{initialize_count[0]}"} - self._response_json = { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "protocolVersion": "2025-06-18", - "capabilities": {"tools": {}}, - "serverInfo": {"name": "test", "version": "1.0"}, - }, - } - self._content = json.dumps(self._response_json).encode() - self.headers["content-type"] = "application/json" - elif self.method == "notifications/initialized": - self.status_code = 204 - self.headers = {} - self._content = b"" - elif self.method == "tools/call": - tool_call_count[0] += 1 - # Always return 404 - self.status_code = 404 - self.headers = {} - self._content = b"" - else: - self.status_code = 200 - self.headers = {} - self._content = b"" - - async def __aenter__(self): - return self - - async def __aexit__(self, *args: Any): - pass - - async def aread(self) -> bytes: - return self._content - - def raise_for_status(self) -> None: - if self.status_code >= 400: - raise Exception(f"HTTP {self.status_code}") - - mock_http_client = self.create_mock_http_client(AlwaysFailMockResponse) - mock_async_client_class.return_value = mock_http_client - - session = McpClient(config=mcp_resource_config, max_retries=1) - - # Should raise McpError after retries exhausted - with patch( - "uipath.platform.UiPath", - return_value=mock_uipath_sdk, - ): - with pytest.raises(McpError): - await session.call_tool("test_tool", {"query": "test"}) - - # Should have reinitialized session (2 initialize calls) - assert initialize_count[0] == 2 - - # Should have attempted tool call twice - assert tool_call_count[0] == 2 - - # HTTP client still created only once - assert mock_async_client_class.call_count == 1 - - await session.dispose() - - @pytest.mark.asyncio - @patch("httpx.AsyncClient") - async def test_dispose_releases_resources( - self, mock_async_client_class, mcp_resource_config, mock_uipath_sdk - ): - """Test that dispose() properly releases session resources.""" - method_call_sequence: list[str] = [] - initialize_count = [0] - tool_call_count = [0] - - MockStreamResponse = self.create_mock_stream_response( - method_call_sequence, initialize_count, tool_call_count - ) - - mock_http_client = self.create_mock_http_client(MockStreamResponse) - mock_async_client_class.return_value = mock_http_client - - session = McpClient(config=mcp_resource_config) - - # Initialize session - with patch( - "uipath.platform.UiPath", - return_value=mock_uipath_sdk, - ): - await session.call_tool("test_tool", {"query": "test"}) - assert await session.get_session_id() is not None - assert session.is_client_initialized - - # Close session - await session.dispose() - - # Verify resources are released - assert await session.get_session_id() is None - assert session._session is None - assert session._stack is None - assert not session.is_client_initialized - - @pytest.mark.asyncio - @patch("httpx.AsyncClient") - async def test_client_initialized_property( - self, mock_async_client_class, mcp_resource_config, mock_uipath_sdk - ): - """Test that is_client_initialized property reflects actual state.""" - method_call_sequence: list[str] = [] - initialize_count = [0] - tool_call_count = [0] - - MockStreamResponse = self.create_mock_stream_response( - method_call_sequence, initialize_count, tool_call_count - ) - - mock_http_client = self.create_mock_http_client(MockStreamResponse) - mock_async_client_class.return_value = mock_http_client - - session = McpClient(config=mcp_resource_config) - - # Before any call - assert not session.is_client_initialized - - with patch( - "uipath.platform.UiPath", - return_value=mock_uipath_sdk, - ): - # After first call - await session.call_tool("test_tool", {"query": "test"}) - assert session.is_client_initialized - - # After dispose - await session.dispose() - assert not session.is_client_initialized - - @pytest.mark.asyncio - @patch("httpx.AsyncClient") - async def test_session_can_be_reused_after_dispose( - self, mock_async_client_class, mcp_resource_config, mock_uipath_sdk - ): - """Test that session can be reinitialized after dispose().""" - method_call_sequence: list[str] = [] - initialize_count = [0] - tool_call_count = [0] - - MockStreamResponse = self.create_mock_stream_response( - method_call_sequence, initialize_count, tool_call_count + try: + yield client + finally: + await client.dispose() + + +@pytest.mark.asyncio +async def test_legacy_httpx_timeout_is_normalized_for_final_client( + mcp_resource_config: AgentMcpResourceConfig, + mock_uipath_sdk: MagicMock, +) -> None: + """The pre-upgrade public timeout type remains accepted by McpClient.""" + endpoint = LegacyMcpEndpoint() + legacy_timeout = httpx.Timeout(20, connect=1, read=2, write=3, pool=4) + + async with configured_client( + mcp_resource_config, + mock_uipath_sdk, + endpoint, + timeout=legacy_timeout, + ) as client: + await client.call_tool("test_tool", {"query": "test"}) + + assert client._http_client is not None + final_timeout = client._http_client.timeout + assert final_timeout.connect == 1 + assert final_timeout.read == 2 + assert final_timeout.write == 3 + assert final_timeout.pool == 4 + + +@pytest.mark.asyncio +async def test_replaces_transport_and_session_after_404( + mcp_resource_config: AgentMcpResourceConfig, + mock_uipath_sdk: MagicMock, +) -> None: + """A terminated session gets a fresh handshake while reusing its HTTP client.""" + endpoint = LegacyMcpEndpoint(failed_tool_calls=1) + async with configured_client( + mcp_resource_config, mock_uipath_sdk, endpoint + ) as client: + result = await client.call_tool("test_tool", {"query": "test"}) + + assert result.structured_content == {"result": "Success from test_tool"} + assert endpoint.initialize_count == 2 + assert endpoint.tool_call_count == 2 + assert endpoint.delete_count == 1 + assert await client.get_session_id() == "session-2" + assert [h["mcp-session-id"] for h in endpoint.headers_for("tools/call")] == [ + "session-1", + "session-2", + ] + + +@pytest.mark.asyncio +async def test_replaces_session_after_official_session_not_found_error( + mcp_resource_config: AgentMcpResourceConfig, + mock_uipath_sdk: MagicMock, +) -> None: + """The SDK server's canonical expired-session response triggers recovery.""" + endpoint = LegacyMcpEndpoint( + failed_tool_calls=1, + failed_tool_message="Session not found", + ) + async with configured_client( + mcp_resource_config, mock_uipath_sdk, endpoint + ) as client: + result = await client.call_tool("test_tool", {"query": "test"}) + + assert result.structured_content == {"result": "Success from test_tool"} + assert endpoint.initialize_count == 2 + assert endpoint.tool_call_count == 2 + assert await client.get_session_id() == "session-2" + + +@pytest.mark.asyncio +async def test_dropped_connection_resumes_the_persisted_session( + mcp_resource_config: AgentMcpResourceConfig, + mock_uipath_sdk: MagicMock, +) -> None: + """``CONNECTION_CLOSED`` reopens the transport but keeps the session. + + A dropped connection is not the server's verdict on the session, so the + retry must resume it. A fresh handshake would start a cold session and, for + a store-backed ``SessionInfo``, throw away the persisted one for nothing. + """ + endpoint = LegacyMcpEndpoint( + failed_tool_calls=1, + failed_tool_message="Connection closed", + failed_tool_code=CONNECTION_CLOSED, + ) + async with configured_client( + mcp_resource_config, mock_uipath_sdk, endpoint + ) as client: + result = await client.call_tool("test_tool", {"query": "test"}) + + assert result.structured_content == {"result": "Success from test_tool"} + # No second handshake: the reconnect adopts the version the first one + # negotiated, so the resumed session costs nothing on the wire. + assert endpoint.initialize_count == 1 + assert endpoint.session_mint_count == 1 + assert await client.get_session_id() == "session-1" + assert [h["mcp-session-id"] for h in endpoint.headers_for("tools/call")] == [ + "session-1", + "session-1", + ] + + +@pytest.mark.asyncio +async def test_auto_mode_does_not_offer_a_minted_id_to_a_legacy_handshake( + mcp_resource_config: AgentMcpResourceConfig, + mock_uipath_sdk: MagicMock, +) -> None: + """The probe is pinned, but a legacy server never issued that ID. + + This endpoint refuses any ``initialize`` naming a session it did not mint, so + a handshake still carrying the affinity ID would be rejected first and the + session established only on the clean retry. + """ + endpoint = LegacyMcpEndpoint() + async with configured_client( + mcp_resource_config, mock_uipath_sdk, endpoint, protocol_mode="auto" + ) as client: + result = await client.call_tool("test_tool", {"query": "test"}) + + assert result.structured_content == {"result": "Success from test_tool"} + assert endpoint.methods[0] == "server/discover" + assert endpoint.headers_for("server/discover")[0].get("mcp-session-id") + # Accepted on the first attempt: the minted ID was withdrawn before it. + assert endpoint.initialize_count == 1 + assert endpoint.headers_for("initialize")[0].get("mcp-session-id") is None + assert await client.get_session_id() == "session-1" + + +@pytest.mark.asyncio +async def test_persisted_session_replaced_when_server_ignores_the_header( + mcp_resource_config: AgentMcpResourceConfig, + mock_uipath_sdk: MagicMock, +) -> None: + """A server that mints on every handshake loses the session but stays usable.""" + endpoint = LegacyMcpEndpoint( + known_session_ids={"persisted-session"}, + mints_new_session_on_initialize=True, + ) + session_info = SessionInfo("persisted-session") + + class PersistedFactory(SessionInfoFactory): + def create_session(self, mcp_server: Any) -> SessionInfo: + return session_info + + async with configured_client( + mcp_resource_config, + mock_uipath_sdk, + endpoint, + session_info_factory=PersistedFactory(), + ) as client: + result = await client.call_tool("test_tool", {"query": "test"}) + + assert result.structured_content == {"result": "Success from test_tool"} + assert endpoint.initialize_count == 1 + assert await client.get_session_id() == "session-1" + assert endpoint.headers_for("tools/call")[0]["mcp-session-id"] == "session-1" + + +@pytest.mark.asyncio +async def test_rejected_persisted_session_is_initialized_and_deleted_once( + mcp_resource_config: AgentMcpResourceConfig, + mock_uipath_sdk: MagicMock, +) -> None: + """A stale restored ID is replaced and only the fresh SDK session is deleted.""" + endpoint = LegacyMcpEndpoint(rejected_session_ids={"expired-session"}) + session_info = SessionInfo("expired-session") + + class PersistedFactory(SessionInfoFactory): + def create_session(self, mcp_server: Any) -> SessionInfo: + return session_info + + async with configured_client( + mcp_resource_config, + mock_uipath_sdk, + endpoint, + session_info_factory=PersistedFactory(), + ) as client: + result = await client.call_tool("test_tool", {"query": "test"}) + + assert result.structured_content == {"result": "Success from test_tool"} + # The refused handshake for the stale ID, then the clean one. + assert endpoint.initialize_count == 2 + assert endpoint.session_mint_count == 1 + assert await client.get_session_id() == "session-1" + + assert endpoint.delete_count == 1 + assert endpoint.headers_for("DELETE")[0]["mcp-session-id"] == "session-1" + + +@pytest.mark.asyncio +async def test_only_an_initialize_response_assigns_a_session_id( + mcp_resource_config: AgentMcpResourceConfig, + mock_uipath_sdk: MagicMock, +) -> None: + """A session ID arriving on any other response must not be adopted. + + The SDK's own transport reads the ID only from the handshake. Persisting it + from any response would let a proxy echoing the header replace a stored ID + mid-connection -- including a client-minted routing key in ``auto`` mode, + whose probe runs on the legacy wire before the era is known. + """ + + class RecordingSessionInfo(SessionInfo): + def __init__(self) -> None: + super().__init__() + self.persisted: list[str | None] = [] + + async def set_session_id(self, session_id: str | None) -> None: + self.persisted.append(session_id) + await super().set_session_id(session_id) + + session_info = RecordingSessionInfo() + + class RecordingFactory(SessionInfoFactory): + def create_session(self, mcp_server: Any) -> SessionInfo: + return session_info + + # The endpoint stamps a *different* session ID onto every non-initialize + # response, the way a session-rewriting proxy would. + endpoint = LegacyMcpEndpoint(repeat_session_header=True) + endpoint._repeated_session_header = lambda: {"mcp-session-id": "proxy-injected"} # type: ignore[method-assign] + + async with configured_client( + mcp_resource_config, + mock_uipath_sdk, + endpoint, + session_info_factory=RecordingFactory(), + ) as client: + await client.call_tool("test_tool", {"query": "test"}) + + assert session_info.persisted == ["session-1"] + assert await client.get_session_id() == "session-1" + + +@pytest.mark.asyncio +async def test_repeated_session_headers_do_not_repeat_external_persistence( + mcp_resource_config: AgentMcpResourceConfig, + mock_uipath_sdk: MagicMock, +) -> None: + """An unchanged response header is not persisted after every MCP call.""" + + class CountingSessionInfo(SessionInfo): + def __init__(self) -> None: + super().__init__() + self.persisted_values: list[str | None] = [] + + async def set_session_id(self, session_id: str | None) -> None: + self.persisted_values.append(session_id) + await super().set_session_id(session_id) + + session_info = CountingSessionInfo() + + class CountingFactory(SessionInfoFactory): + def create_session(self, mcp_server: Any) -> SessionInfo: + return session_info + + endpoint = LegacyMcpEndpoint(repeat_session_header=True) + async with configured_client( + mcp_resource_config, + mock_uipath_sdk, + endpoint, + session_info_factory=CountingFactory(), + ) as client: + await client.list_tools() + await client.call_tool("test_tool", {"query": "first"}) + await client.call_tool("test_tool", {"query": "second"}) + + assert session_info.persisted_values == ["session-1"] + + +@pytest.mark.asyncio +async def test_max_retries_exceeded_raises_mcp_error( + mcp_resource_config: AgentMcpResourceConfig, + mock_uipath_sdk: MagicMock, +) -> None: + """Repeated session termination is surfaced after the configured retry.""" + endpoint = LegacyMcpEndpoint(failed_tool_calls=2) + async with configured_client( + mcp_resource_config, mock_uipath_sdk, endpoint, max_retries=1 + ) as client: + with pytest.raises(MCPError) as exc_info: + await client.call_tool("test_tool", {"query": "test"}) + + assert exc_info.value.code == INVALID_REQUEST + assert endpoint.initialize_count == 2 + assert endpoint.tool_call_count == 2 + + +@pytest.mark.asyncio +async def test_concurrent_recovery_does_not_replace_a_new_session( + mcp_resource_config: AgentMcpResourceConfig, +) -> None: + """A late failure from an old session must not tear down its replacement.""" + client = McpClient(config=mcp_resource_config) + failed_session = MagicMock() + replacement_session = MagicMock() + client._client_initialized = True + client._session = replacement_session + client._session_info = SessionInfo("replacement-id") + open_connection = AsyncMock() + + with patch.object(client, "_open_connection", open_connection): + await client._reinitialize_session(failed_session) + + assert client._session is replacement_session + assert await client.get_session_id() == "replacement-id" + open_connection.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_recovery_continues_when_failed_connection_cleanup_raises( + mcp_resource_config: AgentMcpResourceConfig, +) -> None: + """Closing the failed stack cannot mask recovery of the MCP connection.""" + client = McpClient(config=mcp_resource_config) + failed_session = MagicMock() + failed_stack = MagicMock() + failed_stack.aclose = AsyncMock(side_effect=RuntimeError("cleanup failed")) + client._client_initialized = True + client._session = failed_session + client._connection_stack = failed_stack + client._session_info = SessionInfo("failed-session") + open_connection = AsyncMock() + + with patch.object(client, "_open_connection", open_connection): + await client._reinitialize_session( + failed_session, error=MCPError(INVALID_REQUEST, "Session terminated") ) - mock_http_client = self.create_mock_http_client(MockStreamResponse) - mock_async_client_class.return_value = mock_http_client - - session = McpClient(config=mcp_resource_config) - - with patch( - "uipath.platform.UiPath", - return_value=mock_uipath_sdk, - ): - # First use - await session.call_tool("test_tool", {"query": "first"}) - assert await session.get_session_id() == "test-session-first" - - # Close - await session.dispose() - assert await session.get_session_id() is None - - # Reuse - should create new client and session - # Note: mock returns "test-session-retry" for second initialize - await session.call_tool("test_tool", {"query": "second"}) - assert await session.get_session_id() == "test-session-retry" - assert session.is_client_initialized - - # HTTP client was created twice (once before dispose, once after) - assert mock_async_client_class.call_count == 2 - - await session.dispose() - - @pytest.mark.asyncio - async def test_raises_on_missing_mcp_url(self, mcp_resource_config): - """Test that ValueError is raised when MCP server has no URL configured.""" - mock_sdk = MagicMock() - mock_server = MagicMock() - mock_server.mcp_url = None # No URL configured - mock_sdk.mcp.retrieve_async = AsyncMock(return_value=mock_server) - mock_sdk._config = MagicMock() - mock_sdk._config.secret = "test-token" - - session = McpClient(config=mcp_resource_config) - - with patch( - "uipath.platform.UiPath", - return_value=mock_sdk, - ): - with pytest.raises(ValueError, match="has no URL configured"): - await session.call_tool("test_tool", {"query": "test"}) - - @pytest.mark.asyncio - @patch("httpx.AsyncClient") - async def test_custom_session_info_factory_is_used( - self, mock_async_client_class, mcp_resource_config, mock_uipath_sdk - ): - """Test that a custom SessionInfoFactory is called during initialization.""" - method_call_sequence: list[str] = [] - initialize_count = [0] - tool_call_count = [0] - - MockStreamResponse = self.create_mock_stream_response( - method_call_sequence, initialize_count, tool_call_count + failed_stack.aclose.assert_awaited_once() + open_connection.assert_awaited_once() + assert await client.get_session_id() is None + + +@pytest.mark.asyncio +async def test_concurrent_call_waits_for_recovery_initialization( + mcp_resource_config: AgentMcpResourceConfig, + mock_uipath_sdk: MagicMock, +) -> None: + """A caller cannot use a replacement session before its handshake finishes.""" + endpoint = LegacyMcpEndpoint( + failed_tool_calls=1, + block_initialize_on=2, + ) + + async with configured_client( + mcp_resource_config, mock_uipath_sdk, endpoint + ) as client: + recovery_call = asyncio.create_task( + client.call_tool("test_tool", {"query": "recover"}) ) - mock_http_client = self.create_mock_http_client(MockStreamResponse) - mock_async_client_class.return_value = mock_http_client - - custom_session_info = SessionInfo() - - class TrackingFactory(SessionInfoFactory): - called_with_server = None - - def create_session(self, mcp_server: Any) -> SessionInfo: - TrackingFactory.called_with_server = mcp_server - return custom_session_info - - factory = TrackingFactory() - session = McpClient( - config=mcp_resource_config, - session_info_factory=factory, - ) - - with patch( - "uipath.platform.UiPath", - return_value=mock_uipath_sdk, - ): - await session.call_tool("test_tool", {"query": "test"}) - - # Verify factory was called with the McpServer - assert TrackingFactory.called_with_server is not None - - # Verify our custom SessionInfo instance is used by McpClient - assert session._session_info is custom_session_info - assert await session.get_session_id() == "test-session-first" - - await session.dispose() - - @pytest.mark.asyncio - async def test_skips_initialize_when_session_info_has_id(self, mcp_resource_config): - """Test that _initialize_session skips session.initialize() when SessionInfo has an ID.""" - session = McpClient(config=mcp_resource_config) - - # Simulate already-initialized client with pre-existing session ID - session._session_info = SessionInfo(session_id="pre-existing-id") - session._session = MagicMock() - session._session.initialize = AsyncMock() - - await session._initialize_session() - - # initialize() should NOT be called because session_info already has an ID - session._session.initialize.assert_not_called() - - @pytest.mark.asyncio - async def test_reinitialize_clears_session_info_before_init( - self, mcp_resource_config - ): - """Test that _reinitialize_session clears session info then calls initialize.""" - session = McpClient(config=mcp_resource_config) - - # Simulate already-initialized client with a stale session ID - session._client_initialized = True - session._session_info = SessionInfo(session_id="stale-id") - session._session = MagicMock() - session._session.initialize = AsyncMock() - - await session._reinitialize_session() + await asyncio.wait_for(endpoint.initialize_blocked.wait(), timeout=2) - # Session info should have been cleared before re-initializing - # (set_session_id(None) was called, then _initialize_session ran) - session._session.initialize.assert_called_once() - - # After reinitialize, session_info.session_id is None because - # the mocked initialize() doesn't set a new one - assert await session.get_session_id() is None - - @pytest.mark.asyncio - @patch("httpx.AsyncClient") - async def test_list_tools_initializes_session_and_returns_result( - self, mock_async_client_class, mcp_resource_config, mock_uipath_sdk - ): - """Test that list_tools lazily initializes session and returns tools.""" - method_call_sequence: list[str] = [] - initialize_count = [0] - tool_call_count = [0] + concurrent_list = asyncio.create_task(client.list_tools(force_refresh=True)) + await asyncio.sleep(0) + assert not concurrent_list.done() - MockStreamResponse = self.create_mock_stream_response( - method_call_sequence, initialize_count, tool_call_count + endpoint.release_initialize.set() + call_result, list_result = await asyncio.gather( + recovery_call, + concurrent_list, ) - mock_http_client = self.create_mock_http_client(MockStreamResponse) - mock_async_client_class.return_value = mock_http_client + assert call_result.structured_content == {"result": "Success from test_tool"} + assert list_result.tools[0].name == "test_tool" - client = McpClient(config=mcp_resource_config) - assert not client.is_client_initialized +@pytest.mark.asyncio +async def test_later_call_recovers_after_replacement_initialization_failure( + mcp_resource_config: AgentMcpResourceConfig, + mock_uipath_sdk: MagicMock, +) -> None: + """A failed replacement does not strand the client without a session.""" + endpoint = LegacyMcpEndpoint( + failed_tool_calls=1, + fail_initialize_on={2}, + ) - with patch( - "uipath.platform.UiPath", - return_value=mock_uipath_sdk, - ): - result = await client.list_tools() + async with configured_client( + mcp_resource_config, mock_uipath_sdk, endpoint + ) as client: + with pytest.raises(MCPError, match="Replacement initialization failed"): + await client.call_tool("test_tool", {"query": "first"}) - # Session should have been initialized - assert initialize_count[0] == 1 assert client.is_client_initialized - - # Should return the tools from the mock server - assert result is not None - assert len(result.tools) == 1 - assert result.tools[0].name == "test_tool" - - # Verify protocol flow includes tools/list - assert "initialize" in method_call_sequence - assert "tools/list" in method_call_sequence - # tools/call should NOT have been called - assert "tools/call" not in method_call_sequence - - await client.dispose() - - @pytest.mark.asyncio - @patch("httpx.AsyncClient") - async def test_list_tools_caches_result_across_calls( - self, mock_async_client_class, mcp_resource_config, mock_uipath_sdk - ): - """list_tools caches its result: a second call reuses the session and the - cached tool list, issuing only one tools/list RPC.""" - method_call_sequence: list[str] = [] - initialize_count = [0] - tool_call_count = [0] - - MockStreamResponse = self.create_mock_stream_response( - method_call_sequence, initialize_count, tool_call_count - ) - - mock_http_client = self.create_mock_http_client(MockStreamResponse) - mock_async_client_class.return_value = mock_http_client - - client = McpClient(config=mcp_resource_config) - - with patch( - "uipath.platform.UiPath", - return_value=mock_uipath_sdk, - ): - first = await client.list_tools() - assert initialize_count[0] == 1 - - second = await client.list_tools() - assert initialize_count[0] == 1 # Still only one initialization - - # Fetched once per lifetime: second call returns the cached result, no new RPC. - assert method_call_sequence.count("tools/list") == 1 - assert first is second - - await client.dispose() - - @pytest.mark.asyncio - @patch("httpx.AsyncClient") - async def test_list_tools_force_refresh_bypasses_cache( - self, mock_async_client_class, mcp_resource_config, mock_uipath_sdk - ): - """force_refresh=True re-queries the server instead of returning the cache.""" - method_call_sequence: list[str] = [] - initialize_count = [0] - tool_call_count = [0] - - MockStreamResponse = self.create_mock_stream_response( - method_call_sequence, initialize_count, tool_call_count - ) - - mock_http_client = self.create_mock_http_client(MockStreamResponse) - mock_async_client_class.return_value = mock_http_client - - client = McpClient(config=mcp_resource_config) - - with patch( - "uipath.platform.UiPath", - return_value=mock_uipath_sdk, + assert client._session is None + + result = await client.call_tool("test_tool", {"query": "second"}) + + assert result.structured_content == {"result": "Success from test_tool"} + assert endpoint.initialize_count == 3 + + +@pytest.mark.asyncio +async def test_raises_on_missing_mcp_url( + mcp_resource_config: AgentMcpResourceConfig, +) -> None: + """A server registration without an endpoint fails before allocating HTTP state.""" + sdk = MagicMock() + server = MagicMock() + server.mcp_url = None + sdk.mcp.retrieve_async = AsyncMock(return_value=server) + + client = McpClient(config=mcp_resource_config) + with patch("uipath.platform.UiPath", return_value=sdk): + with pytest.raises(ValueError, match="has no URL configured"): + await client.call_tool("test_tool", {"query": "test"}) + + +@pytest.mark.asyncio +async def test_initialization_failure_cleans_state_and_allows_retry( + mcp_resource_config: AgentMcpResourceConfig, + mock_uipath_sdk: MagicMock, +) -> None: + """A failed handshake releases both stacks and leaves the client reusable.""" + endpoint = LegacyMcpEndpoint() + + async with configured_client( + mcp_resource_config, mock_uipath_sdk, endpoint + ) as client: + with patch.object( + client, + "_initialize_session", + AsyncMock(side_effect=RuntimeError("initialize failed")), ): - await client.list_tools() - await client.list_tools(force_refresh=True) - - # Session reused, but the server is queried twice. - assert initialize_count[0] == 1 - assert method_call_sequence.count("tools/list") == 2 - - await client.dispose() - - @pytest.mark.asyncio - @patch("httpx.AsyncClient") - async def test_dispose_clears_tools_cache( - self, mock_async_client_class, mcp_resource_config, mock_uipath_sdk - ): - """dispose() clears the cached tool list so a reused (or resumed) client - re-fetches it once.""" - method_call_sequence: list[str] = [] - initialize_count = [0] - tool_call_count = [0] - - MockStreamResponse = self.create_mock_stream_response( - method_call_sequence, initialize_count, tool_call_count - ) - - mock_http_client = self.create_mock_http_client(MockStreamResponse) - mock_async_client_class.return_value = mock_http_client - - client = McpClient(config=mcp_resource_config) - - with patch( - "uipath.platform.UiPath", - return_value=mock_uipath_sdk, - ): - await client.list_tools() - assert client._tools_cache is not None - - await client.dispose() - assert client._tools_cache is None - - @pytest.mark.asyncio - @patch.dict(os.environ, {"UIPATH_FOLDER_PATH": "/Shared/TestFolder"}) - @patch("httpx.AsyncClient") - async def test_retrieve_async_uses_name_and_execution_folder_path( - self, mock_async_client_class, mcp_resource_config - ): - """Test that name resolution receives both identities and the execution folder.""" - mock_sdk = MagicMock() - mock_server = MagicMock() - mock_server.mcp_url = "https://test.uipath.com/mcp" - mock_sdk.mcp.retrieve_async = AsyncMock(return_value=mock_server) - mock_sdk._config = MagicMock() - mock_sdk._config.secret = "test-secret-token" - - method_call_sequence: list[str] = [] - initialize_count = [0] - tool_call_count = [0] - - MockStreamResponse = self.create_mock_stream_response( - method_call_sequence, initialize_count, tool_call_count - ) - mock_http_client = self.create_mock_http_client(MockStreamResponse) - mock_async_client_class.return_value = mock_http_client - - session = McpClient(config=mcp_resource_config) - - with patch("uipath.platform.UiPath", return_value=mock_sdk): - await session.call_tool("test_tool", {"query": "test"}) - - mock_sdk.mcp.retrieve_async.assert_called_once_with( - name="test_server", - folder_path="/Shared/TestFolder", - ) + with pytest.raises(RuntimeError, match="initialize failed"): + await client.call_tool("test_tool", {"query": "first"}) + + assert client._stack is None + assert client._connection_stack is None + assert client._http_client is None + assert client._session_info is None + assert client._session is None + assert not client.is_client_initialized - await session.dispose() + result = await client.call_tool("test_tool", {"query": "second"}) + + assert result.structured_content == {"result": "Success from test_tool"} + assert endpoint.initialize_count == 1 + + +def test_only_session_specific_invalid_request_is_retryable() -> None: + """Ordinary INVALID_REQUEST errors must not be mislabeled as disconnects.""" + assert McpClient.is_session_error( + MCPError(code=CONNECTION_CLOSED, message="Connection closed") + ) + assert McpClient.is_session_error( + MCPError(code=INVALID_REQUEST, message="Session terminated") + ) + assert McpClient.is_session_error( + MCPError(code=INVALID_REQUEST, message="Session not found") + ) + assert not McpClient.is_session_error( + MCPError(code=INVALID_REQUEST, message="Invalid request parameters") + ) diff --git a/tests/agent/tools/test_mcp/test_mcp_client_real_http.py b/tests/agent/tools/test_mcp/test_mcp_client_real_http.py new file mode 100644 index 000000000..66fafbfbd --- /dev/null +++ b/tests/agent/tools/test_mcp/test_mcp_client_real_http.py @@ -0,0 +1,666 @@ +"""``McpClient`` driven against real MCP servers over real HTTP. + +Everywhere else the client is exercised through ``httpx2.MockTransport``, which +proves the wire is shaped right but never runs a real server, a real ASGI stack, +or the gateway hop UiPath puts in front of an MCP endpoint. These tests host a +genuine ``MCPServer`` (and hand-written endpoints where a version has to be +pinned) on an ephemeral port and drive the public ``McpClient`` API against it. + +Every assertion is made from the *gateway's* point of view -- what actually +arrived on the wire -- rather than from client internals, so a test fails when +the observable behaviour changes, not when the implementation moves. + +See ``real_server.py`` for the harness. +""" + +from typing import Any + +import pytest +from mcp.shared.exceptions import MCPError + +from uipath_langchain.agent.tools.mcp import SessionInfo, SessionInfoFactory + +from .real_server import ( + HANDSHAKE_VERSIONS, + LEGACY_VERSION, + MODERN_VERSION, + SDK_LOOKUPS, + PinnedVersionServer, + RecordingGateway, + build_sdk_app, + connected_client, + make_client, + make_resource_config, + patched_sdk, + pinned_session_factory, + serve, +) + + +def negotiated_version(gateway: RecordingGateway, rpc_method: str) -> str | None: + """Read the version the SDK stamped on the first request of one method. + + The ``mcp-protocol-version`` header is written by the SDK once negotiation + completes, so it is the wire-visible answer to "what did we settle on". + """ + records = gateway.for_rpc(rpc_method) + return records[0].protocol_version if records else None + + +# --- negotiation per mode --------------------------------------------------- + + +@pytest.mark.asyncio +async def test_legacy_mode_negotiates_the_newest_handshake_version() -> None: + """A real SDK server answers the handshake at 2025-11-25 and mints a session.""" + gateway = RecordingGateway(build_sdk_app()) + async with serve(gateway) as url: + async with connected_client(url, protocol_mode="legacy") as client: + tools = await client.list_tools() + result = await client.call_tool("add", {"a": 2, "b": 3}) + session_id = await client.get_session_id() + + assert sorted(tool.name for tool in tools.tools) == ["add", "multiply"] + assert result.structured_content == {"result": 5} + assert gateway.count("initialize") == 1 + assert gateway.count("server/discover") == 0 + assert negotiated_version(gateway, "tools/call") == LEGACY_VERSION + # The server owns the identity in this era, and it reaches the client on a + # response header rather than being invented locally. + assert session_id is not None + assert gateway.server_session_ids() + assert gateway.for_rpc("tools/call")[0].session_id == session_id + + +@pytest.mark.asyncio +async def test_modern_mode_negotiates_without_any_server_session() -> None: + """2026-07-28 has no session identity, so no response may carry one.""" + gateway = RecordingGateway(build_sdk_app()) + async with serve(gateway) as url: + async with connected_client(url, protocol_mode="modern") as client: + tools = await client.list_tools() + result = await client.call_tool("add", {"a": 2, "b": 3}) + + assert sorted(tool.name for tool in tools.tools) == ["add", "multiply"] + assert result.structured_content == {"result": 5} + assert gateway.count("server/discover") == 1 + assert gateway.count("initialize") == 0 + assert negotiated_version(gateway, "tools/call") == MODERN_VERSION + assert gateway.server_session_ids() == [], ( + "A modern server issued a session ID; the era has no session identity " + "and the client-minted affinity ID must be the only one in play" + ) + + +@pytest.mark.asyncio +async def test_auto_mode_resolves_to_modern_against_a_real_server() -> None: + """A server answering discovery is driven as modern, handshake untouched.""" + gateway = RecordingGateway(build_sdk_app()) + async with serve(gateway) as url: + async with connected_client(url, protocol_mode="auto") as client: + result = await client.call_tool("add", {"a": 4, "b": 5}) + + assert result.structured_content == {"result": 9} + assert gateway.count("server/discover") == 1 + assert gateway.count("initialize") == 0 + assert negotiated_version(gateway, "tools/call") == MODERN_VERSION + assert gateway.server_session_ids() == [] + + +@pytest.mark.asyncio +async def test_auto_mode_falls_back_to_legacy_against_a_handshake_only_server() -> None: + """No discovery endpoint means the probe must fall back, honouring the offer.""" + server = PinnedVersionServer("2025-06-18") + gateway = RecordingGateway(server.build_app()) + async with serve(gateway) as url: + async with connected_client(url, protocol_mode="auto") as client: + result = await client.call_tool("add", {"a": 4, "b": 5}) + session_id = await client.get_session_id() + + assert result.structured_content == {"result": 9} + assert gateway.count("server/discover") == 1 + assert server.discover_count == 0 + assert server.initialize_count == 1 + assert session_id == "session-1" + assert negotiated_version(gateway, "tools/call") == "2025-06-18" + # The probe is pinned like every first request, but the ID it carried was + # minted here and never named a session on this server. The handshake must + # not present it, or a server that routes by the header would refuse it. + assert gateway.for_rpc("server/discover")[0].session_id is not None + assert gateway.for_rpc("initialize")[0].session_id is None + + +@pytest.mark.asyncio +async def test_modern_mode_works_against_a_server_that_refuses_the_handshake() -> None: + """A discover-only server proves modern is not silently falling back. + + This endpoint answers ``server/discover`` and rejects ``initialize`` + outright, so a client that quietly degraded to the legacy handshake could + not complete a single call here. + """ + server = PinnedVersionServer(MODERN_VERSION, modern_only=True) + gateway = RecordingGateway(server.build_app()) + async with serve(gateway) as url: + async with connected_client(url, protocol_mode="modern") as client: + result = await client.call_tool("add", {"a": 2, "b": 3}) + + assert result.structured_content == {"result": 5} + assert server.discover_count == 1 + assert server.initialize_count == 0 + assert gateway.count("initialize") == 0 + assert negotiated_version(gateway, "tools/call") == MODERN_VERSION + + +# --- resume across clients -------------------------------------------------- + + +@pytest.mark.asyncio +async def test_legacy_resume_keeps_the_originally_negotiated_version() -> None: + """A resumed session must keep the version it was negotiated at. + + Two clients share one ``SessionInfo``, standing in for two runs of a + playground agent whose session store outlives the process. The session ID + surviving is only half the contract: the resumed connection must also speak + the version that session was negotiated at. Probing candidate versions + instead -- the pre-existing approach -- always matched the *oldest* + handshake version, silently downgrading every later request. + """ + gateway = RecordingGateway(build_sdk_app()) + shared = SessionInfo() + factory = pinned_session_factory(shared) + + async with serve(gateway) as url: + with patched_sdk(url): + first = make_client(session_info_factory=factory, terminate_on_close=False) + await first.call_tool("add", {"a": 1, "b": 1}) + original_session_id = await first.get_session_id() + await first.dispose() + + resume_boundary = len(gateway.records) + + second = make_client(session_info_factory=factory, terminate_on_close=False) + result = await second.call_tool("add", {"a": 2, "b": 2}) + resumed_session_id = await second.get_session_id() + await second.dispose() + + assert result.structured_content == {"result": 4} + assert original_session_id is not None + assert resumed_session_id == original_session_id + + after_resume = gateway.records[resume_boundary:] + assert after_resume, "The second client sent no requests" + assert all(r.session_id == original_session_id for r in after_resume) + + # Derived from the recording, not asserted against a constant: the contract + # is "the same version as before", so reading it back from the pre-resume + # traffic keeps the guard honest if the server's default ever changes. + before_resume = { + record.protocol_version + for record in gateway.records[:resume_boundary] + if record.protocol_version is not None + } + assert len(before_resume) == 1, ( + f"The first client itself spoke {sorted(before_resume)}; the fixture no " + "longer establishes a single negotiated version to compare against." + ) + versions = { + record.protocol_version + for record in after_resume + if record.protocol_version is not None + } + assert versions == before_resume, ( + f"Requests after resume negotiated {sorted(versions)} instead of " + f"{sorted(before_resume)}. This is the silent-downgrade regression " + "guard: a resumed session that guesses its version settles on the " + "oldest handshake version and downgrades every later request." + ) + resumed_calls = [r for r in after_resume if r.rpc_method == "tools/call"] + assert resumed_calls and resumed_calls[0].protocol_version == LEGACY_VERSION + # And it costs nothing to learn: the version was stored with the ID, so the + # resumed connection adopts it instead of handshaking a second time. + assert gateway.count("initialize") == 1 + assert not [r for r in after_resume if r.rpc_method == "initialize"] + + +@pytest.mark.asyncio +async def test_legacy_resume_survives_a_server_that_refuses_reinitialization() -> None: + """A resumed session must not depend on being allowed to re-handshake. + + The reference TypeScript implementation answers a second ``initialize`` on a + live session with "Server already initialized". A client that resumes by + re-running the handshake loses the persisted session on every run against + such a server -- and with it the gateway affinity the session ID provides. + Adopting the stored version keeps the resume free of wire traffic, so the + server is never asked. + """ + server = PinnedVersionServer("2025-06-18", refuse_reinitialize=True) + gateway = RecordingGateway(server.build_app()) + shared = SessionInfo() + factory = pinned_session_factory(shared) + + async with serve(gateway) as url: + with patched_sdk(url): + first = make_client(session_info_factory=factory, terminate_on_close=False) + await first.call_tool("add", {"a": 1, "b": 1}) + original_session_id = await first.get_session_id() + await first.dispose() + + resume_boundary = len(gateway.records) + + second = make_client(session_info_factory=factory, terminate_on_close=False) + result = await second.call_tool("add", {"a": 2, "b": 2}) + resumed_session_id = await second.get_session_id() + await second.dispose() + + assert result.structured_content == {"result": 4} + assert resumed_session_id == original_session_id == "session-1" + # The handshake was never re-sent, so the server never had to refuse it. + assert server.initialize_count == 1 + assert server.refused_reinitialize_count == 0 + after_resume = gateway.records[resume_boundary:] + assert after_resume and all( + record.session_id == original_session_id for record in after_resume + ) + assert [ + r.protocol_version for r in after_resume if r.rpc_method == "tools/call" + ] == ["2025-06-18"] + + +@pytest.mark.asyncio +async def test_unknown_persisted_session_falls_back_to_a_fresh_session() -> None: + """A stale stored ID is rejected by the server, and the client starts clean.""" + gateway = RecordingGateway(build_sdk_app()) + stored = SessionInfo("never-existed") + + async with serve(gateway) as url: + async with connected_client( + url, session_info_factory=pinned_session_factory(stored) + ) as client: + result = await client.call_tool("add", {"a": 6, "b": 1}) + session_id = await client.get_session_id() + + assert result.structured_content == {"result": 7} + # The refused handshake for the stale ID, then the clean one -- on the same + # transport, because a refused request does not close the connection. + assert gateway.count("initialize") == 2 + handshakes = gateway.for_rpc("initialize") + assert [h.session_id for h in handshakes] == ["never-existed", None] + assert session_id is not None and session_id != "never-existed" + assert gateway.for_rpc("tools/call")[0].session_id == session_id + + +# --- affinity and disposal -------------------------------------------------- + + +@pytest.mark.asyncio +async def test_modern_affinity_pins_one_instance_across_clients() -> None: + """The affinity ID is what replaces gateway routing once sessions are gone. + + ``2026-07-28`` removes ``mcp-session-id`` from the protocol, so UiPath mints + the value itself and keeps sending it on that header purely as a routing + key. Because it is minted *before* negotiating, even ``server/discover`` + is routable -- which a server-assigned session ID never could be. + """ + gateway = RecordingGateway(build_sdk_app()) + shared = SessionInfo() + factory = pinned_session_factory(shared) + + async with serve(gateway) as url: + with patched_sdk(url): + for operands in ((1, 2), (3, 4)): + client = make_client( + session_info_factory=factory, protocol_mode="modern" + ) + try: + await client.call_tool("add", {"a": operands[0], "b": operands[1]}) + finally: + await client.dispose() + + affinity_id = await shared.get_session_id() + assert affinity_id + assert gateway.count("server/discover") == 2 + assert [r.session_id for r in gateway.for_rpc("server/discover")] == [ + affinity_id, + affinity_id, + ] + assert gateway.unpinned() == [], ( + "Requests reached the gateway with no affinity header, so it would have " + "had to route them blind" + ) + assert sorted({record.instance for record in gateway.records}) == ["instance-1"] + # A minted ID is as vulnerable as a restored one: if disposal tore it down, + # the gateway would see a teardown for the instance it is meant to pin. + assert gateway.http_count("DELETE") == 0 + + +@pytest.mark.asyncio +async def test_auto_mode_pins_the_first_request() -> None: + """In ``auto`` the probe must be pinned too, not only the calls after it. + + A serverless gateway routes on the header. An unpinned ``server/discover`` + warms one instance and the first tool call then lands on another -- the + exact scatter the affinity ID exists to prevent. + """ + gateway = RecordingGateway(build_sdk_app()) + shared = SessionInfo() + + async with serve(gateway) as url: + async with connected_client( + url, + protocol_mode="auto", + session_info_factory=pinned_session_factory(shared), + ) as client: + await client.call_tool("add", {"a": 1, "b": 2}) + + affinity_id = await shared.get_session_id() + assert affinity_id + assert gateway.rpc_methods()[0] == "server/discover" + assert gateway.for_rpc("server/discover")[0].session_id == affinity_id + assert gateway.unpinned() == [] + assert sorted({record.instance for record in gateway.records}) == ["instance-1"] + + +@pytest.mark.asyncio +async def test_modern_disposal_does_not_delete_a_restored_affinity_id() -> None: + """A client-minted routing key must never be torn down as if it were a session. + + A restored affinity ID looks exactly like a restored session to the + transport. Deleting it would reach the gateway as a teardown for a live + instance on every run after the first -- precisely the playground case. + """ + gateway = RecordingGateway(build_sdk_app()) + restored = SessionInfo("restored-affinity") + + async with serve(gateway) as url: + async with connected_client( + url, + protocol_mode="modern", + session_info_factory=pinned_session_factory(restored), + terminate_on_close=True, + ) as client: + await client.call_tool("add", {"a": 2, "b": 2}) + + assert gateway.http_count("DELETE") == 0 + # A restored ID is reused rather than replaced, from the very first request. + assert gateway.for_rpc("server/discover")[0].session_id == "restored-affinity" + assert all(record.session_id == "restored-affinity" for record in gateway.records) + # The ID survives disposal, so the next run returns to the same instance. + assert await restored.get_session_id() == "restored-affinity" + + +@pytest.mark.asyncio +async def test_legacy_disposal_deletes_a_restored_session() -> None: + """A restored *server* session is real state, so disposal must terminate it.""" + gateway = RecordingGateway(build_sdk_app()) + shared = SessionInfo() + factory = pinned_session_factory(shared) + + async with serve(gateway) as url: + with patched_sdk(url): + first = make_client(session_info_factory=factory, terminate_on_close=False) + await first.call_tool("add", {"a": 1, "b": 1}) + session_id = await first.get_session_id() + await first.dispose() + + assert gateway.http_count("DELETE") == 0 + resume_boundary = len(gateway.records) + + second = make_client(session_info_factory=factory, terminate_on_close=True) + await second.call_tool("add", {"a": 2, "b": 2}) + await second.dispose() + + deletes = [ + record + for record in gateway.records[resume_boundary:] + if record.http_method == "DELETE" + ] + assert len(deletes) == 1 + assert deletes[0].session_id == session_id + + +# --- retry semantics per era ------------------------------------------------ + + +@pytest.mark.asyncio +async def test_legacy_recovers_from_an_injected_session_termination() -> None: + """A lost legacy session is re-established and the call retried.""" + gateway = RecordingGateway(build_sdk_app(), fault_on_tool_call=1) + async with serve(gateway) as url: + async with connected_client(url, protocol_mode="legacy") as client: + result = await client.call_tool("add", {"a": 3, "b": 4}) + + assert result.structured_content == {"result": 7} + assert gateway.count("initialize") == 2 + assert gateway.count("tools/call") == 2 + calls = gateway.for_rpc("tools/call") + assert calls[0].faulted and not calls[1].faulted + # A fresh handshake means a fresh session, so the retry cannot reuse the + # session the gateway just declared dead. + assert calls[0].session_id != calls[1].session_id + + +@pytest.mark.asyncio +async def test_modern_does_not_retry_an_injected_session_termination() -> None: + """Reconnecting cannot restore state a self-contained request never had. + + The identical response is retried once in legacy mode. Here it must surface + immediately instead of spending the retry budget on something a reconnect + cannot fix. + """ + gateway = RecordingGateway(build_sdk_app(), fault_on_tool_call=1) + async with serve(gateway) as url: + async with connected_client(url, protocol_mode="modern") as client: + with pytest.raises(MCPError): + await client.call_tool("add", {"a": 3, "b": 4}) + + assert gateway.count("tools/call") == 1 + assert gateway.count("server/discover") == 1 + assert gateway.count("initialize") == 0 + + +# --- version breadth -------------------------------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize("protocol_version", HANDSHAKE_VERSIONS) +async def test_legacy_negotiates_every_supported_handshake_version( + protocol_version: str, +) -> None: + """Every handshake version SDK 2 still accepts works through ``McpClient``. + + ``2024-11-05`` and ``2025-03-26`` are covered nowhere else. All four + negotiate identically here -- the server's counter-offer is honoured and + stamped onto every later request -- so no version needs a carve-out. + """ + server = PinnedVersionServer(protocol_version) + gateway = RecordingGateway(server.build_app()) + async with serve(gateway) as url: + async with connected_client(url, protocol_mode="legacy") as client: + result = await client.call_tool("add", {"a": 2, "b": 3}) + session_id = await client.get_session_id() + + assert result.structured_content == {"result": 5} + assert server.initialize_count == 1 + assert session_id == "session-1" + assert negotiated_version(gateway, "tools/call") == protocol_version + assert gateway.for_rpc("tools/call")[0].session_id == "session-1" + + +# --- lifecycle -------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_list_tools_is_cached_until_force_refresh() -> None: + """Discovery is fetched once per client and re-queried only on demand.""" + gateway = RecordingGateway(build_sdk_app()) + async with serve(gateway) as url: + async with connected_client(url) as client: + first = await client.list_tools() + second = await client.list_tools() + assert gateway.count("tools/list") == 1 + + refreshed = await client.list_tools(force_refresh=True) + + assert first is second + assert sorted(tool.name for tool in refreshed.tools) == ["add", "multiply"] + assert gateway.count("tools/list") == 2 + + +@pytest.mark.asyncio +async def test_dispose_then_reuse_reinitializes_the_client() -> None: + """Disposal releases everything, and the next call rebuilds a working client.""" + gateway = RecordingGateway(build_sdk_app()) + async with serve(gateway) as url: + async with connected_client(url) as client: + await client.call_tool("add", {"a": 1, "b": 1}) + await client.dispose() + + assert not client.is_client_initialized + assert await client.get_session_id() is None + + result = await client.call_tool("add", {"a": 5, "b": 5}) + + assert result.structured_content == {"result": 10} + assert client.is_client_initialized + + assert gateway.count("initialize") == 2 + first_call, second_call = gateway.for_rpc("tools/call") + assert first_call.session_id != second_call.session_id + + +@pytest.mark.asyncio +async def test_tool_built_by_the_factory_invokes_over_real_http() -> None: + """Drive the whole seam: factory -> LangChain tool -> McpClient -> the wire. + + Every other tool test substitutes ``MagicMock(spec=McpClient)``, so nothing + in pytest connected the factory to a real server. That gap let a silent + serialization regression ship: ``_normalize_tool_result`` kept using a plain + ``model_dump()`` after SDK 2.0 renamed the model attributes, rewriting + ``mimeType`` to ``mime_type`` for every non-text block handed to the model. + """ + from uipath_langchain.agent.tools.mcp import create_mcp_tools + + gateway = RecordingGateway(build_sdk_app()) + async with serve(gateway) as url: + async with connected_client(url) as client: + tools = await create_mcp_tools(make_resource_config(), client) + add_tool = next(tool for tool in tools if tool.name == "add") + result = await add_tool.ainvoke({"a": 2, "b": 3}) + + blocks = result if isinstance(result, list) else [result] + assert [block.get("text") for block in blocks] == ["5"] + assert gateway.count("tools/call") == 1 + + +@pytest.mark.asyncio +async def test_factory_tool_hands_non_text_blocks_over_in_wire_shape() -> None: + """A non-text block reaches the model camelCased, through the factory path. + + Text blocks serialize identically under either dump mode, so only a + non-text block can catch a snake_case regression. Driving it through + ``build_mcp_tool`` covers the serializer the factory actually installs. + """ + from mcp.types import CallToolResult, ImageContent + + from uipath_langchain.agent.tools.mcp.mcp_tool import _normalize_tool_result + + normalized = _normalize_tool_result( + CallToolResult( + content=[ImageContent(type="image", data="Zm9v", mimeType="image/png")] + ) + ) + + assert normalized == [{"type": "image", "data": "Zm9v", "mimeType": "image/png"}], ( + "non-text blocks must keep their wire spelling for the model" + ) + + +@pytest.mark.asyncio +async def test_dispose_clears_the_tool_cache() -> None: + """A resumed run must not serve a tool list captured before disposal. + + ``dispose()`` clearing ``_tools_cache`` had no assertion anywhere, so a + stale list surviving dispose/reuse would have passed the suite. + """ + gateway = RecordingGateway(build_sdk_app()) + async with serve(gateway) as url: + with patched_sdk(url): + client = make_client() + try: + await client.list_tools() + await client.list_tools() + assert gateway.count("tools/list") == 1, "cache did not hold" + await client.dispose() + await client.list_tools() + assert gateway.count("tools/list") == 2, "cache survived dispose" + finally: + await client.dispose() + + +@pytest.mark.asyncio +async def test_session_info_factory_receives_the_resolved_mcp_server() -> None: + """The factory is handed the resolved server, not just a URL. + + ``SessionInfoDebugStateFactory`` downstream keys its debug-state path on the + server's ``slug``, so losing that argument would break persistence there + while every test here still passed. + """ + seen: list[Any] = [] + + class _RecordingFactory(SessionInfoFactory): + def create_session(self, mcp_server: Any) -> SessionInfo: + seen.append(mcp_server) + return SessionInfo() + + async with serve(build_sdk_app()) as url: + async with connected_client( + url, session_info_factory=_RecordingFactory() + ) as client: + await client.list_tools() + + assert len(seen) == 1 + assert seen[0].slug + assert seen[0].mcp_url + + +@pytest.mark.asyncio +async def test_recovery_reuses_the_same_http_client() -> None: + """Replacing a lost session must not rebuild the authenticated HTTP client. + + Recovery reuses the client precisely so a reconnect costs no TLS handshake + or token resolution. That was previously only a docstring claim. + """ + gateway = RecordingGateway(build_sdk_app(), fault_on_tool_call=1) + async with serve(gateway) as url: + async with connected_client(url, protocol_mode="legacy") as client: + await client.list_tools() + first_http_client = client._http_client + await client.call_tool("add", {"a": 2, "b": 3}) + + assert gateway.count("initialize") == 2, "no session replacement happened" + assert client._http_client is first_http_client + + +@pytest.mark.asyncio +async def test_server_is_resolved_by_name_and_execution_folder( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The lazy SDK lookup must pass the display name and execution folder. + + `uipath debug` applies resource bindings after the graph is built, which is + why this lookup is deferred to the first call. Dropping either argument would + resolve the wrong server -- or the right one in the wrong folder -- while + every wire-level assertion here still passed. The folder comes from + ``UIPATH_FOLDER_PATH``, which the runtime sets per job. + """ + monkeypatch.setenv("UIPATH_FOLDER_PATH", "/Shared/SomeExecutionFolder") + + async with serve(build_sdk_app()) as url: + async with connected_client(url) as client: + await client.list_tools() + + assert len(SDK_LOOKUPS) == 1 + assert SDK_LOOKUPS[0] == { + "name": make_resource_config().name, + "folder_path": "/Shared/SomeExecutionFolder", + } diff --git a/tests/agent/tools/test_mcp/test_mcp_tool.py b/tests/agent/tools/test_mcp/test_mcp_tool.py index b6a18fd82..eaf9ca979 100644 --- a/tests/agent/tools/test_mcp/test_mcp_tool.py +++ b/tests/agent/tools/test_mcp/test_mcp_tool.py @@ -1,14 +1,13 @@ """Tests for mcp_tool.py metadata and functionality.""" -import json import logging -from typing import Any, cast +from typing import cast from unittest.mock import AsyncMock, MagicMock, patch import pytest from langchain_core.tools import BaseTool -from mcp.shared.exceptions import McpError -from mcp.types import ErrorData, ListToolsResult, Tool +from mcp.shared.exceptions import MCPError +from mcp.types import ListToolsResult, Tool from uipath.agent.models.agent import ( AgentMcpResourceConfig, AgentMcpTool, @@ -35,8 +34,6 @@ StructuredToolWithArgumentProperties, ) -logger = logging.getLogger(__name__) - class TestMcpToolMetadata: """Test that MCP tool has correct metadata for observability.""" @@ -330,284 +327,6 @@ async def test_tools_have_correct_metadata(self, mcp_resources): assert "slug" in tool.metadata -class TestMcpToolInvocation: - """Test MCP tool invocation with mocked HTTP. - - This class tests the full flow of tool invocation without mocking the MCP SDK. - Only httpx.AsyncClient is mocked, allowing the real MCP SDK to process messages. - """ - - @pytest.fixture - def mock_uipath_sdk(self): - """Create a mock UiPath SDK for patching.""" - mock_sdk = MagicMock() - mock_server = MagicMock() - mock_server.mcp_url = "https://test.uipath.com/mcp" - mock_sdk.mcp.retrieve_async = AsyncMock(return_value=mock_server) - mock_sdk._config = MagicMock() - mock_sdk._config.secret = "test-secret-token" - return mock_sdk - - def create_mock_stream_response( - self, - method_call_sequence: list[str], - initialize_count: list[int], - tool_call_count: list[int], - session_guid: str = "test-session-12345", - ): - """Create a MockStreamResponse class for testing. - - Reuses the same pattern as test_mcp_client.py. - """ - - class MockStreamResponse: - """Mock HTTP stream response for MCP protocol.""" - - def __init__(self, method: str, url: str, **kwargs: Any): - self.request_method = method - self.url = url - self.kwargs = kwargs - - if method == "GET": - self.status_code = 405 - self.headers = {} - self._content = b"" - return - - json_body = kwargs.get("json", {}) - self.json_body = json_body - self.method = json_body.get("method", "") - - logger.debug(f"Responding to MCP method: {self.method}") - method_call_sequence.append(self.method) - - status_code, response_json, headers = self._build_response() - self.headers = headers or {} - self._response_json = response_json - self.status_code = status_code - - if response_json: - self._content = json.dumps(self._response_json).encode("utf-8") - self.headers["content-type"] = "application/json" - else: - self._content = b"" - - def _build_response(self) -> tuple[int, Any, dict[str, str] | None]: - """Build JSON-RPC response based on method.""" - request_id = self.json_body.get("id") - - if self.method == "initialize": - initialize_count[0] += 1 - return ( - 200, - { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "protocolVersion": "2025-06-18", - "capabilities": {"tools": {}}, - "serverInfo": { - "name": "test-server", - "version": "1.0.0", - }, - }, - }, - {"mcp-session-id": session_guid}, - ) - - elif self.method == "notifications/initialized": - return (204, None, {}) - - elif self.method == "tools/list": - return ( - 200, - { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "tools": [ - { - "name": "search_tool", - "description": "Search for information", - "inputSchema": { - "type": "object", - "properties": {"query": {"type": "string"}}, - "required": ["query"], - }, - "outputSchema": { - "type": "object", - "properties": { - "result": {"type": "string"} - }, - }, - } - ], - }, - }, - {}, - ) - - elif self.method == "tools/call": - tool_call_count[0] += 1 - params = self.json_body.get("params", {}) - tool_name = params.get("name", "unknown") - structured_result = {"result": f"Success from {tool_name}"} - - return ( - 200, - { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "content": [ - { - "type": "text", - "text": json.dumps(structured_result), - } - ], - "structuredContent": structured_result, - "isError": False, - }, - }, - {}, - ) - - else: - if request_id is None: - return (204, None, {}) - return ( - 500, - { - "jsonrpc": "2.0", - "id": request_id, - "error": {"code": -32601, "message": "Method not found"}, - }, - {}, - ) - - async def __aenter__(self): - return self - - async def __aexit__(self, *args: Any, **kwargs: Any): - pass - - async def aread(self) -> bytes: - """Return the response content.""" - return self._content - - def raise_for_status(self) -> None: - """Check response status.""" - if self.status_code >= 400: - raise Exception(f"HTTP {self.status_code}") - - return MockStreamResponse - - def create_mock_http_client(self, mock_stream_response_class: type) -> MagicMock: - """Create a mock HTTP client that uses the given stream response class.""" - mock_client = MagicMock() - mock_client.stream = lambda method, url, **kwargs: mock_stream_response_class( - method, url, **kwargs - ) - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock() - # Mock the delete method for session termination (returns 204 No Content) - mock_delete_response = MagicMock() - mock_delete_response.status_code = 204 - mock_client.delete = AsyncMock(return_value=mock_delete_response) - return mock_client - - @pytest.mark.asyncio - @patch("httpx.AsyncClient") - async def test_tool_invocation_initializes_session_and_returns_result( - self, - mock_async_client_class, - mock_uipath_sdk, - ): - """Smoke test: verify tool invocation initializes MCP session and returns result. - - This test verifies the full integration between create_mcp_tools_from_metadata - and McpClient without mocking any MCP SDK components. - - Expected behavior: - - Session is initialized via MCP protocol (initialize + initialized notification) - - Tool call is sent and result is returned - - Only httpx.AsyncClient is mocked, real MCP SDK processes the messages - """ - # Track MCP method calls - method_call_sequence: list[str] = [] - initialize_count = [0] - tool_call_count = [0] - - # Setup HTTP mock using pattern from test_mcp_client.py - MockStreamResponse = self.create_mock_stream_response( - method_call_sequence, initialize_count, tool_call_count - ) - mock_http_client = self.create_mock_http_client(MockStreamResponse) - mock_async_client_class.return_value = mock_http_client - - # Create resource config - mcp_resource = AgentMcpResourceConfig( - name="test_server", - description="Test server", - folder_path="/Shared/TestFolder", - slug="test-server", - available_tools=[ - AgentMcpTool( - name="search_tool", - description="Search for information", - input_schema={ - "type": "object", - "properties": {"query": {"type": "string"}}, - "required": ["query"], - }, - output_schema={ - "type": "object", - "properties": {"result": {"type": "string"}}, - }, - ) - ], - ) - - # Create McpClient and tools (SDK is called lazily on first tool call) - mcp_client = McpClient(config=mcp_resource) - tools = await create_mcp_tools(mcp_resource, mcp_client) - assert len(tools) == 1 - - tool = tools[0] - assert tool.name == "search_tool" - - # Invoke tool (SDK is called here during initialization) - with patch( - "uipath.platform.UiPath", - return_value=mock_uipath_sdk, - ): - result = await tool.ainvoke({"query": "test query"}) - - # Verify session was initialized - assert initialize_count[0] == 1, ( - f"Expected 1 initialize call, got {initialize_count[0]}" - ) - - # Verify tool was called - assert tool_call_count[0] == 1, ( - f"Expected 1 tool call, got {tool_call_count[0]}" - ) - - # Verify result is returned (content attribute of CallToolResult) - # Result is a list of dicts (model_dump'd TextContent objects) - assert result is not None - assert len(result) == 1 - assert result[0]["type"] == "text" - assert "Success from search_tool" in result[0]["text"] - - # Verify MCP protocol flow - assert "initialize" in method_call_sequence - assert "notifications/initialized" in method_call_sequence - assert "tools/call" in method_call_sequence - - logger.info(f"Method sequence: {method_call_sequence}") - - class TestMcpToolResultSerialization: """Test that tool_fn properly serializes different result types.""" @@ -620,15 +339,22 @@ def mcp_tool(self): ) @pytest.mark.asyncio - async def test_single_object_with_model_dump(self, mcp_tool): - """Test that a single result object with model_dump is serialized.""" - from uipath_langchain.agent.tools.mcp.mcp_tool import build_mcp_tool + async def test_single_object_is_serialized_in_its_wire_shape(self, mcp_tool): + """A lone content block keeps the camelCase spelling the model expects. - mock_content = MagicMock() - mock_content.model_dump.return_value = {"type": "text", "text": "hello"} + Uses a real ``ImageContent`` rather than a mock: SDK 2.0 renamed the + model attributes to snake case while keeping the camelCase names as + serialization aliases, so only a real model exposes the difference. A + mock asserting the call arguments would lock whichever call was written. + """ + from mcp.types import ImageContent + + from uipath_langchain.agent.tools.mcp.mcp_tool import build_mcp_tool mock_result = MagicMock() - mock_result.content = mock_content + mock_result.content = ImageContent( + type="image", data="Zm9v", mimeType="image/png" + ) mock_client = MagicMock(spec=McpClient) mock_client.call_tool = AsyncMock(return_value=mock_result) @@ -636,19 +362,35 @@ async def test_single_object_with_model_dump(self, mcp_tool): tool_fn = build_mcp_tool(mcp_tool, mock_client) result = await tool_fn() - assert result == {"type": "text", "text": "hello"} - mock_content.model_dump.assert_called_once_with(exclude_none=True) + assert result == {"type": "image", "data": "Zm9v", "mimeType": "image/png"} @pytest.mark.asyncio - async def test_list_of_objects_with_model_dump(self, mcp_tool): - """Test that a list of result objects with model_dump are serialized.""" - from uipath_langchain.agent.tools.mcp.mcp_tool import build_mcp_tool + async def test_list_of_blocks_keeps_camel_case_for_non_text_blocks(self, mcp_tool): + """Every block in a list is serialized by alias, not just the first. + + Text blocks are byte-identical either way, so a text-only assertion + cannot catch a snake_case regression -- the list mixes both kinds. + """ + from mcp.types import ( + EmbeddedResource, + ImageContent, + TextContent, + TextResourceContents, + ) - mock_item = MagicMock() - mock_item.model_dump.return_value = {"type": "text", "text": "item1"} + from uipath_langchain.agent.tools.mcp.mcp_tool import build_mcp_tool mock_result = MagicMock() - mock_result.content = [mock_item] + mock_result.content = [ + TextContent(type="text", text="item1"), + ImageContent(type="image", data="Zm9v", mimeType="image/png"), + EmbeddedResource( + type="resource", + resource=TextResourceContents( + uri="file:///x.txt", mimeType="text/plain", text="hi" + ), + ), + ] mock_client = MagicMock(spec=McpClient) mock_client.call_tool = AsyncMock(return_value=mock_result) @@ -656,7 +398,18 @@ async def test_list_of_objects_with_model_dump(self, mcp_tool): tool_fn = build_mcp_tool(mcp_tool, mock_client) result = await tool_fn() - assert result == [{"type": "text", "text": "item1"}] + assert result == [ + {"type": "text", "text": "item1"}, + {"type": "image", "data": "Zm9v", "mimeType": "image/png"}, + { + "type": "resource", + "resource": { + "uri": "file:///x.txt", + "mimeType": "text/plain", + "text": "hi", + }, + }, + ] @pytest.mark.asyncio async def test_plain_value_returned_as_is(self, mcp_tool): @@ -676,7 +429,7 @@ async def test_plain_value_returned_as_is(self, mcp_tool): class TestMcpToolErrorHandling: - """Test that protocol-level McpErrors are mapped to categorized AgentRuntimeErrors.""" + """Test that protocol-level MCPErrors are mapped to categorized AgentRuntimeErrors.""" @pytest.fixture def mcp_tool(self): @@ -686,7 +439,7 @@ def mcp_tool(self): input_schema={"type": "object", "properties": {}}, ) - def _mock_client(self, error: McpError) -> MagicMock: + def _mock_client(self, error: MCPError) -> MagicMock: client = MagicMock(spec=McpClient) client.server_slug = "my-mcp-server" client.call_tool = AsyncMock(side_effect=error) @@ -696,7 +449,7 @@ def _mock_client(self, error: McpError) -> MagicMock: async def test_session_terminated_raises_system_error_with_retry_hint( self, mcp_tool ): - error = McpError(ErrorData(code=32600, message="Session terminated")) + error = MCPError(code=32600, message="Session terminated") client = self._mock_client(error) tool_fn = build_mcp_tool(mcp_tool, client) @@ -714,7 +467,7 @@ async def test_session_terminated_raises_system_error_with_retry_hint( @pytest.mark.asyncio async def test_non_session_mcp_error_includes_server_message(self, mcp_tool): - error = McpError(ErrorData(code=-32601, message="Method not found")) + error = MCPError(code=-32601, message="Method not found") client = self._mock_client(error) tool_fn = build_mcp_tool(mcp_tool, client) @@ -1244,7 +997,7 @@ async def test_breaking_drift_heals_and_asks_retry(self): assert "question (string)" in result client.call_tool.assert_not_awaited() # The schema bound to the model was healed to the live one. - assert tool.args_schema == live_tool.inputSchema + assert tool.args_schema == live_tool.input_schema def test_schema_change_message_lists_param_types(self): """The retry message lists each refreshed param with its type and optionality.""" diff --git a/tests/agent/tools/test_mcp/test_protocol_strategy.py b/tests/agent/tools/test_mcp/test_protocol_strategy.py new file mode 100644 index 000000000..5d8d89b71 --- /dev/null +++ b/tests/agent/tools/test_mcp/test_protocol_strategy.py @@ -0,0 +1,588 @@ +"""Per-era MCP protocol policy, and the two servers a real one cannot imitate. + +Negotiation, affinity, retry semantics and wire identity are all driven against +real servers over real HTTP in ``test_mcp_client_real_http.py``. What is left +here is what a cooperative server cannot express: + +* ``auto`` sending a restored ID before the era is resolved, against a server + with no discovery endpoint. +* A proxy echoing ``mcp-session-id`` back in the modern era, where the client + owns the value and nothing on the wire may overwrite it. + +plus the pure-function policy matrices -- ``is_recoverable``, ``reset``, and +``build_protocol_strategy`` -- which need no server at all. +""" + +import json +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx2 +import pytest +from mcp.shared.exceptions import MCPError +from mcp.types import ( + CONNECTION_CLOSED, + INVALID_REQUEST, + METHOD_NOT_FOUND, + Implementation, + InitializeResult, + ServerCapabilities, +) +from uipath.agent.models.agent import AgentMcpResourceConfig, AgentMcpTool + +from uipath_langchain.agent.tools.mcp import McpClient, SessionInfo, SessionInfoFactory +from uipath_langchain.agent.tools.mcp.protocol_strategy import ( + AutoStrategy, + LegacyHandshakeStrategy, + ModernDiscoveryStrategy, + build_protocol_strategy, + is_session_rejected, +) +from uipath_langchain.agent.tools.mcp.streamable_http import MCP_SESSION_ID + +MODERN_VERSION = "2026-07-28" +LEGACY_VERSION = "2025-11-25" + +TOOL_SCHEMA = { + "name": "test_tool", + "description": "A test tool", + "inputSchema": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + "outputSchema": {"type": "object", "properties": {"result": {"type": "string"}}}, +} + + +class EraMcpEndpoint: + """Streamable HTTP endpoint that can offer either era, or both. + + A modern-only server answers ``server/discover`` and has no ``initialize`` + endpoint at all; a legacy-only server is the reverse. ``auto`` has to pick + correctly against either. + """ + + def __init__( + self, + *, + supports_discover: bool = True, + supports_initialize: bool = False, + echo_session_id: str | None = None, + ) -> None: + self.supports_discover = supports_discover + self.supports_initialize = supports_initialize + # Stands in for a proxy that returns mcp-session-id even in the modern + # era, where the client owns the value. + self.echo_session_id = echo_session_id + self.methods: list[str] = [] + self.request_headers: list[tuple[str, httpx2.Headers]] = [] + self.discover_count = 0 + self.initialize_count = 0 + self.tool_call_count = 0 + self.delete_count = 0 + self.transport = httpx2.MockTransport(self.handle) + + async def handle(self, request: httpx2.Request) -> httpx2.Response: + """Answer the MCP methods under test for whichever era is enabled.""" + if request.method == "GET": + return httpx2.Response(405) + if request.method == "DELETE": + self.delete_count += 1 + self.request_headers.append(("DELETE", request.headers)) + return httpx2.Response(204) + + body = json.loads(request.content) + method = body["method"] + params = body.get("params") or {} + self.methods.append(method) + self.request_headers.append((method, request.headers)) + + if method == "server/discover": + self.discover_count += 1 + if not self.supports_discover: + return self._error(body["id"], METHOD_NOT_FOUND, "Not Found", 404) + return self._result( + body["id"], + { + "supportedVersions": [MODERN_VERSION], + "capabilities": {"tools": {"listChanged": True}}, + "resultType": "complete", + }, + headers=( + {MCP_SESSION_ID: self.echo_session_id} + if self.echo_session_id + else None + ), + ) + if method == "initialize": + self.initialize_count += 1 + if not self.supports_initialize: + return self._error(body["id"], METHOD_NOT_FOUND, "Not Found", 404) + return self._result( + body["id"], + { + "protocolVersion": LEGACY_VERSION, + "capabilities": {"tools": {}}, + "serverInfo": {"name": "test-server", "version": "1.0.0"}, + }, + headers={MCP_SESSION_ID: "server-session-1"}, + ) + if method == "notifications/initialized": + return httpx2.Response(202) + if method == "tools/list": + # ``resultType`` is required on the 2026-07-28 wire, and a cacheable + # result carries its cache directives too. Older peers ignore both. + return self._result( + body["id"], + { + "tools": [TOOL_SCHEMA], + "resultType": "complete", + "ttlMs": 0, + "cacheScope": "private", + }, + ) + if method == "tools/call": + self.tool_call_count += 1 + result = {"result": f"Success from {params['name']}"} + return self._result( + body["id"], + { + "content": [{"type": "text", "text": json.dumps(result)}], + "structuredContent": result, + "isError": False, + "resultType": "complete", + }, + ) + return self._error(body.get("id"), METHOD_NOT_FOUND, "Method not found", 404) + + @staticmethod + def _result( + request_id: Any, + result: dict[str, Any], + *, + headers: dict[str, str] | None = None, + ) -> httpx2.Response: + response_headers = {"content-type": "application/json"} + response_headers.update(headers or {}) + return httpx2.Response( + 200, + headers=response_headers, + json={"jsonrpc": "2.0", "id": request_id, "result": result}, + ) + + @staticmethod + def _error( + request_id: Any, code: int, message: str, status: int + ) -> httpx2.Response: + return httpx2.Response( + status, + headers={"content-type": "application/json"}, + json={ + "jsonrpc": "2.0", + "id": request_id, + "error": {"code": code, "message": message}, + }, + ) + + def headers_for(self, method: str) -> list[httpx2.Headers]: + """Return captured headers for one protocol or HTTP method.""" + return [headers for name, headers in self.request_headers if name == method] + + +@pytest.fixture +def mcp_resource_config() -> AgentMcpResourceConfig: + """Create a minimal MCP resource config for testing.""" + return AgentMcpResourceConfig( + name="test_server", + description="Test MCP server", + folder_path="/Shared/TestFolder", + slug="test-server", + available_tools=[ + AgentMcpTool( + name="test_tool", + description="A test tool", + input_schema={ + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + ) + ], + ) + + +@pytest.fixture +def mock_uipath_sdk() -> MagicMock: + """Create a mock UiPath SDK and resolved MCP server.""" + sdk = MagicMock() + server = MagicMock() + server.mcp_url = "https://test.uipath.com/mcp" + server.slug = "test-server" + server.folder_key = "folder-key" + sdk.mcp.retrieve_async = AsyncMock(return_value=server) + sdk._config.secret = "test-secret-token" + return sdk + + +@asynccontextmanager +async def configured_client( + config: AgentMcpResourceConfig, + sdk: MagicMock, + endpoint: EraMcpEndpoint, + **kwargs: Any, +) -> AsyncIterator[McpClient]: + """Build an McpClient whose real HTTP client uses the mock transport.""" + client = McpClient(config=config, **kwargs) + http_kwargs = { + "headers": {"Authorization": "Bearer test-secret-token"}, + "transport": endpoint.transport, + "follow_redirects": True, + } + with ( + patch("uipath.platform.UiPath", return_value=sdk), + patch( + "uipath_langchain.agent.tools.mcp.mcp_client.get_httpx_client_kwargs", + return_value=http_kwargs, + ), + ): + try: + yield client + finally: + await client.dispose() + + +def _pinned_session_info(session_info: SessionInfo) -> SessionInfoFactory: + class PinnedFactory(SessionInfoFactory): + def create_session(self, mcp_server: Any) -> SessionInfo: + return session_info + + return PinnedFactory() + + +@pytest.mark.asyncio +async def test_auto_mode_sends_a_restored_id_before_the_era_resolves( + mcp_resource_config: AgentMcpResourceConfig, + mock_uipath_sdk: MagicMock, +) -> None: + """A persisted ID goes out on the shared header before the era is known. + + The transport opens before negotiation, and nothing distinguishes a + server-minted session ID from a client-minted affinity ID. Because both eras + use ``mcp-session-id``, no disambiguation is needed. + """ + session_info = SessionInfo("ambiguous-id") + endpoint = EraMcpEndpoint(supports_discover=False, supports_initialize=True) + async with configured_client( + mcp_resource_config, + mock_uipath_sdk, + endpoint, + session_info_factory=_pinned_session_info(session_info), + protocol_mode="auto", + ) as client: + await client.call_tool("test_tool", {"query": "test"}) + + assert endpoint.headers_for("server/discover")[0][MCP_SESSION_ID] == ( + "ambiguous-id" + ) + + +@pytest.mark.asyncio +async def test_modern_mode_ignores_a_server_assigned_session_id( + mcp_resource_config: AgentMcpResourceConfig, + mock_uipath_sdk: MagicMock, +) -> None: + """A response header must not overwrite the client-minted routing key. + + The two eras share ``mcp-session-id``, so a proxy or gateway echoing it back + could otherwise replace the affinity ID mid-connection and scatter the + remaining requests across instances. + """ + session_info = SessionInfo("affinity-keep-me") + endpoint = EraMcpEndpoint(echo_session_id="server-would-assign-this") + async with configured_client( + mcp_resource_config, + mock_uipath_sdk, + endpoint, + session_info_factory=_pinned_session_info(session_info), + protocol_mode="modern", + ) as client: + await client.call_tool("test_tool", {"query": "test"}) + + assert await session_info.get_session_id() == "affinity-keep-me" + sent = [ + headers.get(MCP_SESSION_ID) + for method, headers in endpoint.request_headers + if method != "DELETE" + ] + assert sent and all(value == "affinity-keep-me" for value in sent) + # Filtering DELETE out above would hide a teardown, so count them directly: + # an echoed ID adopted by the transport would surface here. + assert endpoint.delete_count == 0 + + +def _error(code: int, message: str) -> MCPError: + return MCPError(code, message) + + +def test_legacy_recovers_from_session_loss_but_not_from_bad_requests() -> None: + """Only a lost session justifies replacing a legacy connection.""" + strategy = LegacyHandshakeStrategy() + + assert strategy.is_recoverable(_error(CONNECTION_CLOSED, "Connection closed"), None) + assert strategy.is_recoverable(_error(INVALID_REQUEST, "Session terminated"), None) + assert strategy.is_recoverable(_error(INVALID_REQUEST, "Session not found"), None) + assert not strategy.is_recoverable(_error(INVALID_REQUEST, "Bad params"), None) + # The TypeScript SDK's transport refuses an unknown session under -32000, + # which is also CONNECTION_CLOSED -- retryable either way. + assert strategy.is_recoverable( + _error(CONNECTION_CLOSED, "Bad Request: No valid session ID provided"), None + ) + # A bare 404 is only a lost session while a restored ID is still in play. + assert strategy.is_recoverable(_error(METHOD_NOT_FOUND, "Not Found"), "restored") + assert not strategy.is_recoverable(_error(METHOD_NOT_FOUND, "Not Found"), None) + + +def test_modern_recovers_only_from_a_dropped_connection() -> None: + """Every modern request is self-contained, so nothing else is retryable.""" + strategy = ModernDiscoveryStrategy() + + assert strategy.is_recoverable(_error(CONNECTION_CLOSED, "Connection closed"), None) + assert not strategy.is_recoverable( + _error(INVALID_REQUEST, "Session terminated"), "restored" + ) + assert not strategy.is_recoverable( + _error(METHOD_NOT_FOUND, "Not Found"), "restored" + ) + + +def test_a_dropped_transport_is_not_a_verdict_on_the_session() -> None: + """Recovery may only discard the stored ID when the server rejected it. + + ``CONNECTION_CLOSED`` is JSON-RPC's ``-32000``, the same code the + TypeScript SDK's transport uses to refuse a session it does not know, so + the code alone cannot decide. Reading a dropped connection as a verdict + would throw away a live persisted session on a transient failure; reading a + refusal as a drop would resume a session the server has already declared + dead, and the retry would fail identically. + """ + assert not is_session_rejected(_error(CONNECTION_CLOSED, "Connection closed")) + assert is_session_rejected( + _error(CONNECTION_CLOSED, "Bad Request: No valid session ID provided") + ) + assert is_session_rejected(_error(INVALID_REQUEST, "Session terminated")) + # A bare 404 for a restored ID names no session, but it is still the + # server's answer rather than a dead socket. + assert is_session_rejected(_error(METHOD_NOT_FOUND, "Not Found")) + + +@pytest.mark.asyncio +async def test_modern_reset_keeps_the_affinity_id() -> None: + """Discarding the routing ID on failure would abandon the warm instance.""" + strategy = ModernDiscoveryStrategy() + session_info = SessionInfo("affinity-1") + + await strategy.reset(session_info) + + assert await session_info.get_session_id() == "affinity-1" + + +@pytest.mark.asyncio +async def test_legacy_reset_clears_the_stale_session_id() -> None: + """A lost legacy session must not be re-announced on the next handshake.""" + strategy = LegacyHandshakeStrategy() + session_info = SessionInfo("session-1") + await session_info.set_protocol_version(LEGACY_VERSION) + + await strategy.reset(session_info) + + assert await session_info.get_session_id() is None + # The version described the discarded session; keeping it would let the next + # connection adopt a version the replacement never negotiated. + assert await session_info.get_protocol_version() is None + + +def _initialize_result(version: str) -> InitializeResult: + """A handshake result as a server would return it.""" + return InitializeResult( + protocolVersion=version, + capabilities=ServerCapabilities(), + serverInfo=Implementation(name="test-server", version="1.0.0"), + ) + + +@pytest.mark.asyncio +async def test_legacy_resume_adopts_a_stored_version_without_negotiating() -> None: + """A remembered version is all a resumed session needs, so nothing is sent. + + This is the pre-SDK-2 wire behaviour: a stored session is used as-is, with + only the ``mcp-session-id`` header identifying it. Re-running ``initialize`` + inside a live session is refused outright by some servers -- the reference + TypeScript implementation answers "Server already initialized" -- which + would cost the persisted session, and with it the gateway affinity, on every + run against such a server. + """ + strategy = LegacyHandshakeStrategy() + session_info = SessionInfo("persisted-session") + await session_info.set_protocol_version(LEGACY_VERSION) + session = MagicMock() + session.initialize = AsyncMock() + + await strategy.connect(session, session_info) + + session.initialize.assert_not_awaited() + adopted = session.adopt.call_args.args[0] + assert isinstance(adopted, InitializeResult) + assert adopted.protocol_version == LEGACY_VERSION + assert await session_info.get_session_id() == "persisted-session" + assert await session_info.get_protocol_version() == LEGACY_VERSION + + +@pytest.mark.asyncio +async def test_legacy_resume_negotiates_when_the_version_is_unknown() -> None: + """A store written before versions were recorded has nothing to adopt. + + The handshake inside the restored session is the fallback, not the norm: it + is the only way left to learn what that session was negotiated at. + """ + strategy = LegacyHandshakeStrategy() + session_info = SessionInfo("persisted-session") + session = MagicMock() + session.initialize = AsyncMock(return_value=_initialize_result(LEGACY_VERSION)) + + await strategy.connect(session, session_info) + + session.initialize.assert_awaited_once() + session.adopt.assert_not_called() + # Learned now, so the next resume needs no handshake at all. + assert await session_info.get_protocol_version() == LEGACY_VERSION + + +@pytest.mark.asyncio +async def test_legacy_resume_ignores_a_version_it_cannot_speak() -> None: + """A modern version stored against the ID cannot be adopted on this wire.""" + strategy = LegacyHandshakeStrategy() + session_info = SessionInfo("affinity-id") + await session_info.set_protocol_version(MODERN_VERSION) + session = MagicMock() + session.initialize = AsyncMock(return_value=_initialize_result(LEGACY_VERSION)) + + await strategy.connect(session, session_info) + + session.adopt.assert_not_called() + session.initialize.assert_awaited_once() + assert await session_info.get_protocol_version() == LEGACY_VERSION + + +@pytest.mark.asyncio +async def test_legacy_handshake_records_what_it_negotiated() -> None: + """A cold session stores its version, which is what makes resume free.""" + strategy = LegacyHandshakeStrategy() + session_info = SessionInfo() + session = MagicMock() + session.initialize = AsyncMock(return_value=_initialize_result("2025-06-18")) + + await strategy.connect(session, session_info) + + session.initialize.assert_awaited_once() + assert await session_info.get_protocol_version() == "2025-06-18" + + +@pytest.mark.asyncio +async def test_legacy_resume_falls_back_cleanly_when_the_session_is_rejected() -> None: + """A refused handshake still yields a working connection, version recorded.""" + strategy = LegacyHandshakeStrategy() + session_info = SessionInfo("stale-session") + session = MagicMock() + session.initialize = AsyncMock( + side_effect=[ + MCPError(INVALID_REQUEST, "Invalid Request: Server already initialized"), + _initialize_result(LEGACY_VERSION), + ] + ) + + await strategy.connect(session, session_info) + + assert session.initialize.await_count == 2 + assert await session_info.get_protocol_version() == LEGACY_VERSION + + +def test_auto_applies_the_legacy_policy_before_an_era_is_resolved() -> None: + """A failure during the very first probe is judged conservatively.""" + strategy = AutoStrategy() + + assert strategy.is_recoverable(_error(INVALID_REQUEST, "Session terminated"), None) + + +@pytest.mark.asyncio +async def test_legacy_keeps_a_persisted_session_when_the_connection_drops() -> None: + """A dead transport says nothing about whether the session is still valid. + + Clearing the ID here would destroy an externally persisted session -- + permanently, for a store-backed SessionInfo -- over a transient failure, and + the retry would start a cold session instead of resuming the warm one. + """ + strategy = LegacyHandshakeStrategy() + session_info = SessionInfo("persisted-session") + session = MagicMock() + session.initialize = AsyncMock( + side_effect=MCPError(CONNECTION_CLOSED, "Connection closed") + ) + + with pytest.raises(MCPError): + await strategy.connect(session, session_info) + + assert await session_info.get_session_id() == "persisted-session" + # One attempt only: no clean-session fallback on a dead transport. + assert session.initialize.await_count == 1 + + +@pytest.mark.asyncio +async def test_auto_does_not_carry_a_stale_era_through_a_failed_probe() -> None: + """A failed negotiation must not leave the previous era deciding recovery. + + After resolving modern, a later probe that raises would otherwise keep the + modern policy, which refuses to retry session errors -- so a legacy server + reached on the retry would never recover. + """ + strategy = AutoStrategy() + session = MagicMock() + + with patch( + "uipath_langchain.agent.tools.mcp.protocol_strategy.probe_modern_era", + AsyncMock(return_value=True), + ): + await strategy.connect(session, SessionInfo()) + assert not strategy.is_recoverable( + _error(INVALID_REQUEST, "Session terminated"), None + ) + + with patch( + "uipath_langchain.agent.tools.mcp.protocol_strategy.probe_modern_era", + AsyncMock(side_effect=MCPError(INVALID_REQUEST, "probe blew up")), + ): + with pytest.raises(MCPError): + await strategy.connect(session, SessionInfo()) + + assert strategy.is_recoverable(_error(INVALID_REQUEST, "Session terminated"), None) + + +def test_build_protocol_strategy_maps_every_mode() -> None: + """The public ``protocol_mode`` values are the only accepted ones.""" + assert isinstance(build_protocol_strategy("legacy"), LegacyHandshakeStrategy) + assert isinstance(build_protocol_strategy("modern"), ModernDiscoveryStrategy) + assert isinstance(build_protocol_strategy("auto"), AutoStrategy) + + with pytest.raises(ValueError, match="Unknown MCP protocol mode"): + build_protocol_strategy("2026-07-28") # type: ignore[arg-type] + + +def test_legacy_is_the_default_mode( + mcp_resource_config: AgentMcpResourceConfig, +) -> None: + """Existing callers must keep the pre-2026 wire behavior untouched.""" + client = McpClient(config=mcp_resource_config) + + assert isinstance(client._strategy, LegacyHandshakeStrategy) diff --git a/tests/agent/tools/test_mcp/test_protocol_version_support.py b/tests/agent/tools/test_mcp/test_protocol_version_support.py new file mode 100644 index 000000000..bc6b67a7c --- /dev/null +++ b/tests/agent/tools/test_mcp/test_protocol_version_support.py @@ -0,0 +1,90 @@ +"""Guards on the SDK facts the MCP protocol strategies are built on. + +These assert facts about the MCP SDK rather than about UiPath code, which is +unusual for a unit test. They earn their place because each one pins an external +constraint that dictates how the MCP integration is built, and each says what to +change when it flips. They are deliberately cheap: no sockets, no servers. +""" + +import inspect + +from mcp import ClientSession +from mcp.types import UnsupportedProtocolVersionErrorData +from mcp.types.version import ( + HANDSHAKE_PROTOCOL_VERSIONS, + LATEST_HANDSHAKE_VERSION, + LATEST_MODERN_VERSION, + MODERN_PROTOCOL_VERSIONS, +) + + +def test_the_auto_probe_builds_on_public_session_methods() -> None: + """``probe_modern_era`` owns the ``auto`` policy on public ``ClientSession`` seams. + + It sends the probe through ``send_discover(version)`` and installs the + result through ``adopt(result)`` -- the same two calls the SDK's private + ``mode="auto"`` helper is built from, which is deliberately not imported. + A ``-32022`` is read through ``UnsupportedProtocolVersionErrorData.supported``. + If any of these change shape, the probe needs the matching change. + """ + assert list(inspect.signature(ClientSession.send_discover).parameters) == [ + "self", + "version", + ], "ClientSession.send_discover() changed shape; update probe_modern_era" + assert list(inspect.signature(ClientSession.adopt).parameters) == [ + "self", + "result", + ], "ClientSession.adopt() changed shape; update probe_modern_era" + assert "supported" in UnsupportedProtocolVersionErrorData.model_fields, ( + "-32022 error data lost its 'supported' list; probe_modern_era can no " + "longer tell a modern-only server from a legacy one" + ) + + +def test_the_low_level_session_reaches_the_modern_era() -> None: + """``ClientSession.discover()`` is why ``McpClient`` needs no high-level client. + + ``discover`` takes no version argument: it always proposes + ``LATEST_MODERN_VERSION`` and the SDK owns the ``-32022`` retry at a mutual + version. If it gains parameters, ``ModernDiscoveryStrategy`` may be able to + pin a version explicitly. + """ + assert hasattr(ClientSession, "discover"), ( + "ClientSession lost discover(); the modern era is no longer reachable " + "from the low-level session and ModernDiscoveryStrategy needs rework" + ) + parameters = inspect.signature(ClientSession.discover).parameters + assert list(parameters) == ["self"], ( + "ClientSession.discover() gained parameters; the modern strategy may now " + "be able to request a specific protocol version" + ) + assert LATEST_MODERN_VERSION in MODERN_PROTOCOL_VERSIONS + + +def test_initialize_cannot_choose_a_protocol_version() -> None: + """``ClientSession.initialize()`` takes no version argument. + + The legacy strategy always offers ``LATEST_HANDSHAKE_VERSION`` and the server + counters with what it supports, which is why a resumed session re-runs the + handshake to learn its version rather than guessing at one. + """ + parameters = inspect.signature(ClientSession.initialize).parameters + assert list(parameters) == ["self"], ( + "ClientSession.initialize() gained parameters; the legacy strategy may " + "now be able to request a specific protocol version" + ) + assert LATEST_HANDSHAKE_VERSION in HANDSHAKE_PROTOCOL_VERSIONS + + +def test_the_two_eras_share_no_protocol_version() -> None: + """Disjoint version sets are why there are two strategies rather than one. + + A version reachable through both ``initialize`` and ``server/discover`` would + make the era a property of the server rather than of the negotiation, and + ``AutoStrategy`` could stop probing. + """ + overlap = set(MODERN_PROTOCOL_VERSIONS) & set(HANDSHAKE_PROTOCOL_VERSIONS) + assert not overlap, ( + f"Version(s) {sorted(overlap)} are in both eras; the handshake may now " + "reach the modern protocol and AutoStrategy's probe may be redundant" + ) diff --git a/tests/agent/tools/test_mcp/test_session_tools.py b/tests/agent/tools/test_mcp/test_session_tools.py new file mode 100644 index 000000000..273ea89e9 --- /dev/null +++ b/tests/agent/tools/test_mcp/test_session_tools.py @@ -0,0 +1,71 @@ +"""Tests for binding active MCP sessions to LangChain tools.""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from langchain_core.tools import ToolException +from mcp.types import CallToolResult, ListToolsResult, TextContent, Tool + +from uipath_langchain.agent.tools.mcp import load_mcp_tools + + +@pytest.mark.asyncio +async def test_load_mcp_tools_binds_discovery_and_invocation() -> None: + """Discovered MCP 2 schemas and results are usable as LangChain tools.""" + session = MagicMock() + session.list_tools = AsyncMock( + side_effect=[ + ListToolsResult( + nextCursor="next-page", + tools=[ + Tool( + name="echo", + description="Echo a value", + inputSchema={ + "type": "object", + "properties": {"value": {"type": "string"}}, + "required": ["value"], + }, + ) + ], + ), + ListToolsResult(tools=[]), + ] + ) + session.call_tool = AsyncMock( + return_value=CallToolResult( + content=[TextContent(type="text", text="hello")], + isError=False, + ) + ) + + tools = await load_mcp_tools(session) + result = await tools[0].ainvoke({"value": "hello"}) + + assert result == [{"type": "text", "text": "hello"}] + assert session.list_tools.await_count == 2 + assert session.list_tools.await_args_list[0].kwargs == {"params": None} + assert session.list_tools.await_args_list[1].kwargs["params"].cursor == "next-page" + session.call_tool.assert_awaited_once_with("echo", arguments={"value": "hello"}) + + +@pytest.mark.asyncio +async def test_load_mcp_tools_maps_mcp_failures_to_tool_errors() -> None: + """Protocol-level tool failures retain their server-provided message.""" + session = MagicMock() + session.list_tools = AsyncMock( + return_value=ListToolsResult( + tools=[Tool(name="fail", inputSchema={"type": "object"})] + ) + ) + session.call_tool = AsyncMock( + return_value=CallToolResult( + content=[TextContent(type="text", text="server rejected the call")], + isError=True, + ) + ) + + tools = await load_mcp_tools(session) + + with pytest.raises(ToolException, match="server rejected the call"): + await tools[0].ainvoke({}) diff --git a/uv.lock b/uv.lock index 6a95cfa9a..e13b3da0c 100644 --- a/uv.lock +++ b/uv.lock @@ -1567,6 +1567,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55", size = 78784, upload-time = "2025-04-24T22:06:20.566Z" }, ] +[[package]] +name = "httpcore2" +version = "2.9.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "h11" }, + { name = "truststore" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/39/a8/20ed1ed79cbc2ecdf5301c0968ab7c85547212e2a7bd126ddd2d986e206e/httpcore2-2.9.1.tar.gz", hash = "sha256:4d8acbf8b306f48c9d6046591fd5ba4037d1b1b1000d140fc2c3eab1e9a0c0e2", size = 67089, upload-time = "2026-07-24T09:21:03.867Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9f/fb/46c52b781975c335a2bcf1072c7bbc007cbdc8d674217f5ee1daba2c848b/httpcore2-2.9.1-py3-none-any.whl", hash = "sha256:6182472379e855fe4221246a2bb7ecede403bc61c6798062ae1787d051ccde26", size = 82809, upload-time = "2026-07-24T09:21:01.178Z" }, +] + [[package]] name = "httpx" version = "0.28.1" @@ -1606,6 +1619,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/98/f8/a6bc80313a9e93c888fa10534dfce2ad76ff86911b6f485777ce6de6a073/httpx_ws-0.9.0-py3-none-any.whl", hash = "sha256:71640d2fb1bf9a225775015b33cd755cfd4c5f7e21c885192fe3adc4c387b248", size = 15759, upload-time = "2026-03-28T14:11:11.887Z" }, ] +[[package]] +name = "httpx2" +version = "2.9.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "httpcore2" }, + { name = "idna" }, + { name = "truststore" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/21/14/38128fbafd7e0ed41d874df6c9a653d47c2d111cfe59e2b4ac95161b4abd/httpx2-2.9.1.tar.gz", hash = "sha256:1932a768737e3666291582833da748cc4e563c337cf96706fccc04fa6e58764a", size = 95458, upload-time = "2026-07-24T09:21:04.972Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/13/b8/cfd91c4ab9134d386d48f0b6ac662ff3d4be6efdee59ee1c67ebc3c0487c/httpx2-2.9.1-py3-none-any.whl", hash = "sha256:1820fe14a9ab1107bfeff39259987429450b070ec0ff38cc87eb0d8c97fdc71a", size = 91191, upload-time = "2026-07-24T09:21:02.6Z" }, +] + [[package]] name = "huggingface-hub" version = "1.20.1" @@ -2024,20 +2053,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/43/e8/25c50bbad7a05106c7af65557e165d6cb6159c90854dae61de59debe735d/langchain_litellm-0.6.4-py3-none-any.whl", hash = "sha256:60f4e37be1a47dc88f94fac7085675ef8fa04bba92f48735792d82f492120744", size = 26360, upload-time = "2026-04-03T16:56:46.76Z" }, ] -[[package]] -name = "langchain-mcp-adapters" -version = "0.2.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "langchain-core" }, - { name = "mcp" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/d9/52/cebf0ef5b1acef6cbc63d671171d43af70f12d19f55577909c7afa79fb6e/langchain_mcp_adapters-0.2.1.tar.gz", hash = "sha256:58e64c44e8df29ca7eb3b656cf8c9931ef64386534d7ca261982e3bdc63f3176", size = 36394, upload-time = "2025-12-09T16:28:38.98Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/03/81/b2479eb26861ab36be851026d004b2d391d789b7856e44c272b12828ece0/langchain_mcp_adapters-0.2.1-py3-none-any.whl", hash = "sha256:9f96ad4c64230f6757297fec06fde19d772c99dbdfbca987f7b7cfd51ff77240", size = 22708, upload-time = "2025-12-09T16:28:37.877Z" }, -] - [[package]] name = "langchain-openai" version = "1.3.2" @@ -2342,15 +2357,15 @@ wheels = [ [[package]] name = "mcp" -version = "1.26.0" +version = "2.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, - { name = "httpx" }, - { name = "httpx-sse" }, + { name = "httpx2" }, { name = "jsonschema" }, + { name = "mcp-types" }, + { name = "opentelemetry-api" }, { name = "pydantic" }, - { name = "pydantic-settings" }, { name = "pyjwt", extra = ["crypto"] }, { name = "python-multipart" }, { name = "pywin32", marker = "sys_platform == 'win32'" }, @@ -2360,9 +2375,22 @@ dependencies = [ { name = "typing-inspection" }, { name = "uvicorn", marker = "sys_platform != 'emscripten'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/fc/6d/62e76bbb8144d6ed86e202b5edd8a4cb631e7c8130f3f4893c3f90262b10/mcp-1.26.0.tar.gz", hash = "sha256:db6e2ef491eecc1a0d93711a76f28dec2e05999f93afd48795da1c1137142c66", size = 608005, upload-time = "2026-01-24T19:40:32.468Z" } +sdist = { url = "https://files.pythonhosted.org/packages/74/33/32d4dff2c95bb5d897c3ef4c83649a08996b17b58f0a326d2495d4c81179/mcp-2.0.0.tar.gz", hash = "sha256:0f440e735c13ece8bb19bc62cf0b86f4313448432fbb77d35e14034f4e050728", size = 1662284, upload-time = "2026-07-28T13:45:32.346Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/fd/d9/eaa1f80170d2b7c5ba23f3b59f766f3a0bb41155fbc32a69adfa1adaaef9/mcp-1.26.0-py3-none-any.whl", hash = "sha256:904a21c33c25aa98ddbeb47273033c435e595bbacfdb177f4bd87f6dceebe1ca", size = 233615, upload-time = "2026-01-24T19:40:30.652Z" }, + { url = "https://files.pythonhosted.org/packages/67/72/7d7897418912c1d12e87556630dfb7bf0eac71160e9bef8b447960804ee3/mcp-2.0.0-py3-none-any.whl", hash = "sha256:1cb4c75d2d2c7b8c1d756355e5d82a39f2822cc7f13e22a2051d7ca3592349d6", size = 349980, upload-time = "2026-07-28T13:45:28.853Z" }, +] + +[[package]] +name = "mcp-types" +version = "2.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pydantic" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/bb/56/9b8e1c152f61f6c6b07c4b5896c88c7d0ae90bac6ee6306f852fcc5c1eb0/mcp_types-2.0.0.tar.gz", hash = "sha256:d7d939b9285c9961ae8866ba75ef85da34d12bafe276efbf4eb6a131786d8379", size = 66632, upload-time = "2026-07-28T13:45:33.804Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f5/4c/c78d78c3d52b0ac594ad7cc8ef5972adfe070e3597a8a4c6ce0cd39196ea/mcp_types-2.0.0-py3-none-any.whl", hash = "sha256:6b2de797ca2797f568b79529e1b25948e34de511bcc0bd82fef1039a6d1b8eb0", size = 69649, upload-time = "2026-07-28T13:45:30.713Z" }, ] [[package]] @@ -4546,17 +4574,17 @@ wheels = [ [[package]] name = "uipath-langchain" -version = "0.16.16" +version = "0.17.0" source = { editable = "." } dependencies = [ { name = "a2a-sdk" }, { name = "deepagents" }, { name = "httpx" }, + { name = "httpx2" }, { name = "jsonpath-ng" }, { name = "jsonschema-pydantic-converter" }, { name = "langchain" }, { name = "langchain-core" }, - { name = "langchain-mcp-adapters" }, { name = "langgraph" }, { name = "langgraph-checkpoint-sqlite" }, { name = "mcp" }, @@ -4605,7 +4633,9 @@ dev = [ { name = "pytest-mock" }, { name = "ruff" }, { name = "rust-just" }, + { name = "starlette" }, { name = "types-protobuf" }, + { name = "uvicorn" }, { name = "virtualenv" }, ] @@ -4615,14 +4645,14 @@ requires-dist = [ { name = "boto3-stubs", marker = "extra == 'bedrock'", specifier = ">=1.41.4" }, { name = "deepagents", specifier = ">=0.5.9,<0.6.0" }, { name = "httpx", specifier = ">=0.27.0" }, + { name = "httpx2", specifier = ">=2.5.0,<2.10.0" }, { name = "jsonpath-ng", specifier = ">=1.7.0" }, { name = "jsonschema-pydantic-converter", specifier = ">=0.4.0" }, { name = "langchain", specifier = ">=1.2.15,<2.0.0" }, { name = "langchain-core", specifier = ">=1.2.27,<2.0.0" }, - { name = "langchain-mcp-adapters", specifier = "==0.2.1" }, { name = "langgraph", specifier = ">=1.1.8,<2.0.0" }, { name = "langgraph-checkpoint-sqlite", specifier = ">=3.0.3,<4.0.0" }, - { name = "mcp", specifier = "==1.26.0" }, + { name = "mcp", specifier = "==2.0.0" }, { name = "openinference-instrumentation-langchain", specifier = ">=0.1.69,<0.2.0" }, { name = "pillow", specifier = ">=12.1.1" }, { name = "pydantic-settings", specifier = ">=2.6.0" }, @@ -4657,7 +4687,9 @@ dev = [ { name = "pytest-mock", specifier = ">=3.11.1" }, { name = "ruff", specifier = ">=0.9.4" }, { name = "rust-just", specifier = ">=1.39.0" }, + { name = "starlette", specifier = ">=0.41.3" }, { name = "types-protobuf", specifier = "<7" }, + { name = "uvicorn", specifier = ">=0.30.0" }, { name = "virtualenv", specifier = ">=20.36.1" }, ]