diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index 913328e255e..f992d3b091d 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -107,11 +107,13 @@ state. This is not an exactly-once guarantee for external tools: make side-effecting operations idempotent. Legacy unscoped workflow state is not read; migration starts a fresh Responses chain. -For Responses integrations, use a factory when an MCP connection, provider, tool -cache or client carries request identity. A Toolbox's streamable-HTTP writer -inherits the context of the request that **connects** it; sharing that connection -across callers can retain the first call ID. Create the Toolbox and its skills -provider inside the factory, not at process startup. +For Responses integrations, use a factory when a provider, tool cache, client or +credential carries request identity or needs request-owned cleanup. +`FoundryToolbox` resolves platform headers at each operation boundary and +rebinds its MCP session when that identity changes, so a long-lived Toolbox does +not retain the first request's call ID. The integration samples still construct +it inside the factory because they also own request-scoped clients, credentials +and providers. Factory agents are entered/exited for each request, including failed entry and cancellation. `Agent` manages context-managed clients and MCP tools, but it does 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 e58daaab844..564b86fe42e 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py @@ -2,12 +2,14 @@ from __future__ import annotations +import asyncio import inspect import logging import os -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from contextlib import _AsyncGeneratorContextManager # pyright: ignore[reportPrivateUsage] -from typing import TYPE_CHECKING, Any +from contextvars import ContextVar +from typing import TYPE_CHECKING, Any, cast from urllib.parse import urlsplit import httpx @@ -21,8 +23,8 @@ SkillsSourceContext, ) from agent_framework._telemetry import mark_feature_used +from agent_framework.exceptions import ToolExecutionException from azure.ai.agentserver.core import get_request_context -from azure.ai.agentserver.core.platform_headers import FOUNDRY_CALL_ID from typing_extensions import override from ._feature_usage import FeatureIndex @@ -31,10 +33,12 @@ from collections.abc import AsyncGenerator, Generator from datetime import timedelta - from agent_framework import Skill + from agent_framework import Content, Skill from azure.core.credentials import AccessToken, TokenCredential from azure.core.credentials_async import AsyncTokenCredential from mcp.client.session import ClientSession + from mcp.types import ReadResourceResult + from pydantic import AnyUrl AzureCredentialTypes = TokenCredential | AsyncTokenCredential @@ -48,6 +52,10 @@ _TOOLSET_FEATURES_ENV_VAR = "FOUNDRY_AGENT_TOOLSET_FEATURES" # Mandatory preview feature flag for Foundry toolbox requests. _MANDATORY_TOOLBOX_FEATURE = "Toolboxes=V1Preview" +_strict_toolbox_header_provider: ContextVar[bool] = ContextVar( + "strict_toolbox_header_provider", + default=False, +) def _build_toolbox_features_header(additional_features: str | None) -> str: @@ -99,7 +107,7 @@ def _toolbox_name_from_endpoint(endpoint: str) -> str: class _ToolboxAuth(httpx.Auth): - """Injects a fresh bearer token and the platform call-id on every request. + """Inject a fresh bearer token and Toolbox feature flags on every request. Both the synchronous (``sync_auth_flow``) and asynchronous (``async_auth_flow``) httpx auth hooks are implemented, so the same auth works regardless of which @@ -108,11 +116,11 @@ class _ToolboxAuth(httpx.Auth): present. Both synchronous :class:`~azure.core.credentials.TokenCredential` and asynchronous :class:`~azure.core.credentials_async.AsyncTokenCredential` credentials are supported: the async flow awaits an async credential's - ``get_token``, while the sync flow requires a synchronous credential. The - ``x-agent-foundry-call-id`` is read from the hosting context inherited by the - MCP transport task. That task captures context at connection time, so hosted - callers must use a request-owned toolbox/connection rather than sharing one - across requests. The header is absent when no call ID is supplied. + ``get_token``, while the sync flow requires a synchronous credential. + + Platform request headers are injected separately through the MCP dynamic-header + provider so they are resolved in the invoking request task rather than the + connection-time transport task. """ def __init__(self, credential: AzureCredentialTypes, scope: str) -> None: @@ -123,10 +131,6 @@ def __init__(self, credential: AzureCredentialTypes, scope: str) -> None: def _apply_headers(self, request: httpx.Request, token: AccessToken) -> None: request.headers["Authorization"] = f"Bearer {token.token}" - # A Request may be retried or reused after an earlier auth flow, so absence must clear prior caller context. - request.headers.pop(FOUNDRY_CALL_ID, None) - for key, value in get_request_context().platform_headers().items(): - request.headers[key] = value request.headers["Foundry-Features"] = self._features_header def sync_auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]: @@ -167,17 +171,22 @@ class FoundryToolbox(MCPStreamableHTTPTool): - forwards the platform per-request call-id (``x-agent-foundry-call-id``) so the Foundry MCP proxy can resolve the caller context server-side. - The call-id is read from the context inherited by the MCP connection's writer. - Construct this toolbox inside a request-scoped agent factory so it captures - the current caller's context, not a preceding request's. + 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. + Because the toolbox endpoint is a first-party Foundry service, forwarding the opaque caller token to it is safe. Like any MCP tool, the connection lifecycle is driven by the agent: the hosting server enters a factory-created agent for its request and closes its toolbox - and owned HTTP client afterward. An instance-owned agent instead keeps that - connection until shutdown and is not appropriate for differing caller contexts. - Using it as an ``async with`` context manager directly is also supported. + and owned HTTP client afterward. A long-lived Toolbox is also supported: it + keeps its configuration and HTTP client while rebinding its MCP session to + changing platform headers. Using it as an ``async with`` context manager + directly is also supported. Examples: .. code-block:: python @@ -239,6 +248,9 @@ def __init__( mapping configures names per remote function, with ``"*"`` as the global key. timeout: Request timeout in seconds for the underlying HTTP client. kwargs: Additional options forwarded to :class:`~agent_framework.MCPStreamableHTTPTool`. + A custom ``header_provider`` used with :meth:`as_skills_provider` + must resolve without runtime kwargs; skill/resource reads do not + have a function invocation context. """ endpoint = url or _resolve_toolbox_endpoint() tool_name = name or os.environ.get("TOOLBOX_NAME") or _toolbox_name_from_endpoint(endpoint) @@ -250,11 +262,32 @@ def __init__( self._credential = credential self._token_scope = token_scope self._timeout = timeout + caller_header_provider = cast( + Callable[[dict[str, Any]], dict[str, str]] | None, + kwargs.pop("header_provider", None), + ) + + def toolbox_header_provider(runtime_kwargs: dict[str, Any]) -> dict[str, str]: + headers: dict[str, str] = {} + if caller_header_provider is not None: + try: + headers.update(caller_header_provider(runtime_kwargs)) + except KeyError: + if _strict_toolbox_header_provider.get() or self._connection_kwargs is not None: + raise + logger.debug( + "Caller header_provider needs runtime values unavailable to an ambient Toolbox request; " + "applying platform headers only.", + exc_info=True, + ) + headers.update(get_request_context().platform_headers()) + return headers super().__init__( name=tool_name, url=endpoint, http_client=http_client, + header_provider=toolbox_header_provider, load_prompts=load_prompts, load_tools=load_tools, additional_tool_argument_names=additional_tool_argument_names, @@ -267,6 +300,39 @@ async def connect(self, *, reset: bool = False) -> None: await super().connect(reset=reset) mark_feature_used(FeatureIndex.FOUNDRY_TOOLBOX) + @override + async def _prepare_for_run(self, kwargs: Mapping[str, Any]) -> None: + token = _strict_toolbox_header_provider.set(True) + try: + await super()._prepare_for_run(kwargs) + finally: + _strict_toolbox_header_provider.reset(token) + + @override + async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]: + token = _strict_toolbox_header_provider.set(True) + try: + return await super().call_tool(tool_name, **kwargs) + finally: + _strict_toolbox_header_provider.reset(token) + + @override + async def _call_prompt_with_runtime_kwargs( + self, + prompt_name: str, + prompt_arguments: Mapping[str, Any], + runtime_kwargs: Mapping[str, Any], + ) -> str: + token = _strict_toolbox_header_provider.set(True) + try: + return await super()._call_prompt_with_runtime_kwargs( + prompt_name, + prompt_arguments, + runtime_kwargs, + ) + finally: + _strict_toolbox_header_provider.reset(token) + @override def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: """Get an authenticated MCP HTTP client. @@ -291,6 +357,39 @@ async def close(self) -> 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: + await self._ensure_session_identity(headers, {}) + session = self.session + if session is None: + raise RuntimeError( + "FoundryToolbox is not connected, so its resources cannot be read. " + "Pass the toolbox to the agent (tools=...) or enter it as an async " + "context manager before reading resources." + ) + self._active_call_headers = headers + try: + return await session.read_resource(uri) + finally: + self._active_call_headers = None + + def _effective_resource_headers(self) -> dict[str, str]: + """Resolve headers for a resource operation without caller runtime kwargs.""" + token = _strict_toolbox_header_provider.set(True) + try: + try: + return self._effective_headers({}) + except KeyError as ex: + raise ToolExecutionException( + "FoundryToolbox skill and resource reads cannot use a header_provider " + "that requires runtime kwargs. Read required values from a closure or " + "ContextVar, or use a request-owned Toolbox." + ) from ex + finally: + _strict_toolbox_header_provider.reset(token) + def as_skills_provider( self, *, @@ -314,6 +413,11 @@ def as_skills_provider( session and exposes them through a provider you can pass to an agent via ``context_providers=[...]``. + A custom Toolbox ``header_provider`` used with skills must resolve from + ambient state such as a closure or :class:`~contextvars.ContextVar`. + Skill/resource reads do not receive function runtime kwargs, so a provider + that requires them fails explicitly instead of sending incomplete headers. + The toolbox must be **connected** before its skills are discovered (which happens lazily on the first agent run). Connect it by passing the toolbox to the agent via ``tools=`` -- set ``load_tools=False`` if you want skills only @@ -331,11 +435,10 @@ def as_skills_provider( skills; see :class:`~agent_framework.SkillsProvider`. disable_caching: When ``True``, re-query the toolbox on every agent run, re-reading ``skill://index.json`` each time. When ``False`` (the - default), the toolbox's skill discovery is cached after the first run - so the index is read once. The toolbox's advertised skill set is the - same for every caller (the per-request call-id governs execution/ - authorization, not which skills are listed), so a single shared cache - is safe. + default), the toolbox's skill discovery is cached while the effective + platform-header identity is unchanged. Changing the call ID discards + the previous cache before discovery so skill bodies and archive + resources loaded for one caller are not reused by another. cache_refresh_interval: Optional duration after which the cached skill discovery is considered stale and re-read from the toolbox on the next agent run. Useful when a toolbox's attached skills change over the @@ -401,13 +504,20 @@ def as_skills_provider( if archive_max_uncompressed_size_bytes is not None: archive_options["archive_max_uncompressed_size_bytes"] = archive_max_uncompressed_size_bytes - # The toolbox advertises the same skill set to every caller (the per-request - # call-id governs execution/authorization, not which skills are listed), so a - # single shared cache is safe. SkillsProvider won't auto-cache a caller source, - # so we compose the caching ourselves. + # SkillsProvider won't auto-cache a caller source, so compose a bounded + # single-identity cache here. Replacing (rather than keying) the cache on + # identity changes avoids retaining one bucket per high-cardinality call ID. source: SkillsSource = _FoundryToolboxSkillsSource(self, archive_options=archive_options) if not disable_caching: - source = DeduplicatingSkillsSource(CachingSkillsSource(source, refresh_interval=cache_refresh_interval)) + source = DeduplicatingSkillsSource( + _FoundryToolboxCachingSkillsSource( + source, + identity_provider=lambda: tuple( + sorted((name.lower(), value) for name, value in self._effective_resource_headers().items()) + ), + refresh_interval=cache_refresh_interval, + ) + ) return SkillsProvider( source, source_id=source_id, @@ -418,36 +528,68 @@ def as_skills_provider( ) +class _FoundryToolboxResourceSession: + """Expose identity-scoped resource reads through a session-like adapter.""" + + def __init__(self, toolbox: FoundryToolbox) -> None: + self._toolbox = toolbox + + async def read_resource(self, uri: AnyUrl) -> ReadResourceResult: + return await self._toolbox._read_resource(uri) # pyright: ignore[reportPrivateUsage] + + +class _FoundryToolboxCachingSkillsSource(SkillsSource): + """Cache skills only while the effective Toolbox header identity is unchanged.""" + + def __init__( + self, + source: SkillsSource, + *, + identity_provider: Callable[[], tuple[tuple[str, str], ...]], + refresh_interval: timedelta | None, + ) -> None: + self._source = source + self._identity_provider = identity_provider + self._refresh_interval = refresh_interval + self._lock = asyncio.Lock() + self._identity: tuple[tuple[str, str], ...] | None = None + self._cache: CachingSkillsSource | None = None + + async def get_skills(self, context: SkillsSourceContext) -> list[Skill]: + identity = self._identity_provider() + async with self._lock: + if self._cache is None or identity != self._identity: + self._identity = identity + self._cache = CachingSkillsSource(self._source, refresh_interval=self._refresh_interval) + return await self._cache.get_skills(context) + + class _FoundryToolboxSkillsSource(SkillsSource): """Discovers skills from a connected :class:`FoundryToolbox` MCP session. The toolbox's MCP ``session`` is established lazily when the toolbox connects (via the agent or an ``async with`` block) and is **replaced** with a new - object whenever the toolbox reconnects. Skills are therefore bound to a - ``session_provider`` that resolves the toolbox's current session on every - fetch, so cached skills keep using the live session instead of a closed one. + 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 + session instead of a closed one. """ def __init__(self, toolbox: FoundryToolbox, *, archive_options: dict[str, Any] | None = None) -> None: self._toolbox = toolbox + self._resource_session = _FoundryToolboxResourceSession(toolbox) # Explicitly-set MCPSkillsSource archive kwargs; empty means use its defaults. self._archive_options: dict[str, Any] = archive_options or {} - def _require_session(self) -> ClientSession: - """Return the toolbox's current MCP session, or raise if not connected.""" - session = self._toolbox.session - if session is None: - raise RuntimeError( - "FoundryToolbox is not connected, so its skills cannot be discovered. " - "Pass the toolbox to the agent (tools=...) or enter it as an async " - "context manager before the agent runs." - ) - return session + def _resource_session_provider(self) -> ClientSession: + """Return a session-compatible adapter that scopes each resource read.""" + return cast("ClientSession", self._resource_session) async def get_skills(self, context: SkillsSourceContext) -> list[Skill]: - # Fail fast at discovery if not connected, then hand the source a provider - # (not a fixed session) so skills survive a reconnect that swaps the session. - self._require_session() - return await MCPSkillsSource(session_provider=self._require_session, **self._archive_options).get_skills( - context - ) + # Hand the source a provider (not a fixed session) so every resource read + # resolves the current caller identity and survives a reconnect that swaps + # the underlying session. + return await MCPSkillsSource( + session_provider=self._resource_session_provider, + **self._archive_options, + ).get_skills(context) diff --git a/python/packages/foundry_hosting/tests/test_toolbox.py b/python/packages/foundry_hosting/tests/test_toolbox.py index bd35992e397..1a4b586ce0b 100644 --- a/python/packages/foundry_hosting/tests/test_toolbox.py +++ b/python/packages/foundry_hosting/tests/test_toolbox.py @@ -5,20 +5,25 @@ from __future__ import annotations -from collections.abc import Callable +import asyncio +import json +from collections.abc import Awaitable, Callable from datetime import datetime, timedelta, timezone from types import SimpleNamespace -from typing import cast +from typing import Any, cast from unittest.mock import AsyncMock, patch import httpx import pytest from agent_framework import MCPStreamableHTTPTool, SkillsProvider, SkillsSourceContext, SupportsAgentRun +from agent_framework.exceptions import ToolException, ToolExecutionException from azure.ai.agentserver.core import ( FoundryAgentRequestContext, reset_request_context, set_request_context, ) +from mcp.client.session import ClientSession +from pydantic import AnyUrl from agent_framework_foundry_hosting import FoundryToolbox from agent_framework_foundry_hosting._feature_usage import FeatureIndex @@ -71,6 +76,100 @@ async def get_token(self, *scopes: str, **kwargs: object) -> _FakeAccessToken: return _FakeAccessToken(self._token) +class _ToolboxLoopbackServer: + """Minimal MCP server that records platform identity on every request.""" + + def __init__(self) -> None: + self.requests: list[httpx.Request] = [] + self._session_count = 0 + + 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) + + self.requests.append(request) + payload = json.loads(request.content) + method = payload.get("method") + call_id = request.headers.get("x-agent-foundry-call-id") + headers: dict[str, str] = {} + result: dict[str, Any] + + if method == "initialize": + self._session_count += 1 + headers["mcp-session-id"] = f"session-{self._session_count}-{call_id or 'none'}" + result = { + "protocolVersion": payload["params"]["protocolVersion"], + "capabilities": {"tools": {}, "resources": {}}, + "serverInfo": {"name": "toolbox-loopback", "version": "1"}, + } + elif method == "tools/list": + result = { + "tools": [ + { + "name": "record", + "inputSchema": { + "type": "object", + "properties": {"marker": {"type": "string"}}, + }, + } + ] + } + elif method == "tools/call": + result = {"content": [{"type": "text", "text": call_id or "none"}]} + elif method == "resources/read": + uri = payload["params"]["uri"] + if uri == "skill://index.json": + text = json.dumps({ + "skills": [ + { + "name": "loopback-skill", + "description": "Records the current Toolbox request identity.", + "type": "skill-md", + "url": "skill://loopback-skill/SKILL.md", + } + ] + }) + mime_type = "application/json" + else: + text = f"# Loopback skill\n\nCall ID: {call_id or 'none'}" + mime_type = "text/markdown" + result = {"contents": [{"uri": uri, "mimeType": mime_type, "text": text}]} + else: + result = {} + + if "id" not in payload: + return httpx.Response(202) + return httpx.Response( + 200, + headers=headers, + json={"jsonrpc": "2.0", "id": payload["id"], "result": result}, + ) + + def requests_for(self, method: str) -> list[httpx.Request]: + return [request for request in self.requests if json.loads(request.content).get("method") == method] + + +async def _attach_loopback_transport(toolbox: FoundryToolbox, server: _ToolboxLoopbackServer) -> None: + client = toolbox._httpx_client + assert client is not None + toolbox._httpx_client = httpx.AsyncClient( + auth=client.auth, + transport=httpx.MockTransport(server.handle), + timeout=30, + ) + await client.aclose() + + +async def _run_with_request_context(call_id: str, action: Callable[[], Awaitable[Any]]) -> Any: + token = set_request_context(FoundryAgentRequestContext(call_id=call_id, user_id="same-user")) + try: + return await action() + finally: + reset_request_context(token) + + 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" @@ -231,58 +330,232 @@ def test_sync_auth_flow_rejects_async_credential() -> None: next(auth.sync_auth_flow(request)) -async def test_auth_flow_forwards_call_id_when_present() -> None: - auth = _ToolboxAuth(_FakeCredential(), "scope") # type: ignore - request = httpx.Request("POST", "https://h/toolboxes/tb/mcp") +async def test_toolbox_header_provider_forwards_platform_headers_and_preserves_caller_headers() -> None: + toolbox = FoundryToolbox( + _FakeCredential(), # type: ignore + url="https://h/toolboxes/tb/mcp", + header_provider=lambda kwargs: { + "X-Custom-Header": kwargs["custom"], + "x-agent-foundry-call-id": "untrusted-override", + }, + ) token = set_request_context(FoundryAgentRequestContext(call_id="call-xyz")) try: - prepared = await anext(auth.async_auth_flow(request)) + headers = toolbox._effective_headers({"custom": "custom-value"}) finally: reset_request_context(token) + await toolbox.close() - assert prepared.headers["x-agent-foundry-call-id"] == "call-xyz" + assert headers["X-Custom-Header"] == "custom-value" + assert headers["x-agent-foundry-call-id"] == "call-xyz" -async def test_auth_flow_omits_call_id_when_absent() -> None: - auth = _ToolboxAuth(_FakeCredential(), "scope") # type: ignore - request = httpx.Request("POST", "https://h/toolboxes/tb/mcp") +async def test_toolbox_header_provider_does_not_retain_prior_platform_context() -> None: + toolbox = FoundryToolbox( + _FakeCredential(), # type: ignore + url="https://h/toolboxes/tb/mcp", + ) - prepared = await anext(auth.async_auth_flow(request)) + token = set_request_context(FoundryAgentRequestContext(call_id="canary-call-id")) + try: + headers = toolbox._effective_headers({}) + assert headers["x-agent-foundry-call-id"] == "canary-call-id" + finally: + reset_request_context(token) + + try: + headers = toolbox._effective_headers({}) + finally: + await toolbox.close() + + assert "x-agent-foundry-call-id" not in headers - assert "x-agent-foundry-call-id" not in prepared.headers +async def test_custom_header_provider_key_error_does_not_drop_platform_headers() -> None: + server = _ToolboxLoopbackServer() + toolbox = FoundryToolbox( + _FakeCredential(), # type: ignore + url="https://loopback.invalid/mcp", + header_provider=lambda kwargs: {"X-Custom-Header": kwargs["custom"]}, + ) + await _attach_loopback_transport(toolbox, server) + + try: + await _run_with_request_context("CALL-X", toolbox.connect) + await _run_with_request_context( + "CALL-X", + lambda: toolbox.call_tool("record", marker="with-custom", custom="custom-value"), + ) + with pytest.raises(KeyError, match="custom"): + await _run_with_request_context( + "CALL-X", + lambda: toolbox.call_tool("record", marker="missing-custom"), + ) + finally: + await toolbox.close() -async def test_async_auth_flow_removes_call_id_from_reused_request_when_context_absent() -> None: + first_initialize = server.requests_for("initialize")[0] + assert first_initialize.headers["x-agent-foundry-call-id"] == "CALL-X" + assert "X-Custom-Header" not in first_initialize.headers + + call = server.requests_for("tools/call")[-1] + assert call.headers["x-agent-foundry-call-id"] == "CALL-X" + assert call.headers["X-Custom-Header"] == "custom-value" + + +async def test_custom_header_provider_missing_seeded_value_still_fails_connect() -> None: + server = _ToolboxLoopbackServer() + toolbox = FoundryToolbox( + _FakeCredential(), # type: ignore + url="https://loopback.invalid/mcp", + header_provider=lambda kwargs: {"X-Custom-Header": kwargs["custom"]}, + ) + await _attach_loopback_transport(toolbox, server) + + token = set_request_context(FoundryAgentRequestContext(call_id="CALL-X")) + try: + await toolbox._prepare_for_run({}) + with pytest.raises(ToolException): + await toolbox.connect() + finally: + reset_request_context(token) + await toolbox.close() + + assert server.requests_for("initialize") == [] + + +async def test_auth_flow_does_not_resolve_platform_context() -> None: auth = _ToolboxAuth(_FakeCredential(), "scope") # type: ignore request = httpx.Request("POST", "https://h/toolboxes/tb/mcp") - token = set_request_context(FoundryAgentRequestContext(call_id="canary-call-id")) + token = set_request_context(FoundryAgentRequestContext(call_id="call-xyz")) try: prepared = await anext(auth.async_auth_flow(request)) - assert prepared.headers["x-agent-foundry-call-id"] == "canary-call-id" finally: reset_request_context(token) - prepared = await anext(auth.async_auth_flow(request)) - assert "x-agent-foundry-call-id" not in prepared.headers -def test_sync_auth_flow_removes_call_id_from_reused_request_when_context_absent() -> None: - auth = _ToolboxAuth(_FakeCredential(), "scope") # type: ignore - request = httpx.Request("POST", "https://h/toolboxes/tb/mcp") +async def test_toolbox_rebinds_after_connection_without_request_context() -> None: + server = _ToolboxLoopbackServer() + toolbox = FoundryToolbox( + _FakeCredential(), # type: ignore + url="https://loopback.invalid/mcp", + ) + await _attach_loopback_transport(toolbox, server) - token = set_request_context(FoundryAgentRequestContext(call_id="canary-call-id")) try: - prepared = next(auth.sync_auth_flow(request)) - assert prepared.headers["x-agent-foundry-call-id"] == "canary-call-id" + await toolbox.connect() + await _run_with_request_context( + "CALL-B", + lambda: toolbox.call_tool("record", marker="second-request"), + ) finally: - reset_request_context(token) + await toolbox.close() - prepared = next(auth.sync_auth_flow(request)) + initializes = server.requests_for("initialize") + assert [request.headers.get("x-agent-foundry-call-id") for request in initializes] == [None, "CALL-B"] + call = server.requests_for("tools/call")[-1] + assert call.headers["x-agent-foundry-call-id"] == "CALL-B" + assert call.headers["mcp-session-id"].endswith("-CALL-B") - assert "x-agent-foundry-call-id" not in prepared.headers + +async def test_toolbox_concurrent_callers_use_their_own_call_id() -> None: + server = _ToolboxLoopbackServer() + toolbox = FoundryToolbox( + _FakeCredential(), # type: ignore + url="https://loopback.invalid/mcp", + ) + await _attach_loopback_transport(toolbox, server) + + try: + await _run_with_request_context("CALL-A", toolbox.connect) + await asyncio.gather( + _run_with_request_context( + "CALL-B", + lambda: toolbox.call_tool("record", marker="request-b"), + ), + _run_with_request_context( + "CALL-C", + lambda: toolbox.call_tool("record", marker="request-c"), + ), + ) + finally: + await toolbox.close() + + calls = server.requests_for("tools/call") + calls_by_marker = {json.loads(request.content)["params"]["arguments"]["marker"]: request for request in calls} + assert calls_by_marker["request-b"].headers["x-agent-foundry-call-id"] == "CALL-B" + assert calls_by_marker["request-b"].headers["mcp-session-id"].endswith("-CALL-B") + assert calls_by_marker["request-c"].headers["x-agent-foundry-call-id"] == "CALL-C" + assert calls_by_marker["request-c"].headers["mcp-session-id"].endswith("-CALL-C") + + +async def test_toolbox_concurrent_skill_reads_use_their_own_call_id() -> None: + server = _ToolboxLoopbackServer() + toolbox = FoundryToolbox( + _FakeCredential(), # type: ignore + url="https://loopback.invalid/mcp", + load_tools=False, + header_provider=lambda _kwargs: {"X-Custom-Header": "resource-value"}, + ) + await _attach_loopback_transport(toolbox, server) + source = toolbox.as_skills_provider()._source + + async def read_skill() -> tuple[str, str]: + skills = await source.get_skills(_source_context()) + assert len(skills) == 1 + content = await skills[0].get_content() + resource = await skills[0].get_resource("reference.md") + assert resource is not None + return content, await resource.read() + + try: + await _run_with_request_context("CALL-A", toolbox.connect) + results = await asyncio.gather( + _run_with_request_context("CALL-B", read_skill), + _run_with_request_context("CALL-C", read_skill), + ) + finally: + await toolbox.close() + + assert "CALL-B" in results[0][0] + assert "CALL-B" in results[0][1] + assert "CALL-C" in results[1][0] + assert "CALL-C" in results[1][1] + + resource_requests = server.requests_for("resources/read") + assert len(resource_requests) == 6 + assert [request.headers["x-agent-foundry-call-id"] for request in resource_requests].count("CALL-B") == 3 + assert [request.headers["x-agent-foundry-call-id"] for request in resource_requests].count("CALL-C") == 3 + assert all( + request.headers["mcp-session-id"].endswith(f"-{request.headers['x-agent-foundry-call-id']}") + for request in resource_requests + ) + assert all(request.headers["X-Custom-Header"] == "resource-value" for request in resource_requests) + + +async def test_toolbox_skills_reject_header_provider_that_requires_runtime_kwargs() -> None: + server = _ToolboxLoopbackServer() + toolbox = FoundryToolbox( + _FakeCredential(), # type: ignore + url="https://loopback.invalid/mcp", + load_tools=False, + header_provider=lambda kwargs: {"X-Custom-Header": kwargs["custom"]}, + ) + await _attach_loopback_transport(toolbox, server) + source = toolbox.as_skills_provider()._source + + try: + await _run_with_request_context("CALL-X", toolbox.connect) + with pytest.raises(ToolExecutionException, match="cannot use a header_provider"): + await _run_with_request_context("CALL-X", lambda: source.get_skills(_source_context())) + finally: + await toolbox.close() + + assert server.requests_for("resources/read") == [] async def test_close_closes_owned_http_client(monkeypatch: pytest.MonkeyPatch) -> None: @@ -359,14 +632,14 @@ async def test_skills_source_uses_connected_session(monkeypatch: pytest.MonkeyPa _FakeCredential(), # type: ignore url="https://h/toolboxes/tb/mcp", ) - sentinel_session = object() + sentinel_session = AsyncMock() toolbox.session = sentinel_session # type: ignore - captured: dict[str, Callable[[], object]] = {} + captured: dict[str, Callable[[], ClientSession]] = {} captured_kwargs: dict[str, object] = {} class _StubSkillsSource: - def __init__(self, *, session_provider: Callable[[], object], **kwargs: object) -> None: + def __init__(self, *, session_provider: Callable[[], ClientSession], **kwargs: object) -> None: captured["session_provider"] = session_provider captured_kwargs.update(kwargs) @@ -378,13 +651,20 @@ async def get_skills(self, context: SkillsSourceContext) -> list[str]: result = await _FoundryToolboxSkillsSource(toolbox).get_skills(_source_context()) assert result == ["skill-a"] - # The source hands MCPSkillsSource a provider (not a fixed session) that resolves - # the toolbox's current session, so it survives a reconnect that swaps it. + # The source hands MCPSkillsSource a stable adapter that resolves the toolbox's + # current session for every identity-scoped resource read. provider = captured["session_provider"] - assert provider() is sentinel_session - new_session = object() + resource_session = provider() + assert provider() is resource_session + first_uri = AnyUrl("skill://first") + await resource_session.read_resource(first_uri) + sentinel_session.read_resource.assert_awaited_once_with(first_uri) + + new_session = AsyncMock() toolbox.session = new_session # type: ignore - assert provider() is new_session + second_uri = AnyUrl("skill://second") + await resource_session.read_resource(second_uri) + new_session.read_resource.assert_awaited_once_with(second_uri) # No archive options set -> MCPSkillsSource is constructed with defaults. assert captured_kwargs == {} @@ -424,11 +704,12 @@ async def test_skills_source_requires_connection_via_provider() -> None: ) toolbox.session = object() # type: ignore source = _FoundryToolboxSkillsSource(toolbox) - # Discovery captures the bound provider; a later reconnect gap (session is None) - # surfaces the same clear error when the provider is resolved. + resource_session = source._resource_session_provider() + # A cached skill can retain the adapter across reconnects; if the Toolbox is + # actually closed, the read still surfaces the same clear error. toolbox.session = None with pytest.raises(RuntimeError, match="not connected"): - source._require_session() + await resource_session.read_resource(AnyUrl("skill://closed")) class _FakeSkill: @@ -475,6 +756,25 @@ async def test_as_skills_provider_caches_by_default(monkeypatch: pytest.MonkeyPa assert read_count[0] == 1 +async def test_as_skills_provider_rebuilds_cache_when_platform_identity_changes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + toolbox = FoundryToolbox( + _FakeCredential(), # type: ignore + url="https://h/toolboxes/tb/mcp", + ) + toolbox.session = object() # type: ignore + read_count = _patch_counting_mcp_source(monkeypatch) + source = toolbox.as_skills_provider()._source + context = _source_context() + + await _run_with_request_context("CALL-A", lambda: source.get_skills(context)) + await _run_with_request_context("CALL-A", lambda: source.get_skills(context)) + await _run_with_request_context("CALL-B", lambda: source.get_skills(context)) + + assert read_count[0] == 2 + + async def test_as_skills_provider_disable_caching_rereads_every_run(monkeypatch: pytest.MonkeyPatch) -> None: toolbox = FoundryToolbox( _FakeCredential(), # type: ignore diff --git a/python/samples/04-hosting/foundry-hosted-agents/responses/files/README.md b/python/samples/04-hosting/foundry-hosted-agents/responses/files/README.md index dfb81edf78e..9725654ee80 100644 --- a/python/samples/04-hosting/foundry-hosted-agents/responses/files/README.md +++ b/python/samples/04-hosting/foundry-hosted-agents/responses/files/README.md @@ -26,11 +26,12 @@ see the [Toolbox sample](../foundry_toolbox/). Authenticate local runs with `ResponsesHostServer(agent=create_agent, history_source="agent_server")` creates fresh clients and a Toolbox MCP connection for each request and closes their -transports afterward. The MCP writer therefore inherits **this** request's -platform call ID, not an earlier caller's. The outer Responses service supplies -history; the host disables downstream model storage. An outer `store=false` -request writes no host-managed state, but does not undo external tool side -effects or delete uploads. +transports afterward. This request-owned lifecycle gives the sample a clear +cleanup boundary; `FoundryToolbox` also rebinds its MCP session to current +platform headers if an instance is reused across requests. The outer Responses +service supplies history; the host disables downstream model storage. An outer +`store=false` request writes no host-managed state, but does not undo external +tool side effects or delete uploads. Foundry session files and Toolbox code-interpreter container files are different resources. Reading an upload returns its text; it does not mount that upload into diff --git a/python/samples/04-hosting/foundry-hosted-agents/responses/files/main.py b/python/samples/04-hosting/foundry-hosted-agents/responses/files/main.py index 2eb7a213423..4434dcab8d0 100644 --- a/python/samples/04-hosting/foundry-hosted-agents/responses/files/main.py +++ b/python/samples/04-hosting/foundry-hosted-agents/responses/files/main.py @@ -12,7 +12,7 @@ # Copyright (c) Microsoft. All rights reserved. -"""Expose only explicitly uploaded session files, with request-owned Toolbox connections.""" +"""Expose only explicitly uploaded session files with request-owned SDK resources.""" from __future__ import annotations @@ -46,7 +46,7 @@ def read_file(filename: str) -> str: def create_agent() -> Agent: - """Create the client and MCP connection inside the current request context.""" + """Create the client and Toolbox as request-owned resources.""" endpoint = os.environ["FOUNDRY_PROJECT_ENDPOINT"] model = os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"] credential = ( diff --git a/python/samples/04-hosting/foundry-hosted-agents/responses/foundry_toolbox/README.md b/python/samples/04-hosting/foundry-hosted-agents/responses/foundry_toolbox/README.md index fa8f9925ee8..ffd619ea9a5 100644 --- a/python/samples/04-hosting/foundry-hosted-agents/responses/foundry_toolbox/README.md +++ b/python/samples/04-hosting/foundry-hosted-agents/responses/foundry_toolbox/README.md @@ -95,15 +95,15 @@ The agent uses `FoundryChatClient` from the Agent Framework to create an OpenAI- See [main.py](main.py) for the full implementation. -### Request-owned connections +### Request-owned resources The host receives `agent=create_agent`, not a process-wide agent instance. -Every request creates a fresh client and Toolbox MCP connection **under that -request's platform context**, then closes them afterward. The MCP HTTP writer -captures context when connecting; reusing one connection across hosted requests -would retain an earlier caller's `x-agent-foundry-call-id`. This pattern avoids -that reuse without changing the semantics of long-lived Toolbox instances in -other applications. +Every request creates a fresh client and Toolbox MCP connection, then closes +them afterward. This gives the sample one clear owner for its model transports, +credential and Toolbox lifecycle. `FoundryToolbox` also supports long-lived +instances: it resolves the current platform headers before each operation and +reconnects the MCP session when that identity changes, so a shared instance does +not retain an earlier caller's `x-agent-foundry-call-id`. The sample client closes only its own model/project transports and credential, including when MCP entry fails or the request is cancelled. Local authentication diff --git a/python/samples/04-hosting/foundry-hosted-agents/responses/foundry_toolbox/main.py b/python/samples/04-hosting/foundry-hosted-agents/responses/foundry_toolbox/main.py index 0af3721f6e3..671d77726a7 100644 --- a/python/samples/04-hosting/foundry-hosted-agents/responses/foundry_toolbox/main.py +++ b/python/samples/04-hosting/foundry-hosted-agents/responses/foundry_toolbox/main.py @@ -30,7 +30,7 @@ def create_agent() -> Agent: - """Create tools inside this request so the MCP writer captures its call ID.""" + """Create request-owned SDK resources and close them with the agent.""" endpoint = os.environ["FOUNDRY_PROJECT_ENDPOINT"] model = os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"] credential = ( diff --git a/python/samples/04-hosting/foundry-hosted-agents/responses/foundry_toolbox_mcp_skills/README.md b/python/samples/04-hosting/foundry-hosted-agents/responses/foundry_toolbox_mcp_skills/README.md index e6443b7ab48..1e288a9af3d 100644 --- a/python/samples/04-hosting/foundry-hosted-agents/responses/foundry_toolbox_mcp_skills/README.md +++ b/python/samples/04-hosting/foundry-hosted-agents/responses/foundry_toolbox_mcp_skills/README.md @@ -32,11 +32,12 @@ The `FoundryToolbox` is attached to the agent and its skills are exposed through The agent is hosted with the `ResponsesHostServer`, which provisions a REST API endpoint compatible with the OpenAI Responses protocol on `http://localhost:8088`. `agent=create_agent` builds a new Toolbox, skills provider/cache, credential and -client for **each request**. The host enters and exits that agent, so the -streamable-HTTP writer and skill-resource reads inherit the current platform -call ID instead of the first request's context. `tools=` still connects the -Toolbox and `context_providers=` still reads skills from that same session; both -are required. Only loading the vetted sample skill body is auto-approved, not +client for **each request**, giving those resources one deterministic owner and +cleanup boundary. `FoundryToolbox` also scopes tool and skill-resource +operations to the current platform call ID when an instance is reused, rebinding +its MCP session when that identity changes. `tools=` still connects the Toolbox +and `context_providers=` still reads skills from that same session; both are +required. Only loading the vetted sample skill body is auto-approved, not arbitrary skill scripts or external actions. The factory's client context closes its owned SDK transports and credential diff --git a/python/samples/04-hosting/foundry-hosted-agents/responses/foundry_toolbox_mcp_skills/main.py b/python/samples/04-hosting/foundry-hosted-agents/responses/foundry_toolbox_mcp_skills/main.py index b2384b0c3d4..c9b3535ea9f 100644 --- a/python/samples/04-hosting/foundry-hosted-agents/responses/foundry_toolbox_mcp_skills/main.py +++ b/python/samples/04-hosting/foundry-hosted-agents/responses/foundry_toolbox_mcp_skills/main.py @@ -30,7 +30,7 @@ def create_agent() -> Agent: - """Keep skill caches, credentials and the MCP writer within this request.""" + """Keep skill caches, credentials and the MCP lifecycle request-owned.""" endpoint = os.environ["FOUNDRY_PROJECT_ENDPOINT"] model = os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"] credential = (