mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
feat(mcp): v1-backed OAuth token source for authorization_code
V1PerUserTokenStore reads the user's stored access token through v1's mcp_per_user_token_cache (Redis-backed, encrypted) and wraps it in an OAuthToken. v1 holds only the access token (its cache TTL is the lifetime), so no expires_at/refresh_token yet; the v2 cache holds it for its default TTL and the OAuth challenge drives re-auth once v1's cache drops it. Additive: nothing wires it yet, so no behavior change. Step 1b swaps it for a v2-native token store behind the OAuthTokenStore seam.
This commit is contained in:
parent
e284981c95
commit
6710a8e49c
2 changed files with 75 additions and 0 deletions
|
|
@ -0,0 +1,38 @@
|
|||
"""v1-backed ``OAuthTokenStore`` source for the ``authorization_code`` mode.
|
||||
|
||||
Reads the user's stored access token through v1's ``mcp_per_user_token_cache`` (a Redis-backed,
|
||||
encrypted-at-rest per-user cache). This is a temporary adapter: step 1b replaces it with a
|
||||
v2-native token store that also tracks expiry and refresh, behind the same ``OAuthTokenStore`` seam.
|
||||
It imports v1, so it is kept out of the package ``__init__`` like the rest of the adapter layer.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
|
||||
OAuthToken,
|
||||
)
|
||||
|
||||
|
||||
class V1PerUserTokenStore:
|
||||
"""``OAuthTokenStore`` backed by v1's per-user token cache.
|
||||
|
||||
v1 stores only the access token (the cache TTL is its lifetime), so the ``OAuthToken`` carries
|
||||
no ``expires_at`` or ``refresh_token``: the v2 cache holds it for its default TTL, and the OAuth
|
||||
challenge drives re-auth once v1's cache drops it. v1's ``get`` swallows errors as a miss, so
|
||||
this never raises ``TokenStoreUnavailable``.
|
||||
"""
|
||||
|
||||
async def fetch(self, user_id: str, server_id: str) -> Optional[OAuthToken]:
|
||||
if not user_id:
|
||||
return None
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import (
|
||||
mcp_per_user_token_cache,
|
||||
)
|
||||
|
||||
access_token = await mcp_per_user_token_cache.get(user_id, server_id)
|
||||
if not access_token:
|
||||
return None
|
||||
return OAuthToken(access_token=access_token)
|
||||
|
|
@ -0,0 +1,37 @@
|
|||
"""Tests for the v1-backed per-user OAuth token source (V1PerUserTokenStore)."""
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.v1_token_store import (
|
||||
V1PerUserTokenStore,
|
||||
)
|
||||
|
||||
_GET = (
|
||||
"litellm.proxy._experimental.mcp_server.oauth2_token_cache."
|
||||
"mcp_per_user_token_cache.get"
|
||||
)
|
||||
|
||||
|
||||
async def test_wraps_the_v1_access_token():
|
||||
with patch(_GET, new=AsyncMock(return_value="at-123")):
|
||||
token = await V1PerUserTokenStore().fetch("alice", "s")
|
||||
assert token is not None and token.access_token == "at-123"
|
||||
|
||||
|
||||
async def test_missing_token_is_none():
|
||||
with patch(_GET, new=AsyncMock(return_value=None)):
|
||||
assert await V1PerUserTokenStore().fetch("alice", "s") is None
|
||||
|
||||
|
||||
async def test_empty_user_short_circuits_without_hitting_v1():
|
||||
get = AsyncMock(return_value="at")
|
||||
with patch(_GET, new=get):
|
||||
assert await V1PerUserTokenStore().fetch("", "s") is None
|
||||
get.assert_not_called()
|
||||
|
||||
|
||||
async def test_passes_user_and_server_through_to_v1():
|
||||
get = AsyncMock(return_value="at")
|
||||
with patch(_GET, new=get):
|
||||
await V1PerUserTokenStore().fetch("alice", "srv-1")
|
||||
get.assert_awaited_once_with("alice", "srv-1")
|
||||
Loading…
Add table
Reference in a new issue