feat(mcp): generalize the bridge envelope identity to a key_hash or user_id subject

The scripted two-header client mints under a virtual key it presents at the token
endpoint (key_hash), but the interactive DCR client authenticates via SSO at the
bridged authorize, which yields a user, not a key. Make EnvelopeIdentity a
discriminated subject (subject_type key_hash | user_id) with key_hash_identity /
user_identity constructors, and dispatch admission on it: a key_hash reloads the
key, a user_id reloads the user and admits them as themselves (user-level budget
and SCIM enforced via the same centralized gate; no team bound, since a user
belongs to many teams or none). The interactive producer that mints a user_id
envelope lands in the follow-up commit.
This commit is contained in:
Tin Chi Lo 2026-07-11 14:44:10 -07:00
parent 78e5c43301
commit 45fed6a50a
7 changed files with 230 additions and 29 deletions

View file

@ -18,6 +18,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credenti
is_bridge_envelope_shaped,
resolve_bridge_envelope,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import (
EnvelopeIdentity,
)
from litellm.proxy._types import (
UI_TEAM_ID,
LiteLLM_TeamTable,
@ -543,7 +546,7 @@ class MCPRequestHandler:
header_key = server.alias or server.server_name
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)
admitted = await MCPRequestHandler._reload_admitted_principal(result.identity)
await MCPRequestHandler._enforce_admitted_live_policy(admitted=admitted, request=request, route=route)
injected = {header_key: {"Authorization": result.upstream_authorization.get_secret_value()}}
new_headers = {**(mcp_server_auth_headers or {}), **injected}
@ -572,6 +575,58 @@ class MCPRequestHandler:
route=route,
)
@staticmethod
async def _reload_admitted_principal(identity: EnvelopeIdentity) -> UserAPIKeyAuth:
"""Reload the live litellm record the envelope's subject references.
Dispatches on the sealed subject type: a ``key_hash`` reloads the virtual key that
minted the envelope (the scripted two-header client that presents a litellm key at the
token endpoint), a ``user_id`` reloads the user that authenticated interactively (the
DCR client, whose SSO login at the bridged authorize yields a user, not a key). Both
return a ``UserAPIKeyAuth`` the caller runs through the centralized policy gate, so
team/project/org/budget/SCIM enforcement is identical to the principal presenting
itself directly."""
match identity.subject_type:
case "key_hash":
return await MCPRequestHandler._reload_admitted_key(identity.subject)
case "user_id":
return await MCPRequestHandler._reload_admitted_user(identity.subject)
case _:
assert_never(identity.subject_type)
@staticmethod
async def _reload_admitted_user(user_id: str) -> UserAPIKeyAuth:
"""Reload the live user an interactively-minted envelope references and admit them as
themselves.
The DCR client authenticates via SSO at the bridged authorize, which yields a user
subject rather than a virtual key, so the envelope admits under the user's own
identity: the reloaded ``user_id`` rides on the returned ``UserAPIKeyAuth`` and the
caller's centralized policy gate then enforces the user's live budget and org state,
and a SCIM-deactivated owner fails closed here exactly as the key path enforces it. No
team is bound; a user may belong to many teams or none, so the envelope grants the
user's own access rather than silently selecting one team's scope. A missing user
fails closed with a 401 rather than admitting an unresolved identity."""
from litellm.proxy.auth.auth_checks import get_user_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:
user_object = await get_user_object(
user_id=user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
)
except (ProxyException, HTTPException):
raise HTTPException(status_code=401, detail="Invalid or expired credential") from None
if user_object is None:
raise HTTPException(status_code=401, detail="Invalid or expired credential")
if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False:
raise HTTPException(status_code=401, detail="Invalid or expired credential")
return UserAPIKeyAuth(user_id=user_object.user_id)
@staticmethod
async def _reload_admitted_key(key_hash: str) -> UserAPIKeyAuth:
"""Reload the live key record an admitted envelope references and re-check live policy.

View file

@ -961,15 +961,15 @@ def _finish_bridge_mint(
build_bridge_token_response,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import
EnvelopeIdentity,
SealedEnvelope,
UpstreamTokenGrant,
key_hash_identity,
)
grant = _bridge_grant_from_token_response(token_response)
if not isinstance(grant, UpstreamTokenGrant):
return _upstream_rejection_to_mint_error(grant)
identity = EnvelopeIdentity(server_id=mcp_server.server_id, key_hash=ready.key_hash)
identity = key_hash_identity(server_id=mcp_server.server_id, key_hash=ready.key_hash)
sealed = build_bridge_token_response(identity, grant, ready.keys, now)
if not isinstance(sealed, SealedEnvelope):
return "too_large"

View file

@ -67,20 +67,42 @@ typed error, never truncated."""
_ENVELOPE_JWT_ALGORITHM = "HS256"
class EnvelopeIdentity(BaseModel):
"""The litellm identity the envelope binds the inner grant to.
EnvelopeSubjectType: TypeAlias = Literal["key_hash", "user_id"]
"""Discriminator for what litellm principal the envelope binds the 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.
``key_hash`` is a hashed virtual key (the scripted two-header client mints under the key it
presents at the token endpoint); ``user_id`` is a litellm user subject (the interactive DCR
client mints under the SSO-authenticated user, which is the only identity that browser login
yields). Admission reloads a key record for the first and a user record for the second, then
runs both through the same live-policy gate, so team/org/budget/revocation enforcement is
identical either way."""
class EnvelopeIdentity(BaseModel):
"""The litellm principal the envelope binds the inner grant to.
``subject`` is the principal identifier and ``subject_type`` says how to resolve it: a
hashed litellm key (``key_hash``) or a litellm user id (``user_id``), never a raw
credential (and the edge rejects a bare hash or id presented as a bearer). Admission
reloads the live record by it, so the principal's current team/org 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)
server_id: str = Field(min_length=1)
key_hash: str = Field(min_length=1)
subject_type: EnvelopeSubjectType
subject: str = Field(min_length=1)
def key_hash_identity(server_id: str, key_hash: str) -> EnvelopeIdentity:
"""The identity for the scripted client that mints under a presented virtual key."""
return EnvelopeIdentity(server_id=server_id, subject_type="key_hash", subject=key_hash)
def user_identity(server_id: str, user_id: str) -> EnvelopeIdentity:
"""The identity for the interactive DCR client that mints under its SSO user subject."""
return EnvelopeIdentity(server_id=server_id, subject_type="user_id", subject=user_id)
class UpstreamTokenGrant(BaseModel):
@ -200,7 +222,8 @@ class _EnvelopeClaims(BaseModel):
iat: int
exp: int
server_id: str = Field(min_length=1)
key_hash: str = Field(min_length=1)
subject_type: EnvelopeSubjectType
subject: str = Field(min_length=1)
grant: str = Field(min_length=1)
@ -236,7 +259,8 @@ def mint_envelope(
iat=int(now.timestamp()),
exp=int(expires_at.timestamp()),
server_id=identity.server_id,
key_hash=identity.key_hash,
subject_type=identity.subject_type,
subject=identity.subject,
grant=_encrypt_grant_blob(_grant_plaintext(grant), keys.encryption_key),
)
token = ENVELOPE_PREFIX + jwt.encode(
@ -281,7 +305,7 @@ def open_envelope(
if not isinstance(grant, UpstreamTokenGrant):
return grant
return OpenedEnvelope(
identity=EnvelopeIdentity(server_id=claims.server_id, key_hash=claims.key_hash),
identity=EnvelopeIdentity(server_id=claims.server_id, subject_type=claims.subject_type, subject=claims.subject),
grant=grant,
)

View file

@ -4910,6 +4910,7 @@ class TestMCPDcrBridgeDelegateAdmission:
cls,
*,
key_hash=None,
user_id=None,
server_id="bridge-server-id",
access_token="inner-upstream-access-token",
token_type="Bearer",
@ -4921,17 +4922,23 @@ class TestMCPDcrBridgeDelegateAdmission:
envelope_keys_from_master_key,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import (
EnvelopeIdentity,
SealedEnvelope,
UpstreamTokenGrant,
key_hash_identity,
mint_envelope,
user_identity,
)
from pydantic import SecretStr
identity = (
user_identity(server_id=server_id, user_id=user_id)
if user_id is not None
else key_hash_identity(server_id=server_id, key_hash=key_hash or cls._KEY_HASH)
)
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(server_id=server_id, key_hash=key_hash or cls._KEY_HASH),
identity=identity,
grant=UpstreamTokenGrant(
access_token=SecretStr(access_token),
token_type=token_type,
@ -4999,6 +5006,22 @@ class TestMCPDcrBridgeDelegateAdmission:
stack.enter_context(patcher)
yield get_key_object
@staticmethod
@contextlib.contextmanager
def _patch_user_reload(*, return_value=None, side_effect=None):
"""Patch the user-subject reload path an interactively-minted envelope takes: the
``get_user_object`` lookup ``_reload_admitted_user`` runs (which also drives the SCIM gate),
plus the ``prisma_client`` / ``user_api_key_cache`` globals. The centralized gate's own
fetches fail-safe to None under the MagicMock prisma, so an unblocked user admits. Yields the
``get_user_object`` mock so a caller can assert the sealed user_id was the reload key."""
get_user_object = AsyncMock(return_value=return_value, side_effect=side_effect)
with (
patch("litellm.proxy.auth.auth_checks.get_user_object", get_user_object),
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
):
yield get_user_object
@staticmethod
def _mcp_request(path="/mcp/bridge_delegate_server"):
"""A minimal ``Request`` for direct ``_admit_dcr_bridge_delegate`` calls, mirroring how
@ -5060,6 +5083,84 @@ class TestMCPDcrBridgeDelegateAdmission:
"bridge_delegate_server": {"Authorization": "Bearer inner-upstream-access-token"}
}
async def test_user_subject_envelope_admits_under_the_reloaded_user(self):
"""An interactively-minted (user_id) envelope admits under the reloaded USER, not a key: the
reload is keyed by the sealed user_id, the admitted auth carries that user_id, the raw-key
pipeline is never invoked, and the inner upstream token is injected for egress. This is the
interactive-DCR admission the whole flow exists for."""
envelope = self._mint_bridge_envelope(user_id="sso-user-7")
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",
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),
self._patch_user_reload(
return_value=MagicMock(user_id="sso-user-7", metadata={"scim_active": True})
) as get_user_object,
):
mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server()
(auth_result, _h, _s, mcp_server_auth_headers, _o, _r) = await MCPRequestHandler.process_mcp_request(scope)
assert get_user_object.await_args.kwargs["user_id"] == "sso-user-7"
assert auth_result.user_id == "sso-user-7"
mock_auth.assert_not_called()
assert mcp_server_auth_headers == {
"bridge_delegate_server": {"Authorization": "Bearer inner-upstream-access-token"}
}
async def test_user_subject_envelope_missing_user_fails_closed_401(self):
"""A user_id envelope whose user has since been deleted must fail closed: get_user_object
resolves None, so admission 401s instead of admitting an unresolved identity."""
envelope = self._mint_bridge_envelope(user_id="ghost-user")
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_user_reload(return_value=None),
):
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_user_subject_envelope_scim_deactivated_user_fails_closed_401(self):
"""SCIM-deactivating the envelope's user revokes it immediately: the reloaded user carries
scim_active False, so admission 401s rather than letting an offboarded user keep tool access
until the envelope expires."""
envelope = self._mint_bridge_envelope(user_id="offboarded-user")
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_user_reload(
return_value=MagicMock(user_id="offboarded-user", metadata={"scim_active": False})
),
):
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_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

View file

@ -28,13 +28,14 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import
EnvelopeTooLarge,
SealedEnvelope,
UpstreamTokenGrant,
key_hash_identity,
mint_envelope,
)
_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(server_id="srv-456", key_hash="hashed-key-123")
_IDENTITY = key_hash_identity(server_id="srv-456", key_hash="hashed-key-123")
_SERVER_ID = _IDENTITY.server_id
@ -138,7 +139,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(server_id="srv-OTHER", key_hash=_IDENTITY.key_hash)
other_server_identity = key_hash_identity(server_id="srv-OTHER", key_hash=_IDENTITY.subject)
token = _sealed_token(keys, identity=other_server_identity)
result = resolve_bridge_envelope(token, keys, _NOW, _SERVER_ID)
assert isinstance(result, BridgeEnvelopeInvalid)
@ -155,7 +156,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(server_id="srv-café", key_hash=_IDENTITY.key_hash)
unicode_identity = key_hash_identity(server_id="srv-café", key_hash=_IDENTITY.subject)
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)

View file

@ -36,8 +36,10 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import
SealedEnvelope,
UpstreamTokenGrant,
is_envelope,
key_hash_identity,
mint_envelope,
open_envelope,
user_identity,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value, encrypt_value
@ -51,7 +53,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(server_id="srv-456", key_hash="hashed-key-123")
_IDENTITY = key_hash_identity(server_id="srv-456", key_hash="hashed-key-123")
def _full_grant() -> UpstreamTokenGrant:
@ -137,12 +139,13 @@ 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", "server_id", "key_hash", "grant"}
assert set(claims) == {"iss", "iat", "exp", "server_id", "subject_type", "subject", "grant"}
assert claims["iss"] == ENVELOPE_ISSUER
assert claims["iat"] == int(_NOW.timestamp())
assert claims["exp"] == int(_NOW.timestamp()) + 600
assert claims["server_id"] == "srv-456"
assert claims["key_hash"] == "hashed-key-123"
assert claims["subject_type"] == "key_hash"
assert claims["subject"] == "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 +229,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 != "key_hash"})
forged = _forge({key: value for key, value in claims.items() if key != "subject"})
assert isinstance(open_envelope(forged, _KEYS, _NOW), MalformedPayload)
@pytest.mark.parametrize("identity_claim", ["server_id", "key_hash"])
@pytest.mark.parametrize("identity_claim", ["server_id", "subject"])
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 +466,11 @@ 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(server_id="", key_hash="hashed-key-123")
EnvelopeIdentity(server_id="", subject_type="key_hash", subject="hashed-key-123")
with pytest.raises(ValidationError):
EnvelopeIdentity(server_id="srv-456", key_hash="")
EnvelopeIdentity(server_id="srv-456", subject_type="key_hash", subject="")
with pytest.raises(ValidationError):
EnvelopeIdentity(server_id="srv-456", subject_type="not-a-subject-type", subject="x")
with pytest.raises(ValidationError):
EnvelopeKeys(signing_key=SecretStr(""), encryption_key=SecretStr(_ENCRYPTION_KEY))
with pytest.raises(ValidationError):
@ -474,6 +479,20 @@ def test_empty_identity_and_key_fields_are_rejected_at_construction():
UpstreamTokenGrant(access_token=SecretStr(""), token_type="Bearer")
def test_user_subject_identity_round_trips():
"""The user_id subject variant seals and opens with its discriminator intact, so the edge can
tell an interactively-minted (user) envelope from a scripted (key_hash) one and reload the right
kind of record."""
identity = user_identity(server_id="srv-456", user_id="user-42")
sealed = mint_envelope(identity, _full_grant(), _KEYS, _NOW)
assert isinstance(sealed, SealedEnvelope)
opened = open_envelope(sealed.token.get_secret_value(), _KEYS, _NOW)
assert isinstance(opened, OpenedEnvelope)
assert opened.identity.server_id == "srv-456"
assert opened.identity.subject_type == "user_id"
assert opened.identity.subject == "user-42"
def test_public_models_are_frozen():
sealed = mint_envelope(_IDENTITY, _full_grant(), _KEYS, _NOW)
assert isinstance(sealed, SealedEnvelope)
@ -484,4 +503,4 @@ def test_public_models_are_frozen():
with pytest.raises(ValidationError):
opened.grant = _minimal_grant()
with pytest.raises(ValidationError):
_IDENTITY.key_hash = "someone-elses-hash"
_IDENTITY.subject = "someone-elses-hash"

View file

@ -4441,7 +4441,8 @@ async def test_oauth_delegate_bridge_token_exchange_mints_envelope_not_raw_token
keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY)
opened = resolve_bridge_envelope(token, keys, datetime.now(timezone.utc), server.server_id)
assert isinstance(opened, BridgeEnvelopeAdmitted)
assert opened.identity.key_hash == "hashed-litellm-key-77"
assert opened.identity.subject_type == "key_hash"
assert opened.identity.subject == "hashed-litellm-key-77"
assert opened.upstream_authorization.get_secret_value() == "Bearer UPSTREAM-SECRET-TOKEN"