scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

fix(mcp): green CI for authz_code dispatch (format + UTC expiry + v2-seam tests)

Browse source 
- ruff format per_user_oauth_store.py (clears the lint check)
- v2_token_store._iso_to_epoch: anchor a tz-naive expiry to UTC before
  .timestamp(), matching v1's db.py _remaining_token_seconds (Greptile P1) so a
  non-UTC host doesn't read the expiry as local time and skew refresh timing
- test_mcp_stale_session: repoint the 3 discovery tests off the removed v1
  _get_user_oauth_extra_headers_from_db onto the v2 has_user_oauth_token seam;
  the delegate test now asserts the existence check is never consulted (delegate
  short-circuits to the resource_metadata 401 before any token lookup)
- test_mcp_server_manager: repoint test_deferred_mode_uses_v1_auth_value at M2M
  (oauth2 client_credentials), which is still a deferred mode, since per-user
  oauth2 (authorization_code) now routes to the v2 resolver

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Tin Chi Lo 2026-06-26 17:03:47 -07:00
parent 3bf73b9772
commit c813b594b7
4 changed files with 35 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py
View file

@ -108,9 +108,7 @@ def build_per_user_oauth_token_store(
# Cache and refresh coordinator use the foundation's in-process defaults (a single replica needs
# no shared cache or lock); the cross-replica path is layered on separately.
refreshing = RefreshingTokenStore(V2PerUserTokenStore(_read_credential), refresher)
return CachedOAuthTokenStore(
refreshing, default_ttl_seconds=_DEFAULT_TTL_SECONDS
)
return CachedOAuthTokenStore(refreshing, default_ttl_seconds=_DEFAULT_TTL_SECONDS)
class LazyPerUserOAuthTokenStore:

10
litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py
View file

@ -11,7 +11,7 @@ injected, so the DB/decoding plumbing stays testable and out of this seam.
from __future__ import annotations
from collections.abc import Awaitable, Callable
from datetime import datetime
from datetime import datetime, timezone
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
OAuthToken,
@ -22,9 +22,15 @@ CredentialReader = Callable[[str, str], Awaitable["dict[str, object] | None"]]
def _iso_to_epoch(expires_at: str) -> float | None:
try:
return datetime.fromisoformat(expires_at).timestamp()
dt = datetime.fromisoformat(expires_at)
except ValueError:
return None
# A timezone-naive expiry is stored as UTC (db.py writes ``datetime.now(timezone.utc)``),
# so anchor it to UTC before ``.timestamp()`` - otherwise a non-UTC host would read it as
# local time and skew the expiry, diverging from v1's ``_remaining_token_seconds``.
if dt.tzinfo is None:
dt = dt.replace(tzinfo=timezone.utc)
return dt.timestamp()
def _to_oauth_token(payload: dict[str, object]) -> OAuthToken | None:

14
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -4717,15 +4717,23 @@ class TestCreateMcpClientV2Graft:
client._resolved_auth._header_value.get_secret_value() == f"Basic {encoded}"
)
async def test_deferred_mode_uses_v1_auth_value(self):
async def test_m2m_client_credentials_defers_to_v1(self):
# M2M (oauth2 client_credentials) is not migrated: to_server_spec returns
# None, so the graft sets no resolved auth and leaves v1 in charge (v1
# performs the client_credentials grant itself - the static
# authentication_token is never consumed for oauth2, so it does not flow
# to _mcp_auth_value). Per-user oauth2 (authorization_code) is migrated to
# v2 and is exercised separately.
client = await MCPServerManager()._create_mcp_client(
self._http_server(
auth_type=MCPAuth.oauth2, authentication_token="legacy-token"
auth_type=MCPAuth.oauth2,
oauth2_flow="client_credentials",
authentication_token="legacy-token",
)
)
assert client._resolved_auth is None
assert client._mcp_auth_value == "legacy-token"
assert client._mcp_auth_value is None
async def test_static_token_missing_defers_to_v1(self):
client = await MCPServerManager()._create_mcp_client(

27
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py
View file

@ -628,6 +628,7 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401():
oauth_server = MagicMock()
oauth_server.auth_type = MCPAuth.oauth2
oauth_server.needs_user_oauth_token = True
oauth_server.delegate_auth_to_upstream = False
with (
patch(
@ -648,10 +649,10 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401():
return_value=False,
),
patch(
"litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token",
new_callable=AsyncMock,
return_value=None,
) as mock_get_stored_token,
return_value=False,
) as mock_has_token,
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
return_value=oauth_server,
@ -666,7 +667,7 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401():
await handle_streamable_http_mcp(scope, receive, send)
# Verify a 401 was raised
assert mock_get_stored_token.await_count == 1
assert mock_has_token.await_count == 1
assert mock_handle_request.await_count == 0
assert exc_info.value.status_code == 401
assert "www-authenticate" in exc_info.value.headers
@ -817,6 +818,7 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401():
oauth_server = MagicMock()
oauth_server.auth_type = MCPAuth.oauth2
oauth_server.needs_user_oauth_token = True
oauth_server.delegate_auth_to_upstream = False
with (
patch(
@ -837,10 +839,10 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401():
return_value=False,
),
patch(
"litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token",
new_callable=AsyncMock,
return_value={"Authorization": "Bearer cached-token"},
) as mock_get_stored_token,
return_value=True,
) as mock_has_token,
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
return_value=oauth_server,
@ -858,7 +860,7 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401():
):
await handle_streamable_http_mcp(scope, receive, send)
assert mock_get_stored_token.await_count == 1
assert mock_has_token.await_count == 1
assert mock_handle_request.await_count == 1
@ -941,10 +943,9 @@ async def test_handle_streamable_http_mcp_delegated_server_without_token_returns
return_value=False,
),
patch(
"litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token",
new_callable=AsyncMock,
return_value=None,
) as mock_get_stored_token,
) as mock_has_token,
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
return_value=delegated_server,
@ -958,7 +959,9 @@ async def test_handle_streamable_http_mcp_delegated_server_without_token_returns
with pytest.raises(HTTPException) as exc_info:
await handle_streamable_http_mcp(scope, receive, send)
assert mock_get_stored_token.await_count == 1
# Delegate-auth servers raise the resource_metadata challenge before any
# per-user existence check, so the v2 token store is never consulted.
assert mock_has_token.await_count == 0
assert mock_handle_request.await_count == 0
assert exc_info.value.status_code == 401
challenge = exc_info.value.headers["www-authenticate"]