refactor(mcp): single-assignment kdf fallback helpers

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py

@@ -68,6 +68,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent
     SessionSigningConfigError,
     active_session_signing_keys,
     legacy_session_keys_from_master_key,
+    open_session_credential_with_legacy,
     open_session_refresh_bearer,
 )
 from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
@@ -76,11 +77,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i
     MintedSessionToken,
     OpenedSessionToken,
     SessionAudience,
-    SessionExpired,
     SessionKeys,
     SessionPrincipal,
     SessionSigningKeys,
-    SessionTokenOpenError,
     is_session_refresh_token,
     is_session_token,
     mint_session_refresh_token,
@@ -1483,19 +1482,6 @@ async def revoke_refresh_token(token: str, client_id: str, master_key: str | Non
     return Response(content="{}", media_type="application/json", headers=TOKEN_NO_CACHE_HEADERS)
 
 
-def _open_with_legacy(
-    open_token: Callable[[str, SessionSigningKeys, datetime], OpenedSessionToken | SessionTokenOpenError],
-    token: str,
-    keys: SessionSigningKeys,
-    legacy_keys: SessionKeys | None,
-    now: datetime,
-) -> OpenedSessionToken | SessionTokenOpenError:
-    opened: Final = open_token(token, keys, now)
-    if isinstance(opened, OpenedSessionToken) or isinstance(opened, SessionExpired) or legacy_keys is None:
-        return opened
-    return open_token(token, legacy_keys, now)
-
-
 def _inactive_introspection_response() -> Response:
     """RFC 7662 section 2.2: any token the gateway cannot vouch for, whatever the reason
     (wrong family, bad signature, expired, revoked, or a deactivated user), answers 200
@@ -1556,9 +1542,9 @@ async def introspect_gateway_token(
     now: Final = datetime.now(timezone.utc)
     legacy_keys: Final = legacy_session_keys_from_master_key(master_key, keys)
     if is_session_token(token):
-        opened = _open_with_legacy(open_session_token, token, keys, legacy_keys, now)
+        opened = open_session_credential_with_legacy(open_session_token, token, keys, legacy_keys, now)
     elif is_session_refresh_token(token):
-        opened = _open_with_legacy(open_session_refresh_token, token, keys, legacy_keys, now)
+        opened = open_session_credential_with_legacy(open_session_refresh_token, token, keys, legacy_keys, now)
     else:
         return _inactive_introspection_response()
     if not isinstance(opened, OpenedSessionToken):

litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py

@@ -21,6 +21,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import
     EnvelopeIdentity,
     EnvelopeKeys,
     EnvelopeMintError,
+    EnvelopeOpenError,
     OpenedEnvelope,
     OpenedRefreshEnvelope,
     RefreshCredential,
@@ -131,6 +132,34 @@ class BridgeRefreshInvalid(BaseModel):
 BridgeRefreshResult: TypeAlias = BridgeRefreshOpened | BridgeRefreshInvalid
 
 
+def _open_envelope_with_fallback(
+    candidate: str,
+    keys: EnvelopeKeys,
+    legacy_keys: EnvelopeKeys | None,
+    now: datetime,
+) -> OpenedEnvelope | EnvelopeOpenError:
+    """Open an access envelope under the active keys, then under the legacy keys when
+    grace supplied them."""
+    primary: Final = open_envelope(candidate, keys, now)
+    if isinstance(primary, OpenedEnvelope) or legacy_keys is None:
+        return primary
+    return open_envelope(candidate, legacy_keys, now)
+
+
+def _open_refresh_envelope_with_fallback(
+    candidate: str,
+    keys: EnvelopeKeys,
+    legacy_keys: EnvelopeKeys | None,
+    now: datetime,
+) -> OpenedRefreshEnvelope | EnvelopeOpenError:
+    """Open a refresh envelope under the active keys, then under the legacy keys when
+    grace supplied them."""
+    primary: Final = open_refresh_envelope(candidate, keys, now)
+    if isinstance(primary, OpenedRefreshEnvelope) or legacy_keys is None:
+        return primary
+    return open_refresh_envelope(candidate, legacy_keys, now)
+
+
 def open_bridge_refresh_envelope(
     refresh_value: str,
     keys: EnvelopeKeys,
@@ -151,9 +180,7 @@ def open_bridge_refresh_envelope(
     candidate: Final = _strip_bearer(refresh_value)
     if not is_refresh_envelope(candidate):
         return BridgeRefreshInvalid()
-    opened = open_refresh_envelope(candidate, keys, now)
-    if not isinstance(opened, OpenedRefreshEnvelope) and legacy_keys is not None:
-        opened = open_refresh_envelope(candidate, legacy_keys, now)
+    opened: Final = _open_refresh_envelope_with_fallback(candidate, keys, legacy_keys, now)
     if not isinstance(opened, OpenedRefreshEnvelope):
         return BridgeRefreshInvalid()
     if opened.identity.server_id != expected_server_id:
@@ -237,9 +264,7 @@ def resolve_bridge_envelope(
         return BridgeEnvelopeInvalid()
     if not is_envelope(candidate):
         return NotBridgeEnvelope()
-    opened = open_envelope(candidate, keys, now)
-    if not isinstance(opened, OpenedEnvelope) and legacy_keys is not None:
-        opened = open_envelope(candidate, legacy_keys, now)
+    opened: Final = _open_envelope_with_fallback(candidate, keys, legacy_keys, now)
     if not isinstance(opened, OpenedEnvelope):
         return BridgeEnvelopeInvalid()
     if opened.identity.server_id != expected_server_id:

litellm/proxy/_experimental/mcp_server/outbound_credentials/key_derivation.py

@@ -19,12 +19,10 @@ from litellm.proxy.common_utils.fips import is_fips_mode
 LEGACY_KDF_GRACE_ENV_VAR: Final = "LITELLM_MCP_LEGACY_KDF_GRACE"
 DERIVED_KEY_BYTES: Final = 32
 
-# scrypt work factors (RFC 7914). n=2**15 with r=8/p=1 costs ~50ms and ~32MB per derivation.
+# RFC 7914 work factors of the pre-rotation derivation; maxmem is twice the working set.
 _SCRYPT_N: Final = 2**15
 _SCRYPT_R: Final = 8
 _SCRYPT_P: Final = 1
-# scrypt's working-set is ~128 * N * r * p bytes; cap at twice that so the maxmem ceiling scales
-# with every work factor and a future p or r bump does not trip "memory limit exceeded".
 _SCRYPT_MAXMEM: Final = 128 * _SCRYPT_N * _SCRYPT_R * _SCRYPT_P * 2
 
 

litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py

@@ -36,6 +36,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i
     SessionPrincipal,
     SessionRotatedPublicKey,
     SessionSigningKeys,
+    SessionTokenOpenError,
     is_session_refresh_token,
     is_session_token,
     open_session_refresh_token,
@@ -249,13 +250,7 @@ def resolve_session_bearer(
         return SessionBearerInvalid()
     if not is_session_token(candidate):
         return NotSessionBearer()
-    opened = open_session_token(candidate, keys, now)
-    if (
-        not isinstance(opened, OpenedSessionToken)
-        and not isinstance(opened, SessionExpired)
-        and legacy_keys is not None
-    ):
-        opened = open_session_token(candidate, legacy_keys, now)
+    opened: Final = open_session_credential_with_legacy(open_session_token, candidate, keys, legacy_keys, now)
     if isinstance(opened, OpenedSessionToken):
         return SessionBearerAdmitted(principal=opened.principal)
     return SessionBearerInvalid(expired=isinstance(opened, SessionExpired))
@@ -283,6 +278,23 @@ class SessionRefreshInvalid(BaseModel):
 SessionRefreshResult: TypeAlias = SessionRefreshOpened | SessionRefreshInvalid
 
 
+def open_session_credential_with_legacy(
+    open_token: Callable[[str, SessionSigningKeys, datetime], OpenedSessionToken | SessionTokenOpenError],
+    token: str,
+    keys: SessionSigningKeys,
+    legacy_keys: SessionKeys | None,
+    now: datetime,
+) -> OpenedSessionToken | SessionTokenOpenError:
+    """Open a session credential under the active keys, then under the legacy keys when
+    grace supplied them. An :class:`OpenedSessionToken` or :class:`SessionExpired` under
+    the active keys is final: expiry must still report as expiry, so only a signature or
+    format failure falls through to the legacy derivation."""
+    primary: Final = open_token(token, keys, now)
+    if isinstance(primary, (OpenedSessionToken, SessionExpired)) or legacy_keys is None:
+        return primary
+    return open_token(token, legacy_keys, now)
+
+
 def open_session_refresh_bearer(
     refresh_value: str,
     keys: SessionSigningKeys,
@@ -303,13 +315,7 @@ def open_session_refresh_bearer(
     candidate: Final = _strip_bearer(refresh_value)
     if not is_session_refresh_token(candidate):
         return SessionRefreshInvalid()
-    opened = open_session_refresh_token(candidate, keys, now)
-    if (
-        not isinstance(opened, OpenedSessionToken)
-        and not isinstance(opened, SessionExpired)
-        and legacy_keys is not None
-    ):
-        opened = open_session_refresh_token(candidate, legacy_keys, now)
+    opened: Final = open_session_credential_with_legacy(open_session_refresh_token, candidate, keys, legacy_keys, now)
     if not isinstance(opened, OpenedSessionToken):
         return SessionRefreshInvalid()
     if opened.principal.client_id != expected_client_id:

tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py

@@ -10,7 +10,10 @@ through the consumer; and no path leaks the upstream token in a repr.
 
 import hashlib
 import hmac
+from collections.abc import Callable, Mapping
 from datetime import datetime, timedelta, timezone
+from types import MappingProxyType
+from typing import Final
 
 import pytest
 from pydantic import SecretStr
@@ -29,13 +32,6 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credenti
     open_bridge_refresh_envelope,
     resolve_bridge_envelope,
 )
-from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import (
-    LEGACY_KDF_GRACE_ENV_VAR,
-)
-from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import (
-    session_keys_from_master_key,
-)
-from litellm.proxy.common_utils.fips import FIPS_MODE_ENV_VAR
 from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import (
     ENVELOPE_PREFIX,
     EnvelopeIdentity,
@@ -47,6 +43,13 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import
     key_hash_identity,
     mint_envelope,
 )
+from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import (
+    LEGACY_KDF_GRACE_ENV_VAR,
+)
+from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import (
+    session_keys_from_master_key,
+)
+from litellm.proxy.common_utils.fips import FIPS_MODE_ENV_VAR
 
 _NOW = datetime(2026, 7, 9, 12, 0, 0, tzinfo=timezone.utc)
 _MASTER_KEY = "sk-master-key-for-derivation-tests-0123456789"
@@ -332,11 +335,11 @@ def test_is_bridge_envelope_shaped_rejects_non_envelope_bearer():
 def _rfc5869(ikm: bytes, info: bytes) -> bytes:
     """Independent RFC 5869 HKDF-SHA256, hand-written on hmac/hashlib so a wrong KDF
     construction in the product cannot agree with it."""
-    prk = hmac.new(b"\x00" * 32, ikm, hashlib.sha256).digest()
+    prk: Final = hmac.new(b"\x00" * 32, ikm, hashlib.sha256).digest()
     return hmac.new(prk, info + b"\x01", hashlib.sha256).digest()[:32]
 
 
-def _environ(values: dict[str, str]):
+def _environ(values: Mapping[str, str]) -> Callable[[str], str | None]:
     return values.get
 
 
@@ -353,7 +356,7 @@ def _legacy_keys(master_key: str) -> EnvelopeKeys:
 
 
 def test_envelope_key_derivation_matches_an_independent_rfc5869_vector():
-    keys = envelope_keys_from_master_key("sk-1234")
+    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()
     assert (
         keys.encryption_key.get_secret_value() == _rfc5869(b"sk-1234", b"litellm-mcp-bridge:envelope-encryption:").hex()
@@ -363,34 +366,38 @@ def test_envelope_key_derivation_matches_an_independent_rfc5869_vector():
 
 
 def test_session_key_derivation_matches_an_independent_rfc5869_vector():
-    keys = session_keys_from_master_key("sk-1234")
+    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"
 
 
 def test_the_three_derived_keys_are_pairwise_distinct():
-    envelope = envelope_keys_from_master_key(_MASTER_KEY)
-    session = session_keys_from_master_key(_MASTER_KEY)
-    derived = {
-        envelope.signing_key.get_secret_value(),
-        envelope.encryption_key.get_secret_value(),
-        session.signing_key.get_secret_value(),
-    }
+    envelope: Final = envelope_keys_from_master_key(_MASTER_KEY)
+    session: Final = session_keys_from_master_key(_MASTER_KEY)
+    derived: Final = frozenset(
+        (
+            envelope.signing_key.get_secret_value(),
+            envelope.encryption_key.get_secret_value(),
+            session.signing_key.get_secret_value(),
+        )
+    )
     assert len(derived) == 3
 
 
 def test_new_envelopes_open_under_the_hkdf_keys_with_no_legacy_fallback():
-    keys = envelope_keys_from_master_key(_MASTER_KEY)
-    result = resolve_bridge_envelope(_sealed_token(keys), keys, _NOW, _SERVER_ID)
+    keys: Final = envelope_keys_from_master_key(_MASTER_KEY)
+    result: Final = resolve_bridge_envelope(_sealed_token(keys), keys, _NOW, _SERVER_ID)
     assert isinstance(result, BridgeEnvelopeAdmitted)
 
 
 def test_legacy_envelope_opens_during_the_grace_window():
-    legacy = _legacy_keys(_MASTER_KEY)
-    token = _sealed_token(legacy)
-    grace = legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"}))
+    legacy: Final = _legacy_keys(_MASTER_KEY)
+    token: Final = _sealed_token(legacy)
+    grace: Final = legacy_envelope_keys_from_master_key(
+        _MASTER_KEY, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"}))
+    )
     assert grace == legacy
-    result = resolve_bridge_envelope(
+    result: Final = resolve_bridge_envelope(
         token, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID, legacy_keys=grace
     )
     assert isinstance(result, BridgeEnvelopeAdmitted)
@@ -398,32 +405,36 @@ def test_legacy_envelope_opens_during_the_grace_window():
 
 
 def test_legacy_envelope_is_invalid_without_grace():
-    token = _sealed_token(_legacy_keys(_MASTER_KEY))
-    assert legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ({})) is None
-    result = resolve_bridge_envelope(token, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID)
+    token: Final = _sealed_token(_legacy_keys(_MASTER_KEY))
+    assert legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ(MappingProxyType({}))) is None
+    result: Final = resolve_bridge_envelope(token, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID)
     assert isinstance(result, BridgeEnvelopeInvalid)
 
 
 def test_grace_is_disabled_under_fips():
-    grace = legacy_envelope_keys_from_master_key(
+    grace: Final = legacy_envelope_keys_from_master_key(
         _MASTER_KEY,
-        environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true", FIPS_MODE_ENV_VAR: "true"}),
+        environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true", FIPS_MODE_ENV_VAR: "true"})),
     )
     assert grace is None
 
 
-@pytest.mark.parametrize("value", ["yes", "1", "on", "TRUE ", "True"])
+@pytest.mark.parametrize("value", ("yes", "1", "on", "TRUE ", "True"))
 def test_grace_is_enabled_only_by_case_insensitive_true(value: str):
-    result = legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: value}))
-    expected = value.strip().lower() == "true"
+    result: Final = legacy_envelope_keys_from_master_key(
+        _MASTER_KEY, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: value}))
+    )
+    expected: Final = value.strip().lower() == "true"
     assert (result is not None) is expected
 
 
 def test_legacy_refresh_envelope_opens_during_grace():
-    legacy = _legacy_keys(_MASTER_KEY)
-    refresh = _sealed_refresh(legacy)
-    grace = legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"}))
-    result = open_bridge_refresh_envelope(
+    legacy: Final = _legacy_keys(_MASTER_KEY)
+    refresh: Final = _sealed_refresh(legacy)
+    grace: Final = legacy_envelope_keys_from_master_key(
+        _MASTER_KEY, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"}))
+    )
+    result: Final = open_bridge_refresh_envelope(
         refresh, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID, legacy_keys=grace
     )
     assert isinstance(result, BridgeRefreshOpened)
@@ -431,6 +442,6 @@ def test_legacy_refresh_envelope_opens_during_grace():
 
 
 def test_legacy_refresh_envelope_is_invalid_without_grace():
-    refresh = _sealed_refresh(_legacy_keys(_MASTER_KEY))
-    result = open_bridge_refresh_envelope(refresh, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID)
+    refresh: Final = _sealed_refresh(_legacy_keys(_MASTER_KEY))
+    result: Final = open_bridge_refresh_envelope(refresh, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID)
     assert isinstance(result, BridgeRefreshInvalid)

tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py

@@ -1,7 +1,10 @@
 """Tests for the session-token KDF and the edge/token-endpoint resolvers."""
 
 import hashlib
+from collections.abc import Callable, Mapping
 from datetime import datetime, timedelta, timezone
+from types import MappingProxyType
+from typing import Final
 
 import pytest
 from cryptography.hazmat.primitives import serialization
@@ -12,7 +15,9 @@ from pydantic import SecretStr
 from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import (
     envelope_keys_from_master_key,
 )
-from litellm.proxy._experimental.mcp_server.outbound_credentials import session_credentials
+from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import (
+    LEGACY_KDF_GRACE_ENV_VAR,
+)
 from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import (
     NotSessionBearer,
     SessionBearerAdmitted,
@@ -27,9 +32,6 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent
     resolve_session_signing_keys,
     session_keys_from_master_key,
 )
-from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import (
-    LEGACY_KDF_GRACE_ENV_VAR,
-)
 from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
     SESSION_TTL_SECONDS,
     AsymmetricSessionKeys,
@@ -239,12 +241,12 @@ def test_signing_config_error_detail_never_leaks_key_material():
     assert pem.splitlines()[1] not in resolved.detail
 
 
-def _environ(values: dict[str, str]):
+def _environ(values: Mapping[str, str]) -> Callable[[str], str | None]:
     return values.get
 
 
 def _legacy_session_keys(master_key: str) -> SessionKeys:
-    signing = hashlib.scrypt(
+    signing: Final = hashlib.scrypt(
         master_key.encode(),
         salt=b"litellm-mcp-gateway:session-signing:",
         n=2**15,
@@ -257,53 +259,57 @@ def _legacy_session_keys(master_key: str) -> SessionKeys:
 
 
 def test_new_session_tokens_open_under_the_hkdf_key_with_no_legacy_fallback():
-    result = resolve_session_bearer(f"Bearer {_access_token()}", KEYS, NOW)
+    result: Final = resolve_session_bearer(f"Bearer {_access_token()}", KEYS, NOW)
     assert isinstance(result, SessionBearerAdmitted)
 
 
 def test_legacy_session_token_opens_during_the_grace_window():
-    legacy = _legacy_session_keys(MASTER_KEY)
-    minted = mint_session_token(PRINCIPAL, legacy, NOW)
+    legacy: Final = _legacy_session_keys(MASTER_KEY)
+    minted: Final = mint_session_token(PRINCIPAL, legacy, NOW)
     assert isinstance(minted, MintedSessionToken)
-    grace = legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"}))
+    grace: Final = legacy_session_keys_from_master_key(
+        MASTER_KEY, KEYS, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"}))
+    )
     assert grace == legacy
-    result = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW, legacy_keys=grace)
+    result: Final = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW, legacy_keys=grace)
     assert isinstance(result, SessionBearerAdmitted)
     assert result.principal == PRINCIPAL
 
 
 def test_legacy_session_token_is_invalid_without_grace():
-    minted = mint_session_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW)
+    minted: Final = mint_session_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW)
     assert isinstance(minted, MintedSessionToken)
-    assert legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({})) is None
-    result = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW)
+    assert legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ(MappingProxyType({}))) is None
+    result: Final = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW)
     assert isinstance(result, SessionBearerInvalid)
     assert result.expired is False
 
 
 def test_legacy_session_key_is_never_derived_for_configured_rs256_keys():
-    active = AsymmetricSessionKeys(private_key_pem=SecretStr(_rsa_private_pem()), kid="k")
-    grace = legacy_session_keys_from_master_key(
-        MASTER_KEY, active, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})
+    active: Final = AsymmetricSessionKeys(private_key_pem=SecretStr(_rsa_private_pem()), kid="k")
+    grace: Final = legacy_session_keys_from_master_key(
+        MASTER_KEY, active, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"}))
     )
     assert grace is None
 
 
 def test_grace_is_disabled_under_fips():
-    grace = legacy_session_keys_from_master_key(
+    grace: Final = legacy_session_keys_from_master_key(
         MASTER_KEY,
         KEYS,
-        environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true", FIPS_MODE_ENV_VAR: "true"}),
+        environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true", FIPS_MODE_ENV_VAR: "true"})),
     )
     assert grace is None
 
 
 def test_legacy_session_refresh_token_opens_during_grace():
-    legacy = _legacy_session_keys(MASTER_KEY)
-    minted = mint_session_refresh_token(PRINCIPAL, legacy, NOW)
+    legacy: Final = _legacy_session_keys(MASTER_KEY)
+    minted: Final = mint_session_refresh_token(PRINCIPAL, legacy, NOW)
     assert isinstance(minted, MintedSessionToken)
-    grace = legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"}))
-    result = open_session_refresh_bearer(
+    grace: Final = legacy_session_keys_from_master_key(
+        MASTER_KEY, KEYS, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"}))
+    )
+    result: Final = open_session_refresh_bearer(
         minted.token.get_secret_value(),
         KEYS,
         NOW,
@@ -315,37 +321,22 @@ def test_legacy_session_refresh_token_opens_during_grace():
 
 
 def test_legacy_session_refresh_token_is_invalid_without_grace():
-    minted = mint_session_refresh_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW)
+    minted: Final = mint_session_refresh_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW)
     assert isinstance(minted, MintedSessionToken)
-    result = open_session_refresh_bearer(
+    result: Final = open_session_refresh_bearer(
         minted.token.get_secret_value(), KEYS, NOW, expected_client_id="llm_client_abc"
     )
     assert isinstance(result, SessionRefreshInvalid)
 
 
 def test_expired_token_under_the_primary_key_still_reports_expired_during_grace():
-    legacy = _legacy_session_keys(MASTER_KEY)
-    minted = mint_session_token(PRINCIPAL, legacy, NOW)
+    legacy: Final = _legacy_session_keys(MASTER_KEY)
+    minted: Final = mint_session_token(PRINCIPAL, legacy, NOW)
     assert isinstance(minted, MintedSessionToken)
-    grace = legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"}))
-    later = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1)
-    result = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, later, legacy_keys=grace)
+    grace: Final = legacy_session_keys_from_master_key(
+        MASTER_KEY, KEYS, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"}))
+    )
+    later: Final = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1)
+    result: Final = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, later, legacy_keys=grace)
     assert isinstance(result, SessionBearerInvalid)
     assert result.expired is True
-
-
-def test_fallback_tries_the_active_key_before_legacy(monkeypatch):
-    calls: list[object] = []
-    real = session_credentials.open_session_token
-
-    def spy(candidate, keys, now):
-        calls.append(keys)
-        return real(candidate, keys, now)
-
-    monkeypatch.setattr(session_credentials, "open_session_token", spy)
-    grace = legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"}))
-    minted = mint_session_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW)
-    assert isinstance(minted, MintedSessionToken)
-    result = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW, legacy_keys=grace)
-    assert isinstance(result, SessionBearerAdmitted)
-    assert calls[0] is KEYS