Python: Support async credentials in FoundryToolbox

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,

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