From c813b594b749d915a2822d0ef7c154e7abe2884c Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 26 Jun 2026 17:03:47 -0700 Subject: [PATCH] fix(mcp): green CI for authz_code dispatch (format + UTC expiry + v2-seam tests) - 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) --- .../per_user_oauth_store.py | 4 +-- .../outbound_credentials/v2_token_store.py | 10 +++++-- .../mcp_server/test_mcp_server_manager.py | 14 +++++++--- .../mcp_server/test_mcp_stale_session.py | 27 ++++++++++--------- 4 files changed, 35 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py index 2eccc836628..7ffde55c2c1 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py @@ -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: diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py index 458f45367de..ba2b8620254 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py @@ -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: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index e44a345609c..a1ba103edb9 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -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( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py index 67b72dbb695..e4e1890a45a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py @@ -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"]