diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index a5a8275a6ad..adb1812059b 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -252,7 +252,7 @@ class MCPRequestHandler: # Authorization: open the envelope, admit under its recovered identity, and # inject the inner upstream token for egress. A non-envelope bearer on the same # server is NOT admitted here — it falls through to the oauth2 arm, which 401s. - validated_user_api_key_auth, mcp_server_auth_headers = MCPRequestHandler._admit_dcr_bridge_delegate( + validated_user_api_key_auth, mcp_server_auth_headers = await MCPRequestHandler._admit_dcr_bridge_delegate( server=bridge_delegate_target, authorization_value=oauth2_headers["Authorization"], mcp_server_auth_headers=mcp_server_auth_headers, @@ -492,20 +492,23 @@ class MCPRequestHandler: return server @staticmethod - def _admit_dcr_bridge_delegate( + async def _admit_dcr_bridge_delegate( server: MCPServer, authorization_value: str, mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]], ) -> Tuple[UserAPIKeyAuth, Optional[Dict[str, Dict[str, str]]]]: - """Open the bridge envelope and admit the caller under its recovered identity. + """Open the bridge envelope and admit the caller under the live key it references. - The envelope's signature is itself the proof the user authenticated when it was - minted, so the recovered ``user_id`` is admitted without any re-validation. The - inner upstream token is injected under the server's per-server auth-header key so - egress forwards it via the ``PassthroughConfig`` override; the envelope - ``Authorization`` the leak-defense strips never reaches the upstream. A new headers - dict is returned rather than mutating the input. Fails closed with a 401 on an - invalid or expired envelope. + The envelope's signature proves the user authenticated when it was minted, but + authorization is resolved fresh here rather than trusted from the envelope: the + sealed ``key_hash`` reloads the current ``UserAPIKeyAuth`` record, so the key's + present team/org/object-permission restrictions and its revocation state gate the + request instead of a snapshot frozen at mint time. The inner upstream token is + injected under the server's per-server auth-header key so egress forwards it via the + ``PassthroughConfig`` override; the envelope ``Authorization`` the leak-defense + strips never reaches the upstream. A new headers dict is returned rather than + mutating the input. Fails closed with a 401 on an invalid or expired envelope, or + when the referenced key is missing, blocked, or expired. """ from litellm.proxy.proxy_server import master_key @@ -519,14 +522,60 @@ class MCPRequestHandler: header_key = server.server_name or server.alias if header_key is None: raise HTTPException(status_code=500, detail="Server misconfigured: MCP server has no routable name") + admitted = await MCPRequestHandler._reload_admitted_key(result.identity.key_hash) injected = {header_key: {"Authorization": result.upstream_authorization.get_secret_value()}} new_headers = {**(mcp_server_auth_headers or {}), **injected} - return UserAPIKeyAuth(user_id=result.identity.user_id), new_headers + return admitted, new_headers case BridgeEnvelopeInvalid() | NotBridgeEnvelope(): raise HTTPException(status_code=401, detail="Invalid or expired credential") case _: assert_never(result) + @staticmethod + async def _reload_admitted_key(key_hash: str) -> UserAPIKeyAuth: + """Reload the live key record an admitted envelope references. + + Resolving the current ``UserAPIKeyAuth`` (cache first, then DB) is what stops the + envelope from carrying frozen authority: the key's present team/org/object-permission + restrictions ride on the returned object, and a key that has since been deleted, + blocked, or expired fails closed with a 401 here rather than being admitted as an + unrestricted identity. ``get_key_object`` raises for a hash with no key row; a + blocked or expired row is rejected explicitly because ``get_key_object`` resolves a + row without applying those checks (the main ``user_api_key_auth`` pipeline enforces + them downstream, which this admission path bypasses). + """ + from litellm.proxy.auth.auth_checks import get_key_object + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Server misconfigured: no database connection") + try: + key_object = await get_key_object( + hashed_token=key_hash, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + except (ProxyException, HTTPException): + raise HTTPException(status_code=401, detail="Invalid or expired credential") from None + if not MCPRequestHandler._admitted_key_is_active(key_object): + raise HTTPException(status_code=401, detail="Invalid or expired credential") + return key_object + + @staticmethod + def _admitted_key_is_active(key_object: UserAPIKeyAuth) -> bool: + """False when the referenced key is blocked or past its expiry, so a revoked key + cannot be admitted through its still-unexpired envelope. Mirrors the active-key gate + the bridge token endpoint applies at mint time.""" + if key_object.blocked is True: + return False + expires = key_object.expires + if expires is None: + return True + expiry = expires if isinstance(expires, datetime) else datetime.fromisoformat(expires) + if expiry.tzinfo is None or expiry.tzinfo.utcoffset(expiry) is None: + expiry = expiry.replace(tzinfo=timezone.utc) + return expiry >= datetime.now(timezone.utc) + @staticmethod def _resolve_target_server_names(path: str, mcp_servers_header: Optional[List[str]]) -> List[str]: """ diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py index 298bc8d98cc..517c2ef5c8f 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py @@ -10,7 +10,7 @@ as explicit parameters. Wire shape: ``llm_env_`` + an HS256 JWT (same signing approach as the BYOK session bearer in ``byok_oauth_endpoints.py``). Registered claims are ``iss``/``iat``/``exp``; -custom claims are ``user_id``, ``server_id``, and ``grant``, where ``grant`` is the +custom claims are ``server_id``, ``key_hash``, and ``grant``, where ``grant`` is the upstream token grant serialized to JSON, encrypted with the repo's symmetric encryption helpers (``encrypt_value``/``decrypt_value`` from ``encrypt_decrypt_utils`` — the same family ``encrypt_value_helper`` applies to @@ -68,11 +68,19 @@ _ENVELOPE_JWT_ALGORITHM = "HS256" class EnvelopeIdentity(BaseModel): - """The litellm identity the envelope binds the inner grant to.""" + """The litellm identity the envelope binds the inner grant to. + + ``key_hash`` is the hashed litellm key that authorized the mint, never a raw + credential (and the edge rejects a bare hash presented as a bearer). Admission + reloads the live key record by it, so the key's current team/org/object-permission + restrictions and its revocation state are enforced at use time rather than frozen at + mint time. ``server_id`` binds the envelope to one MCP server so it cannot be replayed + across a server boundary. + """ model_config = ConfigDict(frozen=True) - user_id: str = Field(min_length=1) server_id: str = Field(min_length=1) + key_hash: str = Field(min_length=1) class UpstreamTokenGrant(BaseModel): @@ -174,7 +182,7 @@ EnvelopeOpenError: TypeAlias = NotAnEnvelope | BadSignature | Expired | Malforme class _EnvelopeClaims(BaseModel): """Decoded-claims boundary that pins the exact shape :func:`mint_envelope` emits. - ``user_id``/``server_id`` mirror the ``min_length`` constraints of + ``server_id``/``key_hash`` mirror the ``min_length`` constraints of :class:`EnvelopeIdentity` so any claim set that validates here also constructs an identity, keeping :func:`open_envelope` raise-free: a correctly signed JWT with an empty identity claim fails here and maps to ``MalformedPayload``. @@ -191,8 +199,8 @@ class _EnvelopeClaims(BaseModel): iss: str iat: int exp: int - user_id: str = Field(min_length=1) server_id: str = Field(min_length=1) + key_hash: str = Field(min_length=1) grant: str = Field(min_length=1) @@ -227,8 +235,8 @@ def mint_envelope( iss=ENVELOPE_ISSUER, iat=int(now.timestamp()), exp=int(expires_at.timestamp()), - user_id=identity.user_id, server_id=identity.server_id, + key_hash=identity.key_hash, grant=_encrypt_grant_blob(_grant_plaintext(grant), keys.encryption_key), ) token = ENVELOPE_PREFIX + jwt.encode( @@ -273,7 +281,7 @@ def open_envelope( if not isinstance(grant, UpstreamTokenGrant): return grant return OpenedEnvelope( - identity=EnvelopeIdentity(user_id=claims.user_id, server_id=claims.server_id), + identity=EnvelopeIdentity(server_id=claims.server_id, key_hash=claims.key_hash), grant=grant, ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 1002c33783a..9873f3ac7c5 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -1,3 +1,4 @@ +import contextlib import json import os import sys @@ -16,6 +17,8 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + ProxyException, SpecialHeaders, SpecialMCPServerNames, UserAPIKeyAuth, @@ -4875,10 +4878,12 @@ class TestMCPDcrBridgeDelegateAdmission: a single envelope bearer (LIT-4338). The arm fires only for a single ``is_dcr_bridge`` ``is_oauth_delegate`` target carrying an - envelope-shaped Authorization. It opens the litellm-signed envelope, admits under the - recovered identity WITHOUT re-validating (the signature is the proof), and injects the inner - upstream token under the server's per-server auth-header key so egress forwards it. Everything - else must stay on its existing admission path. + envelope-shaped Authorization. It opens the litellm-signed envelope, reloads the live key + record the sealed ``key_hash`` references so the caller is admitted under the key's current + authorization context (team/org/object-permission) and revocation state, and injects the inner + upstream token under the server's per-server auth-header key so egress forwards it. A key that + is missing, blocked, or expired fails closed with a 401. Everything else must stay on its + existing admission path. """ _MASTER_KEY = "sk-bridge-master-key-for-envelope-derivation" @@ -4898,11 +4903,13 @@ class TestMCPDcrBridgeDelegateAdmission: dcr_bridge=dcr_bridge, ) + _KEY_HASH = "hashed-litellm-key-abc123" + @classmethod def _mint_bridge_envelope( cls, *, - user_id="envelope-user-42", + key_hash=None, server_id="bridge-server-id", access_token="inner-upstream-access-token", token_type="Bearer", @@ -4924,7 +4931,7 @@ class TestMCPDcrBridgeDelegateAdmission: keys = envelope_keys_from_master_key(master_key or cls._MASTER_KEY) now = minted_at or datetime.now(timezone.utc) sealed = mint_envelope( - identity=EnvelopeIdentity(user_id=user_id, server_id=server_id), + identity=EnvelopeIdentity(server_id=server_id, key_hash=key_hash or cls._KEY_HASH), grant=UpstreamTokenGrant( access_token=SecretStr(access_token), token_type=token_type, @@ -4936,18 +4943,51 @@ class TestMCPDcrBridgeDelegateAdmission: assert isinstance(sealed, SealedEnvelope), sealed return sealed.token.get_secret_value() - async def test_valid_envelope_admits_identity_and_injects_inner_token(self): - """A valid envelope on a single dcr_bridge oauth_delegate server admits under the envelope's - identity WITHOUT re-validating (user_api_key_auth is never called) and injects the inner - upstream token, keyed by the server name, for egress forwarding.""" - envelope = self._mint_bridge_envelope(user_id="envelope-user-42") + @staticmethod + def _reloaded_key(**overrides): + """A live key record as ``get_key_object`` would return it: carries real authorization + context (key identity, team, org, and an object-permission restricting MCP servers) so a + test can prove admission admits under THAT context rather than a blank identity.""" + defaults = dict( + user_id="envelope-user-42", + api_key=TestMCPDcrBridgeDelegateAdmission._KEY_HASH, + team_id="team-restricted", + org_id="org-restricted", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-1", mcp_servers=["only-this-server"] + ), + ) + defaults.update(overrides) + return UserAPIKeyAuth(**defaults) + + @staticmethod + @contextlib.contextmanager + def _patch_key_reload(*, return_value=None, side_effect=None): + """Patch the live-key reload dependencies used by ``_reload_admitted_key``: the + ``get_key_object`` lookup plus the ``prisma_client`` / ``user_api_key_cache`` globals it + reads. Yields the ``get_key_object`` mock so callers can assert the sealed ``key_hash`` was + the reload key.""" + get_key_object = AsyncMock(return_value=return_value, side_effect=side_effect) + with ( + patch("litellm.proxy.auth.auth_checks.get_key_object", get_key_object), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + ): + yield get_key_object + + async def test_valid_envelope_reloads_live_key_and_admits_its_authorization_context(self): + """A valid envelope admits under the LIVE key record the sealed key_hash references, not a + blank identity: the reload is keyed by that exact hash, and the admitted auth carries the + key's current team/org/object-permission (the MCP tool/server restrictions the finding was + about). The heavyweight ``user_api_key_auth`` pipeline is still never invoked. The inner + upstream token is injected under the per-server key for egress.""" + envelope = self._mint_bridge_envelope(key_hash=self._KEY_HASH) scope = { "type": "http", "method": "POST", "path": "/mcp/bridge_delegate_server", "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], } - with ( patch( "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", @@ -4955,6 +4995,7 @@ class TestMCPDcrBridgeDelegateAdmission: ) as mock_auth, patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_key_reload(return_value=self._reloaded_key()) as get_key_object, ): mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() ( @@ -4966,14 +5007,101 @@ class TestMCPDcrBridgeDelegateAdmission: _raw_headers, ) = await MCPRequestHandler.process_mcp_request(scope) - # Signature is the proof of prior authentication: identity admitted, no re-validation. + # The live key was reloaded by the exact hash the envelope sealed. + assert get_key_object.await_args.kwargs["hashed_token"] == self._KEY_HASH + # Admission carries the reloaded key's authorization context, not a blank UserAPIKeyAuth. assert auth_result.user_id == "envelope-user-42" + assert auth_result.team_id == "team-restricted" + assert auth_result.org_id == "org-restricted" + assert auth_result.object_permission is not None + assert auth_result.object_permission.mcp_servers == ["only-this-server"] + # The full raw-key auth pipeline is still bypassed for the envelope arm. mock_auth.assert_not_called() # Inner upstream token injected under the per-server key so egress forwards it. assert mcp_server_auth_headers == { "bridge_delegate_server": {"Authorization": "Bearer inner-upstream-access-token"} } + async def test_revoked_key_envelope_fails_closed_401(self): + """An envelope whose key has since been deleted must fail closed: ``get_key_object`` raises + for the missing row, so admission 401s instead of admitting the caller as an unrestricted + identity. This is the core regression for the dropped-authorization-context finding.""" + envelope = self._mint_bridge_envelope(key_hash=self._KEY_HASH) + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + revoked = self._patch_key_reload( + side_effect=ProxyException( + message="Authentication Error, Invalid proxy server token passed.", + type="token_not_found_in_db", + param="key", + code=401, + ) + ) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + ) as mock_auth, + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + revoked, + ): + mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + mock_auth.assert_not_called() + + async def test_blocked_key_envelope_fails_closed_401(self): + """A reloaded key that is blocked must fail closed with a 401, so revoking a key by blocking + it takes effect immediately for any envelope still holding its hash.""" + envelope = self._mint_bridge_envelope(key_hash=self._KEY_HASH) + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + with ( + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_key_reload(return_value=self._reloaded_key(blocked=True)), + ): + mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + + async def test_expired_key_record_fails_closed_401(self): + """A reloaded key past its expiry must fail closed with a 401, distinct from an expired + envelope: even a still-valid envelope cannot outlive the key it was minted under.""" + envelope = self._mint_bridge_envelope(key_hash=self._KEY_HASH) + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + expired_at = datetime.now(timezone.utc) - timedelta(hours=1) + + with ( + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_key_reload(return_value=self._reloaded_key(expires=expired_at)), + ): + mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + async def test_alias_only_server_injects_under_alias_egress_can_resolve(self): """When server_name is None, the inner token must be keyed under the alias (which egress resolves), never under server.name (which egress never looks up), so the forwarded token is @@ -4985,7 +5113,6 @@ class TestMCPDcrBridgeDelegateAdmission: "path": "/mcp/bridge_delegate_server", "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], } - with ( patch( "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", @@ -4993,6 +5120,7 @@ class TestMCPDcrBridgeDelegateAdmission: ), patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_key_reload(return_value=self._reloaded_key()), ): mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server( server_name=None, alias="bridge_alias" @@ -5260,11 +5388,14 @@ class TestMCPDcrBridgeDelegateAdmission: """Unit: ``_admit_dcr_bridge_delegate`` must return a NEW headers dict that preserves the caller's existing per-server entries and adds the injected inner token, never mutating the input dict.""" - envelope = self._mint_bridge_envelope(user_id="unit-user") + envelope = self._mint_bridge_envelope() existing = {"other_server": {"Authorization": "Bearer someone-elses-token"}} - with patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY): - auth_result, new_headers = MCPRequestHandler._admit_dcr_bridge_delegate( + with ( + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + self._patch_key_reload(return_value=self._reloaded_key(user_id="unit-user")), + ): + auth_result, new_headers = await MCPRequestHandler._admit_dcr_bridge_delegate( server=self._bridge_delegate_server(), authorization_value=f"Bearer {envelope}", mcp_server_auth_headers=existing, @@ -5286,7 +5417,23 @@ class TestMCPDcrBridgeDelegateAdmission: envelope = self._mint_bridge_envelope() with patch("litellm.proxy.proxy_server.master_key", None): with pytest.raises(HTTPException) as exc_info: - MCPRequestHandler._admit_dcr_bridge_delegate( + await MCPRequestHandler._admit_dcr_bridge_delegate( + server=self._bridge_delegate_server(), + authorization_value=f"Bearer {envelope}", + mcp_server_auth_headers=None, + ) + assert exc_info.value.status_code == 500 + + async def test_admit_helper_raises_500_when_no_db_connection(self): + """Unit: with a valid envelope but no database to reload the key from, admission raises a 500 + rather than admitting on unresolved authorization.""" + envelope = self._mint_bridge_envelope() + with ( + patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), + patch("litellm.proxy.proxy_server.prisma_client", None), + ): + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler._admit_dcr_bridge_delegate( server=self._bridge_delegate_server(), authorization_value=f"Bearer {envelope}", mcp_server_auth_headers=None, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py index dca9dc19314..82e8e2aae89 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py @@ -34,7 +34,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import _NOW = datetime(2026, 7, 9, 12, 0, 0, tzinfo=timezone.utc) _MASTER_KEY = "sk-master-key-for-derivation-tests-0123456789" _ACCESS_TOKEN = "upstream-access-token-do-not-leak-8f14e45fceea" -_IDENTITY = EnvelopeIdentity(user_id="user-123", server_id="srv-456") +_IDENTITY = EnvelopeIdentity(server_id="srv-456", key_hash="hashed-key-123") _SERVER_ID = _IDENTITY.server_id @@ -138,7 +138,7 @@ def test_resolve_envelope_minted_for_another_server_is_invalid(): captured or misrouted envelope cannot forward one server's upstream credential to another. The valid access token stays sealed; the mismatch alone fails the resolve.""" keys = envelope_keys_from_master_key(_MASTER_KEY) - other_server_identity = EnvelopeIdentity(user_id=_IDENTITY.user_id, server_id="srv-OTHER") + other_server_identity = EnvelopeIdentity(server_id="srv-OTHER", key_hash=_IDENTITY.key_hash) token = _sealed_token(keys, identity=other_server_identity) result = resolve_bridge_envelope(token, keys, _NOW, _SERVER_ID) assert isinstance(result, BridgeEnvelopeInvalid) @@ -155,7 +155,7 @@ def test_resolve_non_ascii_server_id_stays_total_and_does_not_raise(): unicode server_id); it stays total and returns a typed result. A matching non-ASCII id admits, a mismatching one is BridgeEnvelopeInvalid, and neither raises.""" keys = envelope_keys_from_master_key(_MASTER_KEY) - unicode_identity = EnvelopeIdentity(user_id=_IDENTITY.user_id, server_id="srv-café") + unicode_identity = EnvelopeIdentity(server_id="srv-café", key_hash=_IDENTITY.key_hash) token = _sealed_token(keys, identity=unicode_identity) assert isinstance(resolve_bridge_envelope(token, keys, _NOW, "srv-café"), BridgeEnvelopeAdmitted) assert isinstance(resolve_bridge_envelope(token, keys, _NOW, "srv-cafe"), BridgeEnvelopeInvalid) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py index 71de206aa1e..b44f3f84cc9 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py @@ -51,7 +51,7 @@ _WRONG_SIGNING = EnvelopeKeys(signing_key=SecretStr(_OTHER_SIGNING_KEY), encrypt _WRONG_ENCRYPTION = EnvelopeKeys(signing_key=SecretStr(_SIGNING_KEY), encryption_key=SecretStr(_OTHER_ENCRYPTION_KEY)) _ACCESS_TOKEN = "upstream-access-token-do-not-leak-8f14e45fceea" _REFRESH_TOKEN = "upstream-refresh-token-do-not-leak-1d0aa4b7" -_IDENTITY = EnvelopeIdentity(user_id="user-123", server_id="srv-456") +_IDENTITY = EnvelopeIdentity(server_id="srv-456", key_hash="hashed-key-123") def _full_grant() -> UpstreamTokenGrant: @@ -137,12 +137,12 @@ def test_minimal_grant_round_trips_without_none_leakage_into_claims(): def test_claim_layout_and_no_plaintext_token_in_envelope(): token = _sealed_token(_full_grant()) claims = _unverified_claims(token) - assert set(claims) == {"iss", "iat", "exp", "user_id", "server_id", "grant"} + assert set(claims) == {"iss", "iat", "exp", "server_id", "key_hash", "grant"} assert claims["iss"] == ENVELOPE_ISSUER assert claims["iat"] == int(_NOW.timestamp()) assert claims["exp"] == int(_NOW.timestamp()) + 600 - assert claims["user_id"] == "user-123" assert claims["server_id"] == "srv-456" + assert claims["key_hash"] == "hashed-key-123" assert _ACCESS_TOKEN not in token assert _ACCESS_TOKEN not in json.dumps(claims) assert _REFRESH_TOKEN not in json.dumps(claims) @@ -226,11 +226,11 @@ def test_wrong_issuer_is_malformed_payload(): def test_missing_identity_claim_is_malformed_payload(): claims = _unverified_claims(_sealed_token(_full_grant())) - forged = _forge({key: value for key, value in claims.items() if key != "user_id"}) + forged = _forge({key: value for key, value in claims.items() if key != "key_hash"}) assert isinstance(open_envelope(forged, _KEYS, _NOW), MalformedPayload) -@pytest.mark.parametrize("identity_claim", ["user_id", "server_id"]) +@pytest.mark.parametrize("identity_claim", ["server_id", "key_hash"]) def test_signed_empty_identity_claim_is_malformed_payload_not_a_raise(identity_claim): claims = _unverified_claims(_sealed_token(_full_grant())) forged = _forge({**claims, identity_claim: ""}) @@ -463,9 +463,9 @@ def test_non_positive_expires_in_is_rejected_at_construction_without_leaking(): def test_empty_identity_and_key_fields_are_rejected_at_construction(): with pytest.raises(ValidationError): - EnvelopeIdentity(user_id="", server_id="srv-456") + EnvelopeIdentity(server_id="", key_hash="hashed-key-123") with pytest.raises(ValidationError): - EnvelopeIdentity(user_id="user-123", server_id="") + EnvelopeIdentity(server_id="srv-456", key_hash="") with pytest.raises(ValidationError): EnvelopeKeys(signing_key=SecretStr(""), encryption_key=SecretStr(_ENCRYPTION_KEY)) with pytest.raises(ValidationError): @@ -484,4 +484,4 @@ def test_public_models_are_frozen(): with pytest.raises(ValidationError): opened.grant = _minimal_grant() with pytest.raises(ValidationError): - _IDENTITY.user_id = "someone-else" + _IDENTITY.key_hash = "someone-elses-hash"