From 61eeee8054d166712f1e46be105145de623d4cf2 Mon Sep 17 00:00:00 2001 From: Ion Mincu Date: Thu, 3 Sep 2026 16:06:40 +0300 Subject: [PATCH 1/2] feat(mcp): upgrade to SDK 2.0 and support the 2026-07-28 protocol Upgrade the MCP Python SDK from 1.26.0 to 2.0.0, replace the copied SDK 1.x Streamable HTTP transport with a thin adapter over the upstream one, and add opt-in support for the 2026-07-28 protocol revision. Cut `streamable_http.py` from roughly 800 lines to a session-aware adapter, so upstream fixes and new protocol behavior arrive without hand-merging a fork. MCP 2 removed the transport's `get_session_id` callback, so two `httpx2` event hooks carry UiPath's externally persisted `SessionInfo` instead: one puts the stored ID on each request, one persists an ID the server returns. Make session recovery correct under concurrency and failure. `initialize()` is idempotent per `ClientSession` in MCP 2, so recovery replaces the transport and session rather than re-initializing, guarded so a late failure from a superseded session cannot tear down its replacement. A failed replacement no longer poisons the client: the next operation reopens. Recovery clears the persisted session ID only on the server's verdict; a plain `CONNECTION_CLOSED` means the transport dropped, so the ID is kept and the reconnect resumes the same session instead of starting a cold one. Reach 2026-07-28 through `ProtocolStrategy`, which abstracts the session lifecycle rather than the transport. Negotiation is one call in either era; what differs is how a connection is negotiated, whether a restored one can be reused, which errors a reconnect can fix, and how the connection is identified on the wire. No high-level client is needed: `ClientSession.discover()` already reaches the modern era, and the `auto` probe is owned here on the public `ClientSession.send_discover` / `adopt` seam rather than importing the SDK's private `mcp.client._probe` helper. Select the era with a keyword-only `protocol_mode`, defaulted to `"legacy"` so every existing caller keeps identical wire behavior. `"auto"` would otherwise move any discovery-capable UiPath MCP server to stateless 2026-07-28 and stop issuing session IDs, breaking the AgentHub playground persistence that `SessionInfoDebugState` exists for. Replace the resumed-session version probe with a second handshake. The probe tried each handshake version with a ping and adopted the first that answered, but servers do not validate that header against what the session negotiated, so the oldest version always won: a session negotiated at 2025-11-25 was adopted as 2025-03-26, silently disabling the server's 2025-11-25 SSE resumability. The version cannot be recovered from the wire -- responses carry only the session ID -- so re-running `initialize` inside the restored session is what learns it. That is safe because the server routes by the session header and mints a new session only when the header is absent. Resume drops from up to four round trips to one. Carry a UiPath-minted affinity ID in the modern era. AgentHub routes serverless MCP instances by `mcp-session-id`, which 2026-07-28 removes, so the modern strategy mints its own value and keeps sending it on that header as an opaque routing key -- ignored by a modern server, and requiring no gateway change. The ID travels on the header only; no request body is rewritten. In both `modern` and `auto` the client mints it before negotiating, so it is present on the very first request, `server/discover` included, and the probe reaches the same instance the tool calls will. When `auto` falls back to the legacy handshake, a freshly minted ID is withdrawn first so a server that routes by the header is not asked to resume a session it never issued; a restored ID is kept for the handshake to resume. `SessionInfo` needed no new API: it stores whichever ID this server uses, so a subclass persists an affinity ID unmodified. Guard that ID against a proxy echoing `mcp-session-id` back, and against disposal sending `DELETE` for a session the server never issued -- which reached the gateway as a teardown for the live instance the ID exists to pin. Scope retries to what each era can recover: a legacy session can be lost and re-established, while every modern request is self-contained, so only a dropped connection is retryable there. Retain compatibility with the previously accepted `httpx.Timeout` API, and replace the incompatible `langchain-mcp-adapters` dependency -- it imports `RequestContext`, which MCP 2 removed -- with a tested first-party session-to-LangChain tool converter. Test the public `McpClient` against a real `MCPServer` over real HTTP rather than a mocked transport: negotiation in every mode, resume asserting the originally negotiated version, affinity pinning across clients, per-era retry and disposal, and all four handshake versions. Add `testcases/simple-http-mcp`, an LLM-free integration testcase hosting MCP over Streamable HTTP on real sockets, covering the same matrix plus the API surface `uipath-agents-python` depends on. Mocked transports are kept only for conditions a cooperative server cannot produce: concurrency races, pathological servers, and pure-function matrices. Move the `uipath new` scaffold pin to the 0.17 minor. The guard from #1052 fails on every minor bump so the scaffold gets reviewed; the templates and hints touch no MCP surface, so only the pin constant changes. BREAKING CHANGE: `McpClient` is a public low-level API that deliberately keeps returning the MCP SDK's raw result models, and SDK 2.0 renamed their Python attributes to snake case. Callers that read them directly must update: CallToolResult.isError -> is_error CallToolResult.structuredContent -> structured_content Tool.inputSchema -> input_schema Tool.outputSchema -> output_schema ListToolsResult.nextCursor -> next_cursor The wire format is unchanged: these became snake_case fields carrying the old names as serialization aliases, so `model_dump(by_alias=True)` still emits camelCase. Note that plain `model_dump()` now emits snake_case keys, which is the one form of this break that fails silently rather than raising. `McpClient.SESSION_ERROR_CODES` is removed; use the `McpClient.is_session_error` static method instead. Removing the vendored transport also drops the unused public names it carried, including `streamablehttp_client`, `StreamableHTTPTransport`, `RequestContext`, `StreamableHTTPError` and `ResumptionError`. Co-Authored-By: Ion Mincu Co-Authored-By: Claude Opus 5 (1M context) Co-Authored-By: Claude Fable 5.1 --- pyproject.toml | 11 +- samples/oauth-external-apps-agent/README.md | 5 +- samples/oauth-external-apps-agent/main.py | 76 +- .../oauth-external-apps-agent/pyproject.toml | 3 +- samples/simple-local-mcp/README.md | 4 +- samples/simple-local-mcp/pyproject.toml | 4 +- .../src/simple-local-mcp/graph.py | 40 +- .../src/simple-local-mcp/math_server.py | 4 +- .../src/simple-local-mcp/weather_server.py | 4 +- samples/simple-remote-mcp/README.md | 3 +- samples/simple-remote-mcp/main.py | 37 +- samples/simple-remote-mcp/pyproject.toml | 3 +- src/uipath_langchain/_cli/cli_new.py | 2 +- .../agent/tools/mcp/__init__.py | 2 + .../agent/tools/mcp/claude.md | 646 +++++-- .../agent/tools/mcp/mcp_client.py | 311 ++-- .../agent/tools/mcp/mcp_tool.py | 58 +- .../agent/tools/mcp/protocol_strategy.py | 384 ++++ .../agent/tools/mcp/session_tools.py | 64 + .../agent/tools/mcp/streamable_http.py | 931 ++-------- testcases/simple-http-mcp/input.json | 1 + testcases/simple-http-mcp/langgraph.json | 6 + testcases/simple-http-mcp/pyproject.toml | 34 + testcases/simple-http-mcp/run.sh | 17 + testcases/simple-http-mcp/src/assert.py | 291 +++ .../src/simple-http-mcp/agents_api.py | 557 ++++++ .../src/simple-http-mcp/graph.py | 437 +++++ .../src/simple-http-mcp/servers.py | 328 ++++ testcases/simple-local-mcp/pyproject.toml | 3 +- .../src/simple-local-mcp/graph.py | 39 +- .../src/simple-local-mcp/math_server.py | 4 +- .../src/simple-local-mcp/weather_server.py | 4 +- tests/agent/tools/test_mcp/claude.md | 639 ++++--- tests/agent/tools/test_mcp/real_server.py | 668 +++++++ tests/agent/tools/test_mcp/test_mcp_client.py | 1567 ++++++++--------- .../test_mcp/test_mcp_client_real_http.py | 618 +++++++ tests/agent/tools/test_mcp/test_mcp_tool.py | 361 +--- .../tools/test_mcp/test_protocol_strategy.py | 446 +++++ .../test_mcp/test_protocol_version_support.py | 90 + .../tools/test_mcp/test_session_tools.py | 71 + uv.lock | 80 +- 41 files changed, 6275 insertions(+), 2578 deletions(-) create mode 100644 src/uipath_langchain/agent/tools/mcp/protocol_strategy.py create mode 100644 src/uipath_langchain/agent/tools/mcp/session_tools.py create mode 100644 testcases/simple-http-mcp/input.json create mode 100644 testcases/simple-http-mcp/langgraph.json create mode 100644 testcases/simple-http-mcp/pyproject.toml create mode 100644 testcases/simple-http-mcp/run.sh create mode 100644 testcases/simple-http-mcp/src/assert.py create mode 100644 testcases/simple-http-mcp/src/simple-http-mcp/agents_api.py create mode 100644 testcases/simple-http-mcp/src/simple-http-mcp/graph.py create mode 100644 testcases/simple-http-mcp/src/simple-http-mcp/servers.py create mode 100644 tests/agent/tools/test_mcp/real_server.py create mode 100644 tests/agent/tools/test_mcp/test_mcp_client_real_http.py create mode 100644 tests/agent/tools/test_mcp/test_protocol_strategy.py create mode 100644 tests/agent/tools/test_mcp/test_protocol_version_support.py create mode 100644 tests/agent/tools/test_mcp/test_session_tools.py 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..2ea0d1639 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,168 @@ 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 # retained, unread 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: ... ``` 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` is retained for compatibility but **nothing reads it**. The +SDK stamps the negotiated version onto every request itself once negotiation +completes, and a resumed legacy session learns its version by re-running the +handshake (see below) rather than by probing. -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. +**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. -Key methods: +#### Upstream StreamableHTTPTransport -| 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` | +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 | re-run the handshake inside the restored session | 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. `CONNECTION_CLOSED` is not -- the transport +dropped, the server never rejected anything -- so the ID is kept and the +reconnect resumes the warm session. Every other recoverable error clears the ID +before the fresh handshake. + +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: re-handshake, do not guess + +A restored session ID arrives without the protocol version it was negotiated at, +and the version cannot be recovered from the wire — server responses carry only +the session ID. The strategy therefore re-runs `initialize` **inside** the +restored session: + +```python +restored = await info.get_session_id() +if restored is None: + await session.initialize() + return +try: + await session.initialize() # lands inside the restored 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 session.initialize() # same transport: a refused request does not close it +``` + +This 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` instead — the pre-existing 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. + +#### 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 +313,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 +323,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 +337,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 +376,21 @@ 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** — `session.initialize()`, whether or not an ID was restored; the + response hook stores the server-assigned ID + - **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 +474,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 +487,20 @@ 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 an ID was restored) └─────────────────┘ ┌─────────────────┐ │ McpClient reads │ ─── await session_info.get_session_id() @@ -348,42 +511,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 +542,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 +559,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 +575,37 @@ 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. + +`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 +642,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 +665,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 error.code != CONNECTION_CLOSED + ): + await self._strategy.reset(self._session_info) + await self._open_connection() ``` ### 4. No `with` Statement for AsyncExitStack @@ -509,17 +698,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 the error was a +plain `CONNECTION_CLOSED` -- 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 +747,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"} +``` -# But session initialized TWICE -assert initialize_count[0] == 2 +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", +] ``` +## 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 +850,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 +904,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 +917,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..b25157c11 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 ( + CONNECTION_CLOSED, + 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, +) +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. ``CONNECTION_CLOSED`` + is not: the transport dropped, so 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 error.code != CONNECTION_CLOSED + ): + 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..75a2233f7 --- /dev/null +++ b/src/uipath_langchain/agent/tools/mcp/protocol_strategy.py @@ -0,0 +1,384 @@ +"""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, + 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 an ``INVALID_REQUEST`` recoverable. +_SESSION_LOST_MARKERS = ("terminated", "expired", "invalid", "not found") + + +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 + message = error.message.lower() + return ( + error.code in (32600, INVALID_REQUEST) + and "session" in message + and any(marker in message for marker in _SESSION_LOST_MARKERS) + ) + + +@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 re-negotiated rather than probed. The server routes + requests purely by the ``mcp-session-id`` header and creates a new session + only when no header is present, so an ``initialize`` carrying a restored ID + lands *inside* that session and returns the version it was negotiated at. + """ + + 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 session.initialize() + logger.info( + "MCP session initialized with session ID: %s", + await info.get_session_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 session.initialize() + logger.info( + "MCP session initialized with session ID: %s", + await info.get_session_id(), + ) + return + + 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, + ) + + 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.""" + await info.set_session_id(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..0a8ccf149 100644 --- a/src/uipath_langchain/agent/tools/mcp/streamable_http.py +++ b/src/uipath_langchain/agent/tools/mcp/streamable_http.py @@ -1,802 +1,225 @@ -"""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. - - 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. - """ + """Store the MCP session ID and allow subclasses to persist it externally.""" def __init__(self, session_id: str | None = None) -> None: self.session_id = session_id + # Exposed for subclasses that want to record what was negotiated. + # Nothing in this package sets or reads it: the SDK stamps the version + # onto every request itself, and a resumed session learns its version by + # re-running the handshake rather than by being told. + 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 -class StreamableHTTPError(Exception): - """Base exception for StreamableHTTP transport errors.""" - +@dataclass(frozen=True) +class SessionIdentityWire: + """How a stored session ID travels on the wire for one protocol era.""" -class ResumptionError(StreamableHTTPError): - """Raised when resumption request is invalid.""" + request_headers: tuple[str, ...] = (MCP_SESSION_ID,) + """Headers the stored ID is sent on. Empty sends none.""" + capture_response_header: str | None = MCP_SESSION_ID + """Header a server-assigned ID is read from, or ``None`` when the client mints it. -@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 - - -class StreamableHTTPTransport: - """StreamableHTTP client transport implementation.""" - - @overload - def __init__( - self, url: str, *, session_info: SessionInfo | None = None - ) -> None: ... - - @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 - - async def _prepare_headers(self) -> dict[str, str]: - """Build MCP-specific request headers. - - 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..dd8ff6775 --- /dev/null +++ b/tests/agent/tools/test_mcp/real_server.py @@ -0,0 +1,668 @@ +"""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``. + """ + + 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 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") + + 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": [_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..d9ccdf1dd 100644 --- a/tests/agent/tools/test_mcp/test_mcp_client.py +++ b/tests/agent/tools/test_mcp/test_mcp_client.py @@ -1,883 +1,768 @@ -"""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"} + # The handshake ran again on the new transport -- inside the same session. + assert endpoint.initialize_count == 2 + 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..01946b0e8 --- /dev/null +++ b/tests/agent/tools/test_mcp/test_mcp_client_real_http.py @@ -0,0 +1,618 @@ +"""``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 + + +@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..3ad5b3d3f --- /dev/null +++ b/tests/agent/tools/test_mcp/test_protocol_strategy.py @@ -0,0 +1,446 @@ +"""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 +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, +) +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) + # 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" + ) + + +@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 strategy.reset(session_info) + + assert await session_info.get_session_id() is None + + +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" }, ] From 9f244cff2885fda4abe53184d66661c0596a59cb Mon Sep 17 00:00:00 2001 From: Ion Mincu Date: Thu, 3 Sep 2026 18:34:12 +0300 Subject: [PATCH 2/2] fix(mcp): resume a legacy session without a second initialize The SDK 2 upgrade made a resumed legacy session re-run `initialize` inside the restored session, because that was the only way left to learn the version the session had been negotiated at. The pre-upgrade client sent no handshake at all on resume, so this changed the wire for every caller of the default `legacy` mode whose session store outlives a connection. It is not only extra traffic. Whether a server accepts a second `initialize` on a live session is implementation-defined: the Python SDK answers it, while the reference TypeScript implementation refuses outright with `-32600 "Invalid Request: Server already initialized"`. Driven against that server, every resume lost the persisted session and started a cold one -- and with it the `mcp-session-id` affinity a gateway routes serverless MCP instances by, which is the whole point of persisting the session. Store the negotiated version alongside the ID instead, and adopt it. The version is the only thing a fresh `ClientSession` lacks, and `ClientSession.adopt()` installs negotiated state locally with no wire traffic, so a resumed run sends only ordinary requests carrying `mcp-session-id`, exactly as it did before the upgrade. `SessionInfo` grows async accessors for the version, mirroring the ID's, so a store that persists one can persist both; `reset` clears them together, so a replacement session can never inherit a version it did not negotiate. The handshake path stays for a store with no version recorded and for a stored version this client cannot speak on a legacy wire. Decide `reset` on the server's verdict rather than on the error code alone. `CONNECTION_CLOSED` is JSON-RPC's implementation-defined `-32000`, which the TypeScript SDK's transport also answers with when it refuses a session it does not know, so comparing the code kept a session the server had just declared dead and spent the retry resuming it. `is_session_rejected` reads the message when the code is ambiguous: a dropped transport still keeps the ID, and a `-32000` naming a lost session now clears it -- as the pre-upgrade client always did. Add a `refuse_reinitialize` mode to the pinned real-HTTP endpoint, so the TypeScript server's refusal is a regression test rather than a manual finding. Co-Authored-By: Claude Opus 5 (1M context) --- .../agent/tools/mcp/claude.md | 110 ++++++++---- .../agent/tools/mcp/mcp_client.py | 12 +- .../agent/tools/mcp/protocol_strategy.py | 159 +++++++++++++++--- .../agent/tools/mcp/streamable_http.py | 26 ++- tests/agent/tools/test_mcp/real_server.py | 23 ++- tests/agent/tools/test_mcp/test_mcp_client.py | 5 +- .../test_mcp/test_mcp_client_real_http.py | 48 ++++++ .../tools/test_mcp/test_protocol_strategy.py | 144 +++++++++++++++- 8 files changed, 454 insertions(+), 73 deletions(-) diff --git a/src/uipath_langchain/agent/tools/mcp/claude.md b/src/uipath_langchain/agent/tools/mcp/claude.md index 2ea0d1639..62a9585a8 100644 --- a/src/uipath_langchain/agent/tools/mcp/claude.md +++ b/src/uipath_langchain/agent/tools/mcp/claude.md @@ -123,10 +123,12 @@ 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 # retained, unread + self.protocol_version: str | None = None async def get_session_id(self) -> 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 @@ -139,14 +141,21 @@ 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. -`protocol_version` is retained for compatibility but **nothing reads it**. The -SDK stamps the negotiated version onto every request itself once negotiation -completes, and a resumed legacy session learns its version by re-running the -handshake (see below) rather than by probing. +`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. + +A subclass persisting externally should override all four accessors, so the ID +and its version are written and cleared together. **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. +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. #### Upstream StreamableHTTPTransport @@ -181,7 +190,7 @@ Negotiation itself is one call either way; what genuinely differs is the session |---------|------------------------------------|------------------------| | Negotiate | `ClientSession.initialize()` | `ClientSession.discover()` | | Identity | `mcp-session-id`, server-minted | none in-protocol; UiPath-minted affinity ID | -| Resume | re-run the handshake inside the restored session | nothing to resume | +| 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` | @@ -195,10 +204,18 @@ class ProtocolStrategy(Protocol): ``` `McpClient` calls `reset` only when the error that triggered recovery is the -server's verdict on the session. `CONNECTION_CLOSED` is not -- the transport -dropped, the server never rejected anything -- so the ID is kept and the -reconnect resumes the warm session. Every other recoverable error clears the ID -before the fresh handshake. +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`: @@ -216,38 +233,53 @@ 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: re-handshake, do not guess +#### Resuming a legacy session: adopt the version, do not renegotiate -A restored session ID arrives without the protocol version it was negotiated at, -and the version cannot be recovered from the wire — server responses carry only -the session ID. The strategy therefore re-runs `initialize` **inside** the -restored session: +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 session.initialize() + 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: - await session.initialize() # lands inside the restored session + 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 session.initialize() # same transport: a refused request does not close it + 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 ``` -This is safe because the server routes purely by the session header and mints a -new session only when the header is **absent** +**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` instead — the pre-existing 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. +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 @@ -379,8 +411,9 @@ During client initialization, `McpClient`: 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** — `session.initialize()`, whether or not an ID was restored; the - response hook stores the server-assigned ID + - **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`; @@ -500,7 +533,8 @@ Phase 2: Connection Initialization (repeated after session loss) ┌─────────────────┐ │ session. │ ─── Sends initialize request │ initialize() │ Response hook calls set_session_id() -│ │ (skipped when an ID was restored) +│ │ (skipped when a stored ID *and* version +│ │ are adopted instead) └─────────────────┘ ┌─────────────────┐ │ McpClient reads │ ─── await session_info.get_session_id() @@ -592,6 +626,10 @@ 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 @@ -683,7 +721,7 @@ async def _reinitialize_session( if ( self._session_info is not None and error is not None - and error.code != CONNECTION_CLOSED + and is_session_rejected(error) ): await self._strategy.reset(self._session_info) await self._open_connection() @@ -708,10 +746,10 @@ async with AsyncExitStack() as stack: ### 5. Reinitialization Reuses the HTTP Client On a recoverable session error, `_reinitialize_session()` closes the old -connection stack, calls `strategy.reset` to clear the ID unless the error was a -plain `CONNECTION_CLOSED` -- 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` +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. diff --git a/src/uipath_langchain/agent/tools/mcp/mcp_client.py b/src/uipath_langchain/agent/tools/mcp/mcp_client.py index b25157c11..425b2b393 100644 --- a/src/uipath_langchain/agent/tools/mcp/mcp_client.py +++ b/src/uipath_langchain/agent/tools/mcp/mcp_client.py @@ -15,7 +15,6 @@ from mcp import ClientSession from mcp.shared.exceptions import MCPError from mcp.types import ( - CONNECTION_CLOSED, CallToolResult, ListToolsResult, ) @@ -28,6 +27,7 @@ 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 @@ -332,10 +332,10 @@ async def _reinitialize_session( The HTTP client and external ``SessionInfo`` object are reused. The strategy is asked to discard persisted session state only when - ``error`` is the server's verdict on the session. ``CONNECTION_CLOSED`` - is not: the transport dropped, so 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. + ``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: @@ -351,7 +351,7 @@ async def _reinitialize_session( if ( self._session_info is not None and error is not None - and error.code != CONNECTION_CLOSED + and is_session_rejected(error) ): await self._strategy.reset(self._session_info) await self._open_connection() diff --git a/src/uipath_langchain/agent/tools/mcp/protocol_strategy.py b/src/uipath_langchain/agent/tools/mcp/protocol_strategy.py index 75a2233f7..1ff152606 100644 --- a/src/uipath_langchain/agent/tools/mcp/protocol_strategy.py +++ b/src/uipath_langchain/agent/tools/mcp/protocol_strategy.py @@ -26,6 +26,9 @@ METHOD_NOT_FOUND, UNSUPPORTED_PROTOCOL_VERSION, DiscoverResult, + Implementation, + InitializeResult, + ServerCapabilities, UnsupportedProtocolVersionErrorData, ) from mcp.types.version import ( @@ -46,8 +49,23 @@ ProtocolMode = Literal["legacy", "auto", "modern"] -#: Session-disconnect markers that make an ``INVALID_REQUEST`` recoverable. -_SESSION_LOST_MARKERS = ("terminated", "expired", "invalid", "not found") +#: 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: @@ -62,12 +80,31 @@ def is_session_error(error: MCPError) -> bool: """ if error.code == CONNECTION_CLOSED: return True - message = error.message.lower() - return ( - error.code in (32600, INVALID_REQUEST) - and "session" in message - and any(marker in message for marker in _SESSION_LOST_MARKERS) - ) + # ``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 @@ -93,10 +130,22 @@ async def reset(self, info: SessionInfo) -> None: class LegacyHandshakeStrategy: """Negotiate through ``initialize`` and reuse server-minted sessions. - A restored session is re-negotiated rather than probed. The server routes - requests purely by the ``mcp-session-id`` header and creates a new session - only when no header is present, so an ``initialize`` carrying a restored ID - lands *inside* that session and returns the version it was negotiated at. + 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: @@ -106,11 +155,10 @@ 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 session.initialize() - logger.info( - "MCP session initialized with session ID: %s", - await info.get_session_id(), - ) + await self._handshake(session, info) + return + + if await self._adopt_restored_session(session, info, restored_id): return try: @@ -131,13 +179,10 @@ async def connect(self, session: ClientSession, info: SessionInfo) -> None: error.code, ) await self.reset(info) - await session.initialize() - logger.info( - "MCP session initialized with session ID: %s", - await info.get_session_id(), - ) + 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( @@ -155,6 +200,68 @@ async def connect(self, session: ClientSession, info: SessionInfo) -> None: 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. @@ -176,8 +283,14 @@ def is_recoverable(self, error: MCPError, restored_id: str | None) -> bool: return restored_id is not None async def reset(self, info: SessionInfo) -> None: - """Clear the stale session ID so the next handshake starts clean.""" + """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: diff --git a/src/uipath_langchain/agent/tools/mcp/streamable_http.py b/src/uipath_langchain/agent/tools/mcp/streamable_http.py index 0a8ccf149..ae15cbfe6 100644 --- a/src/uipath_langchain/agent/tools/mcp/streamable_http.py +++ b/src/uipath_langchain/agent/tools/mcp/streamable_http.py @@ -41,12 +41,17 @@ class SessionInfo: """Store the MCP session ID and allow subclasses to persist it externally.""" + #: 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 - # Exposed for subclasses that want to record what was negotiated. - # Nothing in this package sets or reads it: the SDK stamps the version - # onto every request itself, and a resumed session learns its version by - # re-running the handshake rather than by being told. + # 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: @@ -57,6 +62,19 @@ async def set_session_id(self, session_id: str | None) -> None: """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. + + ``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 + @dataclass(frozen=True) class SessionIdentityWire: diff --git a/tests/agent/tools/test_mcp/real_server.py b/tests/agent/tools/test_mcp/real_server.py index dd8ff6775..832624979 100644 --- a/tests/agent/tools/test_mcp/real_server.py +++ b/tests/agent/tools/test_mcp/real_server.py @@ -112,11 +112,24 @@ class PinnedVersionServer: 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) -> None: + 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 @@ -138,10 +151,18 @@ async def _handle(self, request: Request) -> Response: 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) diff --git a/tests/agent/tools/test_mcp/test_mcp_client.py b/tests/agent/tools/test_mcp/test_mcp_client.py index d9ccdf1dd..0917ac1c0 100644 --- a/tests/agent/tools/test_mcp/test_mcp_client.py +++ b/tests/agent/tools/test_mcp/test_mcp_client.py @@ -403,8 +403,9 @@ async def test_dropped_connection_resumes_the_persisted_session( result = await client.call_tool("test_tool", {"query": "test"}) assert result.structured_content == {"result": "Success from test_tool"} - # The handshake ran again on the new transport -- inside the same session. - assert endpoint.initialize_count == 2 + # 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")] == [ 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 index 01946b0e8..66fafbfbd 100644 --- a/tests/agent/tools/test_mcp/test_mcp_client_real_http.py +++ b/tests/agent/tools/test_mcp/test_mcp_client_real_http.py @@ -216,6 +216,54 @@ async def test_legacy_resume_keeps_the_originally_negotiated_version() -> None: ) 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 diff --git a/tests/agent/tools/test_mcp/test_protocol_strategy.py b/tests/agent/tools/test_mcp/test_protocol_strategy.py index 3ad5b3d3f..5d8d89b71 100644 --- a/tests/agent/tools/test_mcp/test_protocol_strategy.py +++ b/tests/agent/tools/test_mcp/test_protocol_strategy.py @@ -22,7 +22,14 @@ import httpx2 import pytest from mcp.shared.exceptions import MCPError -from mcp.types import CONNECTION_CLOSED, INVALID_REQUEST, METHOD_NOT_FOUND +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 @@ -31,6 +38,7 @@ LegacyHandshakeStrategy, ModernDiscoveryStrategy, build_protocol_strategy, + is_session_rejected, ) from uipath_langchain.agent.tools.mcp.streamable_http import MCP_SESSION_ID @@ -327,6 +335,11 @@ def test_legacy_recovers_from_session_loss_but_not_from_bad_requests() -> 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) @@ -345,6 +358,26 @@ def test_modern_recovers_only_from_a_dropped_connection() -> None: ) +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.""" @@ -361,10 +394,119 @@ 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: