diff --git a/docs/python-sdk/fastmcp-server-auth-providers-azure.mdx b/docs/python-sdk/fastmcp-server-auth-providers-azure.mdx index 89d3c69c4..7c9a982e0 100644 --- a/docs/python-sdk/fastmcp-server-auth-providers-azure.mdx +++ b/docs/python-sdk/fastmcp-server-auth-providers-azure.mdx @@ -14,7 +14,7 @@ using the OAuth Proxy pattern for non-DCR OAuth flows. ## Functions -### `EntraOBOToken` +### `EntraOBOToken` ```python EntraOBOToken(scopes: list[str]) -> str @@ -43,7 +43,7 @@ or OBO exchange fails ## Classes -### `AzureProvider` +### `AzureProvider` Azure (Microsoft Entra) OAuth provider for FastMCP. @@ -78,7 +78,7 @@ Setup: **Methods:** -#### `authorize` +#### `authorize` ```python authorize(self, client: OAuthClientInformationFull, params: AuthorizationParams) -> str @@ -98,17 +98,17 @@ scopes to determine the resource/audience instead of a separate parameter. - Authorization URL to redirect the user to Azure AD -#### `create_obo_credential` +#### `get_obo_credential` ```python -create_obo_credential(self, user_assertion: str) -> OnBehalfOfCredential +get_obo_credential(self, user_assertion: str) -> OnBehalfOfCredential ``` -Create an OnBehalfOfCredential for OBO token exchange. +Get a cached or new OnBehalfOfCredential for OBO token exchange. -Uses the AzureProvider's configuration (client_id, client_secret, -tenant_id, authority) to create a credential that can exchange the -user's token for downstream API tokens. +Credentials are cached by user assertion so the Azure SDK's internal +token cache can avoid redundant OBO exchanges when the same user +calls multiple tools with the same scopes. **Args:** - `user_assertion`: The user's access token to exchange via OBO. @@ -120,7 +120,16 @@ user's token for downstream API tokens. - `ImportError`: If azure-identity is not installed (requires fastmcp[azure]). -### `AzureJWTVerifier` +#### `close_obo_credentials` + +```python +close_obo_credentials(self) -> None +``` + +Close all cached OBO credentials. + + +### `AzureJWTVerifier` JWT verifier pre-configured for Azure AD / Microsoft Entra ID. @@ -157,7 +166,7 @@ Example:: **Methods:** -#### `scopes_supported` +#### `scopes_supported` ```python scopes_supported(self) -> list[str] diff --git a/src/fastmcp/server/auth/providers/azure.py b/src/fastmcp/server/auth/providers/azure.py index b651b74ed..868631d13 100644 --- a/src/fastmcp/server/auth/providers/azure.py +++ b/src/fastmcp/server/auth/providers/azure.py @@ -6,6 +6,8 @@ using the OAuth Proxy pattern for non-DCR OAuth flows. from __future__ import annotations +import hashlib +from collections import OrderedDict from typing import TYPE_CHECKING, Any, cast from key_value.aio.protocols import AsyncKeyValue @@ -166,6 +168,12 @@ class AzureProvider(OAuthProxy): self._tenant_id = tenant_id self._base_authority = base_authority + # Cache of OBO credentials keyed by hash of user assertion token. + # Reusing credentials allows the Azure SDK's internal token cache + # to avoid redundant OBO exchanges for the same user + scopes. + self._obo_credentials: OrderedDict[str, OnBehalfOfCredential] = OrderedDict() + self._obo_max_credentials: int = 128 + # Apply defaults self.identifier_uri = identifier_uri or f"api://{client_id}" self.additional_authorize_scopes: list[str] = parsed_additional_scopes @@ -458,12 +466,12 @@ class AzureProvider(OAuthProxy): logger.debug("Failed to extract Azure claims: %s", e) return None - def create_obo_credential(self, user_assertion: str) -> OnBehalfOfCredential: - """Create an OnBehalfOfCredential for OBO token exchange. + async def get_obo_credential(self, user_assertion: str) -> OnBehalfOfCredential: + """Get a cached or new OnBehalfOfCredential for OBO token exchange. - Uses the AzureProvider's configuration (client_id, client_secret, - tenant_id, authority) to create a credential that can exchange the - user's token for downstream API tokens. + Credentials are cached by user assertion so the Azure SDK's internal + token cache can avoid redundant OBO exchanges when the same user + calls multiple tools with the same scopes. Args: user_assertion: The user's access token to exchange via OBO. @@ -477,13 +485,37 @@ class AzureProvider(OAuthProxy): _require_azure_identity("OBO token exchange") from azure.identity.aio import OnBehalfOfCredential - return OnBehalfOfCredential( + key = hashlib.sha256(user_assertion.encode()).hexdigest() + + if key in self._obo_credentials: + self._obo_credentials.move_to_end(key) + return self._obo_credentials[key] + + credential = OnBehalfOfCredential( tenant_id=self._tenant_id, client_id=self._upstream_client_id, client_secret=self._upstream_client_secret.get_secret_value(), user_assertion=user_assertion, authority=f"https://{self._base_authority}", ) + self._obo_credentials[key] = credential + + # Evict oldest if over capacity + while len(self._obo_credentials) > self._obo_max_credentials: + _, evicted = self._obo_credentials.popitem(last=False) + await evicted.close() + + return credential + + async def close_obo_credentials(self) -> None: + """Close all cached OBO credentials.""" + credentials = list(self._obo_credentials.values()) + self._obo_credentials.clear() + for credential in credentials: + try: + await credential.close() + except Exception: + logger.debug("Error closing OBO credential", exc_info=True) class AzureJWTVerifier(JWTVerifier): @@ -611,12 +643,13 @@ class _EntraOBOToken(Dependency): # type: ignore[misc] """Dependency that performs OBO token exchange for Microsoft Entra. Uses azure.identity's OnBehalfOfCredential for async-native OBO, - with automatic token caching and refresh. + with automatic token caching and refresh. Credentials are cached on + the AzureProvider so repeated tool calls reuse existing credentials + and benefit from the Azure SDK's internal token cache. """ def __init__(self, scopes: list[str]): self.scopes = scopes - self._credential: OnBehalfOfCredential | None = None async def __aenter__(self) -> str: _require_azure_identity("EntraOBOToken") @@ -636,24 +669,13 @@ class _EntraOBOToken(Dependency): # type: ignore[misc] f"Current provider: {type(server.auth).__name__}" ) - self._credential = server.auth.create_obo_credential( + credential = await server.auth.get_obo_credential( user_assertion=access_token.token, ) - try: - result = await self._credential.get_token(*self.scopes) - except BaseException: - await self._credential.close() - self._credential = None - raise - + result = await credential.get_token(*self.scopes) return result.token - async def __aexit__(self, *args: object) -> None: - if self._credential is not None: - await self._credential.close() - self._credential = None - def EntraOBOToken(scopes: list[str]) -> str: """Exchange the user's Entra token for a downstream API token via OBO. diff --git a/tests/server/auth/providers/test_azure.py b/tests/server/auth/providers/test_azure.py index 0ea6166bf..c64902aab 100644 --- a/tests/server/auth/providers/test_azure.py +++ b/tests/server/auth/providers/test_azure.py @@ -1219,10 +1219,10 @@ class TestAzureJWTVerifier: class TestAzureOBOIntegration: - """Tests for azure.identity OBO integration (create_obo_credential, EntraOBOToken).""" + """Tests for azure.identity OBO integration (get_obo_credential, EntraOBOToken).""" - def test_create_obo_credential_returns_configured_credential(self): - """Test that create_obo_credential returns a properly configured credential.""" + async def test_get_obo_credential_returns_configured_credential(self): + """Test that get_obo_credential returns a properly configured credential.""" from unittest.mock import MagicMock, patch provider = AzureProvider( @@ -1238,7 +1238,9 @@ class TestAzureOBOIntegration: with patch( "azure.identity.aio.OnBehalfOfCredential", return_value=mock_credential ) as mock_class: - credential = provider.create_obo_credential(user_assertion="user-token-123") + credential = await provider.get_obo_credential( + user_assertion="user-token-123" + ) mock_class.assert_called_once_with( tenant_id="test-tenant-id", @@ -1249,8 +1251,109 @@ class TestAzureOBOIntegration: ) assert credential is mock_credential - def test_create_obo_credential_with_custom_authority(self): - """Test that create_obo_credential uses custom base_authority.""" + async def test_get_obo_credential_caches_by_assertion(self): + """Test that the same assertion returns the cached credential.""" + from unittest.mock import MagicMock, patch + + provider = AzureProvider( + client_id="test-client-id", + client_secret="test-client-secret", + tenant_id="test-tenant-id", + base_url="https://myserver.com", + required_scopes=["read"], + jwt_signing_key="test-secret", + ) + + mock_credential = MagicMock() + with patch( + "azure.identity.aio.OnBehalfOfCredential", return_value=mock_credential + ) as mock_class: + first = await provider.get_obo_credential(user_assertion="same-token") + second = await provider.get_obo_credential(user_assertion="same-token") + + assert first is second + mock_class.assert_called_once() + + async def test_get_obo_credential_different_assertions_get_different_credentials( + self, + ): + """Test that different assertions produce different credentials.""" + from unittest.mock import MagicMock, patch + + provider = AzureProvider( + client_id="test-client-id", + client_secret="test-client-secret", + tenant_id="test-tenant-id", + base_url="https://myserver.com", + required_scopes=["read"], + jwt_signing_key="test-secret", + ) + + creds = [MagicMock(), MagicMock()] + with patch("azure.identity.aio.OnBehalfOfCredential", side_effect=creds): + first = await provider.get_obo_credential(user_assertion="token-a") + second = await provider.get_obo_credential(user_assertion="token-b") + + assert first is not second + assert first is creds[0] + assert second is creds[1] + + async def test_get_obo_credential_evicts_oldest_when_over_capacity(self): + """Test that credentials are evicted LRU-style when cache is full.""" + from unittest.mock import AsyncMock, MagicMock, patch + + provider = AzureProvider( + client_id="test-client-id", + client_secret="test-client-secret", + tenant_id="test-tenant-id", + base_url="https://myserver.com", + required_scopes=["read"], + jwt_signing_key="test-secret", + ) + provider._obo_max_credentials = 2 + + creds = [MagicMock(close=AsyncMock()) for _ in range(3)] + with patch("azure.identity.aio.OnBehalfOfCredential", side_effect=creds): + await provider.get_obo_credential(user_assertion="token-1") + await provider.get_obo_credential(user_assertion="token-2") + await provider.get_obo_credential(user_assertion="token-3") + + assert len(provider._obo_credentials) == 2 + creds[0].close.assert_awaited_once() + # token-1's credential was evicted + assert ( + await provider.get_obo_credential(user_assertion="token-2") is creds[1] + ) + assert ( + await provider.get_obo_credential(user_assertion="token-3") is creds[2] + ) + + async def test_close_obo_credentials(self): + """Test that close_obo_credentials closes all cached credentials.""" + from unittest.mock import AsyncMock, MagicMock, patch + + provider = AzureProvider( + client_id="test-client-id", + client_secret="test-client-secret", + tenant_id="test-tenant-id", + base_url="https://myserver.com", + required_scopes=["read"], + jwt_signing_key="test-secret", + ) + + creds = [MagicMock(close=AsyncMock()) for _ in range(2)] + with patch("azure.identity.aio.OnBehalfOfCredential", side_effect=creds): + await provider.get_obo_credential(user_assertion="token-a") + await provider.get_obo_credential(user_assertion="token-b") + + await provider.close_obo_credentials() + + assert len(provider._obo_credentials) == 0 + for cred in creds: + cred.close.assert_awaited_once() + + async def test_get_obo_credential_with_custom_authority(self): + """Test that get_obo_credential uses custom base_authority.""" from unittest.mock import MagicMock, patch provider = AzureProvider( @@ -1267,7 +1370,7 @@ class TestAzureOBOIntegration: with patch( "azure.identity.aio.OnBehalfOfCredential", return_value=mock_credential ) as mock_class: - provider.create_obo_credential(user_assertion="user-token") + await provider.get_obo_credential(user_assertion="user-token") call_kwargs = mock_class.call_args[1] assert call_kwargs["authority"] == "https://login.microsoftonline.us"