From 475ed87b12c0fa3a6506b3ea208136c644b9fd58 Mon Sep 17 00:00:00 2001 From: Chinedum Echeta <60179183+cecheta@users.noreply.github.com> Date: Mon, 20 Jul 2026 10:53:04 +0000 Subject: [PATCH 1/4] Python: Support async credentials in `FoundryToolbox` --- .../_toolbox.py | 62 ++++++++++++++----- .../foundry_hosting/tests/test_toolbox.py | 55 +++++++++++++--- 2 files changed, 93 insertions(+), 24 deletions(-) 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..29aae30053a 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,11 +21,12 @@ 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 logger = logging.getLogger(__name__) @@ -74,24 +76,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: TokenCredential | AsyncTokenCredential, 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 +166,7 @@ class FoundryToolbox(MCPStreamableHTTPTool): def __init__( self, - credential: TokenCredential, + credential: TokenCredential | AsyncTokenCredential, *, url: str | None = None, name: str | None = None, @@ -150,9 +178,9 @@ def __init__( """Initialize a Foundry toolbox tool. Args: - credential: A Microsoft Entra credential used to obtain bearer tokens for - the toolbox endpoint. Tokens are requested per outbound request and - cached by the credential. + credential: A synchronous or asynchronous Microsoft Entra credential used + to obtain bearer tokens for the toolbox endpoint. Tokens are requested + per outbound request and cached by the credential. Keyword Args: url: The toolbox MCP endpoint URL. When ``None``, it is resolved from 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 From 8170dcd0290c4b2278e67279eb3215144e091288 Mon Sep 17 00:00:00 2001 From: Chinedum Echeta <60179183+cecheta@users.noreply.github.com> Date: Mon, 20 Jul 2026 12:01:36 +0100 Subject: [PATCH 2/4] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .../agent_framework_foundry_hosting/_toolbox.py | 3 +++ 1 file changed, 3 insertions(+) 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 29aae30053a..2ed11005649 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py @@ -98,6 +98,9 @@ def _apply_headers(self, request: httpx.Request, token: AccessToken) -> None: for key, value in get_request_context().platform_headers().items(): request.headers[key] = value + def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]: + yield from self.sync_auth_flow(request) + 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. From 57d2b4d3bb49313cc4e0d28d69f431e8df9744a7 Mon Sep 17 00:00:00 2001 From: Chinedum Echeta <60179183+cecheta@users.noreply.github.com> Date: Tue, 21 Jul 2026 15:30:14 +0000 Subject: [PATCH 3/4] Refactor: Use AzureCredentialTypes for credential type annotations in Toolbox classes --- .../agent_framework_foundry_hosting/_toolbox.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) 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 2ed11005649..803ed28caad 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py @@ -29,6 +29,8 @@ 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. @@ -89,7 +91,7 @@ class _ToolboxAuth(httpx.Auth): and is absent (no header) for protocol ``1.0.0`` or local development. """ - def __init__(self, credential: TokenCredential | AsyncTokenCredential, scope: str) -> None: + def __init__(self, credential: AzureCredentialTypes, scope: str) -> None: self._credential = credential self._scope = scope @@ -169,7 +171,7 @@ class FoundryToolbox(MCPStreamableHTTPTool): def __init__( self, - credential: TokenCredential | AsyncTokenCredential, + credential: AzureCredentialTypes, *, url: str | None = None, name: str | None = None, @@ -181,9 +183,9 @@ def __init__( """Initialize a Foundry toolbox tool. Args: - credential: A synchronous or asynchronous Microsoft Entra credential used - to obtain bearer tokens for the toolbox endpoint. Tokens are requested - per outbound request and cached by the credential. + credential: A Microsoft Entra credential used to obtain bearer tokens for + the toolbox endpoint. Tokens are requested per outbound request and + cached by the credential. Keyword Args: url: The toolbox MCP endpoint URL. When ``None``, it is resolved from From 47e05cf0bea1b044773801257116f15a9b16d4d8 Mon Sep 17 00:00:00 2001 From: Chinedum Echeta <60179183+cecheta@users.noreply.github.com> Date: Thu, 23 Jul 2026 13:06:39 +0000 Subject: [PATCH 4/4] Remove auth_flow method from _ToolboxAuth class --- .../agent_framework_foundry_hosting/_toolbox.py | 3 --- 1 file changed, 3 deletions(-) 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 803ed28caad..917faa86a15 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py @@ -100,9 +100,6 @@ def _apply_headers(self, request: httpx.Request, token: AccessToken) -> None: for key, value in get_request_context().platform_headers().items(): request.headers[key] = value - def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]: - yield from self.sync_auth_flow(request) - 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.