From 0693d64379ca30cf3ec12a7db9900d955a3b70e3 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Fri, 6 Nov 2026 13:34:54 +0100 Subject: [PATCH 1/4] Python: Allow concurrent MCP calls with request-owned headers --- python/packages/core/AGENTS.md | 16 + python/packages/core/README.md | 20 + python/packages/core/agent_framework/_mcp.py | 443 ++++++++++-- .../core/tests/core/mcp_http_helpers.py | 25 + python/packages/core/tests/core/test_mcp.py | 104 +-- .../core/tests/core/test_mcp_http_auth.py | 651 +++++++++++++++++- .../core/tests/core/test_mcp_observability.py | 3 +- .../_toolbox.py | 51 +- .../foundry_hosting/tests/test_toolbox.py | 137 +++- .../references/integrations.md | 18 + 10 files changed, 1322 insertions(+), 146 deletions(-) create mode 100644 python/packages/core/tests/core/mcp_http_helpers.py diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index b66484a40ec..49461ca95ca 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -285,6 +285,22 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID - **`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. - **MCP HTTP header request scoping** - `static_headers` supplies fixed, origin-scoped headers without serializing concurrent calls. Generated tool and prompt calls give `header_provider` only host runtime kwargs, separately from model arguments; direct `call_tool` calls give it caller kwargs. Model-over-runtime precedence is confined to outbound tool arguments. Fixed and dynamic headers form the session's complete effective identity (case-insensitive names, case-sensitive values), with dynamic values overriding fixed values of the same name. Agent runs reconcile a connected tool before copying its functions; generated tool and prompt calls perform the same check at invocation. Connected run preparation defers when a provider needs runtime values supplied by invocation middleware; invocation-time resolution remains strict. A different identity waits for in-flight discovery, reconnects framework-created sessions without loading, then refreshes session-derived discovery under the same lock. Because a caller-supplied session's established identity is unknown and the wrapper cannot reconnect it, dynamic header resolution on that wrapper is rejected even while disconnected. Ambient requests retain the bound set. With a shared `http_client`, processing stays scoped to the originating `MCPStreamableHTTPTool`, and every injected header is stripped 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. +- **MCP HTTP concurrency** - MCP >=2.2 propagates the sending operation's ContextVars into each outbound HTTP task. + The owner-scoped header hook copies effective headers onto each request, including redirects; do not restore a + mutable instance-level active-header fallback. Same-identity tool, prompt, and Toolbox resource operations may + overlap on one session. Identity changes, discovery, reset, close, and recovery wait for active network attempts; + retries release their attempt before recovery, re-admit the original identity, and use its refreshed argument + allowlist and tool metadata. Failed attempts from the same replaced framework-owned connection share recovery. + Keep the original trace context when rebuilding metadata. Cancelled lifecycle waiters leave their owner running; + new HTTP operations must wait for that setup or teardown to finish. A same-generation nested callback may finish + ahead of a queued identity change, but exclusive operations from active callbacks fail instead of deadlocking. +- **MCP recovery and deadlines** - Failed tool, prompt, discovery, and liveness operations record their connection + and close generation. Central recovery reuses a replacement established by another operation, releases the + lifecycle request lock before discovery, and rejects recovery from before an explicit close. Advance the close + generation before waiting for HTTP exclusivity, including for Toolbox cleanup. Header admission/reconciliation + and connection recovery have separate `request_timeout` budgets; neither wraps the admitted RPC or cleanup. + `None` disables the budgets. Keep recovery's deadline outside the HTTP barrier so waiting for active attempts + is bounded too. - **Streamable HTTP cookies** - Framework-created clients reject response-cookie persistence with or without a `header_provider`; explicit `Cookie` headers from `static_headers` or a provider remain supported. Caller-provided clients retain their cookie behavior and ownership. Applications requiring cookie persistence must scope clients and MCP sessions to one authenticated principal. Cookie rejection does not isolate MCP protocol sessions or other server-side state. - **`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`. diff --git a/python/packages/core/README.md b/python/packages/core/README.md index b3b9c75bfc2..c876419215d 100644 --- a/python/packages/core/README.md +++ b/python/packages/core/README.md @@ -194,6 +194,26 @@ cookie-bearing clients and MCP tool sessions to one authenticated principal; do share a user-specific cookie jar across users. Rejecting cookies does not partition MCP protocol sessions or other server-side state between principals. +### MCP HTTP concurrency and timeouts + +`MCPStreamableHTTPTool` and `FoundryToolbox` attach a separate header snapshot to each +outgoing HTTP request. Calls with identical effective static and dynamic headers can +overlap on one MCP session. A different identity waits for active operations before +reconnecting and refreshing discovery. Discovery, reset, close, and recovery are +exclusive against active network attempts. Calls still share server-side session +state; applications must order state-dependent operations themselves. + +`request_timeout` bounds header admission and session reconciliation together, and +connection recovery separately from the RPC timeout. These are separate budgets, +not an end-to-end deadline. `None` disables these budgets. An admission timeout does +not cancel another caller's active operation. Cleanup may continue after a timeout +or cancellation so the lifecycle owner can finish releasing resources safely. +Caller-owned MCP clients retain their own RPC timeout policy and are not recreated. + +Concurrent failures of the same connection share its replacement. Calling `close()` +interrupts older recovery so it cannot reopen the connection or retry the failed +operation after shutdown. A later explicit `connect()` remains supported. + ## 5. Multi-Agent Orchestration Coordinate multiple agents to collaborate on complex tasks using orchestration patterns: diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 4a80154f55b..533b4eb1e5e 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -10,8 +10,8 @@ import sys import warnings from abc import abstractmethod -from collections.abc import Awaitable, Callable, Collection, Coroutine, Mapping, Sequence -from contextlib import AsyncExitStack, _AsyncGeneratorContextManager # type: ignore +from collections.abc import AsyncGenerator, Awaitable, Callable, Collection, Coroutine, Generator, Mapping, Sequence +from contextlib import AsyncExitStack, _AsyncGeneratorContextManager, asynccontextmanager, contextmanager # type: ignore from contextvars import ContextVar from copy import copy from dataclasses import dataclass @@ -26,6 +26,7 @@ from typing_extensions import deprecated # pragma: no cover +from opentelemetry import context as otel_context from opentelemetry import propagate from opentelemetry import trace as otel_trace from pydantic_core import to_jsonable_python @@ -247,6 +248,7 @@ class MCPSpecificApproval(TypedDict, total=False): ) _MCP_HEADER_OWNER_EXTENSION = "agent_framework.mcp_header_owner" _MCP_INJECTED_HEADER_KEYS_EXTENSION = "agent_framework.mcp_injected_header_keys" +_MCP_REQUEST_HEADERS_EXTENSION = "agent_framework.mcp_request_headers" _MCPHeaderIdentity: TypeAlias = tuple[tuple[str, str], ...] _MCP_DEFAULT_READ_TIMEOUT_SECONDS = 30 MCP_DEFAULT_SSE_READ_TIMEOUT = 60 * 5 @@ -258,6 +260,20 @@ class MCPSpecificApproval(TypedDict, total=False): _DEFAULT_MCP_HOST_PAYLOAD_SIZE_BYTES = 1024 * 1024 +@dataclass +class _MCPHTTPInvocation: + """Track a live network attempt, including callbacks that inherit its context.""" + + owner: object + parent: _MCPHTTPInvocation | None + runtime_kwargs: Mapping[str, Any] + connection: _MCPConnection | None = None + active: bool = False + + +_mcp_http_invocation: ContextVar[_MCPHTTPInvocation | None] = ContextVar("_mcp_http_invocation", default=None) + + @dataclass class _MCPHostPayloadCapture: """Collect one generated MCP result without affecting its public return shape.""" @@ -1247,6 +1263,10 @@ def __init__( self._exit_stack = AsyncExitStack() self._lifecycle_lock = asyncio.Lock() self._lifecycle_request_lock = asyncio.Lock() + self._close_generation = 0 + # Cancelling a lifecycle waiter does not stop the owner's setup or teardown. + self._lifecycle_idle = asyncio.Event() + self._lifecycle_idle.set() self._function_load_lock = asyncio.Lock() self._lifecycle_queue: ( asyncio.Queue[tuple[str, bool, bool, bool, asyncio.Future[None], asyncio.Future[bool]]] | None @@ -2154,6 +2174,8 @@ async def _run_lifecycle_owner(self) -> None: if action == "connect" and future.cancelled(): if not self.is_connected and queue.empty(): return + if queue.empty(): + self._lifecycle_idle.set() continue try: @@ -2211,6 +2233,9 @@ async def _run_lifecycle_owner(self) -> None: else: if not future.done(): future.set_result(None) + finally: + if queue.empty(): + self._lifecycle_idle.set() if action == "close": return @@ -2228,6 +2253,7 @@ async def _run_lifecycle_owner(self) -> None: self._lifecycle_queue = None self._lifecycle_owner_task = None + self._lifecycle_idle.set() def _is_lifecycle_owner_task(self) -> bool: owner_task = self._lifecycle_owner_task @@ -2262,6 +2288,7 @@ async def _run_on_lifecycle_owner( future = asyncio.get_running_loop().create_future() acknowledged: asyncio.Future[bool] = asyncio.get_running_loop().create_future() + self._lifecycle_idle.clear() await queue.put((action, reset, load_configured, reset_discovery, future, acknowledged)) accepted = False try: @@ -2370,9 +2397,13 @@ async def _recover_connection_after_loss( operation: str, inner_exception: Exception, load_configured: bool, + failed_connection: _MCPConnection | None = None, + close_generation: int | None = None, ) -> None: - """Recreate a framework-owned connection or fail clearly for a caller-owned connection.""" - connection = self._connection + """Bound recovery and prevent failed operations from undoing an explicit close.""" + connection = failed_connection if failed_connection is not None else self._connection + if close_generation is None: + close_generation = self._close_generation if connection is not None and not connection.is_framework_owned: connection_type = "Client" if connection.client is not None else "ClientSession" raise _mcp_tool_execution_exception( @@ -2381,18 +2412,67 @@ async def _recover_connection_after_loss( inner_exception=inner_exception, ) from inner_exception - logger.info("MCP connection was lost during %s. Reconnecting...", operation) try: - if load_configured: - await self.connect(reset=True) - else: - await self._reconnect_without_loading() + async with asyncio.timeout(self.request_timeout): + await self._recover_failed_connection( + operation=operation, + failed_connection=connection, + close_generation=close_generation, + load_configured=load_configured, + ) + except TimeoutError as timeout_ex: + raise _mcp_tool_execution_exception( + f"Timed out recovering the MCP connection after {operation}.", + inner_exception=timeout_ex, + ) from timeout_ex + except ToolExecutionException: + raise except Exception as reconnect_ex: raise _mcp_tool_execution_exception( f"Failed to reconnect to MCP server after {operation}.", inner_exception=reconnect_ex, ) from reconnect_ex + async def _recover_failed_connection( + self, + *, + operation: str, + failed_connection: _MCPConnection | None, + close_generation: int, + load_configured: bool, + ) -> None: + """Reuse a replacement connection or reconnect under the lifecycle lock.""" + if self._is_lifecycle_owner_task(): + self._check_recovery_generation(close_generation) + if not self.is_connected or self._connection is failed_connection: + await self._reconnect_without_loading() + else: + async with self._lifecycle_request_lock: + self._check_recovery_generation(close_generation) + if not self.is_connected or self._connection is failed_connection: + logger.info("MCP connection was lost during %s. Reconnecting...", operation) + await self._reconnect_without_loading() + + # Discovery can recover too; holding the lifecycle request lock here would deadlock. + self._check_recovery_generation(close_generation) + if load_configured: + async with self._function_load_lock: + self._check_recovery_generation(close_generation) + if (self.load_tools_flag and not self._tools_loaded) or ( + self.load_prompts_flag and not self._prompts_loaded + ): + try: + await self._load_configured_discovery_locked() + except (Exception, asyncio.CancelledError): + await self.close() + self._reset_session_discovery_state() + raise + self._check_recovery_generation(close_generation) + + def _check_recovery_generation(self, close_generation: int) -> None: + if close_generation != self._close_generation: + raise _mcp_tool_execution_exception("MCP connection recovery was interrupted by close().") + def _is_connection_loss(self, exception: Exception) -> bool: """Return whether an exception proves that the active MCP connection ended.""" from anyio import ClosedResourceError @@ -3114,6 +3194,8 @@ async def _try_list_page_with_recovery( catalog_name = f"{kind}s" operation = f"{kind} discovery" + failed_connection = self._connection + close_generation = self._close_generation async def recover_connection(retry_state: RetryCallState) -> None: outcome = retry_state.outcome @@ -3126,6 +3208,8 @@ async def recover_connection(retry_state: RetryCallState) -> None: operation=operation, inner_exception=exception, load_configured=False, + failed_connection=failed_connection, + close_generation=close_generation, ) @retry( @@ -3135,7 +3219,10 @@ async def recover_connection(retry_state: RetryCallState) -> None: reraise=True, ) async def request_page_with_retry() -> _MCPListResultT | None: + nonlocal failed_connection, close_generation await self._ensure_connected() + failed_connection = self._connection + close_generation = self._close_generation # on reconnect we may have different capabilities. is_supported = self._supports_tools if kind == "tool" else self._supports_prompts @@ -3377,6 +3464,10 @@ async def close(self) -> None: Closes the connection and cleans up resources. """ + self._close_generation += 1 + await self._close_connection() + + async def _close_connection(self) -> None: if self._is_lifecycle_owner_task(): await self._close_on_owner() return @@ -3414,8 +3505,10 @@ async def _ensure_connected(self) -> None: if not self._ping_available: return + failed_connection = self._connection + close_generation = self._close_generation try: - await self._require_connection().check_connection() + await self._check_connection() except MCPError as mcp_exc: if mcp_exc.error.code == -32601: self._ping_available = False @@ -3426,6 +3519,8 @@ async def _ensure_connected(self) -> None: operation="connection liveness check", inner_exception=mcp_exc, load_configured=False, + failed_connection=failed_connection, + close_generation=close_generation, ) return raise _mcp_tool_execution_exception( @@ -3437,6 +3532,8 @@ async def _ensure_connected(self) -> None: operation="connection liveness check", inner_exception=connection_ex, load_configured=False, + failed_connection=failed_connection, + close_generation=close_generation, ) except Exception as connection_ex: connection = self._connection @@ -3445,6 +3542,8 @@ async def _ensure_connected(self) -> None: operation="connection liveness check", inner_exception=connection_ex, load_configured=False, + failed_connection=failed_connection, + close_generation=close_generation, ) return raise _mcp_tool_execution_exception( @@ -3452,6 +3551,9 @@ async def _ensure_connected(self) -> None: inner_exception=connection_ex, ) from connection_ex + async def _check_connection(self) -> None: + await self._require_connection().check_connection() + async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]: """Call a tool with the given arguments. @@ -3490,9 +3592,8 @@ async def _call_tool( raise _mcp_tool_execution_exception( "Tools are not loaded for this server, please set load_tools=True in the constructor." ) - filtered_kwargs, meta = self._prepare_call_kwargs(tool_name, kwargs) - parser = self.parse_tool_results or self._parse_tool_result_from_mcp + trace_context = otel_context.get_current() # Build MCP span attributes for tools/call mcp_span_attrs = self._mcp_base_span_attributes() @@ -3508,17 +3609,27 @@ async def _call_tool( ) as span: return await self._call_tool_with_retries( tool_name, - filtered_kwargs, - meta, + kwargs, + trace_context, parser, span, ) + async def _invoke_tool( + self, tool_name: str, kwargs: dict[str, Any], trace_context: otel_context.Context + ) -> types.CallToolResult: + token = otel_context.attach(trace_context) + try: + filtered_kwargs, meta = self._prepare_call_kwargs(tool_name, kwargs) + finally: + otel_context.detach(token) + return await self._require_connection().call_tool(tool_name, filtered_kwargs, meta=meta) + async def _call_tool_with_retries( self, tool_name: str, - filtered_kwargs: dict[str, Any], - meta: types.RequestParamsMeta | None, + kwargs: dict[str, Any], + trace_context: otel_context.Context, parser: Callable[..., str | list[Content]], span: otel_trace.Span, ) -> str | list[Content]: @@ -3527,13 +3638,11 @@ async def _call_tool_with_retries( from mcp import MCPError for attempt in range(2): + failed_connection = self._connection + close_generation = self._close_generation span_error_recorded = False try: - result = await self._require_connection().call_tool( - tool_name, - filtered_kwargs, - meta=meta, - ) + result = await self._invoke_tool(tool_name, kwargs, trace_context) _capture_mcp_tool_result(result) if result.is_error: parsed = parser(result) @@ -3565,6 +3674,8 @@ async def _call_tool_with_retries( operation="tool call", inner_exception=call_ex, load_configured=True, + failed_connection=failed_connection, + close_generation=close_generation, ) except ToolExecutionException as recovery_error: if span.is_recording(): @@ -3686,8 +3797,10 @@ async def get_prompt(self, prompt_name: str, **kwargs: Any) -> str: capture_unhandled_exceptions=False, ) as span: for attempt in range(2): + failed_connection = self._connection + close_generation = self._close_generation try: - prompt_result = await self._require_connection().get_prompt(prompt_name, kwargs) + prompt_result = await self._invoke_prompt(prompt_name, kwargs) return parser(prompt_result) except ClosedResourceError as cl_ex: if attempt == 0: @@ -3696,6 +3809,8 @@ async def get_prompt(self, prompt_name: str, **kwargs: Any) -> str: operation="prompt call", inner_exception=cl_ex, load_configured=True, + failed_connection=failed_connection, + close_generation=close_generation, ) except ToolExecutionException as recovery_error: set_mcp_span_error(span, type(recovery_error).__name__, str(recovery_error)) @@ -3727,6 +3842,9 @@ async def get_prompt(self, prompt_name: str, **kwargs: Any) -> str: set_mcp_span_error(span, type(retries_exhausted).__name__, str(retries_exhausted)) raise retries_exhausted + async def _invoke_prompt(self, prompt_name: str, kwargs: dict[str, Any]) -> types.GetPromptResult: + return await self._require_connection().get_prompt(prompt_name, kwargs) + async def __aenter__(self) -> Self: """Enter the async context manager. @@ -3861,7 +3979,8 @@ def __init__( access the ``.functions`` list after connecting and set ``result_parser`` on individual ``FunctionTool`` instances. request_timeout: The timeout in seconds for all requests. Framework-owned clients default to 30 seconds; - ``None`` disables the timeout. Caller-owned clients retain their own timeout policy. + ``None`` disables the timeout. Connection recovery uses the same timeout as a separate budget. + Caller-owned clients retain their own request timeout policy. mcp_client: An already-entered high-level MCP ``Client`` to use. The caller retains ownership; closing or resetting this wrapper does not close or replace that client. Mutually exclusive with ``command`` and ``session``. ``request_timeout``, ``args``, ``env``, ``encoding``, and @@ -4135,7 +4254,9 @@ def __init__( access the ``.functions`` list after connecting and set ``result_parser`` on individual ``FunctionTool`` instances. request_timeout: The timeout in seconds for all requests. Framework-owned clients default to 30 seconds; - ``None`` disables the timeout. Caller-owned clients retain their own timeout policy. + ``None`` disables the timeout. HTTP header admission and session reconciliation share one budget; + connection recovery has a separate budget using the same value. Neither bounds the entire call. + Caller-owned clients retain their own request timeout policy. mcp_client: An already-entered high-level MCP ``Client`` to use. The caller retains ownership; closing or resetting this wrapper does not close or replace that client. Mutually exclusive with ``url`` and ``session``. ``request_timeout``, ``terminate_on_close``, ``http_client``, @@ -4203,8 +4324,7 @@ def __init__( static_headers: Optional fixed HTTP headers to inject into requests to the configured ``url`` origin. The headers are copied at construction, included on connection-lifetime requests and tool calls, retained across same-origin - redirects, and removed on cross-origin redirects. Unlike ``header_provider``, - fixed headers do not serialize concurrent tool calls. Use ``header_provider`` + redirects, and removed on cross-origin redirects. Use ``header_provider`` instead when header values depend on runtime invocation arguments. When both are provided, dynamic headers override fixed headers with the same name. header_provider: Optional callable that receives the runtime keyword arguments @@ -4218,7 +4338,10 @@ def __init__( agent middleware) without creating a separate ``httpx2.AsyncClient``. The complete header set used to initialize a connection becomes that session's immutable effective identity. Header names are compared case-insensitively and - values case-sensitively. Before an Agent exposes a connected tool's functions for + values case-sensitively. Operations with the same effective identity may run + concurrently; an identity change waits for active operations before reconnecting. + Each outbound HTTP request retains its own header snapshot, including redirects. + Before an Agent exposes a connected tool's functions for a run, it reconciles the run's effective headers and reconnects when they differ; generated tool and prompt calls perform the same check before sending. Connection-lifetime requests - including discovery, background pings, and resource and prompt reloads - @@ -4350,13 +4473,6 @@ def __init__( self._httpx_client: httpx2.AsyncClient | None = http_client self._static_headers = dict(static_headers or {}) self._header_provider = header_provider - # Headers for the in-flight call_tool invocation. The streamable HTTP transport - # sends requests from tasks spawned at connect time, whose contexts never observe - # ContextVar values set later inside call_tool, so the request hook needs this - # instance-level snapshot as a cross-task fallback. The lock serializes tool calls - # when a header_provider is set: parallel invocations on the same instance would - # otherwise overwrite each other's snapshot and attach the wrong per-call headers. - self._active_call_headers: dict[str, str] | None = None # None means no run seeded this connection, which an empty mapping cannot express: # a run that supplies no kwargs still expects a missing provider key to be an error. self._connection_kwargs: dict[str, Any] | None = None @@ -4369,6 +4485,10 @@ def __init__( self._pending_session_headers: dict[str, str] | None = None self._pending_connection_kwargs: dict[str, Any] | None = None self._call_headers_lock = asyncio.Lock() + self._header_operations = 0 + self._header_operations_idle = asyncio.Event() + self._header_operations_idle.set() + self._header_exclusive_owner: asyncio.Task[Any] | None = None self._header_request_owner = object() self._header_hook_client: httpx2.AsyncClient | None = None @@ -4437,26 +4557,26 @@ async def _inject_headers(request: Request) -> None: # ruff:ignore[unused-async for key in request.extensions.pop(_MCP_INJECTED_HEADER_KEYS_EXTENSION, ()): request.headers.pop(key, None) return - headers = self._static_headers.copy() - if self._header_provider is not None: - # The transport may send this request from a task whose context was - # captured before call_tool set the ContextVar; fall back to the - # instance-level snapshot of the active call's headers. Context values - # are owner-tagged so a nested tool sharing the client cannot inherit - # another tool's credentials. Both are None only when this is an ambient - # request outside call_tool; an active call that legitimately produced - # no headers yields an empty dict and must not trigger the ambient fallback. + request_header_context = request.extensions.get(_MCP_REQUEST_HEADERS_EXTENSION) + snapshot = ( + request_header_context[1] + if request_header_context is not None + and request_header_context[0] is self._header_request_owner + else None + ) + headers = self._static_headers.copy() if snapshot is None else dict(snapshot) + if snapshot is None and self._header_provider is not None: + # MCP >=2.2 carries the sender's context into each HTTP request task. + # Copy it onto the request once, so redirects cannot resolve a new caller. call_header_context = _mcp_call_headers.get(None) dynamic_headers = ( call_header_context[1] if call_header_context is not None and call_header_context[0] is self._header_request_owner else None ) - if dynamic_headers is None: - dynamic_headers = self._active_call_headers else: dynamic_headers = None - if dynamic_headers is None and self._header_provider is not None: + if snapshot is None and dynamic_headers is None and self._header_provider is not None: # Ambient request made outside call_tool (the initialize handshake, # load_tools/load_prompts discovery, or background pings). Invoke the # provider with the kwargs seeded by the run that established this @@ -4488,6 +4608,7 @@ async def _inject_headers(request: Request) -> None: # ruff:ignore[unused-async dynamic_headers = {} if dynamic_headers is not None: headers.update(dynamic_headers) + request.extensions[_MCP_REQUEST_HEADERS_EXTENSION] = (self._header_request_owner, headers.copy()) if self._session_headers is None: self._bind_session_headers(headers) for key in request.extensions.pop(_MCP_INJECTED_HEADER_KEYS_EXTENSION, ()): @@ -4558,6 +4679,204 @@ def _effective_headers(self, kwargs: Mapping[str, Any]) -> dict[str, str]: headers.update(self._header_provider(dict(kwargs))) return headers + @contextmanager + def _http_headers( + self, headers: Mapping[str, str], runtime_kwargs: Mapping[str, Any] | None = None + ) -> Generator[None]: + token = _mcp_call_headers.set((self._header_request_owner, dict(headers))) + invocation_token = _mcp_http_invocation.set( + _MCPHTTPInvocation(self, _mcp_http_invocation.get(), runtime_kwargs or {}) + ) + try: + yield + finally: + _mcp_http_invocation.reset(invocation_token) + _mcp_call_headers.reset(token) + + def _active_http_invocation(self) -> _MCPHTTPInvocation | None: + invocation = _mcp_http_invocation.get() + while invocation is not None: + if invocation.owner is self and invocation.active: + return invocation + invocation = invocation.parent + return None + + @asynccontextmanager + async def _header_admission(self) -> AsyncGenerator[None]: + """Bound admission and identity reconciliation, not the admitted network operation.""" + try: + async with asyncio.timeout(self.request_timeout): + yield + except TimeoutError as timeout_ex: + raise _mcp_tool_execution_exception( + "Timed out waiting for MCP header admission or session reconciliation.", + inner_exception=timeout_ex, + ) from timeout_ex + + def _check_exclusive_header_access(self) -> None: + if ( + self._header_provider is not None + and not self._is_lifecycle_owner_task() + and self._header_exclusive_owner is not asyncio.current_task() + and self._active_http_invocation() is not None + ): + raise _mcp_tool_execution_exception( + "Cannot reconnect, discover, reset, or close an MCP tool from its active HTTP operation." + ) + + @asynccontextmanager + async def _exclusive_header_access(self) -> AsyncGenerator[None]: + if ( + self._header_provider is None + or self._is_lifecycle_owner_task() + or self._header_exclusive_owner is asyncio.current_task() + ): + yield + return + self._check_exclusive_header_access() + async with self._call_headers_lock: + await self._header_operations_idle.wait() + await self._lifecycle_idle.wait() + self._header_exclusive_owner = asyncio.current_task() + try: + yield + finally: + self._header_exclusive_owner = None + + @asynccontextmanager + async def _http_operation(self) -> AsyncGenerator[None]: + if self._header_provider is None: + yield + return + invocation = _mcp_http_invocation.get() + header_context = _mcp_call_headers.get(None) + if invocation is None or invocation.owner is not self or header_context is None: + raise RuntimeError("MCP HTTP operation requires an invocation header snapshot.") + headers = header_context[1] + ancestor = self._active_http_invocation() + if ancestor is not None: + # A queued identity change must not block a callback needed by an active request. + if ( + ancestor.connection is not self._connection + or _mcp_header_identity(headers) != self._session_header_identity + ): + raise _mcp_tool_execution_exception("Cannot change MCP header identity from an active HTTP operation.") + else: + async with self._header_admission(), self._call_headers_lock: + await self._lifecycle_idle.wait() + if _mcp_header_identity(headers) != self._session_header_identity: + await self._header_operations_idle.wait() + self._header_exclusive_owner = asyncio.current_task() + try: + await self._ensure_session_identity(headers, invocation.runtime_kwargs) + finally: + self._header_exclusive_owner = None + invocation.connection = self._connection + invocation.active = True + self._header_operations += 1 + self._header_operations_idle.clear() + if ancestor is not None: + invocation.connection = self._connection + invocation.active = True + self._header_operations += 1 + self._header_operations_idle.clear() + try: + yield + finally: + invocation.active = False + self._header_operations -= 1 + if not self._header_operations: + self._header_operations_idle.set() + + async def connect(self, *, reset: bool = False) -> None: + async with self._exclusive_header_access(): + await super().connect(reset=reset) + + async def close(self) -> None: + self._check_exclusive_header_access() + await super().close() + + async def _close_connection(self) -> None: + async with self._exclusive_header_access(): + await super()._close_connection() + + async def load_tools(self) -> None: + async with self._exclusive_header_access(): + await super().load_tools() + + async def load_prompts(self) -> None: + async with self._exclusive_header_access(): + await super().load_prompts() + + async def _check_connection(self) -> None: + if ( + self._header_provider is None + or self._is_lifecycle_owner_task() + or self._header_exclusive_owner is asyncio.current_task() + ): + await super()._check_connection() + return + invocation = _mcp_http_invocation.get() + if invocation is not None and invocation.owner is self: + async with self._http_operation(): + await super()._check_connection() + else: + async with self._exclusive_header_access(): + await super()._check_connection() + + async def _recover_failed_connection( + self, + *, + operation: str, + failed_connection: _MCPConnection | None, + close_generation: int, + load_configured: bool, + ) -> None: + invocation = _mcp_http_invocation.get() + if ( + operation in ("tool call", "prompt call", "connection liveness check") + and invocation is not None + and invocation.owner is self + and invocation.connection is not None + and self._header_exclusive_owner is not asyncio.current_task() + and not self._is_lifecycle_owner_task() + ): + # Admission can replace the connection after the retry loop records it. + failed_connection = invocation.connection + async with self._exclusive_header_access(): + await super()._recover_failed_connection( + operation=operation, + failed_connection=failed_connection, + close_generation=close_generation, + load_configured=load_configured, + ) + + async def _invoke_tool( + self, tool_name: str, kwargs: dict[str, Any], trace_context: otel_context.Context + ) -> types.CallToolResult: + async with self._http_operation(): + return await super()._invoke_tool(tool_name, kwargs, trace_context) + + async def _invoke_prompt(self, prompt_name: str, kwargs: dict[str, Any]) -> types.GetPromptResult: + async with self._http_operation(): + return await super()._invoke_prompt(prompt_name, kwargs) + + async def get_prompt(self, prompt_name: str, **kwargs: Any) -> str: + if self._header_provider is None: + return await super().get_prompt(prompt_name, **kwargs) + header_context = _mcp_call_headers.get(None) + headers = ( + header_context[1] + if header_context is not None and header_context[0] is self._header_request_owner + else self._effective_headers({}) + ) + invocation = _mcp_http_invocation.get() + runtime_kwargs: Mapping[str, Any] = ( + invocation.runtime_kwargs if invocation is not None and invocation.owner is self else {} + ) + with self._http_headers(headers, runtime_kwargs): + return await super().get_prompt(prompt_name, **kwargs) + async def _prepare_for_run(self, kwargs: Mapping[str, Any]) -> None: if not self.is_connected: await super()._prepare_for_run(kwargs) @@ -4578,7 +4897,7 @@ async def _prepare_for_run(self, kwargs: Mapping[str, Any]) -> None: sensitive_exc_info=True, ) return - async with self._call_headers_lock: + async with self._header_admission(), self._exclusive_header_access(): await self._ensure_session_identity(headers, kwargs) def _bind_session_headers(self, headers: Mapping[str, str]) -> None: @@ -4675,19 +4994,12 @@ async def _call_prompt_with_runtime_kwargs( return await super()._call_prompt_with_runtime_kwargs(prompt_name, prompt_arguments, runtime_kwargs) headers = self._effective_headers(runtime_kwargs) - async with self._call_headers_lock: - await self._ensure_session_identity(headers, runtime_kwargs) - token = _mcp_call_headers.set((self._header_request_owner, headers)) - self._active_call_headers = headers - try: - return await super()._call_prompt_with_runtime_kwargs( - prompt_name, - prompt_arguments, - runtime_kwargs, - ) - finally: - self._active_call_headers = None - _mcp_call_headers.reset(token) + with self._http_headers(headers, runtime_kwargs): + return await super()._call_prompt_with_runtime_kwargs( + prompt_name, + prompt_arguments, + runtime_kwargs, + ) def _release_connection_kwargs(self) -> None: self._connection_kwargs = None @@ -4728,15 +5040,8 @@ async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]: runtime_context = _mcp_tool_runtime_context.get() header_kwargs = runtime_context[1] if runtime_context is not None and runtime_context[0] is self else kwargs headers = self._effective_headers(header_kwargs) - async with self._call_headers_lock: - await self._ensure_session_identity(headers, header_kwargs) - token = _mcp_call_headers.set((self._header_request_owner, headers)) - self._active_call_headers = headers - try: - return await super().call_tool(tool_name, **kwargs) - finally: - self._active_call_headers = None - _mcp_call_headers.reset(token) + with self._http_headers(headers, header_kwargs): + return await super().call_tool(tool_name, **kwargs) return await super().call_tool(tool_name, **kwargs) diff --git a/python/packages/core/tests/core/mcp_http_helpers.py b/python/packages/core/tests/core/mcp_http_helpers.py new file mode 100644 index 00000000000..82aadd9d4db --- /dev/null +++ b/python/packages/core/tests/core/mcp_http_helpers.py @@ -0,0 +1,25 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Deterministic request barriers for MCP HTTP concurrency tests.""" + +import asyncio +import json + +import httpx2 as httpx + + +class MCPHTTPBarrier: + """Hold selected HTTP requests until a test explicitly releases their responses.""" + + def __init__(self, *methods: str) -> None: + self.methods = set(methods) + self.arrivals: asyncio.Queue[httpx.Request] = asyncio.Queue() + self.release = asyncio.Event() + + async def before_request(self, request: httpx.Request) -> None: + if request.method == "POST" and json.loads(request.content).get("method") in self.methods: + self.arrivals.put_nowait(request) + await self.release.wait() + + async def next_request(self) -> httpx.Request: + return await asyncio.wait_for(self.arrivals.get(), timeout=5) diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 23ee4ea01a1..a3463b26f16 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -3009,7 +3009,7 @@ def __init__(self, **kwargs: Any) -> None: self.connect_count = 0 self.sessions: list[Any] = [] - async def connect(self, *, reset: bool = False) -> None: + async def _reconnect_without_loading(self) -> None: self.connect_count += 1 mock_session = Mock(spec=ClientSession) _set_framework_owned_session(self, mock_session) @@ -3021,12 +3021,13 @@ async def connect(self, *, reset: bool = False) -> None: return_value=types.CallToolResult(content=[types.TextContent(type="text", text="recovered")]) ) self.is_connected = True + self._tools_loaded = self._prompts_loaded = True def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: return None # type: ignore[return-value] # pyrefly: ignore[bad-return] # ty: ignore[invalid-return-type] server = TestServer(name="test_server") - await server.connect() + await server._reconnect_without_loading() result = await server.call_tool("test_tool", param="test_value") @@ -7689,15 +7690,13 @@ async def test_mcp_tool_connection_properly_invalidated_after_closed_resource_er original_exit_stack = tool._exit_stack tool._exit_stack.aclose = AsyncMock() # type: ignore[method-assign] - # Mock connect() to avoid trying to start actual process - with patch.object(tool, "connect", new_callable=AsyncMock) as mock_connect: + with patch.object(tool, "_reconnect_without_loading", new_callable=AsyncMock) as mock_connect: - async def restore_session(*, reset=False): - if reset: - await original_exit_stack.aclose() + async def restore_session(): + await original_exit_stack.aclose() _set_framework_owned_session(tool, mock_session) tool.is_connected = True - tool._tools_loaded = True + tool._tools_loaded = tool._prompts_loaded = True mock_connect.side_effect = restore_session @@ -7727,9 +7726,8 @@ async def call_tool_with_error(*args, **kwargs): # This call should trigger reconnection after ClosedResourceError result = await tool.call_tool("test_tool", arg1="value2") assert result is not None - # Verify reconnect was attempted with reset=True assert mock_connect.call_count >= 1 - mock_connect.assert_called_with(reset=True) + mock_connect.assert_awaited_once_with() # Verify _exit_stack.aclose was called during reconnection original_exit_stack.aclose.assert_called() # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] @@ -7787,15 +7785,13 @@ async def test_mcp_tool_get_prompt_reconnection_on_closed_resource_error(): original_exit_stack = tool._exit_stack tool._exit_stack.aclose = AsyncMock() # type: ignore[method-assign] - # Mock connect() to avoid trying to start actual process - with patch.object(tool, "connect", new_callable=AsyncMock) as mock_connect: + with patch.object(tool, "_reconnect_without_loading", new_callable=AsyncMock) as mock_connect: - async def restore_session(*, reset=False): - if reset: - await original_exit_stack.aclose() + async def restore_session(): + await original_exit_stack.aclose() _set_framework_owned_session(tool, mock_session) tool.is_connected = True - tool._prompts_loaded = True + tool._tools_loaded = tool._prompts_loaded = True mock_connect.side_effect = restore_session @@ -7825,9 +7821,8 @@ async def get_prompt_with_error(*args, **kwargs): # This call should trigger reconnection after ClosedResourceError result = await tool.get_prompt("test_prompt", arg1="value2") assert result is not None - # Verify reconnect was attempted with reset=True assert mock_connect.call_count >= 1 - mock_connect.assert_called_with(reset=True) + mock_connect.assert_awaited_once_with() # Verify _exit_stack.aclose was called during reconnection original_exit_stack.aclose.assert_called() # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] @@ -7925,15 +7920,16 @@ async def test_mcp_tool_call_tool_raises_after_reconnection_still_fails() -> Non tool, Mock(call_tool=AsyncMock(side_effect=[ClosedResourceError(), ClosedResourceError()])), ) + tool._tools_loaded = tool._prompts_loaded = True with ( - patch.object(tool, "connect", AsyncMock()) as mock_connect, + patch.object(tool, "_reconnect_without_loading", AsyncMock()) as mock_connect, patch.object(logger, "error") as mock_error, pytest.raises(ToolExecutionException, match="connection lost"), ): await tool.call_tool("remote_tool") - mock_connect.assert_awaited_once_with(reset=True) + mock_connect.assert_awaited_once_with() mock_error.assert_called_once() @@ -7945,15 +7941,16 @@ async def test_mcp_tool_get_prompt_raises_after_reconnection_still_fails() -> No tool, Mock(get_prompt=AsyncMock(side_effect=[ClosedResourceError(), ClosedResourceError()])), ) + tool._tools_loaded = tool._prompts_loaded = True with ( - patch.object(tool, "connect", AsyncMock()) as mock_connect, + patch.object(tool, "_reconnect_without_loading", AsyncMock()) as mock_connect, patch.object(logger, "error") as mock_error, pytest.raises(ToolExecutionException, match="connection lost"), ): await tool.get_prompt("remote_prompt") - mock_connect.assert_awaited_once_with(reset=True) + mock_connect.assert_awaited_once_with() mock_error.assert_called_once() @@ -10439,6 +10436,21 @@ async def test_mcp_streamable_http_tool_header_provider_with_httpx_event_hook(): assert request.headers.get("X-Custom") == "test-value" finally: _mcp_call_headers.reset(token) + token = _mcp_call_headers.set((tool._header_request_owner, {"X-Custom": "later-caller"})) + try: + redirected = Request( + "POST", + "http://example.com/redirected", + headers=request.headers, + extensions=request.extensions.copy(), + ) + await hooks[0](redirected) + assert redirected.headers["X-Custom"] == "test-value" + fresh = _request_for_mcp_tool(tool) + await hooks[0](fresh) + assert fresh.headers["X-Custom"] == "later-caller" + finally: + _mcp_call_headers.reset(token) finally: # Ensure any created httpx client is properly closed if getattr(tool, "_httpx_client", None) is not None: @@ -10533,9 +10545,8 @@ def provider(kw: dict[str, Any]) -> dict[str, str]: hooks = tool._httpx_client.event_hooks.get("request", []) assert len(hooks) == 1 - # Simulate an active call whose provider returned {} (both ContextVar and snapshot set). + # Simulate an active call whose provider returned {}. token = _mcp_call_headers.set((tool._header_request_owner, {})) - tool._active_call_headers = {} try: call_count = 0 request = _request_for_mcp_tool(tool) @@ -10543,7 +10554,6 @@ def provider(kw: dict[str, Any]) -> dict[str, str]: assert "Authorization" not in request.headers assert call_count == 0, "ambient fallback must not run during a set-but-empty call" finally: - tool._active_call_headers = None _mcp_call_headers.reset(token) finally: if getattr(tool, "_httpx_client", None) is not None: @@ -11735,13 +11745,15 @@ async def handler(request: httpx.Request) -> httpx.Response: async def test_mcp_streamable_http_tool_header_provider_snapshot_restored_after_call(): - """Test that the instance-level header snapshot is set during a call and cleared after.""" + """An invocation's header context is scoped to the call rather than shared on the instance.""" + from agent_framework._mcp import _mcp_call_headers + observed_snapshots: list[dict[str, str] | None] = [] original_call_tool = MCPTool.call_tool async def spy_call_tool(self, tool_name, **kwargs): - # Capture the snapshot value during the super call - observed_snapshots.append(self._active_call_headers) + context = _mcp_call_headers.get(None) + observed_snapshots.append(context[1] if context is not None else None) return await original_call_tool(self, tool_name, **kwargs) class _TestServer(MCPStreamableHTTPTool): @@ -11779,20 +11791,21 @@ def get_mcp_client(self): # pyrefly: ignore[bad-override] await server.call_tool("greet", name="Alice", auth_token="bearer-xyz") assert observed_snapshots == [{"X-Auth": "bearer-xyz"}] - assert server._active_call_headers is None + assert _mcp_call_headers.get(None) is None async def test_mcp_streamable_http_tool_header_provider_serializes_concurrent_calls(): - """Concurrent call_tool invocations on the same tool must not mix per-call headers. - - The framework executes parallel tool invocations concurrently, so a second call - must not overwrite the active header snapshot while the first call's requests - are still in flight (that would attach the wrong credentials). - """ + """Different header identities must wait for the preceding operation to finish.""" release = asyncio.Event() in_call = asyncio.Event() + observed_headers: list[dict[str, str]] = [] async def blocking_call_tool(tool_name, *, arguments=None, meta=None): + from agent_framework._mcp import _mcp_call_headers + + context = _mcp_call_headers.get(None) + assert context is not None + observed_headers.append(context[1]) in_call.set() await release.wait() return types.CallToolResult(content=[types.TextContent(type="text", text="Hello!")]) @@ -11833,20 +11846,17 @@ def get_mcp_client(self): # pyrefly: ignore[bad-override] await server.load_tools() first = asyncio.create_task(server.call_tool("greet", name="A", auth_token="token-a")) - await in_call.wait() - assert server._active_call_headers == {"X-Auth": "token-a"} - - second = asyncio.create_task(server.call_tool("greet", name="B", auth_token="token-b")) - for _ in range(10): - await asyncio.sleep(0) - # The second call must be waiting on the first; its headers must not have - # replaced the snapshot the request hook reads for the in-flight call. - assert server._active_call_headers == {"X-Auth": "token-a"} - - release.set() + try: + await in_call.wait() + second = asyncio.create_task(server.call_tool("greet", name="B", auth_token="token-b")) + for _ in range(10): + await asyncio.sleep(0) + assert observed_headers == [{"X-Auth": "token-a"}] + finally: + release.set() await asyncio.gather(first, second) - assert server._active_call_headers is None + assert observed_headers == [{"X-Auth": "token-a"}, {"X-Auth": "token-b"}] # endregion 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 9b8d28265c5..736f481f771 100644 --- a/python/packages/core/tests/core/test_mcp_http_auth.py +++ b/python/packages/core/tests/core/test_mcp_http_auth.py @@ -5,7 +5,7 @@ import asyncio import contextlib import json -from collections.abc import AsyncGenerator, AsyncIterator, Callable +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable from contextvars import ContextVar from typing import Any, Literal, TypeAlias, cast from unittest.mock import AsyncMock, Mock, patch @@ -17,6 +17,8 @@ from agent_framework import FunctionInvocationContext, MCPStreamableHTTPTool from agent_framework.exceptions import ToolException, ToolExecutionException +from .mcp_http_helpers import MCPHTTPBarrier + MCPHTTPServer: TypeAlias = tuple[httpx.AsyncClient, list[httpx.Request], dict[str, list[str]]] MCPHTTPClientFactory: TypeAlias = tuple[Callable[..., httpx.AsyncClient], list[httpx.Request], dict[str, list[str]]] @@ -29,7 +31,7 @@ async def mcp_http_client_factory() -> AsyncIterator[MCPHTTPClientFactory]: async def record_request(request: httpx.Request) -> None: requests.append(request) - async def handle(request: httpx.Request) -> httpx.Response: + async def handle(request: httpx.Request, protocol_version: str | None) -> httpx.Response: if request.url.path == "/unrelated": return httpx.Response(200) principal = request.headers.get("Authorization", "") @@ -48,10 +50,12 @@ async def handle(request: httpx.Request) -> httpx.Response: ) headers: dict[str, str] = {} result: dict[str, Any] = {} - if method == "initialize": + if method == "server/discover" and protocol_version == "2026-07-28": + result = {"supportedVersions": [protocol_version], "capabilities": {"tools": {}, "prompts": {}}} + elif method == "initialize": headers["mcp-session-id"] = f"session-{principal}" result = { - "protocolVersion": body["params"]["protocolVersion"], + "protocolVersion": protocol_version or body["params"]["protocolVersion"], "capabilities": {"tools": {}, "prompts": {}}, "serverInfo": {"name": "auth-test", "version": "1"}, } @@ -103,6 +107,10 @@ async def handle(request: httpx.Request) -> httpx.Response: "description": principal, "messages": [{"role": "user", "content": {"type": "text", "text": f"{principal}:{topic}"}}], } + if protocol_version == "2026-07-28" and method in ("tools/list", "prompts/list"): + result.update(resultType="complete", cacheScope="private", ttlMs=0) + if protocol_version == "2026-07-28" and method in ("tools/call", "prompts/get"): + result["resultType"] = "complete" if "id" not in body: return httpx.Response(202) return httpx.Response(200, headers=headers, json={"jsonrpc": "2.0", "id": body["id"], "result": result}) @@ -110,9 +118,17 @@ async def handle(request: httpx.Request) -> httpx.Response: async_client = httpx.AsyncClient clients: list[httpx.AsyncClient] = [] - def create_client(*, response_cookies: dict[str, str] | None = None, **kwargs: Any) -> httpx.AsyncClient: + def create_client( + *, + response_cookies: dict[str, str] | None = None, + before_request: Callable[[httpx.Request], Awaitable[None]] | None = None, + protocol_version: str | None = None, + **kwargs: Any, + ) -> httpx.AsyncClient: async def handle_with_cookies(request: httpx.Request) -> httpx.Response: - response = await handle(request) + if before_request is not None: + await before_request(request) + response = await handle(request, protocol_version) if response_cookies and request.method == "POST": method = json.loads(request.content).get("method") if method in response_cookies: @@ -171,6 +187,629 @@ def _requests_for_method(requests: list[httpx.Request], method: str) -> list[htt ] +@pytest.mark.parametrize("protocol_version", ["2025-11-25", "2026-07-28"]) +async def test_same_identity_http_requests_overlap( + mcp_http_client_factory: MCPHTTPClientFactory, protocol_version: str +) -> None: + create_client, requests, writes = mcp_http_client_factory + barrier = MCPHTTPBarrier("tools/call", "prompts/get") + client = create_client(before_request=barrier.before_request, protocol_version=protocol_version) + tool = MCPStreamableHTTPTool( + name="parallel", + url="https://mcp.example/mcp", + http_client=client, + header_provider=lambda _: {"Authorization": "token-a", "X-Test-Prompts": "enabled"}, + ) + async with tool: + first = asyncio.create_task(tool.call_tool("record", marker="first")) + second = asyncio.create_task(tool.call_tool("record", marker="second")) + prompt = asyncio.create_task(tool.get_prompt("principal-prompt", topic="parallel")) + try: + arrivals = [await barrier.next_request() for _ in range(3)] + assert {json.loads(request.content)["method"] for request in arrivals} == {"tools/call", "prompts/get"} + assert all(request.headers["Authorization"] == "token-a" for request in arrivals) + assert all(not task.done() for task in (first, second, prompt)) + if protocol_version == "2025-11-25": + assert {request.headers["mcp-session-id"] for request in arrivals} == {"session-token-a"} + finally: + barrier.release.set() + await asyncio.gather(first, second, prompt) + assert await prompt == "token-a:parallel" + assert sorted(writes["token-a"]) == ["first", "second"] + assert len(_requests_for_method(requests, "initialize")) == (1 if protocol_version == "2025-11-25" else 0) + assert len(_requests_for_method(requests, "server/discover")) == 1 + + +@pytest.mark.parametrize("transition", ["identity", "reset", "close", "tools", "prompts"]) +async def test_http_transition_waits_for_active_requests( + mcp_http_client_factory: MCPHTTPClientFactory, transition: str +) -> None: + create_client, requests, writes = mcp_http_client_factory + barrier = MCPHTTPBarrier("tools/call") + tool = _tool(create_client(before_request=barrier.before_request), "token-a") + async with tool: + active = asyncio.create_task(tool.call_tool("record", marker="active")) + await barrier.next_request() + waiting: asyncio.Task[object] + if transition == "identity": + waiting = asyncio.create_task(tool.call_tool("record", marker="new", credential="token-b")) + elif transition == "reset": + waiting = asyncio.create_task(tool.connect(reset=True)) + elif transition == "close": + waiting = asyncio.create_task(tool.close()) + elif transition == "tools": + waiting = asyncio.create_task(tool.load_tools()) + else: + waiting = asyncio.create_task(tool.load_prompts()) + try: + for _ in range(10): + await asyncio.sleep(0) + assert not waiting.done() + assert len(_requests_for_method(requests, "initialize")) == 1 + assert len(_requests_for_method(requests, "tools/list")) == 1 + assert not any(request.method == "DELETE" for request in requests) + finally: + barrier.release.set() + await asyncio.gather(active, waiting) + if transition == "identity": + assert writes["token-b"] == ["new"] + assert all(function.name != "token-a-only" for function in tool.functions) + assert any(function.name == "token-b-only" for function in tool.functions) + assert any( + request.method == "DELETE" and request.headers["Authorization"] == "token-a" for request in requests + ) + assert writes["token-a"] == ["active"] + + +async def test_cancelled_identity_waiter_does_not_block_same_identity_calls( + mcp_http_client_factory: MCPHTTPClientFactory, +) -> None: + create_client, requests, _ = mcp_http_client_factory + barrier = MCPHTTPBarrier("tools/call") + tool = _tool(create_client(before_request=barrier.before_request), "token-a") + async with tool: + first = asyncio.create_task(tool.call_tool("record", marker="first")) + await barrier.next_request() + changing = asyncio.create_task(tool.call_tool("record", marker="changing", credential="token-b")) + for _ in range(10): + await asyncio.sleep(0) + changing.cancel() + with pytest.raises(asyncio.CancelledError): + await changing + second = asyncio.create_task(tool.call_tool("record", marker="second")) + try: + request = await barrier.next_request() + assert request.headers["Authorization"] == "token-a" + assert not first.done() + assert not second.done() + finally: + barrier.release.set() + await asyncio.gather(first, second) + assert len(_requests_for_method(requests, "initialize")) == 1 + + +@pytest.mark.parametrize("stage", ["lock", "active", "reconcile", "lifecycle"]) +async def test_header_admission_timeout_never_dispatches_the_waiting_call( + mcp_http_client_factory: MCPHTTPClientFactory, stage: str +) -> None: + create_client, requests, _ = mcp_http_client_factory + barrier = MCPHTTPBarrier("tools/call" if stage == "active" else "initialize") + client = create_client() + tool = _tool(client, "token-a") + async with tool, asyncio.timeout(5): + client.event_hooks["request"].append(barrier.before_request) + active = None + if stage == "active": + active = asyncio.create_task(tool.call_tool("record", marker="active")) + await barrier.next_request() + elif stage == "lifecycle": + resetting = asyncio.create_task(tool.connect(reset=True)) + await barrier.next_request() + resetting.cancel() + with pytest.raises(asyncio.CancelledError): + await resetting + try: + with patch.object(tool, "request_timeout", 0.05): + async with contextlib.AsyncExitStack() as stack: + if stage == "lock": + await stack.enter_async_context(tool._call_headers_lock) + with pytest.raises(ToolExecutionException, match="Timed out.*header admission") as error: + await tool.call_tool("record", credential="token-b", marker="must-not-send") + assert isinstance(error.value.__cause__, TimeoutError) + assert tool._header_operations == (1 if stage == "active" else 0) + assert not any( + json.loads(request.content)["params"]["arguments"].get("marker") == "must-not-send" + for request in _calls(requests) + ) + if active is not None: + assert not active.done() + finally: + barrier.release.set() + if active is not None: + await active + await tool._lifecycle_idle.wait() + if not tool.is_connected: + await tool.connect() + await tool.call_tool("record", marker="after-timeout") + assert _calls(requests)[-1].headers["Authorization"] == "token-a" + + +@pytest.mark.parametrize("prepare", [False, True]) +async def test_header_admission_none_waits_until_identity_transition_can_finish( + mcp_http_client_factory: MCPHTTPClientFactory, prepare: bool +) -> None: + create_client, _, writes = mcp_http_client_factory + barrier = MCPHTTPBarrier("tools/call") + tool = _tool(create_client(before_request=barrier.before_request), "token-a") + async with tool, asyncio.timeout(5): + active = asyncio.create_task(tool.call_tool("record", marker="active")) + await barrier.next_request() + with patch.object(tool, "request_timeout", None): + waiting = asyncio.create_task( + tool._prepare_for_run({"credential": "token-b"}) + if prepare + else tool.call_tool("record", credential="token-b", marker="waiting") + ) + try: + for _ in range(10): + await asyncio.sleep(0) + assert not waiting.done() + finally: + barrier.release.set() + await asyncio.gather(active, waiting) + assert writes["token-a"] == ["active"] + assert tool._header_operations == 0 + assert tool._session_header_identity == (("authorization", "token-b"),) + + +async def test_run_preparation_uses_the_header_admission_budget(mcp_http_server: MCPHTTPServer) -> None: + client, requests, _ = mcp_http_server + tool = _tool(client, "token-a") + async with tool: + async with tool._call_headers_lock: + with ( + patch.object(tool, "request_timeout", 0.05), + pytest.raises(ToolExecutionException, match="Timed out.*header admission"), + ): + await tool._prepare_for_run({"credential": "token-b"}) + assert len(_requests_for_method(requests, "initialize")) == 1 + await tool.call_tool("record", marker="after-timeout") + + +async def test_header_admission_budget_does_not_wrap_the_admitted_rpc( + mcp_http_client_factory: MCPHTTPClientFactory, +) -> None: + create_client, _, writes = mcp_http_client_factory + barrier = MCPHTTPBarrier("tools/call") + tool = _tool(create_client(before_request=barrier.before_request), "token-a") + async with tool, asyncio.timeout(5): + with patch.object(tool, "request_timeout", 0.05): + calling = asyncio.create_task(tool.call_tool("record", marker="slow-rpc")) + try: + await barrier.next_request() + with pytest.raises(TimeoutError): + await asyncio.wait_for(asyncio.shield(calling), timeout=0.1) + assert not calling.done() + assert tool._header_operations == 1 + finally: + barrier.release.set() + await calling + assert writes["token-a"] == ["slow-rpc"] + + +@pytest.mark.parametrize("blocked_method", ["initialize", "tools/list"]) +@pytest.mark.parametrize("dynamic_headers", [False, True]) +async def test_close_interrupts_recovery_without_retry_and_explicit_reconnect_still_works( + mcp_http_server: MCPHTTPServer, blocked_method: str, dynamic_headers: bool +) -> None: + from anyio import ClosedResourceError + + client, requests, writes = mcp_http_server + tool = ( + _tool(client, "token-a") + if dynamic_headers + else MCPStreamableHTTPTool( + name="static", + url="https://mcp.example/mcp", + http_client=client, + static_headers={"Authorization": "token-a"}, + load_prompts=False, + ) + ) + barrier = MCPHTTPBarrier(blocked_method) + tasks: list[asyncio.Task[Any]] = [] + try: + async with asyncio.timeout(5): + await tool.connect() + client.event_hooks["request"].append(barrier.before_request) + with patch.object(_connected_client(tool), "call_tool", side_effect=ClosedResourceError()): + call = asyncio.create_task(tool.call_tool("record", marker="must-not-retry")) + tasks.append(call) + await barrier.next_request() + closing = asyncio.create_task(tool.close()) + tasks.append(closing) + await asyncio.sleep(0) + barrier.release.set() + with pytest.raises(ToolExecutionException, match="interrupted by close|Failed to reconnect"): + await call + await closing + assert not tool.is_connected + assert tool._connection is None + assert not tool._tools_loaded + assert tool._header_operations == 0 + assert len(_requests_for_method(requests, "initialize")) == 2 + assert writes["token-a"] == [] + await tool.connect() + await tool.call_tool("record", marker="explicit-reconnect") + assert writes["token-a"] == ["explicit-reconnect"] + finally: + barrier.release.set() + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + await tool.close() + + +async def test_close_queued_behind_a_failed_attempt_prevents_recovery(mcp_http_server: MCPHTTPServer) -> None: + from anyio import ClosedResourceError + + client, requests, writes = mcp_http_server + tool = _tool(client, "token-a") + started = asyncio.Event() + release = asyncio.Event() + + async def fail_call(*args: Any, **kwargs: Any) -> None: + started.set() + await release.wait() + raise ClosedResourceError + + async with tool, asyncio.timeout(5): + with patch.object(_connected_client(tool), "call_tool", side_effect=fail_call): + call = asyncio.create_task(tool.call_tool("record", marker="must-not-retry")) + await started.wait() + closing = asyncio.create_task(tool.close()) + try: + await asyncio.sleep(0) + assert not closing.done() + finally: + release.set() + await closing + with pytest.raises(ToolExecutionException, match="interrupted by close"): + await call + assert len(_requests_for_method(requests, "initialize")) == 1 + assert not tool.is_connected + assert writes["token-a"] == [] + + +@pytest.mark.parametrize("timeout", [0.05, None]) +@pytest.mark.parametrize("stage", ["reconnect", "discovery"]) +async def test_recovery_budget_covers_reconnect_and_discovery( + mcp_http_server: MCPHTTPServer, timeout: float | None, stage: str +) -> None: + from anyio import ClosedResourceError + + client, _, writes = mcp_http_server + tool = _tool(client, "token-a") + started = asyncio.Event() + release = asyncio.Event() + method = "_reconnect_without_loading" if stage == "reconnect" else "_load_configured_discovery_locked" + original = getattr(tool, method) + + async def block_recovery() -> None: + started.set() + await release.wait() + await original() + + async with tool, asyncio.timeout(5): + with ( + patch.object(_connected_client(tool), "call_tool", side_effect=ClosedResourceError()), + patch.object(tool, method, side_effect=block_recovery), + patch.object(tool, "request_timeout", timeout), + ): + call = asyncio.create_task(tool.call_tool("record", marker="retry")) + try: + await started.wait() + if timeout is None: + for _ in range(10): + await asyncio.sleep(0) + assert not call.done() + release.set() + await call + assert writes["token-a"] == ["retry"] + else: + with pytest.raises(ToolExecutionException, match="Timed out recovering") as error: + await call + assert isinstance(error.value.__cause__, TimeoutError) + assert writes["token-a"] == [] + finally: + release.set() + call.cancel() + await asyncio.gather(call, return_exceptions=True) + assert tool._header_operations == 0 + + +async def test_recovery_budget_includes_waiting_for_other_active_attempts( + mcp_http_client_factory: MCPHTTPClientFactory, +) -> None: + from anyio import ClosedResourceError + + create_client, _, writes = mcp_http_client_factory + barrier = MCPHTTPBarrier("tools/call") + tool = _tool(create_client(before_request=barrier.before_request), "token-a") + async with tool, asyncio.timeout(5): + active = asyncio.create_task(tool.call_tool("record", marker="active")) + await barrier.next_request() + try: + with ( + patch.object(_connected_client(tool), "call_tool", side_effect=ClosedResourceError()), + patch.object(tool, "request_timeout", 0.05), + pytest.raises(ToolExecutionException, match="Timed out recovering"), + ): + await tool.call_tool("record", marker="must-not-retry") + assert not active.done() + assert tool._header_operations == 1 + finally: + barrier.release.set() + await active + await tool.call_tool("record", marker="after-timeout") + assert writes["token-a"] == ["active", "after-timeout"] + + +async def test_direct_prompt_resolves_current_identity_instead_of_reusing_previous_caller( + mcp_http_server: MCPHTTPServer, +) -> None: + client, requests, _ = mcp_http_server + principal = ContextVar("prompt_principal", default="token-a") + tool = MCPStreamableHTTPTool( + name="prompt", + url="https://mcp.example/mcp", + http_client=client, + header_provider=lambda _: {"Authorization": principal.get(), "X-Test-Prompts": "enabled"}, + ) + async with tool: + token = principal.set("token-b") + try: + assert await tool.get_prompt("principal-prompt", topic="new-caller") == "token-b:new-caller" + finally: + principal.reset(token) + assert _requests_for_method(requests, "prompts/get")[0].headers["Authorization"] == "token-b" + assert len(_requests_for_method(requests, "initialize")) == 2 + + +async def test_request_header_snapshot_cannot_be_reused_by_another_tool_owner( + mcp_http_server: MCPHTTPServer, +) -> None: + from agent_framework._mcp import _MCP_HEADER_OWNER_EXTENSION + + client, requests, _ = mcp_http_server + first = MCPStreamableHTTPTool( + name="first", + url="https://mcp.example/mcp", + http_client=client, + static_headers={"X-First-Only": "first"}, + header_provider=lambda _: {"Authorization": "token-a"}, + ) + second = _tool(client, "token-b") + async with first, second: + original = _requests_for_method(requests, "tools/list")[0] + reused = httpx.Request( + "POST", + original.url, + headers=original.headers, + content=original.content, + extensions=original.extensions.copy(), + ) + reused.extensions[_MCP_HEADER_OWNER_EXTENSION] = second._header_request_owner + for hook in client.event_hooks["request"]: + await hook(reused) + assert reused.headers["Authorization"] == "token-b" + assert "X-First-Only" not in reused.headers + + +async def test_nested_same_identity_call_can_finish_ahead_of_waiting_identity_change( + mcp_http_client_factory: MCPHTTPClientFactory, +) -> None: + create_client, requests, writes = mcp_http_client_factory + outer_started = asyncio.Event() + start_nested = asyncio.Event() + + async def invoke_nested(request: httpx.Request) -> None: + if request.method != "POST": + return + body = json.loads(request.content) + if body.get("method") == "tools/call" and body["params"]["arguments"].get("marker") == "outer": + outer_started.set() + await start_nested.wait() + await tool.call_tool("record", marker="nested") + + tool = _tool(create_client(before_request=invoke_nested), "token-a") + async with tool: + outer = asyncio.create_task(tool.call_tool("record", marker="outer")) + await asyncio.wait_for(outer_started.wait(), timeout=5) + changing = asyncio.create_task(tool.call_tool("record", marker="changed", credential="token-b")) + try: + for _ in range(10): + await asyncio.sleep(0) + assert not changing.done() + finally: + start_nested.set() + await asyncio.wait_for(asyncio.gather(outer, changing), timeout=5) + assert writes["token-a"] == ["nested", "outer"] + assert writes["token-b"] == ["changed"] + assert len(_requests_for_method(requests, "initialize")) == 2 + + +@pytest.mark.parametrize("operation", ["identity", "reset", "close", "tools", "prompts", "recovery"]) +async def test_nested_exclusive_operation_fails_instead_of_waiting_on_its_parent( + mcp_http_client_factory: MCPHTTPClientFactory, operation: str +) -> None: + from anyio import ClosedResourceError + + create_client, _, writes = mcp_http_client_factory + + async def invoke_nested(request: httpx.Request) -> None: + if request.method != "POST" or json.loads(request.content).get("method") != "tools/call": + return + close_generation = tool._close_generation + with pytest.raises(ToolExecutionException, match="active HTTP operation"): + if operation == "identity": + await tool.call_tool("record", marker="nested", credential="token-b") + elif operation == "reset": + await tool.connect(reset=True) + elif operation == "close": + await tool.close() + elif operation == "tools": + await tool.load_tools() + elif operation == "prompts": + await tool.load_prompts() + else: + await tool._recover_connection_after_loss( + operation="tool call", inner_exception=ClosedResourceError(), load_configured=True + ) + assert tool._close_generation == close_generation + + tool = _tool(create_client(before_request=invoke_nested), "token-a") + async with tool: + await asyncio.wait_for(tool.call_tool("record", marker="outer"), timeout=5) + assert writes["token-a"] == ["outer"] + assert not writes["token-b"] + + +@pytest.mark.parametrize("dynamic_headers", [False, True]) +async def test_simultaneous_connection_failures_share_one_recovery( + mcp_http_server: MCPHTTPServer, dynamic_headers: bool +) -> None: + from anyio import ClosedResourceError + + client, requests, writes = mcp_http_server + tool = ( + _tool(client, "token-a") + if dynamic_headers + else MCPStreamableHTTPTool( + name="static", + url="https://mcp.example/mcp", + http_client=client, + static_headers={"Authorization": "token-a"}, + load_prompts=False, + ) + ) + arrivals: asyncio.Queue[None] = asyncio.Queue() + fail = asyncio.Event() + + async def fail_call(*args: Any, **kwargs: Any) -> None: + arrivals.put_nowait(None) + await fail.wait() + raise ClosedResourceError + + async with tool: + original_connection = tool._connection + with patch.object(_connected_client(tool), "call_tool", side_effect=fail_call): + first = asyncio.create_task(tool.call_tool("record", marker="first")) + second = asyncio.create_task(tool.call_tool("record", marker="second")) + try: + await asyncio.wait_for(arrivals.get(), timeout=5) + await asyncio.wait_for(arrivals.get(), timeout=5) + finally: + fail.set() + await asyncio.wait_for(asyncio.gather(first, second), timeout=5) + assert tool._connection is not original_connection + assert len(_requests_for_method(requests, "initialize")) == 2 + assert sorted(writes["token-a"]) == ["first", "second"] + + +async def test_liveness_failure_during_recovery_discovery_recovers_the_replacement_connection( + mcp_http_server: MCPHTTPServer, +) -> None: + from anyio import ClosedResourceError + + client, requests, writes = mcp_http_server + tool = _tool(client, "token-a") + async with tool, asyncio.timeout(5): + original_connection = tool._connection + original_check = tool._check_connection + replacement_failed = False + + async def fail_replacement_once() -> None: + nonlocal replacement_failed + if tool._connection is not original_connection and not replacement_failed: + replacement_failed = True + raise ClosedResourceError + await original_check() + + with ( + patch.object(_connected_client(tool), "call_tool", side_effect=ClosedResourceError()), + patch.object(tool, "_check_connection", side_effect=fail_replacement_once), + ): + await tool.call_tool("record", marker="recovered") + assert replacement_failed + assert len(_requests_for_method(requests, "initialize")) == 3 + assert writes["token-a"] == ["recovered"] + + +async def test_cancelling_recovery_propagates_and_releases_the_exclusive_barrier( + mcp_http_server: MCPHTTPServer, +) -> None: + from anyio import ClosedResourceError + + client, _, writes = mcp_http_server + tool = _tool(client, "token-a") + started = asyncio.Event() + + async def block_recovery() -> None: + started.set() + await asyncio.Event().wait() + + async with tool, asyncio.timeout(5): + with ( + patch.object(_connected_client(tool), "call_tool", side_effect=ClosedResourceError()), + patch.object(tool, "_reconnect_without_loading", side_effect=block_recovery), + ): + calling = asyncio.create_task(tool.call_tool("record", marker="cancelled")) + await started.wait() + calling.cancel() + with pytest.raises(asyncio.CancelledError): + await calling + assert tool._header_operations == 0 + assert tool._header_exclusive_owner is None + await tool.call_tool("record", marker="after-cancel") + assert writes["token-a"] == ["after-cancel"] + + +async def test_cancelled_close_waiter_still_blocks_new_http_operations_until_teardown( + mcp_http_client_factory: MCPHTTPClientFactory, +) -> None: + create_client, _, _ = mcp_http_client_factory + closing = asyncio.Event() + release_close = asyncio.Event() + + async def block_close(request: httpx.Request) -> None: + if request.method == "DELETE": + closing.set() + await release_close.wait() + + tool = _tool(create_client(before_request=block_close), "token-a") + await tool.connect() + old_connection = tool._connection + close = asyncio.create_task(tool.close()) + await asyncio.wait_for(closing.wait(), timeout=5) + close.cancel() + with pytest.raises(asyncio.CancelledError): + await close + calling = asyncio.create_task(tool.call_tool("record", marker="after-close")) + try: + for _ in range(10): + await asyncio.sleep(0) + assert not calling.done() + finally: + release_close.set() + try: + with pytest.raises(ToolExecutionException): + await asyncio.wait_for(calling, timeout=5) + assert tool._connection is not old_connection + finally: + await tool.close() + + async def test_owned_client_does_not_replay_response_cookie_across_principals( mcp_http_client_factory: MCPHTTPClientFactory, ) -> None: diff --git a/python/packages/core/tests/core/test_mcp_observability.py b/python/packages/core/tests/core/test_mcp_observability.py index 47a28b9d205..449a09873f6 100644 --- a/python/packages/core/tests/core/test_mcp_observability.py +++ b/python/packages/core/tests/core/test_mcp_observability.py @@ -333,7 +333,8 @@ async def test_mcp_tools_call_second_connection_loss_log_respects_sensitive_data MCPError(-32000, f"session terminated: {private_error_detail}"), ] ) - tool.connect = AsyncMock(return_value=None) # type: ignore[method-assign] + tool._reconnect_without_loading = AsyncMock(return_value=None) # type: ignore[method-assign] + tool._tools_loaded = tool._prompts_loaded = True span_exporter.clear() with ( diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py index 294afbbe9db..4882bc83426 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py @@ -175,8 +175,9 @@ class FoundryToolbox(MCPStreamableHTTPTool): into the MCP connection's writer task. The complete effective header set is bound to the MCP session; when the caller context changes, the Toolbox reconnects that session before sending the operation. Tool, prompt and - Toolbox skill-resource operations share the same identity lock so concurrent - callers cannot overwrite one another's headers. + Toolbox skill-resource operations retain their own request headers. Operations + with identical effective headers may overlap; identity changes wait for active + operations before reconnecting. Because the toolbox endpoint is a first-party Foundry service, forwarding the opaque caller token to it is safe. @@ -333,6 +334,14 @@ async def _call_prompt_with_runtime_kwargs( finally: _strict_toolbox_header_provider.reset(token) + @override + async def get_prompt(self, prompt_name: str, **kwargs: Any) -> str: + token = _strict_toolbox_header_provider.set(True) + try: + return await super().get_prompt(prompt_name, **kwargs) + finally: + _strict_toolbox_header_provider.reset(token) + @override def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: """Get an authenticated MCP HTTP client. @@ -347,29 +356,29 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: return super().get_mcp_client() @override - async def close(self) -> None: - """Close the MCP session and toolbox HTTP client while preserving credentials and timeout for reconnection.""" - try: - await super().close() - finally: - client = self._httpx_client - if client is not None: - self._httpx_client = None - await client.aclose() + async def _close_connection(self) -> None: + """Close the MCP session before its HTTP client, including after waiter cancellation.""" + async with self._exclusive_header_access(): + try: + await super()._close_connection() + finally: + # The owner may still need HTTP for teardown after its waiter is cancelled. + await self._lifecycle_idle.wait() + client = self._httpx_client + if client is not None: + self._httpx_client = None + await client.aclose() async def _read_resource(self, uri: AnyUrl) -> ReadResourceResult: """Read a Toolbox resource under the current platform-header identity.""" headers = self._effective_resource_headers() - async with self._call_headers_lock: - connection = self._require_connection() - if headers or connection.is_framework_owned: - await self._ensure_session_identity(headers, {}) + connection = self._require_connection() + if not headers and not connection.is_framework_owned: + return await connection.read_resource(str(uri)) + with self._http_headers(headers): + async with self._http_operation(): connection = self._require_connection() - self._active_call_headers = headers - try: - return await connection.read_resource(cast(str, uri)) - finally: - self._active_call_headers = None + return await connection.read_resource(str(uri)) def _effective_resource_headers(self) -> dict[str, str]: """Resolve headers for a resource operation without caller runtime kwargs.""" @@ -567,7 +576,7 @@ class _FoundryToolboxSkillsSource(SkillsSource): (via the agent or an ``async with`` block) and is **replaced** with a new object whenever the toolbox reconnects. Skills are therefore bound to an identity-scoped resource adapter that resolves the current session while - holding the Toolbox operation lock, so cached skills keep using the live + holding a Toolbox operation lease, so cached skills keep using the live session instead of a closed one. """ diff --git a/python/packages/foundry_hosting/tests/test_toolbox.py b/python/packages/foundry_hosting/tests/test_toolbox.py index 324bc0cb01b..3b9f1d6111a 100644 --- a/python/packages/foundry_hosting/tests/test_toolbox.py +++ b/python/packages/foundry_hosting/tests/test_toolbox.py @@ -82,6 +82,7 @@ class _ToolboxLoopbackServer: def __init__(self) -> None: self.requests: list[httpx.Request] = [] self._session_count = 0 + self.before_request: Callable[[httpx.Request], Awaitable[None]] | None = None async def handle(self, request: httpx.Request) -> httpx.Response: if request.method == "GET": @@ -90,6 +91,8 @@ async def handle(self, request: httpx.Request) -> httpx.Response: return httpx.Response(200) self.requests.append(request) + if self.before_request is not None: + await self.before_request(request) payload = json.loads(request.content) method = payload.get("method") call_id = request.headers.get("x-agent-foundry-call-id") @@ -170,6 +173,136 @@ async def _run_with_request_context(call_id: str, action: Callable[[], Awaitable reset_request_context(token) +async def test_same_call_id_overlaps_tool_and_skill_reads_without_freezing_bearer_tokens() -> None: + server = _ToolboxLoopbackServer() + arrivals: asyncio.Queue[httpx.Request] = asyncio.Queue() + release = asyncio.Event() + + async def hold_response(request: httpx.Request) -> None: + if json.loads(request.content).get("method") in ("tools/call", "resources/read"): + arrivals.put_nowait(request) + await release.wait() + + server.before_request = hold_response + credential = _FakeCredential("before-refresh") + toolbox = FoundryToolbox(credential, url="https://h/toolboxes/tb/mcp") # type: ignore + await _attach_loopback_transport(toolbox, server) + try: + await _run_with_request_context("parallel-call", toolbox.connect) + first = asyncio.create_task( + _run_with_request_context("parallel-call", lambda: toolbox.call_tool("record", marker="first")) + ) + first_request = await asyncio.wait_for(arrivals.get(), timeout=5) + credential._token = "after-refresh" + second = asyncio.create_task( + _run_with_request_context("parallel-call", lambda: toolbox.call_tool("record", marker="second")) + ) + resource = asyncio.create_task( + _run_with_request_context( + "parallel-call", lambda: toolbox._read_resource(AnyUrl("skill://loopback-skill/SKILL.md")) + ) + ) + try: + next_requests = [await asyncio.wait_for(arrivals.get(), timeout=5) for _ in range(2)] + assert not first.done() and not second.done() and not resource.done() + assert first_request.headers["Authorization"] == "Bearer before-refresh" + assert all(request.headers["Authorization"] == "Bearer after-refresh" for request in next_requests) + assert all( + request.headers["x-agent-foundry-call-id"] == "parallel-call" + for request in [first_request, *next_requests] + ) + finally: + release.set() + await asyncio.gather(first, second, resource) + for task in (first, second): + result = await task + assert isinstance(result, list) + assert result[0].text == "parallel-call" + assert len(server.requests_for("initialize")) == 1 + finally: + release.set() + await toolbox.close() + + +async def test_cancelled_close_keeps_http_client_alive_until_owner_finishes() -> None: + toolbox = FoundryToolbox(_FakeCredential(), url="https://h/toolboxes/tb/mcp") # type: ignore + await _attach_loopback_transport(toolbox, _ToolboxLoopbackServer()) + await _run_with_request_context("closing-call", toolbox.connect) + client = toolbox._httpx_client + assert client is not None + closing = asyncio.Event() + release = asyncio.Event() + close_on_owner = toolbox._close_on_owner + + async def hold_teardown() -> None: + closing.set() + await release.wait() + await close_on_owner() + + with patch.object(toolbox, "_close_on_owner", side_effect=hold_teardown): + task = asyncio.create_task(toolbox.close()) + try: + await asyncio.wait_for(closing.wait(), timeout=5) + task.cancel() + for _ in range(10): + await asyncio.sleep(0) + assert not task.done() + assert not client.is_closed + finally: + release.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=5) + assert client.is_closed + assert toolbox._httpx_client is None + assert not toolbox.is_connected + + +async def test_toolbox_close_fences_recovery_before_waiting_for_its_barrier() -> None: + from anyio import ClosedResourceError + + toolbox = FoundryToolbox(_FakeCredential(), url="https://h/toolboxes/tb/mcp") # type: ignore + server = _ToolboxLoopbackServer() + await _attach_loopback_transport(toolbox, server) + started = asyncio.Event() + release = asyncio.Event() + + async def fail_call(*args: Any, **kwargs: Any) -> None: + started.set() + await release.wait() + raise ClosedResourceError + + tasks: list[asyncio.Task[Any]] = [] + try: + async with asyncio.timeout(5): + await _run_with_request_context("closing-call", toolbox.connect) + connection = toolbox._require_connection() + assert connection.client is not None + with patch.object(connection.client, "call_tool", side_effect=fail_call): + calling = asyncio.create_task( + _run_with_request_context("closing-call", lambda: toolbox.call_tool("record", marker="no-retry")) + ) + tasks.append(calling) + await started.wait() + closing = asyncio.create_task(toolbox.close()) + tasks.append(closing) + await asyncio.sleep(0) + assert not closing.done() + release.set() + await closing + with pytest.raises(ToolExecutionException, match="interrupted by close"): + await calling + assert not toolbox.is_connected + assert toolbox._httpx_client is None + assert len(server.requests_for("initialize")) == 1 + assert not server.requests_for("tools/call") + finally: + release.set() + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + await toolbox.close() + + def test_resolve_endpoint_prefers_explicit_env(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("TOOLBOX_ENDPOINT", "https://host/toolboxes/tb/mcp?api-version=v1") assert _resolve_toolbox_endpoint() == "https://host/toolboxes/tb/mcp?api-version=v1" @@ -658,13 +791,13 @@ async def get_skills(self, context: SkillsSourceContext) -> list[str]: assert provider() is resource_session first_uri = AnyUrl("skill://first") await resource_session.read_resource(cast(str, first_uri)) - sentinel_session.read_resource.assert_awaited_once_with(first_uri) + sentinel_session.read_resource.assert_awaited_once_with(str(first_uri)) new_session = AsyncMock() toolbox.session = new_session # type: ignore second_uri = AnyUrl("skill://second") await resource_session.read_resource(cast(str, second_uri)) - new_session.read_resource.assert_awaited_once_with(second_uri) + new_session.read_resource.assert_awaited_once_with(str(second_uri)) # No archive options set -> MCPSkillsSource is constructed with defaults. assert captured_kwargs == {} diff --git a/skills/agent-framework-builder-python/references/integrations.md b/skills/agent-framework-builder-python/references/integrations.md index c19bb242bb4..2a2b2c96ff5 100644 --- a/skills/agent-framework-builder-python/references/integrations.md +++ b/skills/agent-framework-builder-python/references/integrations.md @@ -31,6 +31,24 @@ For Foundry, the derived OpenAI connection is owned by the connector even when t Do not share a connector-owned connection across overlapping agent contexts; inject an externally managed SDK client when its lifetime must extend beyond any one agent context. Do not assume a closed SDK connection can reopen. +### MCP HTTP concurrency + +`MCPStreamableHTTPTool` and `FoundryToolbox` keep a separate header snapshot on every outgoing HTTP request. +Calls with identical effective static and dynamic headers can run concurrently on one MCP session. A change to +any effective header, including the Foundry platform call ID, waits for active operations and reconnects before +refreshing discovery. Token refresh performed by HTTP authentication remains request-local. + +Discovery, reset, close, and connection recovery are exclusive. Do not reset, close, refresh discovery, or change +header identity from a callback that must finish an active operation on that same tool; it raises an error rather +than waiting on itself. Concurrent calls still share the server's MCP session state, so applications must order +state-dependent operations themselves. Cancelling a local call does not guarantee that its remote side effects stop. +Caller-owned MCP clients and sessions remain caller-owned; Agent Framework does not recreate them. + +`request_timeout` supplies separate budgets for header admission/session reconciliation and connection recovery, +not a deadline for the entire call. `None` disables those budgets; caller-owned MCP clients keep their own RPC +timeouts. Cleanup may continue after timeout or cancellation. An explicit `close()` interrupts older recovery +without preventing a later explicit `connect()`. + ## Capability categories - Use function tools for application-owned code and MCP for external tool services. From 83cf19af70cad5e1100b1b329d325c1f5934fb15 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Fri, 6 Nov 2026 13:54:45 +0100 Subject: [PATCH 2/4] Python: Address MCP concurrency review findings --- python/packages/core/AGENTS.md | 7 +- python/packages/core/README.md | 8 +- python/packages/core/agent_framework/_mcp.py | 57 +++----- .../core/tests/core/test_mcp_http_auth.py | 124 ++++++++++++++++-- .../foundry_hosting/tests/test_toolbox.py | 65 ++++++++- .../references/integrations.md | 5 +- 6 files changed, 211 insertions(+), 55 deletions(-) diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 49461ca95ca..c0c54780906 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -288,8 +288,11 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID - **MCP HTTP concurrency** - MCP >=2.2 propagates the sending operation's ContextVars into each outbound HTTP task. The owner-scoped header hook copies effective headers onto each request, including redirects; do not restore a mutable instance-level active-header fallback. Same-identity tool, prompt, and Toolbox resource operations may - overlap on one session. Identity changes, discovery, reset, close, and recovery wait for active network attempts; - retries release their attempt before recovery, re-admit the original identity, and use its refreshed argument + overlap on one session. Same-identity run preparation uses shared admission rather than draining active calls. + Direct prompts resolve current headers; generated prompts retain explicitly resolved host runtime headers. + Identity changes, discovery, reset, close, and recovery wait for active network attempts on framework-owned + connections even without a header provider; retries release their attempt before recovery, re-admit the original + identity, and use its refreshed argument allowlist and tool metadata. Failed attempts from the same replaced framework-owned connection share recovery. Keep the original trace context when rebuilding metadata. Cancelled lifecycle waiters leave their owner running; new HTTP operations must wait for that setup or teardown to finish. A same-generation nested callback may finish diff --git a/python/packages/core/README.md b/python/packages/core/README.md index c876419215d..5f3d85ca5a1 100644 --- a/python/packages/core/README.md +++ b/python/packages/core/README.md @@ -200,8 +200,12 @@ MCP protocol sessions or other server-side state between principals. outgoing HTTP request. Calls with identical effective static and dynamic headers can overlap on one MCP session. A different identity waits for active operations before reconnecting and refreshing discovery. Discovery, reset, close, and recovery are -exclusive against active network attempts. Calls still share server-side session -state; applications must order state-dependent operations themselves. +exclusive against active network attempts on framework-owned MCP connections, with +or without a header provider. Same-identity Agent run preparation does not drain +active calls. Direct prompt calls resolve current headers rather than reusing an +inherited invocation snapshot; generated prompts resolve host runtime kwargs +separately from model arguments. Calls still share server-side session state; +applications must order state-dependent operations themselves. `request_timeout` bounds header admission and session reconciliation together, and connection recovery separately from the RPC timeout. These are separate budgets, diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 533b4eb1e5e..9b806498fcb 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -4715,8 +4715,7 @@ async def _header_admission(self) -> AsyncGenerator[None]: def _check_exclusive_header_access(self) -> None: if ( - self._header_provider is not None - and not self._is_lifecycle_owner_task() + not self._is_lifecycle_owner_task() and self._header_exclusive_owner is not asyncio.current_task() and self._active_http_invocation() is not None ): @@ -4726,11 +4725,7 @@ def _check_exclusive_header_access(self) -> None: @asynccontextmanager async def _exclusive_header_access(self) -> AsyncGenerator[None]: - if ( - self._header_provider is None - or self._is_lifecycle_owner_task() - or self._header_exclusive_owner is asyncio.current_task() - ): + if self._is_lifecycle_owner_task() or self._header_exclusive_owner is asyncio.current_task(): yield return self._check_exclusive_header_access() @@ -4745,7 +4740,8 @@ async def _exclusive_header_access(self) -> AsyncGenerator[None]: @asynccontextmanager async def _http_operation(self) -> AsyncGenerator[None]: - if self._header_provider is None: + connection = self._connection + if self._header_provider is None and connection is not None and not connection.is_framework_owned: yield return invocation = _mcp_http_invocation.get() @@ -4809,11 +4805,7 @@ async def load_prompts(self) -> None: await super().load_prompts() async def _check_connection(self) -> None: - if ( - self._header_provider is None - or self._is_lifecycle_owner_task() - or self._header_exclusive_owner is asyncio.current_task() - ): + if self._is_lifecycle_owner_task() or self._header_exclusive_owner is asyncio.current_task(): await super()._check_connection() return invocation = _mcp_http_invocation.get() @@ -4862,19 +4854,8 @@ async def _invoke_prompt(self, prompt_name: str, kwargs: dict[str, Any]) -> type return await super()._invoke_prompt(prompt_name, kwargs) async def get_prompt(self, prompt_name: str, **kwargs: Any) -> str: - if self._header_provider is None: - return await super().get_prompt(prompt_name, **kwargs) - header_context = _mcp_call_headers.get(None) - headers = ( - header_context[1] - if header_context is not None and header_context[0] is self._header_request_owner - else self._effective_headers({}) - ) - invocation = _mcp_http_invocation.get() - runtime_kwargs: Mapping[str, Any] = ( - invocation.runtime_kwargs if invocation is not None and invocation.owner is self else {} - ) - with self._http_headers(headers, runtime_kwargs): + headers = self._effective_headers({}) + with self._http_headers(headers): return await super().get_prompt(prompt_name, **kwargs) async def _prepare_for_run(self, kwargs: Mapping[str, Any]) -> None: @@ -4897,8 +4878,10 @@ async def _prepare_for_run(self, kwargs: Mapping[str, Any]) -> None: sensitive_exc_info=True, ) return - async with self._header_admission(), self._exclusive_header_access(): - await self._ensure_session_identity(headers, kwargs) + with self._http_headers(headers, kwargs): + async with self._http_operation(): + # Reconcile through admission without draining same-identity active calls. + pass def _bind_session_headers(self, headers: Mapping[str, str]) -> None: self._session_headers = dict(headers) @@ -4995,11 +4978,7 @@ async def _call_prompt_with_runtime_kwargs( headers = self._effective_headers(runtime_kwargs) with self._http_headers(headers, runtime_kwargs): - return await super()._call_prompt_with_runtime_kwargs( - prompt_name, - prompt_arguments, - runtime_kwargs, - ) + return await super().get_prompt(prompt_name, **prompt_arguments) def _release_connection_kwargs(self) -> None: self._connection_kwargs = None @@ -5036,13 +5015,11 @@ async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]: Returns: A list of Content items representing the tool output. """ - if self._header_provider is not None: - runtime_context = _mcp_tool_runtime_context.get() - header_kwargs = runtime_context[1] if runtime_context is not None and runtime_context[0] is self else kwargs - headers = self._effective_headers(header_kwargs) - with self._http_headers(headers, header_kwargs): - return await super().call_tool(tool_name, **kwargs) - return await super().call_tool(tool_name, **kwargs) + runtime_context = _mcp_tool_runtime_context.get() + header_kwargs = runtime_context[1] if runtime_context is not None and runtime_context[0] is self else kwargs + headers = self._effective_headers(header_kwargs) + with self._http_headers(headers, header_kwargs): + return await super().call_tool(tool_name, **kwargs) @deprecated("MCP WebSocket transport was removed in MCP v2. Use MCPStreamableHTTPTool instead.") 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 736f481f771..6311e85c0be 100644 --- a/python/packages/core/tests/core/test_mcp_http_auth.py +++ b/python/packages/core/tests/core/test_mcp_http_auth.py @@ -14,7 +14,7 @@ import pytest from mcp.client.session import ClientSession -from agent_framework import FunctionInvocationContext, MCPStreamableHTTPTool +from agent_framework import Agent, FunctionInvocationContext, MCPStreamableHTTPTool from agent_framework.exceptions import ToolException, ToolExecutionException from .mcp_http_helpers import MCPHTTPBarrier @@ -188,17 +188,24 @@ def _requests_for_method(requests: list[httpx.Request], method: str) -> list[htt @pytest.mark.parametrize("protocol_version", ["2025-11-25", "2026-07-28"]) +@pytest.mark.parametrize("header_source", ["dynamic", "static", "client"]) async def test_same_identity_http_requests_overlap( - mcp_http_client_factory: MCPHTTPClientFactory, protocol_version: str + mcp_http_client_factory: MCPHTTPClientFactory, protocol_version: str, header_source: str ) -> None: create_client, requests, writes = mcp_http_client_factory barrier = MCPHTTPBarrier("tools/call", "prompts/get") - client = create_client(before_request=barrier.before_request, protocol_version=protocol_version) + headers = {"Authorization": "token-a", "X-Test-Prompts": "enabled"} + client = create_client( + before_request=barrier.before_request, + protocol_version=protocol_version, + headers=headers if header_source == "client" else None, + ) tool = MCPStreamableHTTPTool( name="parallel", url="https://mcp.example/mcp", http_client=client, - header_provider=lambda _: {"Authorization": "token-a", "X-Test-Prompts": "enabled"}, + static_headers=headers if header_source == "static" else None, + header_provider=(lambda _: headers) if header_source == "dynamic" else None, ) async with tool: first = asyncio.create_task(tool.call_tool("record", marker="first")) @@ -221,12 +228,27 @@ async def test_same_identity_http_requests_overlap( @pytest.mark.parametrize("transition", ["identity", "reset", "close", "tools", "prompts"]) +@pytest.mark.parametrize("header_source", ["dynamic", "static", "client"]) async def test_http_transition_waits_for_active_requests( - mcp_http_client_factory: MCPHTTPClientFactory, transition: str + mcp_http_client_factory: MCPHTTPClientFactory, transition: str, header_source: str ) -> None: + if transition == "identity" and header_source != "dynamic": + pytest.skip("Fixed headers do not have a runtime identity transition.") create_client, requests, writes = mcp_http_client_factory barrier = MCPHTTPBarrier("tools/call") - tool = _tool(create_client(before_request=barrier.before_request), "token-a") + tool = MCPStreamableHTTPTool( + name="transition", + url="https://mcp.example/mcp", + http_client=create_client( + before_request=barrier.before_request, + headers={"Authorization": "token-a"} if header_source == "client" else None, + ), + static_headers={"Authorization": "token-a"} if header_source == "static" else None, + header_provider=(lambda kwargs: {"Authorization": kwargs.get("credential", "token-a")}) + if header_source == "dynamic" + else None, + load_prompts=False, + ) async with tool: active = asyncio.create_task(tool.call_tool("record", marker="active")) await barrier.next_request() @@ -376,6 +398,38 @@ async def test_run_preparation_uses_the_header_admission_budget(mcp_http_server: await tool.call_tool("record", marker="after-timeout") +@pytest.mark.parametrize("stream", [False, True]) +async def test_same_identity_agent_run_prepares_while_another_rpc_is_active( + mcp_http_client_factory: MCPHTTPClientFactory, chat_client_base: Any, stream: bool +) -> None: + create_client, requests, writes = mcp_http_client_factory + barrier = MCPHTTPBarrier("tools/call") + tool = _tool(create_client(before_request=barrier.before_request), "token-a") + agent = Agent(client=chat_client_base, tools=tool) + + async def run_agent() -> str: + if stream: + return "".join([ + update.text + async for update in agent.run( + "prepare", stream=True, function_invocation_kwargs={"credential": "token-a"} + ) + ]) + return (await agent.run("prepare", function_invocation_kwargs={"credential": "token-a"})).text + + async with tool, asyncio.timeout(5): + active = asyncio.create_task(tool.call_tool("record", marker="active")) + await barrier.next_request() + try: + assert await asyncio.wait_for(run_agent(), timeout=1) + assert not active.done() + assert len(_requests_for_method(requests, "initialize")) == 1 + finally: + barrier.release.set() + await active + assert writes["token-a"] == ["active"] + + async def test_header_admission_budget_does_not_wrap_the_admitted_rpc( mcp_http_client_factory: MCPHTTPClientFactory, ) -> None: @@ -576,6 +630,47 @@ async def test_direct_prompt_resolves_current_identity_instead_of_reusing_previo assert len(_requests_for_method(requests, "initialize")) == 2 +async def test_direct_prompt_in_detached_child_does_not_reuse_completed_invocation_headers( + mcp_http_client_factory: MCPHTTPClientFactory, +) -> None: + create_client, requests, _ = mcp_http_client_factory + principal = ContextVar("child_prompt_principal", default="token-a") + start_child = asyncio.Event() + children: list[asyncio.Task[str]] = [] + + async def call_prompt() -> str: + await start_child.wait() + token = principal.set("token-b") + try: + return await tool.get_prompt("principal-prompt", topic="child") + finally: + principal.reset(token) + + async def spawn_child(request: httpx.Request) -> None: + if request.method == "POST" and json.loads(request.content).get("method") == "tools/call": + children.append(asyncio.create_task(call_prompt())) + + tool = MCPStreamableHTTPTool( + name="child-prompt", + url="https://mcp.example/mcp", + http_client=create_client(before_request=spawn_child), + header_provider=lambda _: {"Authorization": principal.get(), "X-Test-Prompts": "enabled"}, + ) + async with tool, asyncio.timeout(5): + try: + await tool.call_tool("record", marker="parent") + assert len(children) == 1 + start_child.set() + assert await children[0] == "token-b:child" + assert _requests_for_method(requests, "prompts/get")[-1].headers["Authorization"] == "token-b" + assert len(_requests_for_method(requests, "initialize")) == 2 + finally: + start_child.set() + for child in children: + child.cancel() + await asyncio.gather(*children, return_exceptions=True) + + async def test_request_header_snapshot_cannot_be_reused_by_another_tool_owner( mcp_http_server: MCPHTTPServer, ) -> None: @@ -640,11 +735,14 @@ async def invoke_nested(request: httpx.Request) -> None: @pytest.mark.parametrize("operation", ["identity", "reset", "close", "tools", "prompts", "recovery"]) +@pytest.mark.parametrize("dynamic_headers", [False, True]) async def test_nested_exclusive_operation_fails_instead_of_waiting_on_its_parent( - mcp_http_client_factory: MCPHTTPClientFactory, operation: str + mcp_http_client_factory: MCPHTTPClientFactory, operation: str, dynamic_headers: bool ) -> None: from anyio import ClosedResourceError + if operation == "identity" and not dynamic_headers: + pytest.skip("Fixed headers do not have a runtime identity transition.") create_client, _, writes = mcp_http_client_factory async def invoke_nested(request: httpx.Request) -> None: @@ -668,7 +766,17 @@ async def invoke_nested(request: httpx.Request) -> None: ) assert tool._close_generation == close_generation - tool = _tool(create_client(before_request=invoke_nested), "token-a") + tool = ( + _tool(create_client(before_request=invoke_nested), "token-a") + if dynamic_headers + else MCPStreamableHTTPTool( + name="static-nested", + url="https://mcp.example/mcp", + http_client=create_client(before_request=invoke_nested), + static_headers={"Authorization": "token-a"}, + load_prompts=False, + ) + ) async with tool: await asyncio.wait_for(tool.call_tool("record", marker="outer"), timeout=5) assert writes["token-a"] == ["outer"] diff --git a/python/packages/foundry_hosting/tests/test_toolbox.py b/python/packages/foundry_hosting/tests/test_toolbox.py index 3b9f1d6111a..124fdc1deb3 100644 --- a/python/packages/foundry_hosting/tests/test_toolbox.py +++ b/python/packages/foundry_hosting/tests/test_toolbox.py @@ -15,7 +15,13 @@ import httpx2 as httpx import pytest -from agent_framework import MCPStreamableHTTPTool, SkillsProvider, SkillsSourceContext, SupportsAgentRun +from agent_framework import ( + FunctionInvocationContext, + MCPStreamableHTTPTool, + SkillsProvider, + SkillsSourceContext, + SupportsAgentRun, +) from agent_framework.exceptions import ToolException, ToolExecutionException from azure.ai.agentserver.core import ( FoundryAgentRequestContext, @@ -104,7 +110,7 @@ async def handle(self, request: httpx.Request) -> httpx.Response: headers["mcp-session-id"] = f"session-{self._session_count}-{call_id or 'none'}" result = { "protocolVersion": payload["params"]["protocolVersion"], - "capabilities": {"tools": {}, "resources": {}}, + "capabilities": {"tools": {}, "prompts": {}, "resources": {}}, "serverInfo": {"name": "toolbox-loopback", "version": "1"}, } elif method == "tools/list": @@ -121,6 +127,10 @@ async def handle(self, request: httpx.Request) -> httpx.Response: } elif method == "tools/call": result = {"content": [{"type": "text", "text": call_id or "none"}]} + elif method == "prompts/list": + result = {"prompts": [{"name": "request-prompt"}]} + elif method == "prompts/get": + result = {"messages": [{"role": "user", "content": {"type": "text", "text": call_id or "none"}}]} elif method == "resources/read": uri = payload["params"]["uri"] if uri == "skill://index.json": @@ -505,6 +515,57 @@ async def test_toolbox_header_provider_does_not_retain_prior_platform_context() assert "x-agent-foundry-call-id" not in headers +@pytest.mark.parametrize("generated", [False, True]) +async def test_prompt_in_detached_child_uses_current_platform_context(generated: bool) -> None: + server = _ToolboxLoopbackServer() + start_child = asyncio.Event() + children: list[asyncio.Task[str]] = [] + toolbox = FoundryToolbox( + _FakeCredential(), # type: ignore + url="https://loopback.invalid/mcp", + load_prompts=True, + ) + await _attach_loopback_transport(toolbox, server) + + async def call_prompt() -> str: + await start_child.wait() + token = set_request_context(FoundryAgentRequestContext(call_id="CALL-B", user_id="same-user")) + try: + if generated: + prompt = next(function for function in toolbox.functions if function.name == "request-prompt") + context = FunctionInvocationContext(function=prompt, arguments={}, kwargs={}) + contents = await prompt.invoke(arguments={}, context=context) + assert len(contents) == 1 + assert contents[0].text is not None + return contents[0].text + return await toolbox.get_prompt("request-prompt") + finally: + reset_request_context(token) + + async def spawn_child(request: httpx.Request) -> None: + if json.loads(request.content).get("method") == "tools/call": + children.append(asyncio.create_task(call_prompt())) + + server.before_request = spawn_child + try: + async with asyncio.timeout(5): + await _run_with_request_context("CALL-A", toolbox.connect) + await _run_with_request_context("CALL-A", lambda: toolbox.call_tool("record", marker="parent")) + assert len(children) == 1 + start_child.set() + assert await children[0] == "CALL-B" + request = server.requests_for("prompts/get")[-1] + assert request.headers["x-agent-foundry-call-id"] == "CALL-B" + assert request.headers["mcp-session-id"].endswith("-CALL-B") + assert len(server.requests_for("initialize")) == 2 + finally: + start_child.set() + for child in children: + child.cancel() + await asyncio.gather(*children, return_exceptions=True) + await toolbox.close() + + async def test_custom_header_provider_key_error_does_not_drop_platform_headers() -> None: server = _ToolboxLoopbackServer() toolbox = FoundryToolbox( diff --git a/skills/agent-framework-builder-python/references/integrations.md b/skills/agent-framework-builder-python/references/integrations.md index 2a2b2c96ff5..f474f09c960 100644 --- a/skills/agent-framework-builder-python/references/integrations.md +++ b/skills/agent-framework-builder-python/references/integrations.md @@ -38,7 +38,10 @@ Calls with identical effective static and dynamic headers can run concurrently o any effective header, including the Foundry platform call ID, waits for active operations and reconnects before refreshing discovery. Token refresh performed by HTTP authentication remains request-local. -Discovery, reset, close, and connection recovery are exclusive. Do not reset, close, refresh discovery, or change +Discovery, reset, close, and connection recovery wait for active attempts on framework-owned MCP connections, +including connections without a header provider. Same-identity Agent run preparation can overlap active calls. +Direct prompts resolve current headers; generated prompts resolve host runtime kwargs separately from model arguments. +Do not reset, close, refresh discovery, or change header identity from a callback that must finish an active operation on that same tool; it raises an error rather than waiting on itself. Concurrent calls still share the server's MCP session state, so applications must order state-dependent operations themselves. Cancelling a local call does not guarantee that its remote side effects stop. From 9f5489d599636cb7c6223c13fdf2e57f91e955e0 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Fri, 6 Nov 2026 15:38:07 +0100 Subject: [PATCH 3/4] Python: Harden MCP request ownership and discovery recovery --- .../0050-python-mcp-v2-client-lifecycle.md | 27 +++ python/packages/core/AGENTS.md | 4 + python/packages/core/README.md | 4 + python/packages/core/agent_framework/_mcp.py | 65 +++++-- .../core/tests/core/mcp_http_helpers.py | 64 +++++++ python/packages/core/tests/core/test_mcp.py | 34 ++-- .../core/tests/core/test_mcp_http_auth.py | 171 +++++++++++++++++- .../_toolbox.py | 4 +- .../foundry_hosting/tests/test_toolbox.py | 52 ++++++ .../references/integrations.md | 3 + 10 files changed, 394 insertions(+), 34 deletions(-) diff --git a/docs/decisions/0050-python-mcp-v2-client-lifecycle.md b/docs/decisions/0050-python-mcp-v2-client-lifecycle.md index 279c2e40653..d411398efa5 100644 --- a/docs/decisions/0050-python-mcp-v2-client-lifecycle.md +++ b/docs/decisions/0050-python-mcp-v2-client-lifecycle.md @@ -128,6 +128,33 @@ Caller-supplied `ClientSession` connections always call the server directly. Age shared cache store or cache-mode configuration. Reconnect or an effective header-identity change replaces the framework-owned Client and its per-client cache. +### Concurrent HTTP operations and identity transitions + +Concurrent operations on a framework-owned HTTP connection may share a session only when their complete effective +header identity matches. Agent Framework copies owner-tagged header snapshots onto individual outgoing requests, +including redirects, rather than retaining a mutable active-call fallback. The supported MCP SDK propagates the +sending operation's ContextVars into its HTTP tasks. Transport request entry points must retain the owner tag; +unsupported untagged sending methods fail explicitly instead of silently dropping dynamic headers. + +Admission tracks network attempts, not complete logical calls. Retries release the failed attempt before recovery, +then reconcile the original identity and rebuild arguments and metadata from the replacement catalog. Identity +changes, explicit discovery, reset, close, and recovery require exclusive access and wait for active attempts. +A queued exclusive operation blocks unrelated arrivals, including new calls with the old identity. Same-identity +callbacks needed by an active attempt may finish first; invocation depth keeps that permission live through nested +operation scopes. Exclusive operations from active callbacks fail explicitly rather than waiting on their parent. +Resource reads enter admission before requiring a ready connection, including during replacement initialization. + +This request-owned approach is preferred over a lease spanning an entire logical call: it avoids shared mutable +request headers, permits recovery after releasing the failed attempt, and distinguishes required nested callbacks +from unrelated arrivals. Its cost is explicit admission and callback-lifetime tracking. It does not create independent +server-side session state, so applications must still order state-dependent operations. + +Catalog cursors and staged pages belong to one connection. Replacement restarts the complete snapshot once from +its first page; a second replacement fails explicitly. Configured discovery also reloads tools if prompt discovery +replaces their connection. Caller-owned Clients and ClientSessions remain caller-owned and are never recreated. +Admission/reconciliation and recovery have separate `request_timeout` budgets, not one end-to-end deadline; +`None` disables those budgets. Cleanup remains owned by the lifecycle task after waiter cancellation. + ### MRTR For both framework-owned and caller-owned high-level connections, `tools/call`, `prompts/get`, and `resources/read` diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index e8934a4d963..5f356cfb094 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -303,6 +303,8 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID Keep the original trace context when rebuilding metadata. Cancelled lifecycle waiters leave their owner running; new HTTP operations must wait for that setup or teardown to finish. A same-generation nested callback may finish ahead of a queued identity change, but exclusive operations from active callbacks fail instead of deadlocking. + Track invocation activity by nesting depth. Resource reads enter admission before requiring the connection. + Transport sending methods must tag request ownership or fail explicitly; do not delegate them unscoped. - **MCP recovery and deadlines** - Failed tool, prompt, discovery, and liveness operations record their connection and close generation. Central recovery reuses a replacement established by another operation, releases the lifecycle request lock before discovery, and rejects recovery from before an explicit close. Advance the close @@ -310,6 +312,8 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID and connection recovery have separate `request_timeout` budgets; neither wraps the admitted RPC or cleanup. `None` disables the budgets. Keep recovery's deadline outside the HTTP barrier so waiting for active attempts is bounded too. + Replacement during discovery discards staged pages and continuation cursors and restarts the complete catalog + once. Another replacement fails explicitly. Configured prompt recovery also refreshes tools from the new session. - **Streamable HTTP cookies** - Framework-created clients reject response-cookie persistence with or without a `header_provider`; explicit `Cookie` headers from `static_headers` or a provider remain supported. Caller-provided clients retain their cookie behavior and ownership. Applications requiring cookie persistence must scope clients and MCP sessions to one authenticated principal. Cookie rejection does not isolate MCP protocol sessions or other server-side state. - **`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`. diff --git a/python/packages/core/README.md b/python/packages/core/README.md index 5f3d85ca5a1..6095f24a7ed 100644 --- a/python/packages/core/README.md +++ b/python/packages/core/README.md @@ -207,6 +207,10 @@ inherited invocation snapshot; generated prompts resolve host runtime kwargs separately from model arguments. Calls still share server-side session state; applications must order state-dependent operations themselves. +Resource reads wait for replacement initialization before requiring the connection. +Catalog discovery restarts the complete snapshot once after a connection replacement, +without reusing old pages or cursors. A second replacement fails explicitly. + `request_timeout` bounds header admission and session reconciliation together, and connection recovery separately from the RPC timeout. These are separate budgets, not an end-to-end deadline. `None` disables these budgets. An admission timeout does diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 01df2e6bf03..41e37a0d9ff 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -268,12 +268,16 @@ class _MCPHTTPInvocation: parent: _MCPHTTPInvocation | None runtime_kwargs: Mapping[str, Any] connection: _MCPConnection | None = None - active: bool = False + active_depth: int = 0 _mcp_http_invocation: ContextVar[_MCPHTTPInvocation | None] = ContextVar("_mcp_http_invocation", default=None) +class _MCPDiscoveryRestart(Exception): + """Restart staged discovery when its connection is replaced.""" + + @dataclass class _MCPHostPayloadCapture: """Collect one generated MCP result without affecting its public return shape.""" @@ -551,14 +555,12 @@ def __init__(self, client: httpx2.AsyncClient, owner: object) -> None: self._owner = owner def __getattr__(self, name: str) -> Any: - # Delegate the rest of the httpx client surface so this wrapper stays a - # drop-in for the MCP transport. Only the request-sending methods below - # are wrapped; anything the transport reads (timeouts, headers, ...) - # comes straight from the caller's client. ``_client`` itself is always a - # real instance attribute, so guard against recursing on a partially - # initialized wrapper. if name == "_client": raise AttributeError(name) + if name in {"get", "post", "put", "patch", "head", "options"}: + raise RuntimeError( + f"MCP HTTP transport cannot use unscoped {name}(); use request(), stream(), send(), or delete()." + ) return getattr(self._client, name) def _tagged_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: @@ -570,6 +572,9 @@ def _tagged_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: def stream(self, *args: Any, **kwargs: Any) -> Any: return self._client.stream(*args, **self._tagged_kwargs(kwargs)) + async def request(self, *args: Any, **kwargs: Any) -> httpx2.Response: + return await self._client.request(*args, **self._tagged_kwargs(kwargs)) + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: request.extensions[_MCP_HEADER_OWNER_EXTENSION] = self._owner return await self._client.send(request, **kwargs) @@ -3070,14 +3075,34 @@ def _determine_approval_mode( async def _load_configured_discovery_locked(self) -> None: """Load configured session discovery while the caller holds the discovery lock.""" + await self._load_discovery_snapshot(self._load_configured_discovery_once) + + async def _load_configured_discovery_once(self) -> None: if self.load_tools_flag: if self._supports_tools: await self._load_tools_locked() self._tools_loaded = True + connection = self._connection if self.load_prompts_flag: if self._supports_prompts: await self._load_prompts_locked() self._prompts_loaded = True + if connection is not self._connection: + # Prompt recovery also invalidates any tools loaded before the replacement. + raise _MCPDiscoveryRestart + + async def _load_discovery_snapshot(self, load_snapshot: Callable[[], Awaitable[None]]) -> None: + """Discard staged pages and restart once after connection replacement.""" + for attempt in range(2): + try: + await load_snapshot() + return + except _MCPDiscoveryRestart as exc: + if attempt: + raise _mcp_tool_execution_exception( + "Failed to load MCP discovery - connection changed again after restarting the catalog.", + inner_exception=exc, + ) from exc async def load_prompts(self) -> None: """Load prompts from the MCP server. @@ -3093,6 +3118,9 @@ async def load_prompts(self) -> None: await self._load_prompts_locked() async def _load_prompts_locked(self) -> None: + await self._load_discovery_snapshot(self._load_prompts_snapshot) + + async def _load_prompts_snapshot(self) -> None: from mcp import types if not self._supports_prompts: @@ -3223,7 +3251,7 @@ async def _try_list_page_with_recovery( kind: Literal["tool", "prompt"], request_page: Callable[[types.PaginatedRequestParams | None], Awaitable[_MCPListResultT]], ) -> _MCPListResultT | None: - """Fetch one catalog page, retrying after connection loss.""" + """Fetch one catalog page and restart its snapshot after connection replacement.""" from anyio import ClosedResourceError from tenacity import RetryCallState, retry, retry_if_exception_type, stop_after_attempt @@ -3246,6 +3274,8 @@ async def recover_connection(retry_state: RetryCallState) -> None: failed_connection=failed_connection, close_generation=close_generation, ) + if self._connection is not failed_connection: + raise _MCPDiscoveryRestart from exception @retry( retry=retry_if_exception_type(ClosedResourceError), @@ -3255,7 +3285,10 @@ async def recover_connection(retry_state: RetryCallState) -> None: ) async def request_page_with_retry() -> _MCPListResultT | None: nonlocal failed_connection, close_generation + previous_connection = self._connection await self._ensure_connected() + if previous_connection is not None and self._connection is not previous_connection: + raise _MCPDiscoveryRestart failed_connection = self._connection close_generation = self._close_generation @@ -3275,7 +3308,10 @@ async def request_page_with_retry() -> _MCPListResultT | None: capture_unhandled_exceptions=False, ) as list_span: try: - return await request_page(params) + result = await request_page(params) + if self._connection is not failed_connection: + raise _MCPDiscoveryRestart + return result except Exception as exc: set_mcp_span_error(list_span, type(exc).__name__, str(exc)) raise @@ -3310,6 +3346,9 @@ async def load_tools(self) -> None: await self._load_tools_locked() async def _load_tools_locked(self) -> None: + await self._load_discovery_snapshot(self._load_tools_snapshot) + + async def _load_tools_snapshot(self) -> None: from mcp import types if not self._supports_tools: @@ -4762,7 +4801,7 @@ def _http_headers( def _active_http_invocation(self) -> _MCPHTTPInvocation | None: invocation = _mcp_http_invocation.get() while invocation is not None: - if invocation.owner is self and invocation.active: + if invocation.owner is self and invocation.active_depth: return invocation invocation = invocation.parent return None @@ -4834,18 +4873,18 @@ async def _http_operation(self) -> AsyncGenerator[None]: finally: self._header_exclusive_owner = None invocation.connection = self._connection - invocation.active = True + invocation.active_depth += 1 self._header_operations += 1 self._header_operations_idle.clear() if ancestor is not None: invocation.connection = self._connection - invocation.active = True + invocation.active_depth += 1 self._header_operations += 1 self._header_operations_idle.clear() try: yield finally: - invocation.active = False + invocation.active_depth -= 1 self._header_operations -= 1 if not self._header_operations: self._header_operations_idle.set() diff --git a/python/packages/core/tests/core/mcp_http_helpers.py b/python/packages/core/tests/core/mcp_http_helpers.py index 82aadd9d4db..a5f5595cc78 100644 --- a/python/packages/core/tests/core/mcp_http_helpers.py +++ b/python/packages/core/tests/core/mcp_http_helpers.py @@ -2,8 +2,11 @@ """Deterministic request barriers for MCP HTTP concurrency tests.""" +from __future__ import annotations + import asyncio import json +from typing import Any import httpx2 as httpx @@ -23,3 +26,64 @@ async def before_request(self, request: httpx.Request) -> None: async def next_request(self) -> httpx.Request: return await asyncio.wait_for(self.arrivals.get(), timeout=5) + + +class MCPPaginatedHTTPServer: + """Serve distinct catalogs and cursors for each initialized connection.""" + + def __init__(self, *, reject_old_cursor: bool = False) -> None: + self.generation = 0 + self.reject_old_cursor = reject_old_cursor + self.pages: list[tuple[str, int, str | None]] = [] + + async def handle(self, request: httpx.Request) -> httpx.Response: + 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 "id" not in body: + return httpx.Response(202) + result: dict[str, Any] = {} + headers: dict[str, str] = {} + if method == "initialize": + self.generation += 1 + headers["mcp-session-id"] = f"session-{self.generation}" + result = { + "protocolVersion": body["params"]["protocolVersion"], + "capabilities": {"tools": {}, "prompts": {}}, + "serverInfo": {"name": "paginated-test", "version": "1"}, + } + elif method in {"tools/list", "prompts/list"}: + cursor = body.get("params", {}).get("cursor") + self.pages.append((method, self.generation, cursor)) + if cursor == "old-cursor" and self.generation > 1 and self.reject_old_cursor: + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": body["id"], + "error": {"code": -32602, "message": "Cursor belongs to an old session"}, + }, + ) + suffix = "tail" if cursor else "head" + kind = "tool" if method == "tools/list" else "prompt" + item: dict[str, Any] = {"name": f"session_{self.generation}_{kind}_{suffix}"} + if kind == "tool": + item.update( + inputSchema={"type": "object", "properties": {f"arg{self.generation}": {"type": "string"}}}, + _meta={"catalog": self.generation}, + ) + else: + item["arguments"] = [{"name": f"arg{self.generation}"}] + result = {"tools" if kind == "tool" else "prompts": [item]} + if cursor is None: + result["nextCursor"] = "old-cursor" if self.generation == 1 else f"cursor-{self.generation}" + elif method == "tools/call": + result = {"content": [{"type": "text", "text": "ok"}]} + return httpx.Response( + 200, + headers=headers, + json={"jsonrpc": "2.0", "id": body["id"], "result": result}, + ) diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index a3463b26f16..9362851adb1 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -9621,7 +9621,7 @@ async def reconnect() -> None: second_session.list_tools.assert_awaited_once() -async def test_load_tools_retries_continuation_page_after_closed_resource() -> None: +async def test_load_tools_restarts_snapshot_after_closed_resource_on_continuation_page() -> None: from anyio import ClosedResourceError tool = MCPTool(name="test_tool", load_tools=True) # type: ignore[abstract] # ty: ignore[call-non-callable] @@ -9655,13 +9655,12 @@ async def reconnect() -> None: with patch.object(tool, "_reconnect_without_loading", AsyncMock(side_effect=reconnect)) as mock_reconnect: await tool.load_tools() - assert [function.name for function in tool.functions] == ["first", "second"] + assert [function.name for function in tool.functions] == ["second"] assert first_session.list_tools.await_count == 2 mock_reconnect.assert_awaited_once_with() - continuation_call = second_session.list_tools.await_args - assert continuation_call is not None - continuation_params = continuation_call.kwargs["params"] - assert continuation_params.cursor == "second" + restarted_call = second_session.list_tools.await_args + assert restarted_call is not None + assert restarted_call.kwargs.get("params") is None async def test_load_tools_fails_after_closed_resource_retry_is_exhausted() -> None: @@ -9678,17 +9677,17 @@ async def test_load_tools_fails_after_closed_resource_retry_is_exhausted() -> No second_session.list_tools = AsyncMock(side_effect=ClosedResourceError()) async def reconnect() -> None: - tool.session = second_session + _set_framework_owned_session(tool, second_session) tool._supports_tools = True with ( patch.object(tool, "_reconnect_without_loading", AsyncMock(side_effect=reconnect)) as mock_reconnect, - pytest.raises(ToolExecutionException, match="Failed to load tools - connection lost"), + pytest.raises(ToolExecutionException, match="connection changed again"), ): await tool.load_tools() first_session.list_tools.assert_awaited_once() - mock_reconnect.assert_awaited_once_with() + assert mock_reconnect.await_count == 2 second_session.list_tools.assert_awaited_once() assert tool.functions == [] @@ -9722,7 +9721,7 @@ async def reconnect() -> None: second_session.list_prompts.assert_awaited_once() -async def test_load_prompts_retries_continuation_page_after_closed_resource() -> None: +async def test_load_prompts_restarts_snapshot_after_closed_resource_on_continuation_page() -> None: from anyio import ClosedResourceError tool = MCPTool(name="test_tool", load_prompts=True) # type: ignore[abstract] # ty: ignore[call-non-callable] @@ -9750,13 +9749,12 @@ async def reconnect() -> None: with patch.object(tool, "_reconnect_without_loading", AsyncMock(side_effect=reconnect)) as mock_reconnect: await tool.load_prompts() - assert [function.name for function in tool.functions] == ["first", "second"] + assert [function.name for function in tool.functions] == ["second"] assert first_session.list_prompts.await_count == 2 mock_reconnect.assert_awaited_once_with() - continuation_call = second_session.list_prompts.await_args - assert continuation_call is not None - continuation_params = continuation_call.kwargs["params"] - assert continuation_params.cursor == "second" + restarted_call = second_session.list_prompts.await_args + assert restarted_call is not None + assert restarted_call.kwargs.get("params") is None async def test_load_prompts_fails_after_closed_resource_retry_is_exhausted() -> None: @@ -9773,17 +9771,17 @@ async def test_load_prompts_fails_after_closed_resource_retry_is_exhausted() -> second_session.list_prompts = AsyncMock(side_effect=ClosedResourceError()) async def reconnect() -> None: - tool.session = second_session + _set_framework_owned_session(tool, second_session) tool._supports_prompts = True with ( patch.object(tool, "_reconnect_without_loading", AsyncMock(side_effect=reconnect)) as mock_reconnect, - pytest.raises(ToolExecutionException, match="Failed to load prompts - connection lost"), + pytest.raises(ToolExecutionException, match="connection changed again"), ): await tool.load_prompts() first_session.list_prompts.assert_awaited_once() - mock_reconnect.assert_awaited_once_with() + assert mock_reconnect.await_count == 2 second_session.list_prompts.assert_awaited_once() assert tool.functions == [] 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 6311e85c0be..f48fc114b2c 100644 --- a/python/packages/core/tests/core/test_mcp_http_auth.py +++ b/python/packages/core/tests/core/test_mcp_http_auth.py @@ -17,7 +17,7 @@ from agent_framework import Agent, FunctionInvocationContext, MCPStreamableHTTPTool from agent_framework.exceptions import ToolException, ToolExecutionException -from .mcp_http_helpers import MCPHTTPBarrier +from .mcp_http_helpers import MCPHTTPBarrier, MCPPaginatedHTTPServer MCPHTTPServer: TypeAlias = tuple[httpx.AsyncClient, list[httpx.Request], dict[str, list[str]]] MCPHTTPClientFactory: TypeAlias = tuple[Callable[..., httpx.AsyncClient], list[httpx.Request], dict[str, list[str]]] @@ -701,6 +701,175 @@ async def test_request_header_snapshot_cannot_be_reused_by_another_tool_owner( assert "X-First-Only" not in reused.headers +async def test_alternate_http_request_entry_point_keeps_owner_and_header_snapshot( + mcp_http_server: MCPHTTPServer, +) -> None: + from agent_framework._mcp import _MCP_HEADER_OWNER_EXTENSION, _MCPHeaderScopedClient + + client, requests, _ = mcp_http_server + tool = _tool(client, "token-a") + extensions = {"test-extension": "retained"} + async with tool: + scoped = _MCPHeaderScopedClient(client, tool._header_request_owner) + with tool._http_headers({"Authorization": "token-b"}): + response = await scoped.request("POST", "https://mcp.example/unrelated", json={}, extensions=extensions) + assert response.status_code == 200 + assert requests[-1].headers["Authorization"] == "token-b" + assert requests[-1].extensions[_MCP_HEADER_OWNER_EXTENSION] is tool._header_request_owner + assert requests[-1].extensions["test-extension"] == "retained" + assert extensions == {"test-extension": "retained"} + assert scoped.headers is client.headers + + +@pytest.mark.parametrize("method", ["get", "post", "put", "patch", "head", "options"]) +async def test_unwrapped_http_sending_methods_fail_before_sending(mcp_http_server: MCPHTTPServer, method: str) -> None: + from agent_framework._mcp import _MCPHeaderScopedClient + + client, requests, _ = mcp_http_server + scoped = _MCPHeaderScopedClient(client, object()) + with pytest.raises(RuntimeError, match=f"cannot use unscoped {method}"): + getattr(scoped, method) + assert requests == [] + + +async def test_same_invocation_reentry_keeps_outer_operation_active(mcp_http_server: MCPHTTPServer) -> None: + client, _, writes = mcp_http_server + tool = _tool(client, "token-a") + async with tool, asyncio.timeout(5): + with tool._http_headers({"Authorization": "token-a"}): + async with tool._http_operation(): + invocation = tool._active_http_invocation() + assert invocation is not None + async with tool._http_operation(): + assert tool._header_operations == 2 + assert tool._active_http_invocation() is invocation + assert tool._header_operations == 1 + with pytest.raises(ToolExecutionException, match="from its active HTTP operation"): + await tool.close() + await tool.call_tool("record", marker="nested-after-inner") + assert tool._active_http_invocation() is None + assert tool._header_operations == 0 + assert tool._header_operations_idle.is_set() + assert writes["token-a"] == ["nested-after-inner"] + + +@pytest.mark.parametrize("kind", ["tool", "prompt"]) +@pytest.mark.parametrize("reject_old_cursor", [False, True]) +async def test_discovery_restarts_all_pages_after_connection_replacement(kind: str, reject_old_cursor: bool) -> None: + from anyio import ClosedResourceError + from mcp import Client + + server = MCPPaginatedHTTPServer(reject_old_cursor=reject_old_cursor) + published: list[str] = [] + async with httpx.AsyncClient(transport=httpx.MockTransport(server.handle)) as client: + tool = MCPStreamableHTTPTool( + name="paginated", + url="https://mcp.example/mcp", + http_client=client, + load_tools=False, + load_prompts=False, + ) + load_catalog = tool.load_tools if kind == "tool" else tool.load_prompts + method = "list_tools" if kind == "tool" else "list_prompts" + original = getattr(Client, method) + + async def fail_old_page(sdk_client: Client, *, cursor: str | None = None, **kwargs: Any) -> Any: + if server.generation == 1 and cursor is not None: + raise ClosedResourceError + return await original(sdk_client, cursor=cursor, **kwargs) + + async with tool, asyncio.timeout(10): + tool._function_load_callback = lambda function, _: published.append(function.name) + with patch.object(Client, method, new=fail_old_page): + await load_catalog() + expected = [f"session_2_{kind}_head", f"session_2_{kind}_tail"] + assert [function.name for function in tool.functions] == expected + assert published == expected + assert server.pages == [ + (f"{kind}s/list", 1, None), + (f"{kind}s/list", 2, None), + (f"{kind}s/list", 2, "cursor-2"), + ] + assert server.generation == 2 + if kind == "tool": + assert tool._tool_call_meta_by_name == {name: {"catalog": 2} for name in expected} + assert tool._tool_param_names_by_name == {name: {"arg2"} for name in expected} + + +@pytest.mark.parametrize("kind", ["tool", "prompt"]) +async def test_discovery_replacement_restart_is_bounded_and_keeps_previous_catalog(kind: str) -> None: + from anyio import ClosedResourceError + from mcp import Client + + server = MCPPaginatedHTTPServer() + async with httpx.AsyncClient(transport=httpx.MockTransport(server.handle)) as client: + tool = MCPStreamableHTTPTool( + name="paginated", + url="https://mcp.example/mcp", + http_client=client, + load_tools=False, + load_prompts=False, + ) + load_catalog = tool.load_tools if kind == "tool" else tool.load_prompts + method = "list_tools" if kind == "tool" else "list_prompts" + original = getattr(Client, method) + + async def fail_every_tail(sdk_client: Client, *, cursor: str | None = None, **kwargs: Any) -> Any: + if cursor is not None: + raise ClosedResourceError + return await original(sdk_client, cursor=cursor, **kwargs) + + async with tool, asyncio.timeout(10): + await load_catalog() + previous_functions = list(tool.functions) + previous_meta = dict(tool._tool_call_meta_by_name) + with ( + patch.object(Client, method, new=fail_every_tail), + pytest.raises(ToolExecutionException, match="connection changed again"), + ): + await load_catalog() + assert tool.functions == previous_functions + assert tool._tool_call_meta_by_name == previous_meta + assert server.generation == 3 + + +async def test_configured_discovery_reloads_tools_when_prompt_recovery_replaces_connection() -> None: + from anyio import ClosedResourceError + from mcp import Client + + server = MCPPaginatedHTTPServer() + original = Client.list_prompts + + async def fail_replacement_prompt_page(sdk_client: Client, *, cursor: str | None = None, **kwargs: Any) -> Any: + if server.generation == 2 and cursor is not None: + raise ClosedResourceError + return await original(sdk_client, cursor=cursor, **kwargs) + + async with httpx.AsyncClient(transport=httpx.MockTransport(server.handle)) as client: + tool = MCPStreamableHTTPTool(name="configured", url="https://mcp.example/mcp", http_client=client) + async with tool, asyncio.timeout(10): + connection = tool._require_connection() + assert connection.client is not None + failed_call = AsyncMock(side_effect=ClosedResourceError) + with ( + patch.object(connection.client, "call_tool", new=failed_call), + patch.object(Client, "list_prompts", new=fail_replacement_prompt_page), + ): + result = await tool.call_tool("session_1_tool_head") + assert isinstance(result, list) + assert len(result) == 1 and result[0].text == "ok" + failed_call.assert_awaited_once() + assert server.generation == 3 + assert {function.name for function in tool.functions} == { + "session_3_tool_head", + "session_3_tool_tail", + "session_3_prompt_head", + "session_3_prompt_tail", + } + assert set(tool._tool_call_meta_by_name) == {"session_3_tool_head", "session_3_tool_tail"} + assert tool._tools_loaded and tool._prompts_loaded + + async def test_nested_same_identity_call_can_finish_ahead_of_waiting_identity_change( mcp_http_client_factory: MCPHTTPClientFactory, ) -> None: diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py index 65fe307a44d..d9a02501c56 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py @@ -473,8 +473,8 @@ async def _close_connection(self) -> None: async def _read_resource(self, uri: AnyUrl) -> ReadResourceResult: """Read a Toolbox resource under the current platform-header identity.""" headers = self._effective_resource_headers() - connection = self._require_connection() - if not headers and not connection.is_framework_owned: + connection = self._connection + if not headers and connection is not None and not connection.is_framework_owned: return await connection.read_resource(str(uri)) with self._http_headers(headers): async with self._http_operation(): diff --git a/python/packages/foundry_hosting/tests/test_toolbox.py b/python/packages/foundry_hosting/tests/test_toolbox.py index 77e5bdb46a7..6396df53df1 100644 --- a/python/packages/foundry_hosting/tests/test_toolbox.py +++ b/python/packages/foundry_hosting/tests/test_toolbox.py @@ -742,6 +742,58 @@ async def read_skill() -> tuple[str, str]: assert all(request.headers["X-Custom-Header"] == "resource-value" for request in resource_requests) +async def test_retained_skill_read_waits_for_replacement_initialization() -> None: + server = _ToolboxLoopbackServer() + toolbox = FoundryToolbox(_FakeCredential(), url="https://loopback.invalid/mcp") # type: ignore + await _attach_loopback_transport(toolbox, server) + source = toolbox.as_skills_provider()._source + initializing = asyncio.Event() + release = asyncio.Event() + tasks: list[asyncio.Task[Any]] = [] + + async def block_replacement(request: httpx.Request) -> None: + if ( + json.loads(request.content).get("method") == "initialize" + and request.headers.get("x-agent-foundry-call-id") == "CALL-B" + ): + initializing.set() + await release.wait() + + try: + async with asyncio.timeout(10): + await _run_with_request_context("CALL-A", toolbox.connect) + skills = await _run_with_request_context("CALL-A", lambda: source.get_skills(_source_context())) + assert len(skills) == 1 + server.before_request = block_replacement + transition = asyncio.create_task( + _run_with_request_context("CALL-B", lambda: toolbox.call_tool("record", marker="transition")) + ) + tasks.append(transition) + await initializing.wait() + read = asyncio.create_task(_run_with_request_context("CALL-B", skills[0].get_content)) + tasks.append(read) + await asyncio.sleep(0) + assert not read.done() + assert not any( + request.headers.get("x-agent-foundry-call-id") == "CALL-B" + for request in server.requests_for("resources/read") + ) + release.set() + result = await transition + assert isinstance(result, list) + assert len(result) == 1 and result[0].text == "CALL-B" + assert "Call ID: CALL-B" in await read + request = server.requests_for("resources/read")[-1] + assert request.headers["x-agent-foundry-call-id"] == "CALL-B" + assert request.headers["mcp-session-id"].endswith("-CALL-B") + finally: + release.set() + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + await toolbox.close() + + async def test_toolbox_skills_reject_header_provider_that_requires_runtime_kwargs() -> None: server = _ToolboxLoopbackServer() toolbox = FoundryToolbox( diff --git a/skills/agent-framework-builder-python/references/integrations.md b/skills/agent-framework-builder-python/references/integrations.md index f474f09c960..6341b23f00e 100644 --- a/skills/agent-framework-builder-python/references/integrations.md +++ b/skills/agent-framework-builder-python/references/integrations.md @@ -47,6 +47,9 @@ than waiting on itself. Concurrent calls still share the server's MCP session st state-dependent operations themselves. Cancelling a local call does not guarantee that its remote side effects stop. Caller-owned MCP clients and sessions remain caller-owned; Agent Framework does not recreate them. +Resource reads wait for replacement initialization. Tool and prompt discovery restarts its complete snapshot once +after connection replacement, without retaining old pages or cursors. Another replacement fails explicitly. + `request_timeout` supplies separate budgets for header admission/session reconciliation and connection recovery, not a deadline for the entire call. `None` disables those budgets; caller-owned MCP clients keep their own RPC timeouts. Cleanup may continue after timeout or cancellation. An explicit `close()` interrupts older recovery From 67d1f4158d1d1a0828bcf0a8fb8b62a3745be5fd Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Fri, 6 Nov 2026 15:46:08 +0100 Subject: [PATCH 4/4] Docs: Remove internal MCP notes from consumer skill --- .../references/integrations.md | 24 ------------------- 1 file changed, 24 deletions(-) diff --git a/skills/agent-framework-builder-python/references/integrations.md b/skills/agent-framework-builder-python/references/integrations.md index 6341b23f00e..c19bb242bb4 100644 --- a/skills/agent-framework-builder-python/references/integrations.md +++ b/skills/agent-framework-builder-python/references/integrations.md @@ -31,30 +31,6 @@ For Foundry, the derived OpenAI connection is owned by the connector even when t Do not share a connector-owned connection across overlapping agent contexts; inject an externally managed SDK client when its lifetime must extend beyond any one agent context. Do not assume a closed SDK connection can reopen. -### MCP HTTP concurrency - -`MCPStreamableHTTPTool` and `FoundryToolbox` keep a separate header snapshot on every outgoing HTTP request. -Calls with identical effective static and dynamic headers can run concurrently on one MCP session. A change to -any effective header, including the Foundry platform call ID, waits for active operations and reconnects before -refreshing discovery. Token refresh performed by HTTP authentication remains request-local. - -Discovery, reset, close, and connection recovery wait for active attempts on framework-owned MCP connections, -including connections without a header provider. Same-identity Agent run preparation can overlap active calls. -Direct prompts resolve current headers; generated prompts resolve host runtime kwargs separately from model arguments. -Do not reset, close, refresh discovery, or change -header identity from a callback that must finish an active operation on that same tool; it raises an error rather -than waiting on itself. Concurrent calls still share the server's MCP session state, so applications must order -state-dependent operations themselves. Cancelling a local call does not guarantee that its remote side effects stop. -Caller-owned MCP clients and sessions remain caller-owned; Agent Framework does not recreate them. - -Resource reads wait for replacement initialization. Tool and prompt discovery restarts its complete snapshot once -after connection replacement, without retaining old pages or cursors. Another replacement fails explicitly. - -`request_timeout` supplies separate budgets for header admission/session reconciliation and connection recovery, -not a deadline for the entire call. `None` disables those budgets; caller-owned MCP clients keep their own RPC -timeouts. Cleanup may continue after timeout or cancellation. An explicit `close()` interrupts older recovery -without preventing a later explicit `connect()`. - ## Capability categories - Use function tools for application-owned code and MCP for external tool services.