mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
refactor(mcp): make the bridge delegate mint a phased failures-as-values pipeline
The dcr_bridge oauth_delegate token mint validated its preconditions in two places: a pre-exchange guard inside exchange_token_with_server (master_key set, resolvable litellm identity) and an authoritative re-check inside the post-exchange _mint_bridge_delegate_token_response. Keeping the two in step by hand is what kept producing the same class of finding: a precondition guarded on one grant branch but not the other, master_key checked after the exchange on one path, identity resolved twice, and each failure raising an ad-hoc HTTPException with its own status and body shape. Model the mint as three phases whose failures are values. _prepare_bridge_mint runs before the exchange, checks every precondition once (master_key, then identity), and returns either a frozen _BridgeMintReady carrying the resolved key hash and the master-key-derived envelope keys, or a _BridgeMintError literal. Because every precondition lives in prepare, and prepare runs before the upstream POST, no failure can burn the single-use code or rotate a refresh token, for either grant type, by construction rather than by a guard we have to remember to keep in sync. _finish_bridge_mint runs after the exchange and has no preconditions left that can fail; its only failure values are properties of the upstream response itself (no usable access_token, or a token too large to seal). One mapper, _bridge_mint_error_response, turns each _BridgeMintError into an RFC 6749 section 5.2-shaped body with a status truthful about where the failure is (400 for the caller, 500 for gateway config, 502 for the upstream), with an exhaustive match plus assert_never so a new failure mode cannot be added without a matching status. Behavior is unchanged for the client. Every failure that previously raised now returns the same status as an OAuth error body, which is the correct token-endpoint contract; the three tests that asserted a raised HTTPException now assert the returned response. _exchange_for_bridge_server additionally asserts the identity resolver is awaited exactly once for a bridge server and never for a non-bridge one.
This commit is contained in:
Tin Chi Lo
parent
e16ad044c3
commit
2f0ddc82f7
2 changed files with 140 additions and 90 deletions
|
|
@ -4,6 +4,7 @@ import json
|
|||
import math
|
||||
import secrets
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple
|
||||
from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
|
||||
|
|
@ -12,6 +13,7 @@ import httpx
|
|||
from fastapi import APIRouter, Form, HTTPException, Request
|
||||
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response
|
||||
from pydantic import BaseModel, SecretStr, ValidationError
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
|
|
@ -39,6 +41,7 @@ from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import (
|
||||
EnvelopeKeys,
|
||||
UpstreamTokenGrant,
|
||||
)
|
||||
from litellm.proxy._types import LiteLLM_MCPServerTable, UserAPIKeyAuth
|
||||
|
|
@ -757,73 +760,112 @@ def _bridge_grant_from_token_response(token_response: object) -> Optional["Upstr
|
|||
)
|
||||
|
||||
|
||||
def _bridge_invalid_request_response() -> JSONResponse:
|
||||
"""RFC 6749 §5.2-shaped ``invalid_request`` for a bridge token exchange that carries no resolvable
|
||||
litellm identity. Returned (not raised) so the OAuth error members sit at the top level rather than
|
||||
wrapped in FastAPI's ``detail``, with the no-store token-endpoint headers, matching the BYOK OAuth
|
||||
endpoint and what a strict DCR client parses per RFC 6749 §5.2."""
|
||||
return JSONResponse(
|
||||
status_code=400,
|
||||
content={
|
||||
"error": "invalid_request",
|
||||
"error_description": (
|
||||
# ---------------------------------------------------------------------------
|
||||
# DCR-bridge oauth_delegate mint: a three-phase pipeline whose failures are values.
|
||||
#
|
||||
# prepare (before the upstream exchange) -> validate every precondition and resolve identity+keys
|
||||
# exchange (the single-use upstream code is consumed here, in exchange_token_with_server)
|
||||
# finish (after the exchange) -> seal the upstream grant into the client-held envelope
|
||||
#
|
||||
# Every precondition lives in ``prepare``, which runs BEFORE the exchange, so no failure can burn the
|
||||
# single-use code or rotate a refresh token, for either grant type -- that whole class of bug is gone
|
||||
# by construction rather than guarded case by case. Failures are values mapped to an OAuth-shaped
|
||||
# response in one place (``_bridge_mint_error_response``), so status codes and the RFC 6749 §5.2 body
|
||||
# shape are uniform. Adding a failure mode is a new literal plus a match arm the type checker forces.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_BridgeMintError = Literal["not_configured", "no_identity", "no_upstream_token", "too_large"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _BridgeMintReady:
|
||||
"""Everything the seal needs, resolved once before the exchange: the authorizing key hash and the
|
||||
master-key-derived envelope keys. Passing this forward means identity resolution and key derivation
|
||||
happen exactly once, and ``_finish_bridge_mint`` has no preconditions left that could fail."""
|
||||
|
||||
key_hash: str
|
||||
keys: "EnvelopeKeys"
|
||||
|
||||
|
||||
def _bridge_mint_error_response(error: _BridgeMintError) -> JSONResponse:
|
||||
"""Map a bridge-mint failure value to its token-endpoint response. One place, RFC 6749 §5.2 shape
|
||||
(top-level ``error``, no-store) for every case, with a status truthful about where the failure is:
|
||||
the caller's request (400), the gateway config (500), or the upstream (502)."""
|
||||
if error == "no_identity":
|
||||
status, code, desc = (
|
||||
400,
|
||||
"invalid_request",
|
||||
(
|
||||
"this server issues a gateway-bound credential; send a litellm credential "
|
||||
"(x-litellm-api-key or Authorization) on the token request"
|
||||
),
|
||||
},
|
||||
headers=TOKEN_NO_CACHE_HEADERS,
|
||||
)
|
||||
elif error == "not_configured":
|
||||
status, code, desc = (
|
||||
500,
|
||||
"server_error",
|
||||
("the gateway is not configured to mint a gateway-bound credential (master_key is not set)"),
|
||||
)
|
||||
elif error == "no_upstream_token":
|
||||
status, code, desc = 502, "server_error", "the upstream token response has no usable access_token"
|
||||
elif error == "too_large":
|
||||
status, code, desc = (
|
||||
502,
|
||||
"server_error",
|
||||
("the upstream token is too large to seal into a gateway-bound credential"),
|
||||
)
|
||||
else:
|
||||
assert_never(error)
|
||||
return JSONResponse(
|
||||
status_code=status, content={"error": code, "error_description": desc}, headers=TOKEN_NO_CACHE_HEADERS
|
||||
)
|
||||
|
||||
|
||||
async def _mint_bridge_delegate_token_response(
|
||||
request: Request, mcp_server: MCPServer, token_response: object
|
||||
) -> JSONResponse:
|
||||
"""Return the client-held envelope bearer for a DCR-bridge ``oauth_delegate`` token exchange.
|
||||
async def _prepare_bridge_mint(request: Request, mcp_server: MCPServer) -> "_BridgeMintReady | _BridgeMintError":
|
||||
"""Phase 1, BEFORE the upstream exchange: validate that the gateway can mint (master_key set) and
|
||||
that the request carries a resolvable litellm identity, and derive the envelope keys. Returns a
|
||||
ready context or a failure value. Running before the exchange is what makes a missing master_key or
|
||||
an unresolvable identity fail closed without consuming the single-use code / rotating a refresh
|
||||
token, for both grant types."""
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
envelope_keys_from_master_key,
|
||||
)
|
||||
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
master_key,
|
||||
)
|
||||
|
||||
The envelope binds the authorizing litellm key (its hash, resolved from the token request) to the
|
||||
upstream grant, so the client holds one bearer that later admits it and forwards the upstream
|
||||
token, with nothing stored server-side. Admission reloads the live key by that hash, so the key's
|
||||
current restrictions and revocation gate the request. Fails closed with an OAuth
|
||||
``invalid_request`` when no active litellm key accompanies the token request rather than minting
|
||||
an unbound credential.
|
||||
"""
|
||||
if not master_key:
|
||||
return "not_configured"
|
||||
key_hash = await _extract_active_key_hash_from_request(request)
|
||||
if not key_hash:
|
||||
return "no_identity"
|
||||
return _BridgeMintReady(key_hash=key_hash, keys=envelope_keys_from_master_key(master_key))
|
||||
|
||||
|
||||
def _finish_bridge_mint(
|
||||
ready: "_BridgeMintReady", mcp_server: MCPServer, token_response: object, now: datetime
|
||||
) -> "JSONResponse | _BridgeMintError":
|
||||
"""Phase 3, AFTER the upstream exchange: seal the upstream grant into the client-held envelope using
|
||||
the pre-resolved identity and keys, so the client holds one bearer that admits it and forwards the
|
||||
upstream token with nothing stored server-side. The only failures here are properties of the
|
||||
upstream response (no usable token, or a token too large to seal), returned as values."""
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
build_bridge_token_response,
|
||||
envelope_keys_from_master_key,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
EnvelopeIdentity,
|
||||
SealedEnvelope,
|
||||
)
|
||||
from litellm.proxy.proxy_server import (
|
||||
master_key, # noqa: PLC0415 # inline import avoids a module-load circular import
|
||||
)
|
||||
|
||||
if not master_key:
|
||||
raise HTTPException(status_code=500, detail="Server misconfigured: master_key is not set")
|
||||
|
||||
key_hash = await _extract_active_key_hash_from_request(request)
|
||||
if not key_hash:
|
||||
return _bridge_invalid_request_response()
|
||||
|
||||
grant = _bridge_grant_from_token_response(token_response)
|
||||
if grant is None:
|
||||
raise HTTPException(status_code=502, detail="Upstream token response has no usable access_token")
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
keys = envelope_keys_from_master_key(master_key)
|
||||
identity = EnvelopeIdentity(server_id=mcp_server.server_id, key_hash=key_hash)
|
||||
sealed = build_bridge_token_response(identity, grant, keys, now)
|
||||
return "no_upstream_token"
|
||||
identity = EnvelopeIdentity(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):
|
||||
# build_bridge_token_response returns EnvelopeTooLarge as a value when the upstream token is
|
||||
# too large to seal; that is an upstream-payload condition, so surface a 502, not a 500.
|
||||
raise HTTPException(
|
||||
status_code=502, detail="Upstream token is too large to seal into a gateway-bound credential"
|
||||
)
|
||||
|
||||
# The JWT exp is int(expires_at.timestamp()) (second-truncated), and admission expires the envelope
|
||||
# against that exp. Report expires_in from the same truncated exp, rounding the elapsed portion up,
|
||||
# so the client is never told the bearer lives past the point admission already rejects it.
|
||||
return "too_large"
|
||||
# Report expires_in from the JWT's own second-truncated exp, rounding the elapsed portion up, so the
|
||||
# client is never told the bearer lives past the point admission (which uses that exp) rejects it.
|
||||
expires_in = max(1, int(sealed.expires_at.timestamp()) - math.ceil(now.timestamp()))
|
||||
body = {"access_token": sealed.token.get_secret_value(), "token_type": "Bearer", "expires_in": expires_in}
|
||||
return JSONResponse(body, headers=TOKEN_NO_CACHE_HEADERS)
|
||||
|
|
@ -902,19 +944,15 @@ async def exchange_token_with_server(
|
|||
if code_verifier:
|
||||
token_data["code_verifier"] = code_verifier
|
||||
|
||||
# A bridge oauth_delegate mint must fail closed BEFORE the upstream exchange consumes or rotates the
|
||||
# single-use code (or refresh token): confirm the gateway can mint at all (master_key set) and that
|
||||
# the request carries a resolvable litellm identity. Applies to both grant types, so an invalid key
|
||||
# or a misconfigured gateway never burns the upstream credential. The mint below re-checks
|
||||
# authoritatively; get_key_object is cache-first, so the identity re-resolution is a cache hit and
|
||||
# adds no extra database round-trip.
|
||||
# Phase 1: for a bridge oauth_delegate mint, validate all preconditions and resolve identity+keys
|
||||
# BEFORE the exchange below consumes the single-use upstream code, and carry the ready context to
|
||||
# phase 3. A failure here returns without ever touching the upstream credential.
|
||||
bridge_mint_ready: _BridgeMintReady | None = None
|
||||
if mcp_server.is_oauth_delegate and mcp_server.is_dcr_bridge:
|
||||
from litellm.proxy.proxy_server import master_key as _bridge_master_key # noqa: PLC0415
|
||||
|
||||
if not _bridge_master_key:
|
||||
raise HTTPException(status_code=500, detail="Server misconfigured: master_key is not set")
|
||||
if not await _extract_active_key_hash_from_request(request):
|
||||
return _bridge_invalid_request_response()
|
||||
prepared = await _prepare_bridge_mint(request, mcp_server)
|
||||
if not isinstance(prepared, _BridgeMintReady):
|
||||
return _bridge_mint_error_response(prepared)
|
||||
bridge_mint_ready = prepared
|
||||
|
||||
async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)
|
||||
response = await async_client.post(
|
||||
|
|
@ -982,8 +1020,11 @@ async def exchange_token_with_server(
|
|||
# A DCR-bridge oauth_delegate server hands the client a gateway-bound envelope (identity plus the
|
||||
# upstream token) instead of the raw upstream token, so the one bearer both admits the caller and
|
||||
# forwards the upstream credential. Only this mode mints; every other server returns the raw token.
|
||||
if mcp_server.is_oauth_delegate and mcp_server.is_dcr_bridge:
|
||||
return await _mint_bridge_delegate_token_response(request, mcp_server, token_response)
|
||||
if bridge_mint_ready is not None:
|
||||
# Phase 3: seal the upstream grant into the client-held envelope; failures map through the same
|
||||
# OAuth-shaped response as the phase-1 preconditions.
|
||||
minted = _finish_bridge_mint(bridge_mint_ready, mcp_server, token_response, datetime.now(timezone.utc))
|
||||
return minted if isinstance(minted, JSONResponse) else _bridge_mint_error_response(minted)
|
||||
|
||||
result = {
|
||||
"access_token": token_response["access_token"],
|
||||
|
|
|
|||
|
|
@ -4375,6 +4375,7 @@ async def _exchange_for_bridge_server(server, upstream_body, key_hash, fake_clie
|
|||
fake_http_response.raise_for_status = MagicMock()
|
||||
fake_http_client = MagicMock()
|
||||
fake_http_client.post = AsyncMock(return_value=fake_http_response)
|
||||
key_resolver = AsyncMock(return_value=key_hash)
|
||||
if fake_client_out is not None:
|
||||
fake_client_out["client"] = fake_http_client
|
||||
|
||||
|
|
@ -4385,11 +4386,11 @@ async def _exchange_for_bridge_server(server, upstream_body, key_hash, fake_clie
|
|||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_active_key_hash_from_request",
|
||||
new=AsyncMock(return_value=key_hash),
|
||||
new=key_resolver,
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY),
|
||||
):
|
||||
return await exchange_token_with_server(
|
||||
response = await exchange_token_with_server(
|
||||
request=_bridge_mock_request(),
|
||||
mcp_server=server,
|
||||
grant_type="authorization_code",
|
||||
|
|
@ -4399,6 +4400,11 @@ async def _exchange_for_bridge_server(server, upstream_body, key_hash, fake_clie
|
|||
client_secret=None,
|
||||
code_verifier="verifier",
|
||||
)
|
||||
if server.is_oauth_delegate and server.is_dcr_bridge:
|
||||
key_resolver.assert_awaited_once()
|
||||
else:
|
||||
key_resolver.assert_not_awaited()
|
||||
return response
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -4457,15 +4463,16 @@ async def test_oauth_delegate_bridge_token_exchange_fails_closed_without_litellm
|
|||
@pytest.mark.asyncio
|
||||
async def test_bridge_envelope_too_large_upstream_token_is_502():
|
||||
"""An upstream token too large to seal into the envelope is an upstream-payload condition, so the
|
||||
mint surfaces a 502 rather than a 500 (build_bridge_token_response returns EnvelopeTooLarge as a
|
||||
value, and the caller maps it to a truthful status)."""
|
||||
mint surfaces a 502 (as an RFC 6749 §5.2 error body, not a raised HTTPException) rather than a 500:
|
||||
build_bridge_token_response returns EnvelopeTooLarge as a value, _finish_bridge_mint returns the
|
||||
"too_large" failure, and _bridge_mint_error_response maps it to a truthful status."""
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
server = _bridge_server(auth_type=MCPAuth.oauth_delegate)
|
||||
upstream = {"access_token": "x" * 40000, "token_type": "Bearer", "expires_in": 3600}
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77")
|
||||
assert exc.value.status_code == 502
|
||||
response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77")
|
||||
assert response.status_code == 502
|
||||
assert json.loads(response.body)["error"] == "server_error"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -4544,8 +4551,10 @@ async def test_bridge_refresh_grant_fails_closed_before_upstream_when_no_identit
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_mint_fails_closed_before_upstream_when_master_key_unset():
|
||||
"""master_key is validated BEFORE the upstream exchange, so a misconfigured gateway 500s without
|
||||
consuming the single-use code, avoiding the burn-then-fail the pre-exchange gate exists to prevent."""
|
||||
"""master_key is validated BEFORE the upstream exchange (in _prepare_bridge_mint), so a
|
||||
misconfigured gateway returns a 500 server_error without consuming the single-use code, avoiding
|
||||
the burn-then-fail the pre-exchange phase exists to prevent. The failure is returned as an RFC 6749
|
||||
error body, not raised."""
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
|
|
@ -4563,19 +4572,19 @@ async def test_bridge_mint_fails_closed_before_upstream_when_master_key_unset():
|
|||
),
|
||||
patch("litellm.proxy.proxy_server.master_key", None),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await exchange_token_with_server(
|
||||
request=_bridge_mock_request(),
|
||||
mcp_server=server,
|
||||
grant_type="authorization_code",
|
||||
code="auth-code",
|
||||
redirect_uri="https://claude.ai/api/mcp/auth_callback",
|
||||
client_id="dcr-client-123",
|
||||
client_secret=None,
|
||||
code_verifier="verifier",
|
||||
)
|
||||
response = await exchange_token_with_server(
|
||||
request=_bridge_mock_request(),
|
||||
mcp_server=server,
|
||||
grant_type="authorization_code",
|
||||
code="auth-code",
|
||||
redirect_uri="https://claude.ai/api/mcp/auth_callback",
|
||||
client_id="dcr-client-123",
|
||||
client_secret=None,
|
||||
code_verifier="verifier",
|
||||
)
|
||||
|
||||
assert exc.value.status_code == 500
|
||||
assert response.status_code == 500
|
||||
assert json.loads(response.body)["error"] == "server_error"
|
||||
fake_http_client.post.assert_not_called()
|
||||
|
||||
|
||||
|
|
@ -4647,18 +4656,18 @@ async def test_bridge_token_exchange_honors_short_float_expires_in_ttl():
|
|||
@pytest.mark.asyncio
|
||||
async def test_oauth_delegate_bridge_token_exchange_missing_access_token_is_502_not_keyerror():
|
||||
"""When the upstream token response has no access_token, a dcr_bridge oauth_delegate exchange
|
||||
returns a clean 502 rather than raising a KeyError. The eager access_token extraction used to run
|
||||
before the bridge branch, so a missing token raised KeyError and _bridge_grant_from_token_response's
|
||||
nil guard (which maps to 502) was dead code; the extraction now lives on the non-bridge path only."""
|
||||
returns a clean 502 error body rather than raising a KeyError. _finish_bridge_mint asks
|
||||
_bridge_grant_from_token_response for a typed grant, gets None, and returns the "no_upstream_token"
|
||||
failure, which maps to 502; nothing indexes token_response["access_token"] on the bridge path."""
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
server = _bridge_server(auth_type=MCPAuth.oauth_delegate)
|
||||
upstream = {"token_type": "Bearer", "expires_in": 3600}
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77")
|
||||
response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77")
|
||||
|
||||
assert exc.value.status_code == 502
|
||||
assert response.status_code == 502
|
||||
assert json.loads(response.body)["error"] == "server_error"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue