diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 1e90594f3c7..aa850c1bb5c 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -121,7 +121,8 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID - **`allowed_tools`** (constructor arg on all `MCPTool` subclasses) - Restricts exposed MCP tools by raw remote MCP tool identity. Prefixed local names remain accepted only when the raw remote name already matches its normalized form; normalized/local aliases do not authorize a different raw remote name. If multiple raw remote tool names map to the same local function name, tool loading raises `ToolExecutionException` instead of first-one-wins shadowing. - **Progressive MCP disclosure** (`use_progressive_disclosure`, `always_load`) - When enabled on any `MCPTool` subclass, the initial model-facing surface is loader tools (`list_mcp_tools` / `load_tool` / `unload_tool`, prefixed by `tool_name_prefix` when configured) plus allowed tools selected by `always_load` and tools loaded earlier on the same `MCPTool` instance. `list_mcp_tools` only reports tools that pass `allowed_tools`; filtered tools are not listed or loadable. Loader tool names are reserved in progressive mode: remote MCP tools whose local generated name collides with a loader name are omitted from the initial/listed surface, and explicit `load_tool` calls return a model-visible message pointing callers to `tool_name_prefix` or excluding the colliding tool. `load_tool` accepts one tool name or a list of tool names and uses `FunctionInvocationContext.add_tools(...)` so the selected generated MCP `FunctionTool`s become available on the next function-calling iteration while keeping existing approval mode, argument filtering, header-provider runtime kwargs, result parsing, OTel, and task behavior. `unload_tool` accepts one dynamically loaded tool name or a list of names and removes them from the live tool list and persisted progressive surface, but it does not remove tools configured in `always_load`. Invalid `always_load` entries are ignored like unmatched `allowed_tools` entries. - **`additional_tool_argument_names`** (constructor arg on all `MCPTool` subclasses) - Opt extra argument names back into the allowlist. Accepts a `Sequence[str]` (applied to every tool) or a `Mapping[str, Sequence[str]]` keyed by **remote tool name**, where the reserved key `"*"` denotes global extras. It is configured only in user code at construction; there is **no per-call/runtime override**, so a model-issued tool call cannot change which names pass through — but note this constrains the *model*, not the *server*, which still widens the effective allowlist through its schema. To use a server that accepts `additionalProperties: true`, list the extra names here and then either (1) manually extend that tool's `inputSchema` (via the `.functions` list after connecting) so the model is prompted to supply them, or (2) supply the values yourself via `function_invocation_kwargs`. If a normal forwarded argument name is supplied by both the model and `function_invocation_kwargs`, the model-supplied value wins; `_meta` is the exception and only trusted runtime/caller metadata is used. -- **`header_provider` request scoping** - When sharing an `http_client`, keep provider processing scoped to the originating `MCPStreamableHTTPTool`, remove its request hook on `close()`, and strip injected headers from cross-origin redirects. +- **`header_provider` request scoping** - When sharing an `http_client`, keep provider processing scoped to the originating `MCPStreamableHTTPTool` and strip injected headers from cross-origin redirects. The session exit stack removes its request hook after transport shutdown, including failed initialization or discovery, and closes framework-created HTTP clients. Failed discovery resets the connection and discovery flags and rolls back partial function/metadata additions. Framework-created sessions are discarded; constructor-supplied sessions remain caller-owned and reusable across cleanup, close, and reset. +- **MCP lifecycle caller cancellation** - The lifecycle owner skips cancelled queued connect requests and waits for the caller to acknowledge successful setup. If the caller cancels before accepting a newly established connection, the owner tears it down before processing the next request. Cancelling a close waiter does not interrupt teardown, and cancelling a redundant connect does not discard a previously established session. - **`function_invocation_kwargs` and MCP servers** - That dict is shared across every tool in the run, including every attached `MCPTool`, and any name in it reaches a server that declares a matching `inputSchema` property. `header_provider` does not mitigate this — it reads the kwargs without consuming them. To keep a credential out of tool arguments, source it outside `function_invocation_kwargs`: read a `ContextVar` inside the provider (this still allows a different value per request), configure a custom `http_client`, or use `env` for `MCPStdioTool`. - **Sampling guardrails** (`sampling_callback`) - Passing `client=` advertises `SamplingCapability` so the server can send `sampling/createMessage`. Because remote servers are untrusted (confused-deputy risk), the default `sampling_callback` is **deny-by-default** and applies, in order: a per-session rate limit (`sampling_max_requests`, default `_DEFAULT_SAMPLING_MAX_REQUESTS`), an approval gate (`sampling_approval_callback`), and a `maxTokens` cap (`sampling_max_tokens`, default `_DEFAULT_SAMPLING_MAX_TOKENS`). The approval callback (constructor arg on all subclasses; exported type alias `SamplingApprovalCallback`) receives the raw `CreateMessageRequestParams`, may be sync or async, and must return truthy to approve. When it is `None` (the default) every sampling request is denied; pass `lambda params: True` to restore legacy auto-approve as an explicit opt-in. Requests and denials are logged at WARNING (content is not logged). The per-session counter resets in `_reset_session_state`. - **`MCPTaskOptions`** (experimental, `MCP_LONG_RUNNING_TASKS` feature, **frozen**) - Per-tool-instance options controlling the SEP-2663 long-running task lifecycle. When the server advertises a tool with `execution.taskSupport == "required"`, `MCPTool.call_tool` transparently routes through `call_tool_as_task`, which sends an augmented `tools/call`, polls `tasks/get` until terminal, and reinterprets `tasks/result` as a normal `CallToolResult`. Instances are immutable; replace via `MCPTool.task_options = MCPTaskOptions(...)`. Fields: diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index a1604eab471..f1bf4ea4bc0 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -576,7 +576,8 @@ def __init__( MCP prompt results to a string. If you need per-function result parsing, access the ``.functions`` list after connecting and set ``result_parser`` on individual ``FunctionTool`` instances. - session: An existing MCP client session to use. + session: An existing MCP client session to use. The caller retains ownership; + closing or resetting this wrapper does not close or replace that session. request_timeout: Timeout in seconds for MCP requests. client: A chat client for sampling callbacks. sampling_approval_callback: Optional gate invoked before each server-initiated @@ -634,9 +635,12 @@ def __init__( self._lifecycle_lock = asyncio.Lock() self._lifecycle_request_lock = asyncio.Lock() self._function_load_lock = asyncio.Lock() - self._lifecycle_queue: asyncio.Queue[tuple[str, bool, bool, asyncio.Future[None]]] | None = None + self._lifecycle_queue: ( + asyncio.Queue[tuple[str, bool, bool, asyncio.Future[None], asyncio.Future[bool]]] | None + ) = None self._lifecycle_owner_task: asyncio.Task[None] | None = None self.session = session + self._owns_session = session is None self.request_timeout = request_timeout self.client = client self.sampling_approval_callback = sampling_approval_callback @@ -1248,11 +1252,30 @@ async def _run_lifecycle_owner(self) -> None: stop_error: BaseException | None = None try: while True: - action, reset, load_configured, future = await queue.get() + action, reset, load_configured, future, acknowledged = await queue.get() + if action == "connect" and future.cancelled(): + if not self.is_connected and queue.empty(): + return + continue try: if action == "connect": + previous_session = self.session + previously_connected = self.is_connected await self._connect_on_owner(reset=reset, load_configured=load_configured) + new_connection = not previously_connected or self.session is not previous_session + accepted = False + try: + if not future.done(): + future.set_result(None) + # A completed result future cannot tell us that its waiter was + # cancelled before consuming it. Await explicit caller acceptance. + accepted = await acknowledged + finally: + if not accepted and new_connection: + await self._close_on_owner() + if not accepted and new_connection and queue.empty(): + return elif action == "close": await self._close_on_owner() else: @@ -1265,6 +1288,10 @@ async def _run_lifecycle_owner(self) -> None: except Exception as ex: if not future.done(): future.set_exception(ex) + else: + logger.warning( + "MCP lifecycle action %s failed after its caller stopped waiting.", action, exc_info=ex + ) else: if not future.done(): future.set_result(None) @@ -1277,7 +1304,7 @@ async def _run_lifecycle_owner(self) -> None: finally: while True: try: - _, _, _, future = queue.get_nowait() + _, _, _, future, _ = queue.get_nowait() except asyncio.QueueEmpty: break if not future.done(): @@ -1313,8 +1340,15 @@ async def _run_on_lifecycle_owner( raise RuntimeError("MCP lifecycle owner is not available.") future = asyncio.get_running_loop().create_future() - await queue.put((action, reset, load_configured, future)) - await future + acknowledged: asyncio.Future[bool] = asyncio.get_running_loop().create_future() + await queue.put((action, reset, load_configured, future, acknowledged)) + accepted = False + try: + await future + accepted = True + finally: + if not acknowledged.done(): + acknowledged.set_result(accepted) async def _safe_close_exit_stack(self) -> BaseException | None: """Safely close the exit stack, handling unexpected cleanup failures. @@ -1363,6 +1397,8 @@ async def _close_and_check_cancelled(self, ex: BaseException) -> tuple[bool, Bas def _reset_session_state(self) -> None: self._server_capabilities = None + self._tools_loaded = False + self._prompts_loaded = False self._supports_tools = True self._supports_prompts = True self._supports_logging = None @@ -1412,7 +1448,8 @@ async def _connect_on_owner(self, *, reset: bool = False, load_configured: bool """ if reset: await self._safe_close_exit_stack() - self.session = None + if self._owns_session: + self.session = None self.is_connected = False self._reset_session_state() self._exit_stack = AsyncExitStack() @@ -1499,33 +1536,53 @@ async def _connect_on_owner(self, *, reset: bool = False, load_configured: bool logger.debug(error_msg, exc_info=True) raise ToolException(error_msg, inner_exception=ex if isinstance(ex, Exception) else None) from ex self.session = session - elif self.session._request_id == 0: # type: ignore[attr-defined] - # If the session is not initialized, we need to reinitialize it - with create_mcp_client_span("initialize", attributes=self._mcp_base_span_attributes()) as init_span: - initialize_result = await self.session.initialize() - init_span.set_attribute(OtelAttr.MCP_PROTOCOL_VERSION, initialize_result.protocolVersion) - self._set_server_capabilities(getattr(initialize_result, "capabilities", None)) - elif self._server_capabilities is None: - self._set_server_capabilities(getattr(self.session, "_server_capabilities", None)) - logger.debug("Connected to MCP server: %s", self.session) - self.is_connected = True - if load_configured and self.load_tools_flag: - if self._supports_tools: - await self.load_tools() - self._tools_loaded = True - if load_configured and self.load_prompts_flag: - if self._supports_prompts: - await self.load_prompts() - self._prompts_loaded = True - - if logger.level != logging.NOTSET and self._supports_logging is not False: + self._owns_session = True + else: try: - level_name = cast( - Any, next(level for level, value in LOG_LEVEL_MAPPING.items() if value == logger.level) - ) - await self.session.set_logging_level(level_name) - except Exception as exc: - logger.warning("Failed to set log level to %s", logger.level, exc_info=exc) + if self.session._request_id == 0: # type: ignore[attr-defined] + # If the session is not initialized, we need to reinitialize it + with create_mcp_client_span("initialize", attributes=self._mcp_base_span_attributes()) as init_span: + initialize_result = await self.session.initialize() + init_span.set_attribute(OtelAttr.MCP_PROTOCOL_VERSION, initialize_result.protocolVersion) + self._set_server_capabilities(getattr(initialize_result, "capabilities", None)) + elif self._server_capabilities is None: + self._set_server_capabilities(getattr(self.session, "_server_capabilities", None)) + except (Exception, asyncio.CancelledError): + await self._close_on_owner() + raise + functions_before_discovery = self._functions.copy() + call_meta_before_discovery = self._tool_call_meta_by_name + task_support_before_discovery = self._tool_task_support_by_name + param_names_before_discovery = self._tool_param_names_by_name + try: + logger.debug("Connected to MCP server: %s", self.session) + self.is_connected = True + if load_configured and self.load_tools_flag: + if self._supports_tools: + await self.load_tools() + self._tools_loaded = True + if load_configured and self.load_prompts_flag: + if self._supports_prompts: + await self.load_prompts() + self._prompts_loaded = True + + if logger.level != logging.NOTSET and self._supports_logging is not False: + try: + level_name = cast( + Any, next(level for level, value in LOG_LEVEL_MAPPING.items() if value == logger.level) + ) + await self.session.set_logging_level(level_name) + except Exception as exc: + logger.warning("Failed to set log level to %s", logger.level, exc_info=exc) + except (Exception, asyncio.CancelledError): + try: + await self._close_on_owner() + finally: + self._functions[:] = functions_before_discovery + self._tool_call_meta_by_name = call_meta_before_discovery + self._tool_task_support_by_name = task_support_before_discovery + self._tool_param_names_by_name = param_names_before_discovery + raise async def _sampling_request_approved(self, params: types.CreateMessageRequestParams) -> bool: """Run the configured sampling approval gate. @@ -2052,7 +2109,8 @@ async def _close_on_owner(self) -> None: await self._safe_close_exit_stack() self._exit_stack = AsyncExitStack() - self.session = None + if self._owns_session: + self.session = None self.is_connected = False self._reset_session_state() @@ -2903,7 +2961,8 @@ def __init__( access the ``.functions`` list after connecting and set ``result_parser`` on individual ``FunctionTool`` instances. request_timeout: The default timeout in seconds for all requests. - session: The session to use for the MCP connection. + session: An existing MCP client session to use. The caller retains ownership; + closing or resetting this wrapper does not close or replace that session. description: The description of the tool. approval_mode: The approval mode for the tool. This can be: - "always_require": The tool always requires approval before use. @@ -3101,7 +3160,8 @@ def __init__( access the ``.functions`` list after connecting and set ``result_parser`` on individual ``FunctionTool`` instances. request_timeout: The default timeout in seconds for all requests. - session: The session to use for the MCP connection. + session: An existing MCP client session to use. The caller retains ownership; + closing or resetting this wrapper does not close or replace that session. description: The description of the tool. approval_mode: The approval mode for the tool. This can be: - "always_require": The tool always requires approval before use. @@ -3265,6 +3325,7 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: timeout=Timeout(MCP_DEFAULT_TIMEOUT, read=MCP_DEFAULT_SSE_READ_TIMEOUT), ) self._httpx_client = http_client + self._exit_stack.push_async_callback(self._close_owned_http_client, http_client) if not hasattr(self, "_inject_headers_hook"): @@ -3321,6 +3382,9 @@ async def _inject_headers(request: Request) -> None: # ruff:ignore[unused-async self._header_hook_client = http_client if self._inject_headers_hook not in http_client.event_hooks["request"]: http_client.event_hooks["request"].append(self._inject_headers_hook) + # Register before transport entry so failed connections clean up too, + # while successful sessions keep the hook through transport shutdown. + self._exit_stack.callback(self._remove_header_hook) transport_http_client = ( _MCPHeaderScopedClient(http_client, self._header_request_owner) if http_client is not None else None @@ -3332,6 +3396,14 @@ async def _inject_headers(request: Request) -> None: # ruff:ignore[unused-async terminate_on_close=self.terminate_on_close if self.terminate_on_close is not None else True, ) + async def _close_owned_http_client(self, http_client: AsyncClient) -> None: + """Release a framework-created client without retaining it for reconnect.""" + try: + await http_client.aclose() + finally: + if self._httpx_client is http_client: + self._httpx_client = None + def _remove_header_hook(self) -> None: """Detach this tool's request hook from its HTTP client.""" if self._header_hook_client is None or not hasattr(self, "_inject_headers_hook"): @@ -3459,7 +3531,8 @@ def __init__( access the ``.functions`` list after connecting and set ``result_parser`` on individual ``FunctionTool`` instances. request_timeout: The default timeout in seconds for all requests. - session: The session to use for the MCP connection. + session: An existing MCP client session to use. The caller retains ownership; + closing or resetting this wrapper does not close or replace that session. description: The description of the tool. approval_mode: The approval mode for the tool. This can be: - "always_require": The tool always requires approval before use. diff --git a/python/packages/core/tests/core/test_mcp_http_auth.py b/python/packages/core/tests/core/test_mcp_http_auth.py new file mode 100644 index 00000000000..04c27676060 --- /dev/null +++ b/python/packages/core/tests/core/test_mcp_http_auth.py @@ -0,0 +1,734 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +import asyncio +import contextlib +import json +import sys +from collections.abc import AsyncGenerator, AsyncIterator +from typing import Any, Literal, TypeAlias +from unittest.mock import AsyncMock, Mock, patch + +import httpx +import pytest + +from agent_framework import MCPStreamableHTTPTool +from agent_framework.exceptions import ToolException, ToolExecutionException + +MCPHTTPServer: TypeAlias = tuple[httpx.AsyncClient, list[httpx.Request], dict[str, list[str]]] + + +@pytest.fixture +async def mcp_http_server() -> AsyncIterator[MCPHTTPServer]: + requests: list[httpx.Request] = [] + writes: dict[str, list[str]] = {"token-a": [], "token-b": [], "token-c": []} + + async def record_request(request: httpx.Request) -> None: + requests.append(request) + + async def handle(request: httpx.Request) -> httpx.Response: + if request.url.path == "/unrelated": + return httpx.Response(200) + principal = request.headers.get("Authorization", "") + if principal not in writes: + return httpx.Response(401) + if request.method == "GET": + return httpx.Response(405) + if request.method == "DELETE": + return httpx.Response(200) + body = json.loads(request.content) + method = body.get("method") + headers: dict[str, str] = {} + result: dict[str, Any] = {} + if method == "initialize": + headers["mcp-session-id"] = f"session-{principal}" + result = { + "protocolVersion": body["params"]["protocolVersion"], + "capabilities": {"tools": {}, "prompts": {}}, + "serverInfo": {"name": "auth-test", "version": "1"}, + } + elif method == "tools/list": + result = { + "tools": [ + { + "name": "record", + "inputSchema": {"type": "object", "properties": {"marker": {"type": "string"}}}, + } + ] + } + elif method == "tools/call": + await asyncio.sleep(0) + marker = body["params"].get("arguments", {}).get("marker") + if marker is not None: + writes[principal].append(marker) + result = {"content": [{"type": "text", "text": principal}]} + elif method == "prompts/list": + result = {"prompts": []} + if "id" not in body: + return httpx.Response(202) + return httpx.Response(200, headers=headers, json={"jsonrpc": "2.0", "id": body["id"], "result": result}) + + async with httpx.AsyncClient( + transport=httpx.MockTransport(handle), event_hooks={"request": [record_request]} + ) as client: + yield client, requests, writes + + +def _tool(client: httpx.AsyncClient, principal: str) -> MCPStreamableHTTPTool: + return MCPStreamableHTTPTool( + name=principal, + url="https://mcp.example/mcp", + http_client=client, + load_prompts=False, + header_provider=lambda kwargs: {"Authorization": kwargs.get("credential", principal)}, + ) + + +def _calls(requests: list[httpx.Request]) -> list[httpx.Request]: + return [ + request + for request in requests + if request.method == "POST" and json.loads(request.content).get("method") == "tools/call" + ] + + +@pytest.mark.parametrize("principals", [("token-a", "token-b"), ("token-b", "token-a")]) +async def test_shared_client_keeps_tools_and_caller_requests_isolated( + mcp_http_server: MCPHTTPServer, principals: tuple[str, str] +) -> None: + client, requests, writes = mcp_http_server + original_hooks = list(client.event_hooks["request"]) + first, second = (_tool(client, principal) for principal in principals) + + async with first: + await first.call_tool("record") + async with second: + await second.call_tool("record") + await first.call_tool("record", marker="first-only") + await client.get("https://mcp.example/unrelated") + assert "Authorization" not in requests[-1].headers + await first.call_tool("record") + await first.connect(reset=True) + await first.call_tool("record") + assert len(client.event_hooks["request"]) == len(original_hooks) + 1 + + calls = _calls(requests) + assert [request.headers["Authorization"] for request in calls] == [ + principals[0], + principals[1], + principals[0], + principals[0], + principals[0], + ] + assert calls[2].headers["mcp-session-id"] == f"session-{principals[0]}" + assert all( + request.headers["mcp-session-id"] == f"session-{request.headers['Authorization']}" + for request in requests + if "mcp-session-id" in request.headers + ) + assert writes[principals[0]] == ["first-only"] + assert writes[principals[1]] == [] + assert client.event_hooks["request"] == original_hooks + assert not client.is_closed + await client.get("https://mcp.example/unrelated") + assert "Authorization" not in requests[-1].headers + + +async def test_closing_another_tool_does_not_skip_inflight_request_hooks(mcp_http_server: MCPHTTPServer) -> None: + client, requests, _ = mcp_http_server + started = asyncio.Event() + release = asyncio.Event() + + async def pause_call(request: httpx.Request) -> None: + if request.method == "POST" and json.loads(request.content).get("method") == "tools/call": + started.set() + await release.wait() + + async with _tool(client, "token-a") as first: + client.event_hooks["request"].append(pause_call) + async with _tool(client, "token-b") as second: + call = asyncio.create_task(second.call_tool("record")) + try: + await asyncio.wait_for(started.wait(), timeout=5) + await first.close() + finally: + release.set() + result = await call + assert isinstance(result, list) + assert result[0].text == "token-b" + assert _calls(requests)[-1].headers["Authorization"] == "token-b" + assert pause_call in client.event_hooks["request"] + + +@pytest.mark.parametrize("failure", ["entry", "cancellation", "initialize"]) +@pytest.mark.parametrize("owned_client", [False, True]) +async def test_transport_failure_cleans_up_hooks_and_owned_client( + mcp_http_server: MCPHTTPServer, failure: str, owned_client: bool +) -> None: + client, _, _ = mcp_http_server + original_hooks = list(client.event_hooks["request"]) + tool = _tool(client, "token-a") + if owned_client: + tool = MCPStreamableHTTPTool( + name="owned", url="https://mcp.example/mcp", header_provider=lambda _: {"Authorization": "token-a"} + ) + if failure == "initialize": + tool = MCPStreamableHTTPTool( + name="invalid", + url="https://mcp.example/mcp", + http_client=None if owned_client else client, + header_provider=lambda _: {"Authorization": "invalid-token"}, + ) + + @contextlib.asynccontextmanager + async def transport(**kwargs: Any) -> AsyncGenerator[tuple[()]]: + if failure == "cancellation": + task = asyncio.current_task() + assert task is not None + task.cancel() + await asyncio.sleep(0) + raise RuntimeError("transport entry failed") + yield () + + error = asyncio.CancelledError if failure == "cancellation" and sys.version_info >= (3, 11) else ToolException + transport_patch = ( + contextlib.nullcontext() + if failure == "initialize" + else patch("agent_framework._mcp.streamable_http_client", side_effect=transport) + ) + try: + with transport_patch, patch("httpx.AsyncClient", return_value=client), pytest.raises(error): + await tool.connect() + assert client.event_hooks["request"] == original_hooks + assert client.is_closed is owned_client + finally: + await tool.close() + + +async def test_owned_client_is_closed_after_successful_session(mcp_http_server: MCPHTTPServer) -> None: + client, _, _ = mcp_http_server + original_hooks = list(client.event_hooks["request"]) + tool = MCPStreamableHTTPTool( + name="owned", + url="https://mcp.example/mcp", + load_prompts=False, + header_provider=lambda _: {"Authorization": "token-a"}, + ) + with patch("httpx.AsyncClient", return_value=client): + async with tool: + await tool.call_tool("record") + assert client.is_closed + assert client.event_hooks["request"] == original_hooks + + +async def test_connecting_another_tool_during_a_call_does_not_capture_its_headers( + mcp_http_server: MCPHTTPServer, +) -> None: + client, requests, _ = mcp_http_server + second = _tool(client, "token-b") + connected = False + + async def connect_second(request: httpx.Request) -> None: + nonlocal connected + if not connected and request.method == "POST" and json.loads(request.content).get("method") == "tools/call": + connected = True + await second.connect() + + client.event_hooks["request"].append(connect_second) + try: + async with _tool(client, "token-a") as first: + await first.call_tool("record", credential="token-c") + await second.call_tool("record") + assert [request.headers["Authorization"] for request in _calls(requests)] == ["token-c", "token-b"] + second_initializes = [ + request + for request in requests + if request.method == "POST" + and json.loads(request.content).get("method") == "initialize" + and request.headers["Authorization"] == "token-b" + ] + assert len(second_initializes) == 1 + finally: + await second.close() + assert not client.is_closed + await client.get("https://mcp.example/unrelated") + assert "Authorization" not in requests[-1].headers + + +async def test_shared_client_concurrent_calls_keep_dynamic_headers_isolated(mcp_http_server: MCPHTTPServer) -> None: + client, requests, writes = mcp_http_server + async with _tool(client, "token-a") as first, _tool(client, "token-b") as second: + await asyncio.gather( + first.call_tool("record", credential="token-c", marker="first"), + second.call_tool("record", credential="token-b", marker="second"), + ) + calls = _calls(requests) + assert {request.headers["mcp-session-id"]: request.headers["Authorization"] for request in calls} == { + "session-token-a": "token-c", + "session-token-b": "token-b", + } + assert writes == {"token-a": [], "token-b": ["second"], "token-c": ["first"]} + assert all("credential" not in json.loads(request.content)["params"]["arguments"] for request in calls) + + +async def test_headerless_transport_does_not_inherit_another_transports_credentials( + mcp_http_server: MCPHTTPServer, +) -> None: + client, requests, _ = mcp_http_server + second = MCPStreamableHTTPTool( + name="unauthenticated", url="https://mcp.example/mcp", http_client=client, load_prompts=False + ) + attempted = False + + async def connect_second(request: httpx.Request) -> None: + nonlocal attempted + if not attempted and request.method == "POST" and json.loads(request.content).get("method") == "tools/call": + attempted = True + with pytest.raises(ToolException): + await second.connect() + + client.event_hooks["request"].append(connect_second) + try: + async with _tool(client, "token-a") as first: + await first.call_tool("record") + assert attempted + assert _calls(requests)[-1].headers["Authorization"] == "token-a" + assert any( + request.method == "POST" + and json.loads(request.content).get("method") == "initialize" + and "Authorization" not in request.headers + for request in requests + ) + finally: + await second.close() + + +async def test_failed_connect_removes_only_its_own_authentication_hook(mcp_http_server: MCPHTTPServer) -> None: + client, requests, _ = mcp_http_server + async with _tool(client, "token-a") as first: + original_hooks = list(client.event_hooks["request"]) + failed = _tool(client, "invalid-token") + try: + with pytest.raises(ToolException): + await failed.connect() + assert client.event_hooks["request"] == original_hooks + await first.call_tool("record") + assert _calls(requests)[-1].headers["Authorization"] == "token-a" + finally: + await failed.close() + assert not client.is_closed + + +async def test_prepared_transport_hook_is_removed_on_close(mcp_http_server: MCPHTTPServer) -> None: + client, _, _ = mcp_http_server + original_hooks = list(client.event_hooks["request"]) + tool = _tool(client, "token-a") + tool.get_mcp_client() + tool.get_mcp_client() + await tool.close() + assert client.event_hooks["request"] == original_hooks + + +@pytest.mark.parametrize("discovery_method", ["load_tools", "load_prompts"]) +@pytest.mark.parametrize("entry_method", ["connect", "context_manager"]) +@pytest.mark.parametrize("owned_client", [False, True]) +@pytest.mark.parametrize("failure_type", [ToolExecutionException, asyncio.CancelledError]) +async def test_discovery_failure_cleans_up_resources( + mcp_http_server: MCPHTTPServer, + discovery_method: str, + entry_method: str, + owned_client: bool, + failure_type: type[ToolExecutionException] | type[asyncio.CancelledError], +) -> None: + client, _, _ = mcp_http_server + original_hooks = list(client.event_hooks["request"]) + tool = MCPStreamableHTTPTool( + name="discovery", + url="https://mcp.example/mcp", + http_client=None if owned_client else client, + header_provider=lambda _: {"Authorization": "token-a"}, + ) + failure = failure_type("discovery failed") + try: + with ( + patch("httpx.AsyncClient", return_value=client), + patch.object(tool, discovery_method, new=AsyncMock(side_effect=failure)), + pytest.raises(failure_type, match="discovery failed") as error, + ): + if entry_method == "connect": + await tool.connect() + else: + async with tool: + pytest.fail("Failed discovery must not enter the context manager") + assert error.value is failure + assert client.event_hooks["request"] == original_hooks + assert client.is_closed is owned_client + assert tool.session is None + assert not tool.is_connected + assert not tool._tools_loaded + assert not tool._prompts_loaded + finally: + await tool.close() + + +@pytest.mark.parametrize("discovery_method", ["tools/list", "prompts/list"]) +@pytest.mark.parametrize("owned_client", [False, True]) +async def test_discovery_failure_retry_starts_a_fresh_session( + discovery_method: Literal["tools/list", "prompts/list"], owned_client: bool +) -> None: + initialize_count = 0 + failure_count = 0 + clients: list[httpx.AsyncClient] = [] + + async def handle(request: httpx.Request) -> httpx.Response: + nonlocal initialize_count, failure_count + if request.method == "GET": + return httpx.Response(405) + if request.method == "DELETE": + return httpx.Response(200) + body = json.loads(request.content) + method = body.get("method") + if method == discovery_method and body.get("params", {}).get("cursor") and failure_count < 2: + failure_count += 1 + return httpx.Response( + 200, json={"jsonrpc": "2.0", "id": body["id"], "error": {"code": -32603, "message": "discovery failed"}} + ) + result: dict[str, Any] = {} + headers: dict[str, str] = {} + if method == "initialize": + initialize_count += 1 + headers["mcp-session-id"] = f"session-{initialize_count}" + result = { + "protocolVersion": body["params"]["protocolVersion"], + "capabilities": {"tools": {}, "prompts": {}}, + "serverInfo": {"name": "discovery-test", "version": "1"}, + } + elif method == "tools/list": + result = { + "tools": [ + { + "name": "record" if failure_count == 2 else "partial_record", + "inputSchema": {"type": "object", "properties": {"marker": {"type": "string"}}}, + "_meta": {"source": "discovery"}, + "execution": {"taskSupport": "optional"}, + } + ] + } + elif method == "prompts/list": + result = {"prompts": [{"name": "partial_prompt"}] if failure_count < 2 else []} + if method == discovery_method and failure_count < 2: + result["nextCursor"] = "next-page" + if "id" not in body: + return httpx.Response(202) + return httpx.Response(200, headers=headers, json={"jsonrpc": "2.0", "id": body["id"], "result": result}) + + async_client = httpx.AsyncClient + + def create_client(**kwargs: Any) -> httpx.AsyncClient: + client = async_client(transport=httpx.MockTransport(handle), **kwargs) + clients.append(client) + return client + + tool = MCPStreamableHTTPTool( + name="discovery", + url="https://mcp.example/mcp", + http_client=None if owned_client else create_client(), + header_provider=lambda _: {"Authorization": "token-a"}, + ) + from mcp.shared.exceptions import McpError + + try: + with patch("httpx.AsyncClient", side_effect=create_client): + for _ in range(2): + with pytest.raises(McpError, match="discovery failed"): + await tool.connect() + assert tool.session is None + assert not tool.is_connected + assert not tool._tools_loaded + assert not tool._prompts_loaded + assert tool.functions == [] + assert tool._tool_call_meta_by_name == {} + assert tool._tool_task_support_by_name == {} + assert tool._tool_param_names_by_name == {} + assert all(not client.event_hooks["request"] for client in clients) + assert all(client.is_closed is owned_client for client in clients) + await tool.connect() + assert tool.is_connected + assert [function.name for function in tool.functions] == ["record"] + assert initialize_count == 3 + assert len(clients) == (3 if owned_client else 1) + await tool.close() + assert all(not client.event_hooks["request"] for client in clients) + assert all(client.is_closed is owned_client for client in clients) + finally: + await tool.close() + for client in clients: + await client.aclose() + + +@pytest.mark.parametrize("blocked_method", ["initialize", "tools/list", "prompts/list"]) +@pytest.mark.parametrize("entry_method", ["connect", "context_manager"]) +@pytest.mark.parametrize("owned_client", [False, True]) +async def test_cancelled_connect_caller_releases_abandoned_resources( + mcp_http_server: MCPHTTPServer, blocked_method: str, entry_method: str, owned_client: bool +) -> None: + client, requests, _ = mcp_http_server + setup_started = asyncio.Event() + release_setup = asyncio.Event() + cleanup_finished = asyncio.Event() + + async def block_setup(request: httpx.Request) -> None: + if request.method == "POST" and json.loads(request.content).get("method") == blocked_method: + setup_started.set() + await release_setup.wait() + + client.event_hooks["request"].append(block_setup) + original_hooks = list(client.event_hooks["request"]) + tool = MCPStreamableHTTPTool( + name="cancelled-caller", + url="https://mcp.example/mcp", + http_client=None if owned_client else client, + header_provider=lambda _: {"Authorization": "token-a"}, + ) + close_on_owner = tool._close_on_owner + + async def record_cleanup() -> None: + await close_on_owner() + cleanup_finished.set() + + async def enter() -> None: + if entry_method == "connect": + await tool.connect() + else: + async with tool: + pytest.fail("Cancelled setup must not enter the context manager") + + with patch("httpx.AsyncClient", return_value=client), patch.object(tool, "_close_on_owner", record_cleanup): + caller = asyncio.create_task(enter()) + try: + await asyncio.wait_for(setup_started.wait(), timeout=5) + caller.cancel() + with pytest.raises(asyncio.CancelledError): + await caller + assert not client.is_closed + assert tool._lifecycle_owner_task is not None + assert not tool._lifecycle_owner_task.done() + + release_setup.set() + await asyncio.wait_for(cleanup_finished.wait(), timeout=5) + assert not tool.is_connected + assert tool.session is None + assert tool._lifecycle_owner_task is None + assert client.event_hooks["request"] == original_hooks + assert client.is_closed is owned_client + assert any( + request.method == "DELETE" and request.headers.get("Authorization") == "token-a" for request in requests + ) + finally: + release_setup.set() + await tool.close() + + +async def test_cancelled_queued_connect_does_not_start_transport(mcp_http_server: MCPHTTPServer) -> None: + client, requests, _ = mcp_http_server + tool = _tool(client, "token-a") + caller = asyncio.create_task(tool.connect()) + # This runs before the lifecycle owner created by connect() gets its first turn. + asyncio.get_running_loop().call_soon(caller.cancel, None) + with pytest.raises(asyncio.CancelledError): + await caller + assert tool._lifecycle_owner_task is None + await tool.close() + assert requests == [] + + +async def test_connect_cancelled_during_result_delivery_releases_session(mcp_http_server: MCPHTTPServer) -> None: + client, _, _ = mcp_http_server + original_hooks = list(client.event_hooks["request"]) + tool = _tool(client, "token-a") + cleanup_finished = asyncio.Event() + connect_on_owner = tool._connect_on_owner + close_on_owner = tool._close_on_owner + + async def cancel_before_delivery(*, reset: bool = False, load_configured: bool = True) -> None: + await connect_on_owner(reset=reset, load_configured=load_configured) + # Setup succeeds, then the caller is cancelled before consuming the completed future. + asyncio.get_running_loop().call_soon(caller.cancel, None) + + async def record_cleanup() -> None: + await close_on_owner() + cleanup_finished.set() + + with ( + patch.object(tool, "_connect_on_owner", cancel_before_delivery), + patch.object(tool, "_close_on_owner", record_cleanup), + ): + caller = asyncio.create_task(tool.connect()) + try: + with pytest.raises(asyncio.CancelledError): + await caller + await asyncio.wait_for(cleanup_finished.wait(), timeout=5) + assert tool.session is None + assert not tool.is_connected + assert tool._lifecycle_owner_task is None + assert client.event_hooks["request"] == original_hooks + finally: + await tool.close() + + +async def test_abandoned_connect_is_cleaned_before_next_caller(mcp_http_server: MCPHTTPServer) -> None: + client, requests, _ = mcp_http_server + setup_started = asyncio.Event() + release_setup = asyncio.Event() + + async def block_setup(request: httpx.Request) -> None: + if request.method == "POST" and json.loads(request.content).get("method") == "tools/list": + setup_started.set() + await release_setup.wait() + + client.event_hooks["request"].append(block_setup) + original_hooks = list(client.event_hooks["request"]) + tool = _tool(client, "token-a") + caller = asyncio.create_task(tool.connect()) + try: + await asyncio.wait_for(setup_started.wait(), timeout=5) + abandoned_session = tool.session + caller.cancel() + with pytest.raises(asyncio.CancelledError): + await caller + retry = asyncio.create_task(tool.connect()) + release_setup.set() + await asyncio.wait_for(retry, timeout=5) + assert tool.is_connected + assert tool.session is not abandoned_session + assert ( + sum( + request.method == "POST" and json.loads(request.content).get("method") == "initialize" + for request in requests + ) + == 2 + ) + assert len(client.event_hooks["request"]) == len(original_hooks) + 1 + finally: + release_setup.set() + await tool.close() + + +@pytest.mark.parametrize("discovery_method", ["load_tools", "load_prompts"]) +@pytest.mark.parametrize("entry_method", ["connect", "context_manager"]) +@pytest.mark.parametrize("failure_type", [ToolExecutionException, RuntimeError, asyncio.CancelledError]) +async def test_failed_discovery_preserves_caller_supplied_session( + mcp_http_server: MCPHTTPServer, + discovery_method: str, + entry_method: str, + failure_type: type[Exception] | type[asyncio.CancelledError], +) -> None: + client, _, _ = mcp_http_server + async with _tool(client, "token-a") as source: + supplied_session = source.session + assert supplied_session is not None + original_hooks = list(client.event_hooks["request"]) + borrowed = MCPStreamableHTTPTool( + name="borrowed", url="https://must-not-connect.example/mcp", session=supplied_session + ) + failure = failure_type("discovery failed") + expected_error = ( + ToolExecutionException + if entry_method == "context_manager" and failure_type is RuntimeError + else failure_type + ) + with patch.object( + borrowed, "get_mcp_client", Mock(side_effect=AssertionError("Unexpected transport")) + ) as transport: + try: + with ( + patch.object(borrowed, discovery_method, AsyncMock(side_effect=failure)), + pytest.raises(expected_error), + ): + if entry_method == "connect": + await borrowed.connect() + else: + async with borrowed: + pytest.fail("Failed discovery must not enter the context manager") + assert borrowed.session is supplied_session + assert not borrowed.is_connected + assert client.event_hooks["request"] == original_hooks + assert not client.is_closed + await supplied_session.send_ping() + + await borrowed.connect() + assert borrowed.is_connected + assert borrowed.session is supplied_session + await borrowed.connect(reset=True) + assert borrowed.session is supplied_session + transport.assert_not_called() + finally: + await borrowed.close() + assert borrowed.session is supplied_session + await supplied_session.send_ping() + + +async def test_cancelled_redundant_connect_keeps_existing_session(mcp_http_server: MCPHTTPServer) -> None: + client, _, _ = mcp_http_server + async with _tool(client, "token-a") as tool: + existing_session = tool.session + original_hooks = list(client.event_hooks["request"]) + connect_on_owner = tool._connect_on_owner + + async def cancel_before_delivery(*, reset: bool = False, load_configured: bool = True) -> None: + await connect_on_owner(reset=reset, load_configured=load_configured) + asyncio.get_running_loop().call_soon(caller.cancel, None) + + with patch.object(tool, "_connect_on_owner", cancel_before_delivery): + caller = asyncio.create_task(tool.connect()) + with pytest.raises(asyncio.CancelledError): + await caller + # A subsequent request also confirms that the abandoned action finished processing. + await tool.connect() + assert tool.is_connected + assert tool.session is existing_session + assert client.event_hooks["request"] == original_hooks + result = await tool.call_tool("record") + assert isinstance(result, list) + assert result[0].text == "token-a" + + +async def test_cancelled_borrowed_session_caller_can_retry(mcp_http_server: MCPHTTPServer) -> None: + client, _, _ = mcp_http_server + async with _tool(client, "token-a") as source: + supplied_session = source.session + assert supplied_session is not None + borrowed = MCPStreamableHTTPTool( + name="borrowed", url="https://must-not-connect.example/mcp", session=supplied_session + ) + setup_started = asyncio.Event() + release_setup = asyncio.Event() + load_tools = borrowed.load_tools + + async def block_discovery() -> None: + setup_started.set() + await release_setup.wait() + await load_tools() + + with ( + patch.object(borrowed, "load_tools", block_discovery), + patch.object( + borrowed, "get_mcp_client", Mock(side_effect=AssertionError("Unexpected transport")) + ) as transport, + ): + caller = asyncio.create_task(borrowed.connect()) + try: + await asyncio.wait_for(setup_started.wait(), timeout=5) + caller.cancel() + with pytest.raises(asyncio.CancelledError): + await caller + release_setup.set() + await borrowed.connect() + assert borrowed.is_connected + assert borrowed.session is supplied_session + transport.assert_not_called() + await supplied_session.send_ping() + finally: + release_setup.set() + await borrowed.close() diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index 1523f2261f0..a8329dea278 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -2,12 +2,15 @@ """Tests for ``DefaultMCPToolHandler``. -These tests exercise the real handler against a fake ``MCPStreamableHTTPTool`` +Most tests exercise the real handler against a fake ``MCPStreamableHTTPTool`` (no real MCP server, no real network) to cover the parts of the handler not exercisable through the executor stub: cache hit/miss/eviction, concurrent connect via in-flight futures, header isolation across cache keys, string-result normalisation, ``load_prompts=False`` verification, and owned-vs-caller httpx close semantics. + +The shared-client regression also exercises the real MCP SDK transport against +an in-process HTTPX mock server. """ from __future__ import annotations @@ -147,6 +150,83 @@ def _invocation( ) +@pytest.mark.parametrize("cache_max_size", [1, 2]) +async def test_shared_client_isolates_cached_authentication_and_cleans_up_hooks(cache_max_size: int) -> None: + sessions: dict[str, str] = {} + calls: list[tuple[str, str]] = [] + writes: dict[str, list[str]] = {"token-a": [], "token-b": []} + + async def caller_hook(request: httpx.Request) -> None: + request.headers["X-Caller"] = "preserved" + + async def handle(request: httpx.Request) -> httpx.Response: + assert request.headers["X-Caller"] == "preserved" + principal = request.headers.get("Authorization", "") + if principal not in writes: + return httpx.Response(401) + if request.method == "GET": + return httpx.Response(405) + if request.method == "DELETE": + return httpx.Response(200) + body = json.loads(request.content) + headers: dict[str, str] = {} + result: dict[str, Any] = {} + if body.get("method") == "initialize": + session_id = f"session-{len(sessions)}" + sessions[session_id] = principal + headers["mcp-session-id"] = session_id + result = { + "protocolVersion": body["params"]["protocolVersion"], + "capabilities": {"tools": {}}, + "serverInfo": {"name": "auth-test", "version": "1"}, + } + elif body.get("method") == "tools/list": + result = { + "tools": [ + { + "name": "search", + "inputSchema": {"type": "object", "properties": {"marker": {"type": "string"}}}, + } + ] + } + elif body.get("method") == "tools/call": + calls.append((request.headers["mcp-session-id"], principal)) + if marker := body["params"].get("arguments", {}).get("marker"): + writes[principal].append(marker) + result = {"content": [{"type": "text", "text": principal}]} + if "id" not in body: + return httpx.Response(202) + return httpx.Response(200, headers=headers, json={"jsonrpc": "2.0", "id": body["id"], "result": result}) + + async with httpx.AsyncClient( + transport=httpx.MockTransport(handle), event_hooks={"request": [caller_hook]} + ) as client: + + async def client_provider(invocation: MCPToolInvocation) -> httpx.AsyncClient: + return client + + async with DefaultMCPToolHandler(client_provider=client_provider, cache_max_size=cache_max_size) as handler: + outputs: list[str | None] = [] + for index, principal in enumerate(("token-a", "token-b", "token-a")): + result = await handler.invoke_tool( + _invocation( + headers={"Authorization": principal}, + arguments={"marker": "a-only"} if index == 2 else {}, + ) + ) + assert not result.is_error + outputs.append(result.outputs[0].text) + assert outputs == ["token-a", "token-b", "token-a"] + assert len(client.event_hooks["request"]) == 1 + cache_max_size + + assert [principal for _, principal in calls] == ["token-a", "token-b", "token-a"] + assert all(sessions[session_id] == principal for session_id, principal in calls) + assert len(sessions) == (3 if cache_max_size == 1 else 2) + assert writes == {"token-a": ["a-only"], "token-b": []} + assert client.event_hooks["request"] == [caller_hook] + assert not client.is_closed + + # ---------- Construction ---------------------------------------------------