mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-23 05:54:19 +02:00
Cache OBO credentials on AzureProvider for token reuse (#3212)
* Cache OBO credentials on AzureProvider for token reuse * chore: Update SDK documentation * Close evicted OBO credentials properly * chore: Update SDK documentation --------- Co-authored-by: marvin-context-protocol[bot] <225465937+marvin-context-protocol[bot]@users.noreply.github.com>
This commit is contained in:
Jeremiah Lowin
GitHub
parent
9b248a15b3
commit
7aba0df323
3 changed files with 173 additions and 39 deletions
|
|
@ -14,7 +14,7 @@ using the OAuth Proxy pattern for non-DCR OAuth flows.
|
|||
|
||||
## Functions
|
||||
|
||||
### `EntraOBOToken` <sup><a href="https://github.com/PrefectHQ/fastmcp/blob/main/src/fastmcp/server/auth/providers/azure.py#L658" target="_blank"><Icon icon="github" style="width: 14px; height: 14px;" /></a></sup>
|
||||
### `EntraOBOToken` <sup><a href="https://github.com/PrefectHQ/fastmcp/blob/main/src/fastmcp/server/auth/providers/azure.py#L680" target="_blank"><Icon icon="github" style="width: 14px; height: 14px;" /></a></sup>
|
||||
|
||||
```python
|
||||
EntraOBOToken(scopes: list[str]) -> str
|
||||
|
|
@ -43,7 +43,7 @@ or OBO exchange fails
|
|||
|
||||
## Classes
|
||||
|
||||
### `AzureProvider` <sup><a href="https://github.com/PrefectHQ/fastmcp/blob/main/src/fastmcp/server/auth/providers/azure.py#L31" target="_blank"><Icon icon="github" style="width: 14px; height: 14px;" /></a></sup>
|
||||
### `AzureProvider` <sup><a href="https://github.com/PrefectHQ/fastmcp/blob/main/src/fastmcp/server/auth/providers/azure.py#L33" target="_blank"><Icon icon="github" style="width: 14px; height: 14px;" /></a></sup>
|
||||
|
||||
|
||||
Azure (Microsoft Entra) OAuth provider for FastMCP.
|
||||
|
|
@ -78,7 +78,7 @@ Setup:
|
|||
|
||||
**Methods:**
|
||||
|
||||
#### `authorize` <sup><a href="https://github.com/PrefectHQ/fastmcp/blob/main/src/fastmcp/server/auth/providers/azure.py#L235" target="_blank"><Icon icon="github" style="width: 14px; height: 14px;" /></a></sup>
|
||||
#### `authorize` <sup><a href="https://github.com/PrefectHQ/fastmcp/blob/main/src/fastmcp/server/auth/providers/azure.py#L243" target="_blank"><Icon icon="github" style="width: 14px; height: 14px;" /></a></sup>
|
||||
|
||||
```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` <sup><a href="https://github.com/PrefectHQ/fastmcp/blob/main/src/fastmcp/server/auth/providers/azure.py#L461" target="_blank"><Icon icon="github" style="width: 14px; height: 14px;" /></a></sup>
|
||||
#### `get_obo_credential` <sup><a href="https://github.com/PrefectHQ/fastmcp/blob/main/src/fastmcp/server/auth/providers/azure.py#L469" target="_blank"><Icon icon="github" style="width: 14px; height: 14px;" /></a></sup>
|
||||
|
||||
```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` <sup><a href="https://github.com/PrefectHQ/fastmcp/blob/main/src/fastmcp/server/auth/providers/azure.py#L489" target="_blank"><Icon icon="github" style="width: 14px; height: 14px;" /></a></sup>
|
||||
#### `close_obo_credentials` <sup><a href="https://github.com/PrefectHQ/fastmcp/blob/main/src/fastmcp/server/auth/providers/azure.py#L510" target="_blank"><Icon icon="github" style="width: 14px; height: 14px;" /></a></sup>
|
||||
|
||||
```python
|
||||
close_obo_credentials(self) -> None
|
||||
```
|
||||
|
||||
Close all cached OBO credentials.
|
||||
|
||||
|
||||
### `AzureJWTVerifier` <sup><a href="https://github.com/PrefectHQ/fastmcp/blob/main/src/fastmcp/server/auth/providers/azure.py#L521" target="_blank"><Icon icon="github" style="width: 14px; height: 14px;" /></a></sup>
|
||||
|
||||
|
||||
JWT verifier pre-configured for Azure AD / Microsoft Entra ID.
|
||||
|
|
@ -157,7 +166,7 @@ Example::
|
|||
|
||||
**Methods:**
|
||||
|
||||
#### `scopes_supported` <sup><a href="https://github.com/PrefectHQ/fastmcp/blob/main/src/fastmcp/server/auth/providers/azure.py#L569" target="_blank"><Icon icon="github" style="width: 14px; height: 14px;" /></a></sup>
|
||||
#### `scopes_supported` <sup><a href="https://github.com/PrefectHQ/fastmcp/blob/main/src/fastmcp/server/auth/providers/azure.py#L601" target="_blank"><Icon icon="github" style="width: 14px; height: 14px;" /></a></sup>
|
||||
|
||||
```python
|
||||
scopes_supported(self) -> list[str]
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue