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 8d800adc571..917faa86a15 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py @@ -2,6 +2,7 @@ from __future__ import annotations +import inspect import logging import os from typing import TYPE_CHECKING @@ -20,13 +21,16 @@ from azure.ai.agentserver.core import get_request_context if TYPE_CHECKING: - from collections.abc import Generator + from collections.abc import AsyncGenerator, Generator from datetime import timedelta from agent_framework import Skill - from azure.core.credentials import TokenCredential + from azure.core.credentials import AccessToken, TokenCredential + from azure.core.credentials_async import AsyncTokenCredential from mcp.client.session import ClientSession + AzureCredentialTypes = TokenCredential | AsyncTokenCredential + logger = logging.getLogger(__name__) # Default Microsoft Entra scope for Foundry data-plane access. @@ -74,24 +78,50 @@ 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. - ``auth_flow`` runs for *every* outbound request (connection handshake as well - as tool calls), so the bearer token is always present. The per-request - ``x-agent-foundry-call-id`` is read from the request-scoped context populated - by the hosting endpoint; it resolves to a fresh value on each request and is - absent (no header) for protocol ``1.0.0`` or local development. + Both the synchronous (``sync_auth_flow``) and asynchronous (``async_auth_flow``) + httpx auth hooks are implemented, so the same auth works regardless of which + transport the toolbox client uses. Each runs for *every* outbound request + (connection handshake as well as tool calls), so the bearer token is always + 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 + per-request ``x-agent-foundry-call-id`` is read from the request-scoped context + populated by the hosting endpoint; it resolves to a fresh value on each request + and is absent (no header) for protocol ``1.0.0`` or local development. """ - def __init__(self, credential: TokenCredential, scope: str) -> None: + def __init__(self, credential: AzureCredentialTypes, scope: str) -> None: self._credential = credential self._scope = scope - def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]: - # azure-core credentials cache the token internally and only refresh near - # expiry, so calling get_token per request is cheap. - token = self._credential.get_token(self._scope).token - request.headers["Authorization"] = f"Bearer {token}" + def _apply_headers(self, request: httpx.Request, token: AccessToken) -> None: + request.headers["Authorization"] = f"Bearer {token.token}" for key, value in get_request_context().platform_headers().items(): request.headers[key] = value + + def sync_auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]: + # azure-core credentials cache the token internally and only refresh near + # expiry, so calling get_token per request is cheap. + token = self._credential.get_token(self._scope) + if inspect.isawaitable(token): + close = getattr(token, "close", None) + if callable(close): + close() + raise RuntimeError( + "An async credential cannot be used with the synchronous auth flow; " + "use a synchronous TokenCredential or drive the toolbox with an httpx.AsyncClient." + ) + self._apply_headers(request, token) + yield request + + async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]: + # Sync credentials return the token directly; async credentials return an + # awaitable to await. + token = self._credential.get_token(self._scope) + if inspect.isawaitable(token): + token = await token + self._apply_headers(request, token) yield request @@ -138,7 +168,7 @@ class FoundryToolbox(MCPStreamableHTTPTool): def __init__( self, - credential: TokenCredential, + credential: AzureCredentialTypes, *, url: str | None = None, name: str | None = None, diff --git a/python/packages/foundry_hosting/tests/test_toolbox.py b/python/packages/foundry_hosting/tests/test_toolbox.py index 4990213a1a1..0c8e2c49554 100644 --- a/python/packages/foundry_hosting/tests/test_toolbox.py +++ b/python/packages/foundry_hosting/tests/test_toolbox.py @@ -57,6 +57,18 @@ def get_token(self, *scopes: str, **kwargs: object) -> _FakeAccessToken: return _FakeAccessToken(self._token) +class _FakeAsyncCredential: + """Minimal stand-in for azure.core.credentials_async.AsyncTokenCredential.""" + + def __init__(self, token: str = "fake-token") -> None: + self._token = token + self.scopes: list[str] = [] + + async def get_token(self, *scopes: str, **kwargs: object) -> _FakeAccessToken: + self.scopes.extend(scopes) + return _FakeAccessToken(self._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" @@ -106,36 +118,65 @@ def test_init_derives_name_and_defaults() -> None: assert toolbox.load_prompts_flag is False -def test_auth_flow_injects_bearer_token() -> None: +async def test_auth_flow_injects_bearer_token() -> None: cred = _FakeCredential("abc123") auth = _ToolboxAuth(cred, "https://ai.azure.com/.default") # type: ignore request = httpx.Request("POST", "https://h/toolboxes/tb/mcp") - flow = auth.auth_flow(request) - prepared = next(flow) + prepared = await anext(auth.async_auth_flow(request)) assert prepared.headers["Authorization"] == "Bearer abc123" assert cred.scopes == ["https://ai.azure.com/.default"] -def test_auth_flow_forwards_call_id_when_present() -> None: +async def test_auth_flow_injects_bearer_token_async_credential() -> None: + cred = _FakeAsyncCredential("async123") + auth = _ToolboxAuth(cred, "https://ai.azure.com/.default") # type: ignore + request = httpx.Request("POST", "https://h/toolboxes/tb/mcp") + + prepared = await anext(auth.async_auth_flow(request)) + + assert prepared.headers["Authorization"] == "Bearer async123" + assert cred.scopes == ["https://ai.azure.com/.default"] + + +def test_sync_auth_flow_injects_bearer_token() -> None: + cred = _FakeCredential("sync123") + auth = _ToolboxAuth(cred, "https://ai.azure.com/.default") # type: ignore + request = httpx.Request("POST", "https://h/toolboxes/tb/mcp") + + prepared = next(auth.sync_auth_flow(request)) + + assert prepared.headers["Authorization"] == "Bearer sync123" + assert cred.scopes == ["https://ai.azure.com/.default"] + + +def test_sync_auth_flow_rejects_async_credential() -> None: + auth = _ToolboxAuth(_FakeAsyncCredential(), "scope") # type: ignore + request = httpx.Request("POST", "https://h/toolboxes/tb/mcp") + + with pytest.raises(RuntimeError, match="async credential"): + 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") token = set_request_context(FoundryAgentRequestContext(call_id="call-xyz")) try: - prepared = next(auth.auth_flow(request)) + prepared = await anext(auth.async_auth_flow(request)) finally: reset_request_context(token) assert prepared.headers["x-agent-foundry-call-id"] == "call-xyz" -def test_auth_flow_omits_call_id_when_absent() -> None: +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") - prepared = next(auth.auth_flow(request)) + prepared = await anext(auth.async_auth_flow(request)) assert "x-agent-foundry-call-id" not in prepared.headers