From 3590960149b8a37f1c34173431ea955f9e36d155 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 7 Sep 2026 16:27:35 +0200 Subject: [PATCH 1/3] Python: Clean up MCP HTTP resources after failed connections Register header-hook and owned HTTP client cleanup on the MCP session exit stack so failed initialization, cancellation, reconnects, and shutdown release resources without closing caller-owned clients. Preserve the request-scoping and redirect behavior from #8039 and cover shared-client declarative cache reuse and eviction. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- python/packages/core/AGENTS.md | 2 +- python/packages/core/agent_framework/_mcp.py | 12 + .../core/tests/core/test_mcp_http_auth.py | 315 ++++++++++++++++++ .../tests/test_default_mcp_tool_handler.py | 82 ++++- 4 files changed, 409 insertions(+), 2 deletions(-) create mode 100644 python/packages/core/tests/core/test_mcp_http_auth.py diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 1e90594f3c7..28c76a601c8 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -121,7 +121,7 @@ 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 connections, and closes framework-created HTTP clients. Caller-owned clients and other tools' hooks remain reusable. - **`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..c0ec3df393a 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -3265,6 +3265,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 +3322,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 +3336,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"): 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..6875ea8e9da --- /dev/null +++ b/python/packages/core/tests/core/test_mcp_http_auth.py @@ -0,0 +1,315 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +import asyncio +import contextlib +import json +from collections.abc import AsyncGenerator, AsyncIterator +from typing import Any +from unittest.mock import patch + +import httpx +import pytest + +from agent_framework import MCPStreamableHTTPTool +from agent_framework.exceptions import ToolException + + +@pytest.fixture +async def mcp_http_server() -> AsyncIterator[tuple[httpx.AsyncClient, list[httpx.Request], dict[str, list[str]]]]: + 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": {}}, + "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}]} + 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, principals): + 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): + 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() + await call + 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, failure, owned_client): + 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" 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): + 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): + 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): + 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): + 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): + 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): + 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 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 --------------------------------------------------- From d97c41ee3ffc2f4336526c4f7425085a6e5232da Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 7 Sep 2026 20:42:50 +0200 Subject: [PATCH 2/3] Python: Reset MCP setup after failed discovery Close the lifecycle-owned session when post-transport setup fails, reset discovery flags, and roll back partially loaded functions and metadata before retrying. Preserve caller-owned HTTP clients and the original setup exception. Correct the Python 3.10 cancellation expectation, explicitly type async fixture values for Zuban, and assert the joined task result instead of ignoring it. Add discovery failure, cancellation, context-entry, and paginated retry coverage. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- python/packages/core/AGENTS.md | 2 +- python/packages/core/agent_framework/_mcp.py | 73 ++++--- .../core/tests/core/test_mcp_http_auth.py | 184 ++++++++++++++++-- 3 files changed, 216 insertions(+), 43 deletions(-) diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 28c76a601c8..71bad8637da 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -121,7 +121,7 @@ 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` and strip injected headers from cross-origin redirects. The session exit stack removes its request hook after transport shutdown, including failed connections, and closes framework-created HTTP clients. Caller-owned clients and other tools' hooks remain reusable. +- **`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 also resets the connection and discovery flags and rolls back partial function/metadata additions so retries create a fresh session. Caller-owned clients and other tools' hooks remain reusable. - **`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 c0ec3df393a..b576f2d2def 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -1363,6 +1363,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 @@ -1499,33 +1501,52 @@ 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: + 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. diff --git a/python/packages/core/tests/core/test_mcp_http_auth.py b/python/packages/core/tests/core/test_mcp_http_auth.py index 6875ea8e9da..4f177574008 100644 --- a/python/packages/core/tests/core/test_mcp_http_auth.py +++ b/python/packages/core/tests/core/test_mcp_http_auth.py @@ -5,19 +5,22 @@ import asyncio import contextlib import json +import sys from collections.abc import AsyncGenerator, AsyncIterator -from typing import Any -from unittest.mock import patch +from typing import Any, Literal, TypeAlias +from unittest.mock import AsyncMock, patch import httpx import pytest from agent_framework import MCPStreamableHTTPTool -from agent_framework.exceptions import ToolException +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[tuple[httpx.AsyncClient, list[httpx.Request], dict[str, list[str]]]]: +async def mcp_http_server() -> AsyncIterator[MCPHTTPServer]: requests: list[httpx.Request] = [] writes: dict[str, list[str]] = {"token-a": [], "token-b": [], "token-c": []} @@ -42,7 +45,7 @@ async def handle(request: httpx.Request) -> httpx.Response: headers["mcp-session-id"] = f"session-{principal}" result = { "protocolVersion": body["params"]["protocolVersion"], - "capabilities": {"tools": {}}, + "capabilities": {"tools": {}, "prompts": {}}, "serverInfo": {"name": "auth-test", "version": "1"}, } elif method == "tools/list": @@ -60,6 +63,8 @@ async def handle(request: httpx.Request) -> httpx.Response: 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}) @@ -89,7 +94,9 @@ def _calls(requests: list[httpx.Request]) -> list[httpx.Request]: @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, principals): +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) @@ -128,7 +135,7 @@ async def test_shared_client_keeps_tools_and_caller_requests_isolated(mcp_http_s assert "Authorization" not in requests[-1].headers -async def test_closing_another_tool_does_not_skip_inflight_request_hooks(mcp_http_server): +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() @@ -147,14 +154,18 @@ async def pause_call(request: httpx.Request) -> None: await first.close() finally: release.set() - await call + 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, failure, owned_client): +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") @@ -180,7 +191,7 @@ async def transport(**kwargs: Any) -> AsyncGenerator[tuple[()]]: raise RuntimeError("transport entry failed") yield () - error = asyncio.CancelledError if failure == "cancellation" else ToolException + error = asyncio.CancelledError if failure == "cancellation" and sys.version_info >= (3, 11) else ToolException transport_patch = ( contextlib.nullcontext() if failure == "initialize" @@ -195,7 +206,7 @@ async def transport(**kwargs: Any) -> AsyncGenerator[tuple[()]]: await tool.close() -async def test_owned_client_is_closed_after_successful_session(mcp_http_server): +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( @@ -211,7 +222,9 @@ async def test_owned_client_is_closed_after_successful_session(mcp_http_server): assert client.event_hooks["request"] == original_hooks -async def test_connecting_another_tool_during_a_call_does_not_capture_its_headers(mcp_http_server): +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 @@ -243,7 +256,7 @@ async def connect_second(request: httpx.Request) -> None: assert "Authorization" not in requests[-1].headers -async def test_shared_client_concurrent_calls_keep_dynamic_headers_isolated(mcp_http_server): +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( @@ -259,7 +272,9 @@ async def test_shared_client_concurrent_calls_keep_dynamic_headers_isolated(mcp_ 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): +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 @@ -289,7 +304,7 @@ async def connect_second(request: httpx.Request) -> None: await second.close() -async def test_failed_connect_removes_only_its_own_authentication_hook(mcp_http_server): +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"]) @@ -305,7 +320,7 @@ async def test_failed_connect_removes_only_its_own_authentication_hook(mcp_http_ assert not client.is_closed -async def test_prepared_transport_hook_is_removed_on_close(mcp_http_server): +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") @@ -313,3 +328,140 @@ async def test_prepared_transport_hook_is_removed_on_close(mcp_http_server): 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() From be918b77222325d7029740612e7f94a3020f2f49 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Tue, 8 Sep 2026 08:24:55 +0200 Subject: [PATCH 3/3] Python: Handle cancelled MCP connect callers and borrowed sessions Require caller acknowledgement before retaining a newly established MCP connection. Skip cancelled queued connects, clean up abandoned setup on the lifecycle owner before subsequent requests, and release idle owners without disrupting established sessions or close operations. Track session ownership so failed discovery, close, and reset preserve constructor-supplied sessions instead of opening the configured transport on retry. Cover true caller cancellation during setup/result delivery and real borrowed-session retries. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- python/packages/core/AGENTS.md | 3 +- python/packages/core/agent_framework/_mcp.py | 62 +++- .../core/tests/core/test_mcp_http_auth.py | 269 +++++++++++++++++- 3 files changed, 321 insertions(+), 13 deletions(-) diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 71bad8637da..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` 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 also resets the connection and discovery flags and rolls back partial function/metadata additions so retries create a fresh session. Caller-owned clients and other tools' hooks remain reusable. +- **`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 b576f2d2def..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. @@ -1414,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() @@ -1501,6 +1536,7 @@ 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 + self._owns_session = True else: try: if self.session._request_id == 0: # type: ignore[attr-defined] @@ -2073,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() @@ -2924,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. @@ -3122,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. @@ -3492,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 index 4f177574008..04c27676060 100644 --- a/python/packages/core/tests/core/test_mcp_http_auth.py +++ b/python/packages/core/tests/core/test_mcp_http_auth.py @@ -8,7 +8,7 @@ import sys from collections.abc import AsyncGenerator, AsyncIterator from typing import Any, Literal, TypeAlias -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, Mock, patch import httpx import pytest @@ -465,3 +465,270 @@ def create_client(**kwargs: Any) -> httpx.AsyncClient: 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()