From 48e3d911b58f68304e5da842aec29cdab0e5382b Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 8 Oct 2026 01:06:21 +0000 Subject: [PATCH] fix(mcp): honor the legacy KDF grace window on the Claude Code refresh grant refresh_proxy_credential opened presented session refresh tokens without legacy_keys, so the Claude Code /oauth/token refresh grant skipped the scrypt grace window that the DCR refresh, revoke and introspect paths honor. Also replaces the sk-1234 KAT literal in the outbound credential known-answer tests with sk-kat-master-key and the matching pinned hex. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/gateway_dcr_flow.py | 4 +- .../integration/mcp/test_mcp_kdf_rotation.py | 75 +++++++++++++++++++ .../test_bridge_credentials.py | 21 ++++-- .../mcp_server/test_gateway_dcr_flow.py | 73 +++++++++++++++++- 4 files changed, 163 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 6be3962af8e..2b22b45c6cf 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -1550,7 +1550,9 @@ async def refresh_proxy_credential( signing: Final = resolve_session_signing(master_key, "mcp_gateway refresh grant") if isinstance(signing, Response): return signing - opened: Final = _open_presented_refresh_token(refresh_token, client_id, signing) + opened: Final = _open_presented_refresh_token( + refresh_token, client_id, signing, legacy_keys=legacy_session_keys_from_master_key(master_key, signing.keys) + ) if isinstance(opened, Response): return opened if opened.principal.audience != PROXY_API_AUDIENCE: diff --git a/tests/integration/mcp/test_mcp_kdf_rotation.py b/tests/integration/mcp/test_mcp_kdf_rotation.py index 9fdce17d7d8..8609a1ab9a5 100644 --- a/tests/integration/mcp/test_mcp_kdf_rotation.py +++ b/tests/integration/mcp/test_mcp_kdf_rotation.py @@ -14,6 +14,7 @@ from typing import Final import httpx import pytest +import yaml from cryptography.hazmat.primitives.hashes import SHA256 from cryptography.hazmat.primitives.kdf.hkdf import HKDF from integration._support.client import Gateway @@ -70,6 +71,11 @@ class _IntrospectForm(TypedDict): token: ReadOnly[str] +class _ClaudeCodeTokenForm(TypedDict): + grant_type: ReadOnly[str] + refresh_token: ReadOnly[str] + + class _ClientRegistration(TypedDict): redirect_uris: ReadOnly[tuple[str, ...]] client_name: ReadOnly[str] @@ -139,6 +145,24 @@ def _refresh_envelope(identity_server: str, key: str, keys: EnvelopeKeys, upstre return sealed.token.get_secret_value() +def _proxy_api_session_refresh(user_id: str, keys: SessionKeys) -> str: + minted: Final = mint_session_refresh_token( + SessionPrincipal(user_id=user_id, client_id="claude_code", audience="proxy_api"), + keys, + datetime.now(timezone.utc), + ) + assert isinstance(minted, MintedSessionToken), minted + return minted.token.get_secret_value() + + +def _claude_code_gateway_config(path: Path) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["general_settings"]["enable_claude_code_gateway"] = True + config_path: Final = path / "claude-code-gateway.yaml" + config_path.write_text(yaml.safe_dump(config)) + return config_path + + def _session_refresh(user_id: str, client_id: str, keys: SessionKeys) -> str: minted: Final = mint_session_refresh_token( SessionPrincipal(user_id=user_id, client_id=client_id), keys, datetime.now(timezone.utc) @@ -492,3 +516,54 @@ def test_session_access_token_minted_under_legacy_scrypt_key_introspects_active_ ) assert response.status_code == 200, response.text assert response.json()["active"] is True, response.text + + +def test_claude_code_gateway_refresh_of_an_hkdf_session_token_is_renewed_without_grace( + gateway: Gateway, tmp_path: Path +) -> None: + config: Final = _claude_code_gateway_config(tmp_path) + with owned_proxy(gateway, tmp_path, {}, config=config) as graced, graced.scenario() as scenario: + refresh: Final = _proxy_api_session_refresh(scenario.user(), _hkdf_session_keys(graced.key)) + response: Final = graced.client.post( + "/claude_code_gateway/oauth/token", + data=_ClaudeCodeTokenForm(grant_type="refresh_token", refresh_token=refresh), + ) + assert response.status_code == 200, response.text + renewed: Final = response.json() + opened: Final = open_session_refresh_token( + renewed["refresh_token"], _hkdf_session_keys(graced.key), datetime.now(timezone.utc) + ) + assert isinstance(opened, OpenedSessionToken), opened + + +def test_claude_code_gateway_refresh_of_a_legacy_scrypt_session_token_is_rejected_without_grace( + gateway: Gateway, tmp_path: Path +) -> None: + config: Final = _claude_code_gateway_config(tmp_path) + with owned_proxy(gateway, tmp_path, {}, config=config) as graced, graced.scenario() as scenario: + refresh: Final = _proxy_api_session_refresh(scenario.user(), _scrypt_session_keys(graced.key)) + response: Final = graced.client.post( + "/claude_code_gateway/oauth/token", + data=_ClaudeCodeTokenForm(grant_type="refresh_token", refresh_token=refresh), + ) + assert response.status_code == 400, response.text + assert response.json()["error"] == "invalid_grant", response.text + + +def test_claude_code_gateway_refresh_of_a_legacy_scrypt_session_token_is_renewed_during_the_legacy_grace_window( + gateway: Gateway, tmp_path: Path +) -> None: + config: Final = _claude_code_gateway_config(tmp_path) + with owned_proxy(gateway, tmp_path, GRACE, config=config) as graced, graced.scenario() as scenario: + refresh: Final = _proxy_api_session_refresh(scenario.user(), _scrypt_session_keys(graced.key)) + response: Final = graced.client.post( + "/claude_code_gateway/oauth/token", + data=_ClaudeCodeTokenForm(grant_type="refresh_token", refresh_token=refresh), + ) + assert response.status_code == 200, response.text + renewed: Final = response.json() + assert renewed["refresh_token"] != refresh, response.text + opened: Final = open_session_refresh_token( + renewed["refresh_token"], _hkdf_session_keys(graced.key), datetime.now(timezone.utc) + ) + assert isinstance(opened, OpenedSessionToken), opened diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py index 673c418041b..63d73730950 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py @@ -356,19 +356,26 @@ def _legacy_keys(master_key: str) -> EnvelopeKeys: def test_envelope_key_derivation_matches_an_independent_rfc5869_vector(): - keys: Final = envelope_keys_from_master_key("sk-1234") - assert keys.signing_key.get_secret_value() == _rfc5869(b"sk-1234", b"litellm-mcp-bridge:envelope-signing:").hex() + keys: Final = envelope_keys_from_master_key("sk-kat-master-key") assert ( - keys.encryption_key.get_secret_value() == _rfc5869(b"sk-1234", b"litellm-mcp-bridge:envelope-encryption:").hex() + keys.signing_key.get_secret_value() + == _rfc5869(b"sk-kat-master-key", b"litellm-mcp-bridge:envelope-signing:").hex() + ) + assert ( + keys.encryption_key.get_secret_value() + == _rfc5869(b"sk-kat-master-key", b"litellm-mcp-bridge:envelope-encryption:").hex() ) # literal pin computed from the same hand implementation (RFC 5869, salt = 32 zero bytes) - assert keys.signing_key.get_secret_value() == "331032bfe2b8d86bd00586ea69c0d854140019ca0e0adbe4dd0cc9e06cb36605" + assert keys.signing_key.get_secret_value() == "20386b6579b080c3cc909c641fbf358c2db29869107289d56a823e9394064c3c" def test_session_key_derivation_matches_an_independent_rfc5869_vector(): - keys: Final = session_keys_from_master_key("sk-1234") - assert keys.signing_key.get_secret_value() == _rfc5869(b"sk-1234", b"litellm-mcp-gateway:session-signing:").hex() - assert keys.signing_key.get_secret_value() == "511f07b16ceeb5fba1660033785d4b9a4a6e80c749a51d639a688c287671a0e6" + keys: Final = session_keys_from_master_key("sk-kat-master-key") + assert ( + keys.signing_key.get_secret_value() + == _rfc5869(b"sk-kat-master-key", b"litellm-mcp-gateway:session-signing:").hex() + ) + assert keys.signing_key.get_secret_value() == "3a40b8858630f352a274de0e488b9cb05388346f7bcd4dc0da9cdec480652b44" def test_the_three_derived_keys_are_pairwise_distinct(): diff --git a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 76d071b1852..123b0a0d73c 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -10,6 +10,7 @@ from typing import Final from urllib.parse import parse_qs, urlparse import pytest +from pydantic import SecretStr from starlette.requests import Request from litellm.caching.caching import DualCache @@ -40,9 +41,11 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( native_client_auth_contract, native_client_authorize, open_gateway_dcr_client, + refresh_proxy_credential, register_aggregate_client, revoke_refresh_token, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import legacy_scrypt from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( SessionBearerAdmitted, SessionRefreshOpened, @@ -53,6 +56,8 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( SESSION_ISSUER, SESSION_REFRESH_PREFIX, + MintedSessionToken, + SessionKeys, SessionPrincipal, mint_session_refresh_token, mint_session_token, @@ -1914,7 +1919,9 @@ def _redis_that(async_increment, get=None): cache.redis_cache.async_increment = async_increment cache.redis_cache.check_and_fix_namespace = MagicMock(side_effect=lambda key: key) cache.redis_cache.init_async_client.return_value.get = get or AsyncMock(return_value=None) - cache.redis_cache.async_get_cache = AsyncMock(side_effect=AssertionError("peek must read the client, not the wrapper")) + cache.redis_cache.async_get_cache = AsyncMock( + side_effect=AssertionError("peek must read the client, not the wrapper") + ) cache.async_increment_cache = AsyncMock(side_effect=AssertionError("must not fall back to in-memory")) return cache @@ -2068,7 +2075,11 @@ async def test_refresh_of_a_rotated_token_answers_503_before_minting_while_redis (await _redeem_native(await _native_code(client_id, cache=issued), client_id, _Minter(), cache=issued)).body ) rotated = json.loads( - (await _refresh_native(payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(return_value=1)))).body + ( + await _refresh_native( + payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(return_value=1)) + ) + ).body )["refresh_token"] minter = _Minter() @@ -2578,3 +2589,61 @@ async def test_token_exchange_relays_a_mint_refusal(failure, status, error): response = await _exchange_native(client_id, _Minter(failure), _Exchanger()) assert response.status_code == status assert json.loads(response.body)["error"] == error + + +def _legacy_session_refresh(user_id="u1", client_id="claude_code", team_id="team-b"): + signing: Final = legacy_scrypt(MASTER_KEY, b"litellm-mcp-gateway:session-signing:") + assert signing is not None + minted: Final = mint_session_refresh_token( + SessionPrincipal(user_id=user_id, client_id=client_id, audience="proxy_api", team_id=team_id), + SessionKeys(signing_key=SecretStr(signing)), + datetime.now(timezone.utc), + ) + assert isinstance(minted, MintedSessionToken) + return minted.token.get_secret_value() + + +@pytest.mark.asyncio +async def test_claude_code_refresh_of_a_legacy_scrypt_token_is_renewed_during_the_grace_window(monkeypatch): + """The Claude Code ``/oauth/token`` refresh grant opens a scrypt-minted proxy-API + refresh token while ``LITELLM_MCP_LEGACY_KDF_GRACE=true``, exactly like the DCR + refresh grant, and the rotated token it returns opens under the HKDF key.""" + monkeypatch.setenv("LITELLM_MCP_LEGACY_KDF_GRACE", "true") + monkeypatch.delenv("LITELLM_FIPS_MODE", raising=False) + minter: Final = _Minter() + renewed: Final = await refresh_proxy_credential( + refresh_token=_legacy_session_refresh(), + client_id="claude_code", + master_key=MASTER_KEY, + cache=DualCache(), + mint_proxy_credential=minter, + ) + assert renewed.status_code == 200, renewed.body + body: Final = json.loads(renewed.body) + assert body["access_token"] == "sk-cli-u1" + assert body["user_id"] == "u1" + opened: Final = open_session_refresh_bearer( + body["refresh_token"], + session_keys_from_master_key(MASTER_KEY), + datetime.now(timezone.utc), + expected_client_id="claude_code", + ) + assert isinstance(opened, SessionRefreshOpened) + + +@pytest.mark.asyncio +async def test_claude_code_refresh_of_a_legacy_scrypt_token_is_rejected_without_grace(monkeypatch): + """Without the grace variable the same scrypt-minted token answers invalid_grant and + the minter never runs.""" + monkeypatch.delenv("LITELLM_MCP_LEGACY_KDF_GRACE", raising=False) + minter: Final = _Minter() + refused: Final = await refresh_proxy_credential( + refresh_token=_legacy_session_refresh(), + client_id="claude_code", + master_key=MASTER_KEY, + cache=DualCache(), + mint_proxy_credential=minter, + ) + assert refused.status_code == 400 + assert json.loads(refused.body)["error"] == "invalid_grant" + assert minter.calls == []