diff --git a/docs/mcp-sdk-2-upgrade.md b/docs/mcp-sdk-2-upgrade.md new file mode 100644 index 000000000..9a9f88828 --- /dev/null +++ b/docs/mcp-sdk-2-upgrade.md @@ -0,0 +1,136 @@ +# MCP Python SDK 2.0 upgrade review + +## Outcome + +`uipath-langchain-python` now pins `mcp==2.0.0`, the latest stable MCP Python +SDK in the local upstream checkout. The lockfile resolves its new `mcp-types` +and `httpx2` dependencies. The unused `langchain-mcp-adapters==0.2.1` direct +dependency was removed; the core package does not import it. Standalone samples +that use the adapter declare it in their own `pyproject.toml` files. + +The UiPath client continues to use the SDK's low-level `ClientSession` and +Streamable HTTP transport. Its externally persisted session-ID extension is +now a small adapter around the upstream transport rather than a private copy of +the complete transport. + +## What UiPath changed in the old SDK 1.26 transport copy + +The private `streamable_http.py` was introduced in commit `9c038fa2` and was +based on `mcp.client.streamable_http` from MCP Python SDK 1.26. Compared with +that upstream implementation, UiPath added: + +- An asynchronous `SessionInfo` abstraction whose `get_session_id()` and + `set_session_id()` methods can be overridden to load and save AgentHub debug + state. +- Asynchronous request-header preparation so every request can load the latest + externally stored session ID. +- Persistence of the `mcp-session-id` returned by an initialization response. +- A `session_info` argument on the local context manager, replacing the old + transport's session-ID callback shape. +- Raw response-body logging for HTTP error responses. + +The resulting file duplicated roughly 800 lines of SDK transport code. That +made fixes and new protocol behavior in upstream Streamable HTTP unavailable +without manually merging the copy. + +## How Streamable HTTP evolved in SDK 2.0 + +SDK 2.0's upstream transport now owns substantially more behavior than the +1.26 copy, including: + +- Legacy initialization and 2026 modern-protocol routing. +- `Mcp-Protocol-Version`, `Mcp-Method`, and `Mcp-Name` headers. +- Correct JSON-RPC errors from non-2xx response bodies and request-scoped + fallback errors when a body is absent. +- SSE resumption with `Last-Event-ID` and bounded reconnection. +- 2026 HTTP cancellation by aborting the in-flight request POST. +- GET channel and DELETE session lifecycle handling. +- Per-request error delivery rather than transport-wide exception groups. + +The UiPath adapter now delegates all of this to +`mcp.client.streamable_http.streamable_http_client`. Two `httpx2` event hooks +provide the UiPath-specific behavior: + +1. Before a request, asynchronously load `SessionInfo` and set or remove + `mcp-session-id`. +2. After a response, persist a returned `mcp-session-id` through `SessionInfo`. + +This retains compatibility with `SessionInfoDebugState` in +`uipath-agents-python` without forking the transport again. + +The old raw error-body logging was deliberately not recreated. SDK 2.0 now +parses a JSON-RPC error carried by a non-2xx response and surfaces its message +through `MCPError`; logging an arbitrary raw server body would add payload and +credential-leak risk without improving the structured error path. + +## SDK 2.0 breaking changes relevant here + +| SDK 1.x API | SDK 2.0 API / behavior | Upgrade action | +| --- | --- | --- | +| `McpError` | `MCPError(code, message, data)` | Updated imports, catches, construction, and tests. | +| Python model fields such as `inputSchema` and `outputSchema` | `input_schema` and `output_schema` | Updated all attribute reads. Wire JSON remains camelCase. | +| JSON-RPC root-model wrappers and `.root` | Plain discriminated message unions | The old copied transport was removed, eliminating these accesses locally. | +| `httpx` plus `httpx-sse` | `httpx2`, including SSE support | MCP connection and timeout types now use `httpx2`. | +| `timedelta` session timeout values | Seconds as `float` (or `None`) | The UiPath HTTP timeout is represented by `httpx2.Timeout`. | +| Transport `get_session_id` callback | No callback | Replaced with request/response event hooks. | +| `StreamableHTTPTransport.protocol_version` | Removed | Version handling is left to `ClientSession` and the transport. | +| Transport failures may surface through an `ExceptionGroup` | A request receives an `MCPError` | Retry logic catches `MCPError` directly. | +| Recalling `ClientSession.initialize()` could be used as local recovery logic | Initialization is idempotent per `ClientSession` | Recovery now replaces the transport and `ClientSession`, then performs a fresh handshake. | +| Experimental Tasks APIs | Removed | No UiPath code used them. | + +## Protocol-version and backward-compatibility behavior + +MCP SDK 2.0 declares these legacy handshake versions: + +- `2024-11-05` +- `2025-03-26` +- `2025-06-18` +- `2025-11-25` + +It also declares `2026-07-28` as a modern protocol version. The high-level SDK +`Client(mode="auto")` probes modern discovery and falls back to a legacy +initialization handshake. + +UiPath currently uses low-level `ClientSession.initialize()`. That method sends +the latest legacy version (`2025-11-25`) and accepts any version in the legacy +handshake set returned by the server. Therefore: + +| Server behavior | Current UiPath client | +| --- | --- | +| Negotiates `2025-03-26` | Supported and tested. | +| Negotiates `2025-06-18` | Supported and tested. | +| Negotiates `2025-11-25` | Supported and tested. | +| Supports 2026 but also accepts legacy initialize | Connects in legacy mode. | +| Supports only modern `2026-07-28` discovery | Not supported by the current low-level UiPath connection path. | + +Supporting a 2026-only server would require adopting the high-level auto mode +or reproducing its discover/adopt flow. That is a separate behavior change from +this dependency upgrade. + +## Session recovery details + +For sessions initialized in the current process, SDK 2.0 maps a bare HTTP 404 +to `MCPError(INVALID_REQUEST, "Session terminated")`. UiPath recognizes that +error and `CONNECTION_CLOSED`, closes the old connection stack, clears the +external session ID, and opens a fresh transport and `ClientSession` over the +same authenticated HTTP client. + +An externally restored session ID is not stored inside the new transport; it is +injected by the request hook. Consequently, the transport initially maps a bare +404 to `METHOD_NOT_FOUND`. UiPath disambiguates that exact bare-404 error when +an external session ID was attached, clears the stale ID, initializes a new +session, and retries. JSON-RPC `METHOD_NOT_FOUND` errors with a response body +are not retried. + +## Validation added + +The MCP tests use the real SDK 2.0 `ClientSession` and Streamable HTTP transport +over `httpx2.MockTransport`. They cover: + +- Negotiation with `2025-03-26`, `2025-06-18`, and `2025-11-25` servers. +- Session-header capture and reuse. +- Replacing the transport/session after HTTP 404 while reusing the HTTP client. +- Reuse of an externally persisted session without another initialization. +- Recovery from an expired externally persisted session. +- Retry exhaustion and non-session error classification. +- Tool listing cache/refresh, disposal/reuse, and full tool-call mapping. diff --git a/pyproject.toml b/pyproject.toml index 8d2eebf6f..7c343bd04 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,8 +21,7 @@ dependencies = [ "openinference-instrumentation-langchain>=0.1.56", "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>=0.2.0,<1.0.0", diff --git a/src/uipath_langchain/agent/tools/mcp/claude.md b/src/uipath_langchain/agent/tools/mcp/claude.md index 107c70cbb..36a5f5cb6 100644 --- a/src/uipath_langchain/agent/tools/mcp/claude.md +++ b/src/uipath_langchain/agent/tools/mcp/claude.md @@ -24,7 +24,7 @@ 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) +└── streamable_http.py # SessionInfo + thin adapter over the MCP SDK transport ``` ### Public Exports (`__init__.py`) @@ -44,43 +44,29 @@ transport helper used only by `McpClient`. ## 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 set or remove + the `mcp-session-id` header. +2. After every response, persist a returned `mcp-session-id` 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. #### SessionInfo @@ -92,46 +78,38 @@ class SessionInfo: self.session_id = session_id 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. - -#### StreamableHTTPTransport +**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. -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. +#### Upstream StreamableHTTPTransport -Key methods: - -| 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. + --- ### SessionInfoFactory @@ -160,7 +138,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,7 +148,7 @@ 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 │ ├─────────────────────────────────────────────────────────────┤ @@ -182,34 +161,34 @@ 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 # legacy handshake │ │ - _ensure_session() -> ClientSession │ -│ - _reinitialize_session() -> None │ -│ - _is_session_error(error) -> bool │ +│ - _reinitialize_session(failed_session) -> None │ +│ + is_session_error(error) -> bool │ └─────────────────────────────────────────────────────────────┘ ``` @@ -220,11 +199,16 @@ 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()` +4. Passes the `SessionInfo` to the local adapter, which opens the SDK transport +5. Creates a new `ClientSession` over those streams +6. If no ID was restored, calls `session.initialize()`; the response hook stores + the server-assigned ID +7. Reads the new session ID via `await session_info.get_session_id()` -On session reinitialization (404 retry), only steps 5-6 repeat. +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 +292,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 +305,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 +329,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 +360,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 +377,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 ───────────────────────│ @@ -430,8 +397,16 @@ The following error codes trigger automatic session reinitialization: | Code | Meaning | Source | |------|---------|--------| -| `32600` | Session terminated | HTTP 404 converted by transport | -| `-32000` | Server error | Can indicate session not found | +| `CONNECTION_CLOSED` | 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 | + +`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"`; `McpClient` treats that exact shape as +recoverable only while `SessionInfo` still contains the restored ID. Structured +JSON-RPC method errors are not retried. ## Key Implementation Details @@ -468,13 +443,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 +466,19 @@ 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 +) -> 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._connection_stack.aclose() + await self._session_info.set_session_id(None) + await self._open_connection() ``` ### 4. No `with` Statement for AsyncExitStack @@ -509,17 +490,20 @@ 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, clears the stale ID, 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,30 +537,32 @@ 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.py` | Real SDK 2 transport, legacy negotiation, persisted sessions, recovery, caching, and disposal | +| `test_mcp_tool.py` | Tool factory, schema refresh, result serialization, and error mapping | +| `test_session_info.py` | SessionInfo and SessionInfoFactory contract | ### Key Test Classes | Class | Tests | |-------|-------| -| `TestMcpClient` | Session lifecycle, 404 retry, client reuse | +| Module-level client tests | 2025 negotiation, persisted sessions, 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 most important recovery test verifies a fresh session with HTTP client reuse: ```python -# HTTP client created only ONCE (not recreated on retry) -assert mock_async_client_class.call_count == 1 - -# But session initialized TWICE -assert initialize_count[0] == 2 +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", +] ``` ## Guidelines for Changes @@ -585,10 +571,10 @@ assert initialize_count[0] == 2 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 the legacy-version, persisted-session, 404 recovery, and DELETE tests ### Adding New Factory Functions @@ -601,17 +587,17 @@ 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 +2. It 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 here; `_open_connection()` owns the transport/session stack +4. The response hook handles `set_session_id` — `_initialize_session` only reads via `get_session_id` +5. Verify recovery creates two sessions while retaining one HTTP client ### Adding New Methods to McpClient @@ -638,7 +624,7 @@ 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 + thin SDK transport adapter | | `mcp_client.py` | uipath-langchain | SessionInfoFactory + McpClient | | `mcp_tool.py` | uipath-langchain | Tool factory functions | | `__init__.py` | uipath-langchain | Public exports | @@ -649,30 +635,34 @@ 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`. This low-level UiPath path uses +`ClientSession.initialize()`, so a server that supports only modern +`2026-07-28` discovery is not supported here; the SDK high-level +`Client(mode="auto")` owns that probe/fallback behavior. ## 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..2306120c9 100644 --- a/src/uipath_langchain/agent/tools/mcp/mcp_client.py +++ b/src/uipath_langchain/agent/tools/mcp/mcp_client.py @@ -10,12 +10,16 @@ from contextlib import AsyncExitStack 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, + INVALID_REQUEST, + METHOD_NOT_FOUND, + CallToolResult, + ListToolsResult, +) from uipath._utils._ssl_context import get_httpx_client_kwargs from uipath.runtime.base import UiPathDisposableProtocol @@ -61,20 +65,18 @@ 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: httpx2.Timeout | float | None = None, max_retries: int = 1, session_info_factory: SessionInfoFactory | None = None, terminate_on_close: bool = True, @@ -93,7 +95,7 @@ def __init__( Defaults to ``SessionInfoFactory`` which returns a plain SessionInfo. """ self._config = config - self._timeout = timeout or httpx.Timeout(600) + self._timeout = timeout or httpx2.Timeout(600) self._max_retries = max_retries self._session_info_factory = session_info_factory or SessionInfoFactory() self._terminate_on_close = terminate_on_close @@ -112,15 +114,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 +144,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,59 +174,71 @@ 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, + ) + ) + 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. + """Initialize a newly-created MCP session when no persisted ID exists. 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. + a session ID from the server. MCP 2 makes this method idempotent on a + ``ClientSession``; recovery therefore creates a new session first. Requires: Client must be initialized first (_initialize_client). """ @@ -267,38 +278,74 @@ async def _ensure_session(self) -> ClientSession: if not self._client_initialized: await self._initialize_client() - return self._session # type: ignore[return-value] + if self._session is None: + raise RuntimeError("MCP client initialized without a session") + return self._session - async def _reinitialize_session(self) -> None: - """Reinitialize only the MCP session after a disconnect error. + async def _reinitialize_session( + self, failed_session: ClientSession | None = None + ) -> None: + """Replace the transport/session after a disconnect and initialize again. - 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. + 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. """ 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 failed_session is not None and self._session is not failed_session: + logger.debug( + "MCP session was already replaced by another operation" + ) + return + if self._connection_stack is not None: + await self._connection_stack.aclose() + self._connection_stack = None + self._session = None if self._session_info: await self._session_info.set_session_id(None) - await self._initialize_session() + await self._open_connection() - def _is_session_error(self, error: McpError) -> bool: - """Check if an McpError indicates a session disconnect. + @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. """ + 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 ("terminated", "expired", "invalid") + ) + ) + + async def _is_recoverable_session_error(self, error: MCPError) -> bool: + """Recognize explicit and persisted-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 and a fresh initialization is safe. + """ + if self.is_session_error(error): + return True + if error.code != METHOD_NOT_FOUND or error.message != "Not Found": + return False return ( - hasattr(error, "error") - and hasattr(error.error, "code") - and error.error.code in self.SESSION_ERROR_CODES + self._session_info is not None + and (await self._session_info.get_session_id()) is not None ) async def _execute_with_retry( @@ -309,7 +356,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 +368,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 +381,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: + is_session_error = await self._is_recoverable_session_error(e) + if is_session_error 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) retry_count += 1 continue else: @@ -408,17 +457,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..834f86018 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 — + 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, @@ -359,7 +359,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 +376,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/streamable_http.py b/src/uipath_langchain/agent/tools/mcp/streamable_http.py index fab856a96..f9805b71c 100644 --- a/src/uipath_langchain/agent/tools/mcp/streamable_http.py +++ b/src/uipath_langchain/agent/tools/mcp/streamable_http.py @@ -1,802 +1,87 @@ -"""StreamableHTTP Client Transport Module. +"""Session-aware adapter for the MCP SDK's Streamable HTTP transport. -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. +The MCP SDK owns the transport implementation. UiPath only adds asynchronous, +externally-persistable session ID storage through :class:`SessionInfo`. """ -import contextlib -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 +from typing import Any -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, +import httpx2 +from mcp.client.streamable_http import ( + streamable_http_client as sdk_streamable_http_client, ) -from typing_extensions import deprecated - -logger = logging.getLogger(__name__) - -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" - - -JSON = "application/json" -SSE = "text/event-stream" - -# Sentinel value for detecting unset optional parameters -_UNSET = object() 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 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.""" - - -class ResumptionError(StreamableHTTPError): - """Raised when resumption request is invalid.""" - - -@dataclass -class RequestContext: - """Context for a request operation.""" - - client: httpx.AsyncClient - session_message: SessionMessage - metadata: ClientMessageMetadata | None - read_stream_writer: StreamWriter - headers: dict[str, str] | None = None # Deprecated - no longer used - sse_read_timeout: float | None = None # Deprecated - no longer used - - -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 - - 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) - - headers = await self._prepare_headers() - headers[LAST_EVENT_ID] = last_event_id - - # 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) - - 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}") - - async def get_session_id(self) -> str | None: - """Get the current session ID.""" - return await self._session_info.get_session_id() - - @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. +) -> AsyncGenerator[tuple[Any, Any], None]: + """Open the SDK transport while synchronizing its session header externally. - 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 + 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. """ - 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: - 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, - ) - - try: - yield ( - read_stream, - write_stream, - ) - 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 + info = session_info or SessionInfo() + owns_client = http_client is None + client = http_client or httpx2.AsyncClient( + follow_redirects=True, + timeout=httpx2.Timeout(30, read=300), ) - # 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 + async def apply_session_id(request: httpx2.Request) -> None: + session_id = await info.get_session_id() + if session_id is None: + request.headers.pop(MCP_SESSION_ID, None) + else: + request.headers[MCP_SESSION_ID] = session_id + + async def capture_session_id(response: httpx2.Response) -> None: + session_id = response.headers.get(MCP_SESSION_ID) + if session_id is not None: + await info.set_session_id(session_id) + + 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: + yield streams + else: + async with sdk_streamable_http_client( + url, + http_client=client, + terminate_on_close=terminate_on_close, + ) as streams: + yield streams + finally: + client.event_hooks["request"].remove(apply_session_id) + client.event_hooks["response"].remove(capture_session_id) diff --git a/tests/agent/tools/test_mcp/claude.md b/tests/agent/tools/test_mcp/claude.md index 4822c7110..b4e7cb9f5 100644 --- a/tests/agent/tools/test_mcp/claude.md +++ b/tests/agent/tools/test_mcp/claude.md @@ -4,7 +4,7 @@ > > When you modify `test_mcp_client.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 +17,34 @@ 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: +The client tests run the real MCP SDK 2.0 `ClientSession` and Streamable HTTP +transport over `httpx2.MockTransport`. Only the remote endpoint is simulated. +This approach: - Tests the actual MCP protocol flow -- Validates error handling with real `McpError` exceptions -- Ensures `ClientSession.initialize()` behaves correctly when called multiple times +- Validates error handling with real `MCPError` exceptions +- Verifies that recovery replaces an idempotently initialized `ClientSession` +- Exercises persisted-session request/response hooks - Catches integration issues between our code and the SDK ## 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 +├── test_mcp_client.py # Real SDK 2 transport/session integration +│ ├── LegacyMcpEndpoint # httpx2.MockTransport request handler +│ ├── test_negotiates_supported_legacy_protocol_versions +│ ├── test_replaces_transport_and_session_after_404 ← Key test +│ ├── test_persisted_session_is_reused_without_initialize +│ ├── test_expired_persisted_session_falls_back_to_fresh_initialize +│ ├── test_max_retries_exceeded_raises_mcp_error +│ ├── test_concurrent_recovery_does_not_replace_a_new_session +│ ├── test_list_tools_cache_and_force_refresh +│ ├── test_dispose_allows_client_reuse +│ ├── test_raises_on_missing_mcp_url +│ └── test_only_session_specific_invalid_request_is_retryable +│ +├── test_session_info.py # SessionInfo + SessionInfoFactory contract │ └── test_mcp_tool.py # Tool factory tests (17 tests) ├── TestMcpToolMetadata (class) @@ -64,8 +67,8 @@ tests/agent/tools/test_mcp/ │ ├── 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 @@ -105,29 +108,35 @@ 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`). +(`test_list_tools_cache_and_force_refresh`; disposal/reuse is covered separately). ## 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 +146,7 @@ async def test_something(self, mock_async_client_class): ├─────────────────────────────────────────────────────────────┤ │ │ │ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ │ -│ │ McpClient │ ──► │ MCP SDK │ ──► │ HTTP Mock │ │ +│ │ McpClient │ ──► │ MCP SDK 2 │ ──► │ MockTransport│ │ │ │ │ │ (real) │ │ (mocked) │ │ │ └─────────────┘ └─────────────┘ └─────────────┘ │ │ ▲ │ │ │ @@ -148,62 +157,72 @@ 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, + ) -> None: + self.protocol_version = protocol_version + self.failed_tool_calls = failed_tool_calls + self.methods: list[str] = [] + self.request_headers: list[tuple[str, httpx2.Headers]] = [] + 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: ... ``` ### 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 | Returns selected legacy version and a new ID | +| `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 +231,80 @@ 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_negotiates_supported_legacy_protocol_versions -**Purpose:** Verify lazy initialization on first `call_tool()` +Parameterizes `2025-03-26`, `2025-06-18`, and `2025-11-25`. It verifies the +real SDK accepts each server-selected handshake version and stamps the selected +version plus session ID on the subsequent tool request. -**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 -``` - -#### test_session_reused_across_calls +#### test_replaces_transport_and_session_after_404 ⭐ -**Purpose:** Verify session persists across multiple tool calls +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 -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 +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_reinitializes_on_404_error ⭐ +#### Persisted-session tests -**Purpose:** THE KEY TEST - verify client reuse on session reinit +`test_persisted_session_is_reused_without_initialize` verifies an ID restored +through a custom `SessionInfoFactory` is injected into `tools/call` without a +new handshake. -**Setup:** -```python -MockStreamResponse = self.create_mock_stream_response( - ..., - fail_first_tool_call=True, # First tools/call returns 404 -) -``` - -**Critical Assertions:** -```python -# Session was reinitialized (initialize called twice) -assert initialize_count[0] == 2 - -# Tool call was retried -assert tool_call_count[0] == 2 - -# Session ID changed -assert session.session_id == "test-session-retry" - -# KEY: HTTP client created only ONCE (not recreated) -assert mock_async_client_class.call_count == 1 -``` - -#### test_max_retries_exceeded +`test_expired_persisted_session_falls_back_to_fresh_initialize` verifies the +special bare-404 path: the SDK transport does not know an externally injected +ID, so UiPath recognizes `METHOD_NOT_FOUND`/`"Not Found"` while the ID is still +present, clears it, initializes a new session, and retries. -**Purpose:** Verify `McpError` is raised after max retries +#### Retry and concurrency tests -**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_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. -#### test_dispose_releases_resources +#### Cache, disposal, and configuration tests -**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 -``` - -#### test_client_initialized_property - -**Purpose:** Verify `is_client_initialized` property accuracy - -**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 -``` - -#### test_session_can_be_reused_after_dispose - -**Purpose:** Verify session can be fully reinitialized after `dispose()` - -**Assertions:** -```python -await session.call_tool(...) -await session.dispose() -await session.call_tool(...) # Should work! - -# HTTP client created TWICE (once before dispose, once after) -assert mock_async_client_class.call_count == 2 -``` +- `test_list_tools_cache_and_force_refresh` verifies normal caching and explicit + refresh over the real protocol path. +- `test_dispose_allows_client_reuse` verifies disposal resets the state and a + later call creates another HTTP client/session stack. +- `test_raises_on_missing_mcp_url` verifies endpoint validation happens before + HTTP resources are allocated. ### TestCreateMcpToolsFromAgent Tests @@ -432,47 +382,50 @@ for tool in tools: ## Guidelines for Adding New Tests -### 1. Use the Factory Methods +### 1. Use the Shared Endpoint and Client Context -Always use the provided factory methods: +Use `LegacyMcpEndpoint` and `configured_client` so tests keep the real SDK +transport/session path: ```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 +433,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 +446,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 +489,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 +523,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 +567,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 +576,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 +586,9 @@ def _build_response(self): | File | Purpose | |------|---------| -| `test_mcp_client.py` | McpClient session tests (7 tests) | -| `test_mcp_tool.py` | Tool factory tests (17 tests) | +| `test_mcp_client.py` | SDK 2 transport, legacy versions, session persistence/recovery, caching, disposal | +| `test_session_info.py` | Async session ID store and factory | +| `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/test_mcp_client.py b/tests/agent/tools/test_mcp/test_mcp_client.py index dc3d117a3..51dae67b6 100644 --- a/tests/agent/tools/test_mcp/test_mcp_client.py +++ b/tests/agent/tools/test_mcp/test_mcp_client.py @@ -1,883 +1,395 @@ -"""Tests for McpClient class.""" +"""Tests for the MCP 2 Streamable HTTP client integration.""" 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 httpx2 import pytest +from mcp.shared.exceptions import MCPError +from mcp.types import 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( +class LegacyMcpEndpoint: + """Small Streamable HTTP endpoint used to exercise the real MCP SDK transport.""" + + def __init__( 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 - - 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}" - ) - 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, + protocol_version: str = "2025-11-25", + *, + failed_tool_calls: int = 0, + ) -> None: + self.protocol_version = protocol_version + self.failed_tool_calls = failed_tool_calls + self.methods: list[str] = [] + self.request_headers: list[tuple[str, httpx2.Headers]] = [] + 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: + """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 + return self._json_response( + body["id"], + { + "protocolVersion": self.protocol_version, + "capabilities": {"tools": {}}, + "serverInfo": {"name": "test-server", "version": "1.0.0"}, + }, + headers={"mcp-session-id": f"session-{self.initialize_count}"}, + ) + if method == "notifications/initialized": + return httpx2.Response(202) + if method == "tools/list": + return self._json_response( + body["id"], + { + "tools": [ { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "protocolVersion": "2025-06-18", - "capabilities": {"tools": {}}, - "serverInfo": { - "name": "test-server", - "version": "1.0.0", - }, + "name": "test_tool", + "description": "A test tool", + "inputSchema": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], }, - }, - {"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"} - }, - }, - } - ], + "outputSchema": { + "type": "object", + "properties": {"result": {"type": "string"}}, }, - }, - {}, - ) - - 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, - }, - }, - {}, - ) - - 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() - 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 - ) - - 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() - - @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. - - 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 - """ - 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 - ): - """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 + } + ] + }, + ) + if method == "tools/call": + self.tool_call_count += 1 + if self.tool_call_count <= self.failed_tool_calls: + 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, + }, + ) + return httpx2.Response( + 404, + json={ + "jsonrpc": "2.0", + "id": body.get("id"), + "error": {"code": METHOD_NOT_FOUND, "message": "Method not found"}, + }, ) - 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 + @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 - - 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() - - # 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 + 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: 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 list_tools lazily initializes session and returns tools.""" - 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 +@pytest.mark.parametrize("protocol_version", ["2025-03-26", "2025-06-18", "2025-11-25"]) +async def test_negotiates_supported_legacy_protocol_versions( + protocol_version: str, + mcp_resource_config: AgentMcpResourceConfig, + mock_uipath_sdk: MagicMock, +) -> None: + """MCP 2's low-level initialize handshake remains compatible with 2025 servers.""" + endpoint = LegacyMcpEndpoint(protocol_version) + 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 == 1 + assert endpoint.tool_call_count == 1 + assert await client.get_session_id() == "session-1" + assert endpoint.headers_for("tools/call")[0]["mcp-session-id"] == "session-1" + assert ( + endpoint.headers_for("tools/call")[0]["mcp-protocol-version"] + == protocol_version ) - mock_http_client = self.create_mock_http_client(MockStreamResponse) - mock_async_client_class.return_value = mock_http_client - - client = McpClient(config=mcp_resource_config) - - assert not client.is_client_initialized - - with patch( - "uipath.platform.UiPath", - return_value=mock_uipath_sdk, - ): - result = await client.list_tools() - - # 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 +@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_persisted_session_is_reused_without_initialize( + mcp_resource_config: AgentMcpResourceConfig, + mock_uipath_sdk: MagicMock, +) -> None: + """The UiPath SessionInfo extension injects an externally restored session ID.""" + endpoint = LegacyMcpEndpoint() + 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: + await client.call_tool("test_tool", {"query": "test"}) + + assert endpoint.initialize_count == 0 + assert endpoint.headers_for("tools/call")[0]["mcp-session-id"] == ( + "persisted-session" ) - 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 +@pytest.mark.asyncio +async def test_expired_persisted_session_falls_back_to_fresh_initialize( + mcp_resource_config: AgentMcpResourceConfig, + mock_uipath_sdk: MagicMock, +) -> None: + """A 404 for an externally restored session is treated as session expiry.""" + endpoint = LegacyMcpEndpoint(failed_tool_calls=1) + 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: + await client.call_tool("test_tool", {"query": "test"}) + + assert endpoint.initialize_count == 1 + assert endpoint.tool_call_count == 2 + assert await client.get_session_id() == "session-1" + assert [h["mcp-session-id"] for h in endpoint.headers_for("tools/call")] == [ + "expired-session", + "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") + client._open_connection = AsyncMock() + + await client._reinitialize_session(failed_session) + + assert client._session is replacement_session + assert await client.get_session_id() == "replacement-id" + client._open_connection.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_list_tools_cache_and_force_refresh( + mcp_resource_config: AgentMcpResourceConfig, + mock_uipath_sdk: MagicMock, +) -> None: + """Tool discovery still caches normally across the SDK upgrade.""" + endpoint = LegacyMcpEndpoint() + async with configured_client( + mcp_resource_config, mock_uipath_sdk, endpoint + ) as client: + first = await client.list_tools() + second = await client.list_tools() + refreshed = await client.list_tools(force_refresh=True) - 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, - ): - 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 - + assert refreshed.tools[0].input_schema["required"] == ["query"] + assert endpoint.methods.count("tools/list") == 2 + + +@pytest.mark.asyncio +async def test_dispose_allows_client_reuse( + mcp_resource_config: AgentMcpResourceConfig, + mock_uipath_sdk: MagicMock, +) -> None: + """Disposal closes resources and a later call builds a new client/session.""" + endpoint = LegacyMcpEndpoint() + async with configured_client( + mcp_resource_config, mock_uipath_sdk, endpoint + ) as client: + await client.call_tool("test_tool", {"query": "first"}) 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) + assert not client.is_client_initialized + assert client._session is None + assert await client.get_session_id() is None - with patch("uipath.platform.UiPath", return_value=mock_sdk): - await session.call_tool("test_tool", {"query": "test"}) + await client.call_tool("test_tool", {"query": "second"}) + assert endpoint.initialize_count == 2 + assert client.is_client_initialized - mock_sdk.mcp.retrieve_async.assert_called_once_with( - name="test_server", - folder_path="/Shared/TestFolder", - ) - await session.dispose() +@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"}) + + +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=INVALID_REQUEST, message="Session terminated") + ) + assert not McpClient.is_session_error( + MCPError(code=INVALID_REQUEST, message="Invalid request parameters") + ) diff --git a/tests/agent/tools/test_mcp/test_mcp_tool.py b/tests/agent/tools/test_mcp/test_mcp_tool.py index b6a18fd82..73f002150 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.""" @@ -676,7 +395,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 +405,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 +415,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 +433,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 +963,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/uv.lock b/uv.lock index d47c30bbc..07506d8b8 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. +exclude-newer = "2026-08-04T10:26:42.895483Z" exclude-newer-span = "P2D" [options.exclude-newer-package] @@ -1537,6 +1537,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" @@ -1576,6 +1589,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" @@ -1985,20 +2014,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" @@ -2303,15 +2318,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'" }, @@ -2321,9 +2336,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/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/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/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]] @@ -4508,7 +4536,6 @@ dependencies = [ { name = "jsonschema-pydantic-converter" }, { name = "langchain" }, { name = "langchain-core" }, - { name = "langchain-mcp-adapters" }, { name = "langgraph" }, { name = "langgraph-checkpoint-sqlite" }, { name = "mcp" }, @@ -4569,10 +4596,9 @@ requires-dist = [ { 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.56" }, { name = "pillow", specifier = ">=12.1.1" }, { name = "pydantic-settings", specifier = ">=2.6.0" },