diff --git a/docs/decisions/0050-python-mcp-v2-client-lifecycle.md b/docs/decisions/0050-python-mcp-v2-client-lifecycle.md index 279c2e40653..b9fa52262b3 100644 --- a/docs/decisions/0050-python-mcp-v2-client-lifecycle.md +++ b/docs/decisions/0050-python-mcp-v2-client-lifecycle.md @@ -92,6 +92,27 @@ lifecycle. Constructor combinations that mix those ownership modes or configure owns raise `ValueError` instead of silently ignoring parameters. Temporarily retained legacy constructor and `ClientSession` combinations remain compatible and will be removed in vNext. +### Managed connection readiness + +For framework-created connections, the lifecycle owner completes transport setup, +configured catalog discovery, recovery, and teardown. Operations wait for that +owner before borrowing a connection; transport availability alone is not readiness. +Public managed catalog refreshes also run on the owner, so discovery recovery and +failure cleanup never wait on a lifecycle request held by another task while +holding the discovery lock. A catalog interrupted by connection replacement is +restarted as a complete snapshot rather than resumed using the previous session's cursor. + +Same-header HTTP operations share identity leases and execute concurrently outside +admission synchronization. Identity changes wait for leases to drain. Lease release +is synchronous, so cancellation cannot prevent releasing the last lease. + +Separating transport recovery and discovery across invocation tasks was rejected: +it exposes a replacement before its schemas and metadata are ready and creates +cross-task lock-ordering dependencies. Keeping the complete transition on the +existing owner preserves one authority without depending on transport-specific +ContextVar propagation. Caller-supplied connections do not receive these additional +lifecycle or synchronization guarantees. + Catalog page and item limits apply after the MCP SDK parses each response. They do not establish a universal transport response-size bound; applications must connect to MCP servers trusted to return reasonably sized responses and configure transport-specific limits where the selected transport exposes them. diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index f4e21ced608..fec2dbb6801 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -289,7 +289,8 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID independently listing tools; `SecureMCPToolProxy.refresh_labels()` explicitly performs a bounded tool refresh before reapplying local policy. - **`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 HTTP header request scoping** - 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. For managed connections, same-identity tools, generated prompts, and Foundry Toolbox resources share a header snapshot while running concurrently. An admission lock and an idle event coordinate identity changes; different identities wait outside the lock, and active-operation release never awaits a lock. Agent runs reconcile a connected tool before copying its functions; generated calls perform the same check at invocation. Connected run preparation defers missing runtime values; invocation-time resolution remains strict. The lifecycle owner completes identity replacement and discovery before admitting further operations. Dynamic headers on a caller-supplied session with unknown identity remain rejected. Ambient requests retain the bound set. Shared HTTP-client hooks stay tool-scoped and strip injected headers from cross-origin redirects. Session cleanup removes request hooks after transport shutdown and closes framework-created HTTP clients. Failed discovery rolls back partial state; constructor-supplied connections remain caller-owned and do not gain managed synchronization guarantees. +- **MCP concurrent recovery** - Each failed operation records its connection and close generation. Central recovery submits the complete recovery action to the lifecycle owner, which reuses an already-ready replacement and performs discovery/cleanup on that same task. Public managed catalog discovery is owner-serialized and restarts a complete snapshot if the connection changes. Calls wait for lifecycle readiness and reprepare tool arguments/metadata after recovery. Explicit `close()` interrupts older recovery without preventing subsequent explicit connection or discovery. Header admission and recovery each use `request_timeout` as a separate budget. The prior split recovery path and per-admission synchronize request are not used. - **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. - **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`. diff --git a/python/packages/core/README.md b/python/packages/core/README.md index b3b9c75bfc2..88a02bc45f1 100644 --- a/python/packages/core/README.md +++ b/python/packages/core/README.md @@ -179,6 +179,39 @@ asyncio.run(main()) You can explore additional agent samples [here](https://github.com/microsoft/agent-framework/tree/main/python/samples/02-agents). +### MCP HTTP concurrency and timeouts + +For connections created by `MCPTool`, `MCPStreamableHTTPTool` permits concurrent tool +and prompt calls when their effective HTTP headers match. +Fixed and dynamic headers together identify the +MCP session; header names are compared case-insensitively and values case-sensitively. +A caller with different headers waits for active operations to finish before the +framework reconnects and refreshes session discovery. Foundry Toolbox resource reads +use the same coordination. + +Direct `get_prompt()` calls resolve current ambient headers; generated prompts +use their invocation's runtime headers. Prompt arguments never supply authentication. +Foundry Toolbox closes its owned HTTP client on the lifecycle owner after MCP +teardown, even if the caller waiting for `close()` is cancelled. + +The lifecycle owner completes connection setup and configured discovery before +admitting operations. Concurrent failures reuse a replacement only after that +replacement's catalogs are ready. Tool retries reprepare arguments and metadata +from the refreshed catalog. Catalog discovery restarts from the beginning if its +managed connection changes; it does not combine pages from different sessions. + +Header admission, session reconciliation, and connection recovery use `request_timeout` +as separate budgets (30 seconds by default; `None` disables them). These budgets do +not form an end-to-end deadline. Admission timeout does not cancel another caller's +active tool call. Releasing a call's header lease does not wait on a lock, including +when the call is cancelled. +Calling `close()` interrupts in-flight recovery; it does not permit that recovery +to reopen the session. Subsequent explicit connection or discovery remains supported. + +Caller-supplied `mcp_client` and `session` objects retain their existing ownership +and passthrough behavior. Their lifecycle and synchronization remain the caller's +responsibility; these managed-connection readiness guarantees do not apply to them. + ### MCP HTTP client ownership and cookies **Breaking change:** HTTP clients created by `MCPStreamableHTTPTool` no longer persist diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 32b5bbdc3b2..9091e585030 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, Mapping, Sequence +from contextlib import AsyncExitStack, _AsyncGeneratorContextManager, asynccontextmanager # type: ignore from contextvars import ContextVar from copy import copy from dataclasses import dataclass @@ -609,6 +609,22 @@ async def read_resource(self, uri: str) -> types.ReadResourceResult: ... +@dataclass +class _MCPLifecycleRequest: + action: Literal["connect", "recover", "close", "load_tools", "load_prompts"] + reset: bool + load_configured: bool + reset_discovery: bool + future: asyncio.Future[None] + acknowledged: asyncio.Future[bool] + failed_connection: _MCPConnection | None = None + close_generation: int | None = None + + +class _MCPDiscoveryRestart(Exception): + """Restart a catalog snapshot after its connection is replaced.""" + + class _ClientMCPConnection: """MCP connection backed by a high-level Client.""" @@ -1253,10 +1269,11 @@ def __init__( self._lifecycle_lock = asyncio.Lock() self._lifecycle_request_lock = asyncio.Lock() self._function_load_lock = asyncio.Lock() - self._lifecycle_queue: ( - asyncio.Queue[tuple[str, bool, bool, bool, asyncio.Future[None], asyncio.Future[bool]]] | None - ) = None + self._lifecycle_queue: asyncio.Queue[_MCPLifecycleRequest] | None = None self._lifecycle_owner_task: asyncio.Task[None] | None = None + self._lifecycle_idle = asyncio.Event() + self._lifecycle_idle.set() + self._close_generation = 0 if mcp_client is not None and session is not None: raise ValueError("`mcp_client` and `session` are mutually exclusive. If possible, prefer `mcp_client`.") @@ -2155,21 +2172,20 @@ async def _run_lifecycle_owner(self) -> None: stop_error: BaseException | None = None try: while True: - action, reset, load_configured, reset_discovery, future, acknowledged = await queue.get() - if action == "connect" and future.cancelled(): + request = await queue.get() + action, future = request.action, request.future + if action != "close" and future.cancelled(): if not self.is_connected and queue.empty(): return + if queue.empty(): + self._lifecycle_idle.set() continue + previous_connection = self._connection + previously_connected = self.is_connected try: - if action == "connect": - previous_connection = self._connection - previously_connected = self.is_connected - await self._connect_on_owner( - reset=reset, - load_configured=load_configured, - reset_discovery=reset_discovery, - ) + await self._execute_lifecycle_request(request) + if action != "close": new_connection = not previously_connected or self._connection is not previous_connection accepted = False try: @@ -2177,24 +2193,27 @@ async def _run_lifecycle_owner(self) -> None: future.set_result(None) # A completed result future cannot tell us that its waiter was # cancelled before consuming it. Await explicit caller acceptance. - accepted = await acknowledged + accepted = await request.acknowledged finally: if not accepted and new_connection: await self._close_on_owner() - if reset_discovery: + if request.reset_discovery: self._reset_session_discovery_state() if not accepted and new_connection and queue.empty(): return - elif action == "close": - await self._close_on_owner() - else: - raise RuntimeError(f"Unknown MCP lifecycle action: {action}") except asyncio.CancelledError as ex: stop_error = ex if not future.done(): future.set_exception(ex) raise except Exception as ex: + if ( + self._connection is not previous_connection + and self._connection is not None + and self._connection.is_framework_owned + ): + await self._close_on_owner() + self._reset_session_discovery_state() if not future.done(): future.set_exception(ex) else: @@ -2211,11 +2230,14 @@ async def _run_lifecycle_owner(self) -> None: # cancelled-connect branch above. The connected check matters because # is_connected is set before tools and prompts are loaded: when loading fails # the session is live and still needs this owner to close it later. - if action == "connect" and not self.is_connected and queue.empty(): + if action != "close" and not self.is_connected and queue.empty(): return else: if not future.done(): future.set_result(None) + finally: + if queue.empty(): + self._lifecycle_idle.set() if action == "close": return @@ -2225,14 +2247,58 @@ async def _run_lifecycle_owner(self) -> None: finally: while True: try: - _, _, _, _, future, _ = queue.get_nowait() + request = queue.get_nowait() except asyncio.QueueEmpty: break - if not future.done(): - future.set_exception(stop_error or RuntimeError("MCP lifecycle owner stopped unexpectedly.")) + if not request.future.done(): + request.future.set_exception( + stop_error or RuntimeError("MCP lifecycle owner stopped unexpectedly.") + ) self._lifecycle_queue = None self._lifecycle_owner_task = None + self._lifecycle_idle.set() + + async def _execute_lifecycle_request(self, request: _MCPLifecycleRequest) -> None: + """Keep transport changes and catalog publication on the lifecycle owner.""" + if request.close_generation is not None and request.close_generation != self._close_generation: + raise _mcp_tool_execution_exception("MCP lifecycle operation was interrupted by close().") + if request.action == "recover": + if not self.is_connected or self._connection is request.failed_connection: + await self._reconnect_without_loading() + self._reset_session_discovery_state() + if request.load_configured: + try: + async with self._function_load_lock: + await self._load_configured_discovery_locked() + except (Exception, asyncio.CancelledError): + await self._close_on_owner() + self._reset_session_discovery_state() + raise + elif request.action == "connect": + await self._connect_on_owner( + reset=request.reset, + load_configured=request.load_configured, + reset_discovery=request.reset_discovery, + ) + elif request.action == "close": + await self._close_on_owner() + else: + previous_connection = self._connection + if request.action == "load_tools": + await self.load_tools() + else: + await self.load_prompts() + if ( + self._connection is not previous_connection + and self._connection is not None + and self._connection.client is not None + ): + async with self._function_load_lock: + await self._load_configured_discovery_locked() + if request.close_generation is not None and request.close_generation != self._close_generation: + await self._close_on_owner() + raise _mcp_tool_execution_exception("MCP lifecycle operation was interrupted by close().") def _is_lifecycle_owner_task(self) -> bool: owner_task = self._lifecycle_owner_task @@ -2240,34 +2306,29 @@ def _is_lifecycle_owner_task(self) -> bool: async def _run_on_lifecycle_owner( self, - action: str, + action: Literal["connect", "recover", "close", "load_tools", "load_prompts"], *, reset: bool = False, load_configured: bool = True, reset_discovery: bool = False, + failed_connection: _MCPConnection | None = None, + close_generation: int | None = None, ) -> None: await self._ensure_lifecycle_owner() + future = asyncio.get_running_loop().create_future() + acknowledged: asyncio.Future[bool] = asyncio.get_running_loop().create_future() + request = _MCPLifecycleRequest( + action, reset, load_configured, reset_discovery, future, acknowledged, failed_connection, close_generation + ) if self._is_lifecycle_owner_task(): - if action == "connect": - await self._connect_on_owner( - reset=reset, - load_configured=load_configured, - reset_discovery=reset_discovery, - ) - elif action == "close": - await self._close_on_owner() - else: - raise RuntimeError(f"Unknown MCP lifecycle action: {action}") + await self._execute_lifecycle_request(request) return - queue = self._lifecycle_queue if queue is None: raise RuntimeError("MCP lifecycle owner is not available.") - - future = asyncio.get_running_loop().create_future() - acknowledged: asyncio.Future[bool] = asyncio.get_running_loop().create_future() - await queue.put((action, reset, load_configured, reset_discovery, future, acknowledged)) + self._lifecycle_idle.clear() + await queue.put(request) accepted = False try: await future @@ -2276,6 +2337,18 @@ async def _run_on_lifecycle_owner( if not acknowledged.done(): acknowledged.set_result(accepted) + async def _wait_for_lifecycle(self) -> None: + """Wait for a managed connection's complete lifecycle action, including discovery.""" + connection = self._connection + if self._is_lifecycle_owner_task() or (connection is not None and not connection.is_framework_owned): + return + async with asyncio.timeout(self.request_timeout): + while True: + await self._lifecycle_idle.wait() + if not self._lifecycle_idle.is_set(): + continue + return + async def _safe_close_exit_stack(self) -> BaseException | None: """Safely close the exit stack, handling unexpected cleanup failures. @@ -2375,9 +2448,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 + """Recover the failed connection once, reusing a replacement another caller established.""" + 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( @@ -2386,12 +2463,20 @@ 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): + logger.info("MCP connection was lost during %s. Requesting recovery...", operation) + await self._run_on_lifecycle_owner( + "recover", + load_configured=load_configured, + failed_connection=connection, + close_generation=close_generation, + ) + 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 Exception as reconnect_ex: raise _mcp_tool_execution_exception( f"Failed to reconnect to MCP server after {operation}.", @@ -2990,15 +3075,36 @@ def _determine_approval_mode( async def _load_configured_discovery_locked(self) -> None: """Load configured session discovery while the caller holds the discovery lock.""" - if self.load_tools_flag: + if self.load_tools_flag and not self._tools_loaded: if self._supports_tools: await self._load_tools_locked() self._tools_loaded = True - if self.load_prompts_flag: + if self.load_prompts_flag and not self._prompts_loaded: if self._supports_prompts: await self._load_prompts_locked() self._prompts_loaded = True + async def _load_discovery_snapshot( + self, + kind: Literal["tools", "prompts"], + load_snapshot: Callable[[], Awaitable[None]], + ) -> None: + """Restart complete catalog discovery once when its managed connection changes.""" + from tenacity import AsyncRetrying, retry_if_exception_type, stop_after_attempt + + try: + async for attempt in AsyncRetrying( + retry=retry_if_exception_type(_MCPDiscoveryRestart), + stop=stop_after_attempt(2), + reraise=True, + ): + with attempt: + await load_snapshot() + except _MCPDiscoveryRestart as ex: + raise _mcp_tool_execution_exception( + f"MCP connection changed repeatedly during {kind} discovery.", inner_exception=ex + ) from ex + async def load_prompts(self) -> None: """Load prompts from the MCP server. @@ -3009,10 +3115,17 @@ async def load_prompts(self) -> None: ToolExecutionException: If the MCP server is not connected or a configured allow/approval name matches multiple raw remote names. """ + if not self._is_lifecycle_owner_task() and (self._connection is None or self._connection.is_framework_owned): + await self._run_on_lifecycle_owner("load_prompts", close_generation=self._close_generation) + return async with self._function_load_lock: await self._load_prompts_locked() + self._prompts_loaded = True async def _load_prompts_locked(self) -> None: + await self._load_discovery_snapshot("prompts", self._load_prompts_snapshot) + + async def _load_prompts_snapshot(self) -> None: from mcp import types if not self._supports_prompts: @@ -3149,6 +3262,8 @@ async def _try_list_page_with_recovery( catalog_name = f"{kind}s" operation = f"{kind} discovery" + failed_connection: _MCPConnection | None = None + close_generation = self._close_generation async def recover_connection(retry_state: RetryCallState) -> None: outcome = retry_state.outcome @@ -3161,7 +3276,15 @@ async def recover_connection(retry_state: RetryCallState) -> None: operation=operation, inner_exception=exception, load_configured=False, + failed_connection=failed_connection, + close_generation=close_generation, ) + if ( + failed_connection is not None + and failed_connection.client is not None + and self._connection is not failed_connection + ): + raise _MCPDiscoveryRestart @retry( retry=retry_if_exception_type(ClosedResourceError), @@ -3170,7 +3293,16 @@ async def recover_connection(retry_state: RetryCallState) -> None: reraise=True, ) async def request_page_with_retry() -> _MCPListResultT | None: + nonlocal failed_connection, close_generation + close_generation = self._close_generation + previous_connection = self._connection await self._ensure_connected() + if ( + previous_connection is not None + and previous_connection.client is not None + and self._connection is not previous_connection + ): + raise _MCPDiscoveryRestart # on reconnect we may have different capabilities. is_supported = self._supports_tools if kind == "tool" else self._supports_prompts @@ -3188,6 +3320,7 @@ async def request_page_with_retry() -> _MCPListResultT | None: capture_unhandled_exceptions=False, ) as list_span: try: + failed_connection = self._connection return await request_page(params) except Exception as exc: set_mcp_span_error(list_span, type(exc).__name__, str(exc)) @@ -3219,10 +3352,17 @@ async def load_tools(self) -> None: map to the same local name, or a configured allow/approval name matches multiple raw remote names. """ + if not self._is_lifecycle_owner_task() and (self._connection is None or self._connection.is_framework_owned): + await self._run_on_lifecycle_owner("load_tools", close_generation=self._close_generation) + return async with self._function_load_lock: await self._load_tools_locked() + self._tools_loaded = True async def _load_tools_locked(self) -> None: + await self._load_discovery_snapshot("tools", self._load_tools_snapshot) + + async def _load_tools_snapshot(self) -> None: from mcp import types if not self._supports_tools: @@ -3423,6 +3563,7 @@ async def close(self) -> None: Closes the connection and cleans up resources. """ + self._close_generation += 1 if self._is_lifecycle_owner_task(): await self._close_on_owner() return @@ -3460,6 +3601,8 @@ async def _ensure_connected(self) -> None: if not self._ping_available: return + connection = self._connection + close_generation = self._close_generation try: await self._require_connection().check_connection() except MCPError as mcp_exc: @@ -3467,11 +3610,13 @@ async def _ensure_connected(self) -> None: self._ping_available = False logger.debug("Skipping future MCP pings because the server does not support ping.") return - if self._is_connection_loss(mcp_exc) or self._connection is None or self._connection.is_framework_owned: + if self._is_connection_loss(mcp_exc) or connection is None or connection.is_framework_owned: await self._recover_connection_after_loss( operation="connection liveness check", inner_exception=mcp_exc, load_configured=False, + failed_connection=connection, + close_generation=close_generation, ) return raise _mcp_tool_execution_exception( @@ -3483,14 +3628,17 @@ async def _ensure_connected(self) -> None: operation="connection liveness check", inner_exception=connection_ex, load_configured=False, + failed_connection=connection, + close_generation=close_generation, ) except Exception as connection_ex: - connection = self._connection if connection is None or connection.is_framework_owned: await self._recover_connection_after_loss( operation="connection liveness check", inner_exception=connection_ex, load_configured=False, + failed_connection=connection, + close_generation=close_generation, ) return raise _mcp_tool_execution_exception( @@ -3536,9 +3684,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 + parent_span = otel_trace.get_current_span() # Build MCP span attributes for tools/call mcp_span_attrs = self._mcp_base_span_attributes() @@ -3554,19 +3701,19 @@ async def _call_tool( ) as span: return await self._call_tool_with_retries( tool_name, - filtered_kwargs, - meta, + kwargs, parser, span, + parent_span, ) async def _call_tool_with_retries( self, tool_name: str, - filtered_kwargs: dict[str, Any], - meta: types.RequestParamsMeta | None, + kwargs: dict[str, Any], parser: Callable[..., str | list[Content]], span: otel_trace.Span, + parent_span: otel_trace.Span, ) -> str | list[Content]: """Execute the MCP tools/call RPC with retry logic.""" from anyio import ClosedResourceError @@ -3574,8 +3721,16 @@ async def _call_tool_with_retries( for attempt in range(2): span_error_recorded = False + connection = self._connection + close_generation = self._close_generation try: - result = await self._require_connection().call_tool( + await self._wait_for_lifecycle() + connection = self._require_connection() + with otel_trace.use_span( + parent_span, end_on_exit=False, record_exception=False, set_status_on_exception=False + ): + filtered_kwargs, meta = self._prepare_call_kwargs(tool_name, kwargs) + result = await connection.call_tool( tool_name, filtered_kwargs, meta=meta, @@ -3611,6 +3766,8 @@ async def _call_tool_with_retries( operation="tool call", inner_exception=call_ex, load_configured=True, + failed_connection=connection, + close_generation=close_generation, ) except ToolExecutionException as recovery_error: if span.is_recording(): @@ -3732,8 +3889,12 @@ async def get_prompt(self, prompt_name: str, **kwargs: Any) -> str: capture_unhandled_exceptions=False, ) as span: for attempt in range(2): + connection = self._connection + close_generation = self._close_generation try: - prompt_result = await self._require_connection().get_prompt(prompt_name, kwargs) + await self._wait_for_lifecycle() + connection = self._require_connection() + prompt_result = await connection.get_prompt(prompt_name, kwargs) return parser(prompt_result) except ClosedResourceError as cl_ex: if attempt == 0: @@ -3742,6 +3903,8 @@ async def get_prompt(self, prompt_name: str, **kwargs: Any) -> str: operation="prompt call", inner_exception=cl_ex, load_configured=True, + failed_connection=connection, + close_generation=close_generation, ) except ToolExecutionException as recovery_error: set_mcp_span_error(span, type(recovery_error).__name__, str(recovery_error)) @@ -3912,7 +4075,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 @@ -4196,7 +4360,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. HTTP header admission and connection recovery each use 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 ``url`` and ``session``. ``request_timeout``, ``terminate_on_close``, ``http_client``, @@ -4269,8 +4434,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 @@ -4278,13 +4442,16 @@ def __init__( of HTTP headers to inject into every outbound request to the MCP server. Generated tool calls supply only this host runtime context, never model arguments, even when names collide. Direct ``call_tool`` calls use their - caller-supplied kwargs. Model-over-runtime precedence applies only to - outbound tool arguments, not to authentication. + caller-supplied kwargs. Direct ``get_prompt`` calls resolve ambient headers + without using prompt arguments as authentication inputs. Model-over-runtime + precedence applies only to outbound tool arguments, not to authentication. Use this to forward per-request context (e.g. authentication tokens set in 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. Calls with the same effective headers can run concurrently; + a different header set waits for active operations to finish before replacing the session. + 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 - @@ -4416,12 +4583,8 @@ 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. + # One header identity is leased by all concurrent operations. Keep a snapshot + # for transports that do not propagate the invoking task's ContextVars. 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. @@ -4435,6 +4598,9 @@ 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._active_header_calls = 0 + self._call_headers_idle = asyncio.Event() + self._call_headers_idle.set() self._header_request_owner = object() self._header_hook_client: httpx2.AsyncClient | None = None @@ -4644,8 +4810,8 @@ async def _prepare_for_run(self, kwargs: Mapping[str, Any]) -> None: sensitive_exc_info=True, ) return - async with self._call_headers_lock: - await self._ensure_session_identity(headers, kwargs) + async with self._acquire_call_headers(headers, kwargs): + pass def _bind_session_headers(self, headers: Mapping[str, str]) -> None: self._session_headers = dict(headers) @@ -4671,27 +4837,12 @@ def _discard_pending_session_headers(self) -> None: self._pending_connection_kwargs = None async def _reconnect_for_identity_change(self) -> None: - async with self._function_load_lock: - if self._is_lifecycle_owner_task(): - await self._connect_on_owner(reset=True, load_configured=False, reset_discovery=True) - else: - async with self._lifecycle_request_lock: - await self._run_on_lifecycle_owner( - "connect", - reset=True, - load_configured=False, - reset_discovery=True, - ) - - try: - await self._load_configured_discovery_locked() - except (Exception, asyncio.CancelledError): - self._reset_session_discovery_state() - if self._is_lifecycle_owner_task(): - await self._close_on_owner() - else: - await self.close() - raise + await self._run_on_lifecycle_owner( + "connect", + reset=True, + reset_discovery=True, + close_generation=self._close_generation, + ) async def _ensure_session_identity( self, @@ -4731,6 +4882,66 @@ async def _ensure_session_identity( raise self._promote_pending_session_headers() + @asynccontextmanager + async def _acquire_call_headers( + self, + headers: dict[str, str], + kwargs: Mapping[str, Any], + ) -> AsyncGenerator[None]: + """Wait outside the admission lock until the requested identity can be used.""" + connection = self._connection + if connection is not None and not connection.is_framework_owned: + await self._ensure_session_identity(headers, kwargs) + yield + return + identity = _mcp_header_identity(headers) + try: + async with asyncio.timeout(self.request_timeout): + while True: + await self._wait_for_lifecycle() + async with self._call_headers_lock: + if self._active_call_headers is None or identity == _mcp_header_identity( + self._active_call_headers + ): + await self._ensure_session_identity(headers, kwargs) + yield + return + await self._call_headers_idle.wait() + 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 + + @asynccontextmanager + async def _use_call_headers( + self, + headers: dict[str, str], + kwargs: Mapping[str, Any], + ) -> AsyncGenerator[None]: + """Share a header snapshot until every operation using its identity finishes.""" + connection = self._connection + if connection is not None and not connection.is_framework_owned: + await self._ensure_session_identity(headers, kwargs) + yield + return + async with self._acquire_call_headers(headers, kwargs): + if self._active_header_calls == 0: + self._active_call_headers = headers + self._active_header_calls += 1 + self._call_headers_idle.clear() + + token = _mcp_call_headers.set((self._header_request_owner, headers)) + try: + yield + finally: + _mcp_call_headers.reset(token) + # No suspension point: cancellation cannot strand the last lease. + self._active_header_calls -= 1 + if self._active_header_calls == 0: + self._active_call_headers = None + self._call_headers_idle.set() + async def _call_prompt_with_runtime_kwargs( self, prompt_name: str, @@ -4741,19 +4952,30 @@ 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) + async with self._use_call_headers(headers, runtime_kwargs): + return await super().get_prompt(prompt_name, **prompt_arguments) + + async def get_prompt(self, prompt_name: str, **kwargs: Any) -> str: + """Call a prompt using the current ambient header identity. + + Args: + prompt_name: The remote prompt name. + + Keyword Args: + kwargs: Prompt arguments, not inputs to the authentication header provider. + + Returns: + The parsed prompt result. + + Raises: + ToolExecutionException: If admission, connection, or prompt execution fails. + KeyError: If the header provider requires unavailable runtime values. + """ + if self._header_provider is None: + return await super().get_prompt(prompt_name, **kwargs) + headers = self._effective_headers({}) + async with self._use_call_headers(headers, {}): + return await super().get_prompt(prompt_name, **kwargs) def _release_connection_kwargs(self) -> None: self._connection_kwargs = None @@ -4775,8 +4997,9 @@ async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]: The returned headers are attached to every HTTP request made during this tool call via a request hook on the underlying HTTP client. Fixed and dynamic headers form the effective identity. If they differ from a framework-created - session's identity, the tool reconnects before sending the call; caller-supplied sessions - reject the change. + session's identity, the tool waits for active operations to finish and reconnects before + sending the call; caller-supplied sessions reject the change. Calls with the same effective + headers can run concurrently. The provider does not consume runtime kwargs: their separately merged tool argument mapping continues to :meth:`MCPTool.call_tool` and its outbound filter. @@ -4794,15 +5017,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) + async with self._use_call_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/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 23ee4ea01a1..0d0b8337b40 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,14 @@ 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 = True + 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 +7691,15 @@ 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: + # Mock transport replacement to avoid starting an actual process. + 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._prompts_loaded = True mock_connect.side_effect = restore_session @@ -7727,9 +7729,9 @@ 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 + # Verify the failed transport was replaced. assert mock_connect.call_count >= 1 - mock_connect.assert_called_with(reset=True) + mock_connect.assert_called_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 +7789,15 @@ 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: + # Mock transport replacement to avoid starting an actual process. + 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 = True mock_connect.side_effect = restore_session @@ -7825,9 +7827,9 @@ 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 + # Verify the failed transport was replaced. assert mock_connect.call_count >= 1 - mock_connect.assert_called_with(reset=True) + mock_connect.assert_called_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 +7927,17 @@ 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 = True + 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 +7949,17 @@ 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 = True + 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() @@ -8096,7 +8102,7 @@ async def __aexit__(self, exc_type, exc, tb): assert transport_contexts[1].exit_task is transport_contexts[1].enter_task -async def test_mcp_tool_connect_from_lifecycle_owner_bypasses_request_lock() -> None: +async def test_mcp_tool_connect_from_lifecycle_owner_bypasses_request_queue() -> None: """connect(reset=True) should bypass the request queue when already on the owner task.""" import asyncio @@ -8110,18 +8116,21 @@ async def test_mcp_tool_connect_from_lifecycle_owner_bypasses_request_lock() -> async def connect_from_owner_task() -> None: tool._lifecycle_owner_task = asyncio.current_task() try: - async with tool._lifecycle_request_lock: - await tool.connect(reset=True) + await tool.connect(reset=True) finally: tool._lifecycle_owner_task = None - with patch.object(tool, "_connect_on_owner", AsyncMock()) as mock_connect_on_owner: + with ( + patch.object(tool, "_connect_on_owner", AsyncMock()) as mock_connect_on_owner, + patch.object(tool, "_run_on_lifecycle_owner", AsyncMock()) as queue_request, + ): await asyncio.wait_for(connect_from_owner_task(), timeout=0.1) mock_connect_on_owner.assert_awaited_once_with(reset=True) + queue_request.assert_not_awaited() -async def test_mcp_tool_close_from_lifecycle_owner_bypasses_request_lock() -> None: +async def test_mcp_tool_close_from_lifecycle_owner_bypasses_request_queue() -> None: """close() should bypass the request queue when already on the owner task.""" import asyncio @@ -8135,15 +8144,18 @@ async def test_mcp_tool_close_from_lifecycle_owner_bypasses_request_lock() -> No async def close_from_owner_task() -> None: tool._lifecycle_owner_task = asyncio.current_task() try: - async with tool._lifecycle_request_lock: - await tool.close() + await tool.close() finally: tool._lifecycle_owner_task = None - with patch.object(tool, "_close_on_owner", AsyncMock()) as mock_close_on_owner: + with ( + patch.object(tool, "_close_on_owner", AsyncMock()) as mock_close_on_owner, + patch.object(tool, "_run_on_lifecycle_owner", AsyncMock()) as queue_request, + ): await asyncio.wait_for(close_from_owner_task(), timeout=0.1) mock_close_on_owner.assert_awaited_once_with() + queue_request.assert_not_awaited() async def test_mcp_tool_safe_close_reraises_other_runtime_errors(): 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..87c2e1ccacf 100644 --- a/python/packages/core/tests/core/test_mcp_http_auth.py +++ b/python/packages/core/tests/core/test_mcp_http_auth.py @@ -12,6 +12,7 @@ import httpx2 as httpx import pytest +from anyio import ClosedResourceError from mcp.client.session import ClientSession from agent_framework import FunctionInvocationContext, MCPStreamableHTTPTool @@ -110,13 +111,20 @@ 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, + rewrite_response: Callable[[httpx.Request, httpx.Response], httpx.Response] | None = None, + **kwargs: Any, + ) -> httpx.AsyncClient: async def handle_with_cookies(request: httpx.Request) -> httpx.Response: response = await handle(request) if response_cookies and request.method == "POST": method = json.loads(request.content).get("method") if method in response_cookies: response.headers["Set-Cookie"] = response_cookies[method] + if rewrite_response is not None: + response = rewrite_response(request, response) return response client = async_client(transport=httpx.MockTransport(handle_with_cookies), **kwargs) @@ -525,6 +533,57 @@ async def test_captured_prompt_reconciles_run_identity_at_invocation(mcp_http_se await tool.close() +@pytest.mark.parametrize("detached_child", [False, True]) +async def test_direct_prompt_uses_current_ambient_identity( + mcp_http_server: MCPHTTPServer, detached_child: bool +) -> None: + client, requests, _ = mcp_http_server + principal = ContextVar("direct_prompt_principal", default="token-a") + child_release = asyncio.Event() + children: list[asyncio.Task[str]] = [] + tool = MCPStreamableHTTPTool( + name="direct-prompt", + url="https://mcp.example/mcp", + http_client=client, + static_headers={"X-Test-Prompts": "enabled"}, + header_provider=lambda _: {"Authorization": principal.get()}, + ) + + async def read_prompt() -> str: + if detached_child: + await child_release.wait() + token = principal.set("token-b") + try: + return await tool.get_prompt("principal-prompt", topic="current") + 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(read_prompt())) + + async with tool, asyncio.timeout(5): + try: + if detached_child: + client.event_hooks["request"].append(spawn_child) + await tool.call_tool("record", marker="parent") + assert len(children) == 1 + child_release.set() + result = await children[0] + else: + result = await read_prompt() + assert result == "token-b:current" + request = _requests_for_method(requests, "prompts/get")[-1] + assert request.headers["Authorization"] == "token-b" + assert request.headers["mcp-session-id"] == "session-token-b" + assert len(_requests_for_method(requests, "initialize")) == 2 + finally: + child_release.set() + for child in children: + child.cancel() + await asyncio.gather(*children, return_exceptions=True) + + async def test_run_preparation_reconnects_before_exposing_principal_tools(mcp_http_server: MCPHTTPServer) -> None: client, requests, _ = mcp_http_server tool = _tool(client, "token-a") @@ -563,6 +622,476 @@ async def test_concurrent_callers_use_sessions_bound_to_their_own_headers(mcp_ht await tool.close() +@pytest.mark.parametrize("first_completion", ["success", "cancel", "nested"]) +@pytest.mark.parametrize("transition", ["tool", "run"]) +async def test_same_header_calls_overlap_and_identity_change_waits_for_all_calls( + mcp_http_server: MCPHTTPServer, first_completion: str, transition: str +) -> None: + client, requests, writes = mcp_http_server + started = {marker: asyncio.Event() for marker in ("first", "second")} + release = {marker: asyncio.Event() for marker in started} + switch_requested = asyncio.Event() + + def provide_headers(kwargs: dict[str, Any]) -> dict[str, str]: + principal = kwargs.get("credential", "token-a") + if principal == "token-b": + switch_requested.set() + return {"Authorization": principal} + + async def block_call(request: httpx.Request) -> None: + if request.method != "POST": + return + body = json.loads(request.content) + if body.get("method") == "tools/call": + marker = body["params"].get("arguments", {}).get("marker") + if marker in started: + started[marker].set() + await release[marker].wait() + if marker == "first" and first_completion == "nested": + await tool.call_tool("record", credential="token-a", marker="nested") + + tool = MCPStreamableHTTPTool( + name="concurrent", + url="https://mcp.example/mcp", + http_client=client, + load_prompts=False, + header_provider=provide_headers, + ) + async with tool, asyncio.timeout(5): + client.event_hooks["request"].append(block_call) + async with asyncio.TaskGroup() as tasks: + first = tasks.create_task(tool.call_tool("record", credential="token-a", marker="first")) + await started["first"].wait() + second = tasks.create_task(tool.call_tool("record", credential="token-a", marker="second")) + await started["second"].wait() + switch = tasks.create_task( + tool.call_tool("record", credential="token-b", marker="switch") + if transition == "tool" + else tool._prepare_for_run({"credential": "token-b"}) + ) + await switch_requested.wait() + assert len(_requests_for_method(requests, "initialize")) == 1 + + if first_completion == "cancel": + first.cancel() + with pytest.raises(asyncio.CancelledError): + await first + else: + release["first"].set() + await first + assert not switch.done() + assert len(_requests_for_method(requests, "initialize")) == 1 + + release["second"].set() + await second + await switch + + assert [request.headers["Authorization"] for request in _requests_for_method(requests, "initialize")] == [ + "token-a", + "token-b", + ] + calls = _calls(requests) + assert all( + request.headers["mcp-session-id"] == f"session-{request.headers['Authorization']}" for request in calls + ) + assert ( + writes["token-a"] + == { + "success": ["first", "second"], + "cancel": ["second"], + "nested": ["nested", "first", "second"], + }[first_completion] + ) + assert writes["token-b"] == (["switch"] if transition == "tool" else []) + + +@pytest.mark.parametrize("operation", ["tool", "prompt"]) +async def test_concurrent_calls_share_recovery_of_their_failed_connection( + mcp_http_server: MCPHTTPServer, operation: str +) -> None: + client, requests, _ = mcp_http_server + recovery_started = asyncio.Event() + release_recovery = asyncio.Event() + first_retry_started = asyncio.Event() + second_retry_started = asyncio.Event() + release_first_retry = asyncio.Event() + both_failed_calls_started = asyncio.Event() + failed_calls = 0 + method = "tools/call" if operation == "tool" else "prompts/get" + argument = "marker" if operation == "tool" else "topic" + tool = MCPStreamableHTTPTool( + name="shared-recovery", + url="https://mcp.example/mcp", + http_client=client, + load_prompts=True, + static_headers={"X-Test-Prompts": "enabled"}, + header_provider=lambda kwargs: {"Authorization": kwargs.get("credential", "token-a")}, + ) + + async def lose_connection(name: str, **kwargs: Any) -> Any: + nonlocal failed_calls + failed_calls += 1 + if failed_calls == 2: + both_failed_calls_started.set() + await both_failed_calls_started.wait() + if kwargs["arguments"][argument] == "second": + await first_retry_started.wait() + raise ClosedResourceError + + async def block_recovery_and_retry(request: httpx.Request) -> None: + if request.method != "POST": + return + body = json.loads(request.content) + if body.get("method") == "initialize": + recovery_started.set() + await release_recovery.wait() + elif body.get("method") == method: + marker = body["params"]["arguments"][argument] + if marker == "first": + first_retry_started.set() + await release_first_retry.wait() + elif marker == "second": + second_retry_started.set() + + async with tool, asyncio.timeout(5): + sdk_client = _connected_client(tool) + client.event_hooks["request"].append(block_recovery_and_retry) + with patch.object(sdk_client, "call_tool" if operation == "tool" else "get_prompt", new=lose_connection): + async with asyncio.TaskGroup() as tasks: + calls = [ + tasks.create_task( + tool.call_tool("record", credential="token-a", marker=marker) + if operation == "tool" + else tool._call_prompt_with_runtime_kwargs( + "principal-prompt", {"topic": marker}, {"credential": "token-a"} + ) + ) + for marker in ("first", "second") + ] + await recovery_started.wait() + arriving = tasks.create_task(tool.call_tool("record", credential="token-a", marker="arriving")) + await asyncio.sleep(0) + assert not arriving.done() + release_recovery.set() + await second_retry_started.wait() + release_first_retry.set() + + assert failed_calls == 2 + assert len(_requests_for_method(requests, "initialize")) == 2 + assert all(call.result() is not None for call in calls) + assert arriving.result() is not None + assert all(request.headers["Authorization"] == "token-a" for request in _requests_for_method(requests, method)) + assert sorted( + json.loads(request.content)["params"]["arguments"][argument] + for request in _requests_for_method(requests, method) + ) == (["arriving", "first", "second"] if operation == "tool" else ["first", "second"]) + + +@pytest.mark.parametrize("operation", ["tool", "discovery"]) +@pytest.mark.parametrize("blocked_method", ["initialize", "tools/list"]) +async def test_close_interrupts_recovery_without_reopening_or_retrying( + mcp_http_server: MCPHTTPServer, blocked_method: str, operation: str +) -> None: + client, requests, writes = mcp_http_server + tool = _tool(client, "token-a") + recovery_started = asyncio.Event() + release_recovery = asyncio.Event() + + async def block_recovery(request: httpx.Request) -> None: + if request.method == "POST" and json.loads(request.content).get("method") == blocked_method: + recovery_started.set() + await release_recovery.wait() + + call = closing = None + try: + async with asyncio.timeout(5): + await tool.connect() + client.event_hooks["request"].append(block_recovery) + with patch.object( + _connected_client(tool), + "call_tool" if operation == "tool" else "list_tools", + side_effect=ClosedResourceError(), + ): + call = asyncio.create_task( + tool.call_tool("record", credential="token-a", marker="must-not-retry") + if operation == "tool" + else tool.load_tools() + ) + await recovery_started.wait() + closing = asyncio.create_task(tool.close()) + await asyncio.sleep(0) + release_recovery.set() + with pytest.raises(ToolExecutionException, match=r"Failed to reconnect|interrupted by close\(\)"): + await call + await closing + + assert not tool.is_connected + assert tool._connection is None + assert writes["token-a"] == [] + assert len(_requests_for_method(requests, "initialize")) == 2 + assert not tool._tools_loaded + assert tool._active_header_calls == 0 + finally: + release_recovery.set() + for task in (call, closing): + if task is not None: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await tool.close() + + +async def test_idle_header_admission_does_not_queue_lifecycle_synchronization(mcp_http_server: MCPHTTPServer) -> None: + client, _, writes = mcp_http_server + tool = _tool(client, "token-a") + async with tool: + with patch.object(tool, "_run_on_lifecycle_owner", wraps=tool._run_on_lifecycle_owner) as lifecycle: + await asyncio.gather( + tool.call_tool("record", credential="token-a", marker="first"), + tool.call_tool("record", credential="token-a", marker="second"), + ) + lifecycle.assert_not_called() + assert sorted(writes["token-a"]) == ["first", "second"] + + +async def test_recovery_waits_for_catalog_and_reprepares_arguments_and_metadata( + mcp_http_client_factory: MCPHTTPClientFactory, +) -> None: + create_client, requests, _ = mcp_http_client_factory + session_count = 0 + catalog_started = asyncio.Event() + release_catalog = asyncio.Event() + + def change_catalog(request: httpx.Request, response: httpx.Response) -> httpx.Response: + nonlocal session_count + if request.method != "POST": + return response + method = json.loads(request.content).get("method") + if method == "initialize": + session_count += 1 + if method == "tools/list": + payload = response.json() + remote_tool = payload["result"]["tools"][0] + remote_tool["_meta"] = {"catalog-version": str(session_count)} + if session_count > 1: + remote_tool["inputSchema"]["properties"]["fresh"] = {"type": "string"} + return httpx.Response(200, json=payload) + return response + + async def block_new_catalog(request: httpx.Request) -> None: + if request.method == "POST" and json.loads(request.content).get("method") == "tools/list" and session_count > 1: + catalog_started.set() + await release_catalog.wait() + + async with create_client(rewrite_response=change_catalog) as client: + tool = _tool(client, "token-a") + async with tool, asyncio.timeout(5), asyncio.TaskGroup() as tasks: + client.event_hooks["request"].append(block_new_catalog) + with patch.object(_connected_client(tool), "call_tool", side_effect=ClosedResourceError()): + recovering = tasks.create_task( + tool.call_tool("record", credential="token-a", marker="recovering", fresh="value") + ) + await catalog_started.wait() + arriving = tasks.create_task( + tool.call_tool("record", credential="token-a", marker="arriving", fresh="value") + ) + await asyncio.sleep(0) + assert _calls(requests) == [] + release_catalog.set() + + assert recovering.result() is not None + assert arriving.result() is not None + assert session_count == 2 + for request in _calls(requests): + params = json.loads(request.content)["params"] + assert params["arguments"]["fresh"] == "value" + assert params["_meta"]["catalog-version"] == "2" + + +@pytest.mark.parametrize("cancel_call", [False, True]) +async def test_header_lease_cleanup_does_not_wait_for_admission_lock( + mcp_http_server: MCPHTTPServer, cancel_call: bool +) -> None: + client, _, _ = mcp_http_server + tool = _tool(client, "token-a") + started = asyncio.Event() + release = asyncio.Event() + + async def block_call(request: httpx.Request) -> None: + if request.method == "POST" and json.loads(request.content).get("method") == "tools/call": + started.set() + await release.wait() + + call = None + try: + async with tool, asyncio.timeout(5): + client.event_hooks["request"].append(block_call) + call = asyncio.create_task(tool.call_tool("record", credential="token-a")) + await started.wait() + async with tool._call_headers_lock: + if cancel_call: + call.cancel() + with pytest.raises(asyncio.CancelledError): + _ = await call + else: + release.set() + assert await call is not None + assert tool._active_header_calls == 0 + assert tool._active_call_headers is None + assert tool._call_headers_idle.is_set() + finally: + release.set() + if call is not None: + call.cancel() + await asyncio.gather(call, return_exceptions=True) + + +async def test_recovery_discovery_failure_does_not_deadlock_queued_connect( + mcp_http_client_factory: MCPHTTPClientFactory, +) -> None: + create_client, requests, _ = mcp_http_client_factory + session_count = 0 + discovery_started = asyncio.Event() + release_discovery = asyncio.Event() + + def reject_recovery_catalog(request: httpx.Request, response: httpx.Response) -> httpx.Response: + nonlocal session_count + if request.method != "POST": + return response + payload = json.loads(request.content) + if payload.get("method") == "initialize": + session_count += 1 + if payload.get("method") == "tools/list" and session_count == 2: + return httpx.Response( + 200, json={"jsonrpc": "2.0", "id": payload["id"], "error": {"code": -32001, "message": "rejected"}} + ) + return response + + async def block_discovery(request: httpx.Request) -> None: + if ( + request.method == "POST" + and json.loads(request.content).get("method") == "tools/list" + and session_count == 2 + ): + discovery_started.set() + await release_discovery.wait() + + async with create_client(rewrite_response=reject_recovery_catalog) as client: + tool = _tool(client, "token-a") + async with tool, asyncio.timeout(5): + client.event_hooks["request"].append(block_discovery) + with patch.object(_connected_client(tool), "call_tool", side_effect=ClosedResourceError()): + recovering = asyncio.create_task(tool.call_tool("record", credential="token-a")) + connecting = None + try: + await discovery_started.wait() + connecting = asyncio.create_task(tool.connect()) + await asyncio.sleep(0) + assert not connecting.done() + release_discovery.set() + results = await asyncio.gather(recovering, connecting, return_exceptions=True) + assert isinstance(results[0], ToolExecutionException) + assert results[1] is None + assert tool.is_connected + assert session_count == 3 + finally: + release_discovery.set() + for task in (recovering, connecting): + if task is not None: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await tool.call_tool("record", credential="token-a", marker="after-recovery") + assert len(_calls(requests)) == 1 + + +@pytest.mark.parametrize("stage", ["lock", "active", "reconcile"]) +async def test_header_admission_timeout_does_not_dispatch_the_waiting_call( + mcp_http_server: MCPHTTPServer, stage: str +) -> None: + client, requests, _ = mcp_http_server + blocked = asyncio.Event() + release = asyncio.Event() + tool = _tool(client, "token-a") + + async def block_request(request: httpx.Request) -> None: + if request.method != "POST": + return + method = json.loads(request.content).get("method") + if (stage == "active" and method == "tools/call") or (stage == "reconcile" and method == "initialize"): + blocked.set() + await release.wait() + + async with tool, asyncio.timeout(5), asyncio.TaskGroup() as tasks: + client.event_hooks["request"].append(block_request) + active = None + if stage == "active": + active = tasks.create_task(tool.call_tool("record", credential="token-a", marker="active")) + await blocked.wait() + 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"): + await tool.call_tool("record", credential="token-b", marker="waiting") + assert not any( + json.loads(request.content)["params"]["arguments"].get("marker") == "waiting" + for request in _calls(requests) + ) + assert tool._active_header_calls == (1 if stage == "active" else 0) + if stage == "reconcile": + following = asyncio.create_task(tool.call_tool("record", credential="token-a", marker="after-timeout")) + try: + await asyncio.sleep(0) + assert not following.done() + release.set() + with pytest.raises(ToolExecutionException): + await following + finally: + following.cancel() + await asyncio.gather(following, return_exceptions=True) + finally: + release.set() + if active is not None: + await active + + +@pytest.mark.parametrize("timeout", [1, None]) +async def test_connection_recovery_respects_configured_timeout( + mcp_http_server: MCPHTTPServer, timeout: int | None +) -> None: + client, _, _ = mcp_http_server + tool = _tool(client, "token-a") + recovering = asyncio.Event() + release = asyncio.Event() + + async def block_recovery() -> None: + recovering.set() + await release.wait() + + async with tool, asyncio.timeout(5): + tool.request_timeout = timeout + with ( + patch.object(_connected_client(tool), "call_tool", side_effect=ClosedResourceError()), + patch.object(tool, "_reconnect_without_loading", new=block_recovery), + ): + call = asyncio.create_task(tool.call_tool("record", credential="token-a")) + await recovering.wait() + if timeout is None: + await asyncio.sleep(0.1) + assert not call.done() + call.cancel() + with pytest.raises(asyncio.CancelledError): + await call + else: + with pytest.raises(ToolExecutionException, match="Timed out recovering"): + await call + release.set() + await tool._lifecycle_idle.wait() + assert tool._lifecycle_idle.is_set() + assert tool._active_header_calls == 0 + await tool.call_tool("record", credential="token-a") + + @pytest.mark.parametrize("model_credential", ["token-b", ""]) async def test_generated_tool_keeps_authentication_in_runtime_context( mcp_http_server: MCPHTTPServer, model_credential: str diff --git a/python/packages/core/tests/core/test_mcp_observability.py b/python/packages/core/tests/core/test_mcp_observability.py index 47a28b9d205..7ad75525a73 100644 --- a/python/packages/core/tests/core/test_mcp_observability.py +++ b/python/packages/core/tests/core/test_mcp_observability.py @@ -333,10 +333,12 @@ 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._tools_loaded = True + tool._prompts_loaded = True span_exporter.clear() with ( + patch.object(tool, "_reconnect_without_loading", AsyncMock()), caplog.at_level(logging.ERROR, logger="agent_framework._mcp"), pytest.raises(ToolExecutionException, match="connection lost"), ): 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 51172446aa8..a9c349cd9b2 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py @@ -238,9 +238,10 @@ class FoundryToolbox(MCPStreamableHTTPTool): Platform headers are resolved at each operation boundary before work crosses 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. + waits for active operations to finish and reconnects that session before + sending the operation. Tool, prompt and Toolbox skill-resource operations + with the same effective headers can run concurrently without overwriting + one another's headers. Because the toolbox endpoint is a first-party Foundry service, forwarding the opaque caller token to it is safe. @@ -417,6 +418,29 @@ async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]: finally: _strict_toolbox_header_provider.reset(token) + @override + async def get_prompt(self, prompt_name: str, **kwargs: Any) -> str: + """Retrieve a prompt using current platform headers. + + Args: + prompt_name: The remote prompt name. + + Keyword Args: + kwargs: Remote prompt arguments, not inputs to the authentication header provider. + + Returns: + The parsed prompt result. + + Raises: + ToolExecutionException: If prompts are disabled or the MCP prompt operation fails. + KeyError: If a custom header provider requires unavailable runtime values. + """ + token = _strict_toolbox_header_provider.set(True) + try: + return await super().get_prompt(prompt_name, **kwargs) + finally: + _strict_toolbox_header_provider.reset(token) + @override async def _call_prompt_with_runtime_kwargs( self, @@ -448,10 +472,10 @@ 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.""" + async def _close_on_owner(self) -> None: + """Close the owned HTTP client after the lifecycle owner finishes MCP teardown.""" try: - await super().close() + await super()._close_on_owner() finally: client = self._httpx_client if client is not None: @@ -461,16 +485,11 @@ async def close(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() - 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() - self._active_call_headers = headers - try: - return await connection.read_resource(cast(str, uri)) - finally: - self._active_call_headers = None + connection = self._connection + if not headers and connection is not None and not connection.is_framework_owned: + return await connection.read_resource(str(uri)) + async with self._use_call_headers(headers, {}): + return await self._require_connection().read_resource(str(uri)) def _effective_resource_headers(self) -> dict[str, str]: """Resolve headers for a resource operation without caller runtime kwargs.""" @@ -668,7 +687,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 the Toolbox header identity, so cached skills keep using the live session instead of a closed one. """ diff --git a/python/packages/foundry_hosting/tests/foundry_hosting/toolbox_approval_helpers.py b/python/packages/foundry_hosting/tests/foundry_hosting/toolbox_approval_helpers.py index 81fd0550e76..fc3803544c9 100644 --- a/python/packages/foundry_hosting/tests/foundry_hosting/toolbox_approval_helpers.py +++ b/python/packages/foundry_hosting/tests/foundry_hosting/toolbox_approval_helpers.py @@ -8,6 +8,7 @@ from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Mapping, Sequence from contextlib import asynccontextmanager from typing import Any +from unittest.mock import patch import httpx2 as httpx from agent_framework import ( @@ -85,17 +86,21 @@ async def handle(self, request: httpx.Request) -> httpx.Response: async def toolbox_connection( server: ToolboxApprovalServer, *, connect: bool = True, **kwargs: Any ) -> AsyncGenerator[FoundryToolbox]: - toolbox = FoundryToolbox(ToolboxCredential(), url="https://toolbox.invalid/mcp", **kwargs) - original = toolbox._httpx_client - assert original is not None - await original.aclose() - toolbox._httpx_client = httpx.AsyncClient(transport=httpx.MockTransport(server.handle)) - try: - if connect: - await toolbox.connect() - yield toolbox - finally: - await toolbox.close() + async_client = httpx.AsyncClient + + def create_client(*args: Any, **client_kwargs: Any) -> httpx.AsyncClient: + client_kwargs.setdefault("transport", httpx.MockTransport(server.handle)) + return async_client(*args, **client_kwargs) + + # Failed setup closes the owned client; cover its replacement as well as startup. + with patch.object(httpx, "AsyncClient", new=create_client): + toolbox = FoundryToolbox(ToolboxCredential(), url="https://toolbox.invalid/mcp", **kwargs) + try: + if connect: + await toolbox.connect() + yield toolbox + finally: + await toolbox.close() def toolbox_tool( diff --git a/python/packages/foundry_hosting/tests/test_toolbox.py b/python/packages/foundry_hosting/tests/test_toolbox.py index 52a04d34b1b..d04af369d46 100644 --- a/python/packages/foundry_hosting/tests/test_toolbox.py +++ b/python/packages/foundry_hosting/tests/test_toolbox.py @@ -118,7 +118,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": {}, "resources": {}, "prompts": {}}, "serverInfo": {"name": "toolbox-loopback", "version": "1"}, } elif method == "tools/list": @@ -135,6 +135,21 @@ 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": "record-prompt", + "arguments": [{"name": "marker", "required": True}], + } + ] + } + 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": @@ -510,6 +525,102 @@ async def test_toolbox_concurrent_callers_use_their_own_call_id() -> None: assert calls_by_marker["request-c"].headers["mcp-session-id"].endswith("-CALL-C") +async def test_toolbox_direct_prompt_resolves_current_call_id() -> None: + server = _ToolboxLoopbackServer() + toolbox = FoundryToolbox( + _FakeCredential(), # type: ignore + url="https://loopback.invalid/mcp", + load_prompts=True, + load_tools=False, + ) + await _attach_loopback_transport(toolbox, server) + + try: + await _run_with_request_context("CALL-A", toolbox.connect) + first = await _run_with_request_context( + "CALL-A", + lambda: toolbox.get_prompt("record-prompt", marker="first-request"), + ) + second = await _run_with_request_context( + "CALL-B", + lambda: toolbox.get_prompt("record-prompt", marker="second-request"), + ) + finally: + await toolbox.close() + + assert "CALL-A" in first + assert "CALL-B" in second + prompts = server.requests_for("prompts/get") + assert [request.headers["x-agent-foundry-call-id"] for request in prompts] == ["CALL-A", "CALL-B"] + assert all( + request.headers["mcp-session-id"].endswith(f"-{request.headers['x-agent-foundry-call-id']}") + for request in prompts + ) + assert json.loads(prompts[-1].content)["params"]["arguments"] == {"marker": "second-request"} + + +async def test_toolbox_direct_prompt_preserves_header_provider_errors_without_using_prompt_arguments() -> None: + server = _ToolboxLoopbackServer() + toolbox = FoundryToolbox( + _FakeCredential(), # type: ignore + url="https://loopback.invalid/mcp", + load_prompts=True, + load_tools=False, + header_provider=lambda kwargs: {"X-Custom-Header": kwargs["custom"]}, + ) + await _attach_loopback_transport(toolbox, server) + + try: + await _run_with_request_context("CALL-A", toolbox.connect) + with pytest.raises(KeyError, match="custom"): + await _run_with_request_context( + "CALL-B", + lambda: toolbox.get_prompt("record-prompt", marker="prompt-marker", custom="model-value"), + ) + await _run_with_request_context("CALL-A", lambda: toolbox.connect(reset=True)) + finally: + await toolbox.close() + + assert server.requests_for("prompts/get") == [] + assert all("X-Custom-Header" not in request.headers for request in server.requests_for("initialize")) + + +async def test_toolbox_generated_prompt_uses_runtime_headers_not_prompt_arguments() -> None: + server = _ToolboxLoopbackServer() + toolbox = FoundryToolbox( + _FakeCredential(), # type: ignore + url="https://loopback.invalid/mcp", + load_prompts=True, + load_tools=False, + header_provider=lambda kwargs: {"X-Custom-Header": kwargs["custom"]}, + ) + await _attach_loopback_transport(toolbox, server) + + try: + await toolbox._prepare_for_run({"custom": "setup-value"}) + await _run_with_request_context("CALL-A", toolbox.connect) + prompt = next(function for function in toolbox.functions if function.name == "record-prompt") + context = FunctionInvocationContext( + function=prompt, + arguments={"marker": "prompt-marker"}, + kwargs={"custom": "runtime-value"}, + ) + result = await _run_with_request_context( + "CALL-B", + lambda: prompt.invoke(arguments={"marker": "prompt-marker"}, context=context), + ) + finally: + await toolbox.close() + + assert isinstance(result, list) + assert result[0].text == "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 request.headers["X-Custom-Header"] == "runtime-value" + assert json.loads(request.content)["params"]["arguments"] == {"marker": "prompt-marker"} + + async def test_toolbox_concurrent_skill_reads_use_their_own_call_id() -> None: server = _ToolboxLoopbackServer() toolbox = FoundryToolbox( @@ -554,6 +665,47 @@ async def read_skill() -> tuple[str, str]: assert all(request.headers["X-Custom-Header"] == "resource-value" for request in resource_requests) +async def test_toolbox_skill_read_waits_for_in_progress_identity_reconnect() -> 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 + reconnect_started = asyncio.Event() + release_reconnect = asyncio.Event() + + async def pause_initialize(request: httpx.Request) -> None: + if request.method == "POST" and json.loads(request.content).get("method") == "initialize": + reconnect_started.set() + await release_reconnect.wait() + + try: + await _run_with_request_context("CALL-A", toolbox.connect) + skills = await _run_with_request_context("CALL-A", lambda: source.get_skills(_source_context())) + client = toolbox._httpx_client + assert client is not None + client.event_hooks["request"].append(pause_initialize) + + async with asyncio.timeout(5), asyncio.TaskGroup() as tasks: + call = tasks.create_task( + _run_with_request_context("CALL-B", lambda: toolbox.call_tool("record", marker="reconnect")) + ) + await reconnect_started.wait() + read = tasks.create_task(_run_with_request_context("CALL-B", skills[0].get_content)) + await asyncio.sleep(0) + assert not read.done() + release_reconnect.set() + + assert call.result() is not None + assert "Call ID: CALL-B" in read.result() + assert len(server.requests_for("initialize")) == 2 + 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_reconnect.set() + await toolbox.close() + + async def test_toolbox_skills_reject_header_provider_that_requires_runtime_kwargs() -> None: server = _ToolboxLoopbackServer() toolbox = FoundryToolbox( @@ -593,6 +745,54 @@ async def test_close_closes_owned_http_client(monkeypatch: pytest.MonkeyPatch) - aclose.assert_awaited_once() +async def test_cancelled_close_keeps_http_client_open_until_owner_teardown_finishes( + caplog: pytest.LogCaptureFixture, +) -> None: + server = _ToolboxLoopbackServer() + toolbox = FoundryToolbox(_FakeCredential(), url="https://loopback.invalid/mcp") # type: ignore + await _attach_loopback_transport(toolbox, server) + client = toolbox._httpx_client + assert client is not None + teardown_started = asyncio.Event() + release_teardown = asyncio.Event() + owner: asyncio.Task[None] | None = None + + async def pause_termination(request: httpx.Request) -> None: + if request.method == "DELETE": + teardown_started.set() + await release_teardown.wait() + + try: + await toolbox.connect() + client.event_hooks["request"].append(pause_termination) + + async with asyncio.timeout(5), asyncio.TaskGroup() as tasks: + close = tasks.create_task(toolbox.close()) + await teardown_started.wait() + owner = toolbox._lifecycle_owner_task + assert owner is not None + close.cancel() + with pytest.raises(asyncio.CancelledError): + await close + + assert not client.is_closed + assert toolbox._httpx_client is client + assert not owner.done() + release_teardown.set() + await owner + + assert client.is_closed + assert toolbox._httpx_client is None + assert not toolbox.is_connected + assert "Session termination failed" not in caplog.text + finally: + release_teardown.set() + if owner is not None: + async with asyncio.timeout(5): + await owner + await toolbox.close() + + def test_as_skills_provider_returns_provider() -> None: toolbox = FoundryToolbox( _FakeCredential(), # type: ignore @@ -675,13 +875,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..9707b692b55 100644 --- a/skills/agent-framework-builder-python/references/integrations.md +++ b/skills/agent-framework-builder-python/references/integrations.md @@ -31,6 +31,13 @@ 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. +For framework-created `MCPStreamableHTTPTool` and Foundry Toolbox connections, tool and prompt calls +with the same effective HTTP headers can run concurrently. Different headers require the current operations +to finish before the MCP session +is replaced; do not bypass that boundary to share caller-specific sessions or catalogs. `request_timeout` also +bounds header admission and recovery separately from the RPC timeout. See the +[MCP HTTP concurrency guidance](https://github.com/microsoft/agent-framework/blob/main/python/packages/core/README.md#mcp-http-concurrency-and-timeouts). + ## Capability categories - Use function tools for application-owned code and MCP for external tool services.