scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

feat(mcp): graft token_exchange (RFC 8693 OBO) to the v2 resolver

Browse source 
Completes TokenExchangeConfig with the client authentication the exchange grant requires
(client_id/client_secret) plus scopes -- v1's exchange (and has_token_exchange_config) require
them, so the v2 config was incomplete. Then grafts the arm on the existing header seam, riding the
Subject edge (the inbound token is the Subject's inbound_token):

- request_token(endpoint, data, clock) is factored out of the client_credentials fetcher (both
  hit a token endpoint and parse access_token/expires_in); the fetcher now calls it.
- HttpxTokenExchanger does the RFC 8693 grant via request_token: grant_type=token-exchange,
  subject_token, subject_token_type, client_id/client_secret, audience=resource (RFC 8707), scope.
  The inbound token is sent only to the exchange endpoint, never upstream.
- _to_server_spec maps has_token_exchange_config servers to TokenExchangeConfig (endpoint falling
  back to token_url, resource = audience or url), placed before client_credentials to match v1's
  cascade. _provider injects HttpxTokenExchanger; the per-user exchanged-token cache stays the
  in-memory TokenStore (best-effort, re-exchangeable).

Intended divergence from v1: the exchange always binds to resource (audience or upstream URL,
RFC 8707); v1 sent audience only when explicitly set. Parity for servers with an explicit
audience; a new binding for those without (the resource-binding design).

Tests cover the exchanger grant shape + error mapping and the server->config mapping. 97 tests
pass; types.py adds no new errors, port bodies/bridge typecheck clean. Live e2e is deferred: it
needs an IdP exchange endpoint and an inbound subject_token (JWT), which the config-only local
harness does not set up.
This commit is contained in:
Tin Chi Lo 2026-06-18 18:15:48 -07:00
parent 49e12c872d
commit 2f69b705cc
5 changed files with 203 additions and 50 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

138
litellm/proxy/_experimental/mcp_server/v2_port_bodies.py
View file

@ -27,6 +27,7 @@ from litellm.proxy.gateway.mcp.outbound_credentials.types import (
ClientCredentialsConfig,
CredError,
StaticKeys,
TokenExchangeConfig,
)
from litellm.proxy.gateway.mcp.result import Error, Ok, Result
from litellm.types.llms.custom_http import httpxSpecialProvider
@ -40,6 +41,60 @@ class _TokenResponse(BaseModel):
expires_in: Optional[int] = None
async def request_token(
endpoint: str, data: dict[str, str], clock: Clock
) -> Result[StoredToken, CredError]:
"""POST an OAuth token request and parse the response into a StoredToken.
Shared by the client_credentials grant and the RFC 8693 exchange (both hit a token endpoint
and read access_token / expires_in). Errors-as-values: network / 5xx -> upstream_unavailable,
4xx -> misconfigured, missing access_token -> misconfigured. The raw lifetime is stored; the
resolver arm's _REFRESH_BUFFER (= v1's default buffer) handles proactive re-mint.
"""
client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP)
try:
response = await client.post(endpoint, data=data)
except Exception as e: # network / timeout / DNS
return Error(
CredError.of_upstream_unavailable(f"token endpoint unreachable: {e}")
)
if response is None:
return Error(
CredError.of_upstream_unavailable("token endpoint returned no response")
)
if response.status_code >= 500:
return Error(
CredError.of_upstream_unavailable(
f"token endpoint returned {response.status_code}"
)
)
if response.status_code >= 400:
return Error(
CredError.of_misconfigured(
f"token request rejected ({response.status_code})"
)
)
try:
parsed = _TokenResponse.model_validate(response.json())
except ValidationError:
return Error(
CredError.of_misconfigured("token response missing a valid access_token")
)
ttl = timedelta(
seconds=(
parsed.expires_in
if parsed.expires_in is not None
else MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL
)
)
return Ok(
StoredToken(
access_token=SecretStr(parsed.access_token),
expires_at=clock.now() + ttl,
)
)
class HttpxClientCredentialsFetcher:
"""RFC 6749 client_credentials grant via litellm's configured httpx client.
@ -63,58 +118,47 @@ class HttpxClientCredentialsFetcher:
if config.scopes:
data["scope"] = " ".join(config.scopes)
client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP)
try:
response = await client.post(config.token_url, data=data)
except Exception as e: # network / timeout / DNS
return Error(
CredError.of_upstream_unavailable(
f"client_credentials token endpoint unreachable: {e}"
)
)
return await request_token(config.token_url, data, self._clock)
if response is None:
return Error(
CredError.of_upstream_unavailable(
"client_credentials token endpoint returned no response"
)
)
if response.status_code >= 500:
return Error(
CredError.of_upstream_unavailable(
f"client_credentials token endpoint returned {response.status_code}"
)
)
if response.status_code >= 400:
_TOKEN_EXCHANGE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange"
class HttpxTokenExchanger:
"""RFC 8693 token exchange (OBO) via litellm's configured httpx client.
Swaps the caller's inbound subject_token for a token bound to the upstream (audience=resource,
RFC 8707), mirroring v1's exchange grant: the gateway authenticates as an OAuth client via
client_id / client_secret. The inbound token is sent only to the exchange endpoint, never to
the upstream.
"""
def __init__(self, clock: Optional[Clock] = None) -> None:
self._clock: Clock = clock or SystemClock()
async def exchange(
self, config: TokenExchangeConfig, subject_token: SecretStr, resource: str
) -> Result[StoredToken, CredError]:
if not config.token_exchange_endpoint:
return Error(
CredError.of_misconfigured(
f"client_credentials grant rejected ({response.status_code})"
"token_exchange: no exchange endpoint configured"
)
)
try:
parsed = _TokenResponse.model_validate(response.json())
except ValidationError:
return Error(
CredError.of_misconfigured(
"client_credentials response missing a valid access_token"
)
)
# Store the raw lifetime; the resolver arm's _REFRESH_BUFFER (60s, = v1's default
# buffer) handles proactive re-mint, so no buffer is subtracted here.
ttl = timedelta(
seconds=(
parsed.expires_in
if parsed.expires_in is not None
else MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL
)
)
return Ok(
StoredToken(
access_token=SecretStr(parsed.access_token),
expires_at=self._clock.now() + ttl,
)
)
data = {
"grant_type": _TOKEN_EXCHANGE_GRANT_TYPE,
"subject_token": subject_token.get_secret_value(),
"subject_token_type": config.subject_token_type,
}
if config.client_id:
data["client_id"] = config.client_id
if config.client_secret is not None:
data["client_secret"] = config.client_secret.get_secret_value()
if resource:
data["audience"] = resource
if config.scopes:
data["scope"] = " ".join(config.scopes)
return await request_token(config.token_exchange_endpoint, data, self._clock)
class HttpxSigV4Signer:

20
litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py
View file

@ -23,6 +23,7 @@ from litellm._logging import verbose_logger
from litellm.proxy._experimental.mcp_server.v2_port_bodies import (
HttpxClientCredentialsFetcher,
HttpxSigV4Signer,
HttpxTokenExchanger,
V1ByokCredentialStore,
)
from litellm.proxy.gateway.mcp.outbound_credentials.clock import SystemClock
@ -102,7 +103,7 @@ def _provider() -> UpstreamCredentialProvider:
clock=SystemClock(),
service_token_store=InMemoryServiceTokenStore(),
client_credentials_fetcher=HttpxClientCredentialsFetcher(),
token_exchanger=unwired,
token_exchanger=HttpxTokenExchanger(),
signer_factory=HttpxSigV4Signer(),
)
@ -137,6 +138,23 @@ def _to_server_spec(server: MCPServer) -> Optional[ServerSpec]:
key_source=SharedKey(value=SecretStr(token)),
),
)
if server.has_token_exchange_config:
# token_exchange takes precedence over client_credentials (matches v1's cascade); the arm
# binds the exchanged token to this resource (audience, RFC 8707).
return ServerSpec(
server_id=server.server_id,
resource=server.audience or resource,
config=TokenExchangeConfig(
subject_token_type=server.subject_token_type,
token_exchange_endpoint=server.token_exchange_endpoint
or server.token_url,
client_id=server.client_id,
client_secret=(
SecretStr(server.client_secret) if server.client_secret else None
),
scopes=tuple(server.scopes or ()),
),
)
if (
server.has_client_credentials
and server.client_id

8
litellm/proxy/gateway/mcp/outbound_credentials/types.py
View file

@ -166,14 +166,18 @@ class ClientCredentialsConfig(BaseModel):
class TokenExchangeConfig(BaseModel):
"""RFC 8693 OBO; swap the caller's live subject_token for a token bound to the upstream's
audience (`server.resource`). The exchange runs at the IdP's token endpoint, discovered
(RFC 8414); only the inbound token type and an optional manual endpoint override are config.
audience (`server.resource`, RFC 8707). The gateway authenticates to the exchange endpoint as
an OAuth client (`client_id`/`client_secret`); the inbound token is sent only to that endpoint,
never to the upstream.
"""
model_config = ConfigDict(frozen=True)
kind: Literal[AuthSpecKind.token_exchange] = AuthSpecKind.token_exchange
subject_token_type: str = "urn:ietf:params:oauth:token-type:access_token"
token_exchange_endpoint: str | None = None
client_id: str | None = None
client_secret: SecretStr | None = None
scopes: tuple[str, ...] = ()
class SharedKey(BaseModel):

58
tests/test_litellm/proxy/_experimental/mcp_server/test_v2_port_bodies.py
View file

@ -11,6 +11,7 @@ from litellm.proxy._experimental.mcp_server import v2_port_bodies
from litellm.proxy._experimental.mcp_server.v2_port_bodies import (
HttpxClientCredentialsFetcher,
HttpxSigV4Signer,
HttpxTokenExchanger,
V1ByokCredentialStore,
_classify_sigv4_error,
)
@ -21,6 +22,7 @@ from litellm.proxy.gateway.mcp.outbound_credentials.types import (
AwsSigV4Config,
ClientCredentialsConfig,
StaticKeys,
TokenExchangeConfig,
)
from litellm.proxy.gateway.mcp.result import Error, Ok
@ -207,3 +209,59 @@ async def test_byok_store_db_error_is_upstream_unavailable():
result = await V1ByokCredentialStore(reader=reader).get(_cred_key())
assert isinstance(result, Error)
assert result.error.tag == "upstream_unavailable"
def _te_config(endpoint="https://idp/exchange", scopes=(), client_secret="csecret"):
return TokenExchangeConfig(
token_exchange_endpoint=endpoint,
client_id="cid",
client_secret=SecretStr(client_secret) if client_secret else None,
scopes=tuple(scopes),
)
async def test_exchange_builds_rfc8693_grant(monkeypatch):
fake = _patch(
monkeypatch,
_FakeResponse(200, {"access_token": "exch-tok", "expires_in": 3600}),
)
result = await HttpxTokenExchanger().exchange(
_te_config(scopes=["a"]), SecretStr("inbound-jwt"), "https://up.example/mcp"
)
assert isinstance(result, Ok)
assert result.ok.access_token.get_secret_value() == "exch-tok"
d = fake.posted["data"]
assert d["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange"
assert d["subject_token"] == "inbound-jwt"
assert d["subject_token_type"] == "urn:ietf:params:oauth:token-type:access_token"
assert d["client_id"] == "cid"
assert d["client_secret"] == "csecret"
assert d["audience"] == "https://up.example/mcp" # bound to the resource (RFC 8707)
assert d["scope"] == "a"
assert fake.posted["url"] == "https://idp/exchange"
async def test_exchange_no_endpoint_is_misconfigured():
result = await HttpxTokenExchanger().exchange(
_te_config(endpoint=None), SecretStr("jwt"), "https://up/mcp"
)
assert isinstance(result, Error)
assert result.error.tag == "misconfigured"
async def test_exchange_rejected_is_misconfigured(monkeypatch):
_patch(monkeypatch, _FakeResponse(400, {}))
result = await HttpxTokenExchanger().exchange(
_te_config(), SecretStr("jwt"), "https://up/mcp"
)
assert isinstance(result, Error)
assert result.error.tag == "misconfigured"
async def test_exchange_endpoint_down_is_upstream_unavailable(monkeypatch):
_patch(monkeypatch, _FakeResponse(503, {}))
result = await HttpxTokenExchanger().exchange(
_te_config(), SecretStr("jwt"), "https://up/mcp"
)
assert isinstance(result, Error)
assert result.error.tag == "upstream_unavailable"

29
tests/test_litellm/proxy/_experimental/mcp_server/test_v2_resolver_bridge.py
View file

@ -295,3 +295,32 @@ async def test_byok_server_maps_to_byok_key_source():
assert isinstance(spec.config, ApiKeyConfig)
assert isinstance(spec.config.key_source, Byok)
assert spec.config.header_name == "X-API-Key"
async def test_token_exchange_server_maps_to_config():
from litellm.proxy._experimental.mcp_server.v2_resolver_bridge import (
_to_server_spec,
)
from litellm.proxy.gateway.mcp.outbound_credentials.types import TokenExchangeConfig
server = MCPServer(
server_id="obo",
name="obo",
transport=MCPTransport.http,
url="https://up.example/mcp",
auth_type=MCPAuth.oauth2_token_exchange,
client_id="cid",
client_secret="csecret",
token_exchange_endpoint="https://idp/exchange",
audience="https://aud.example",
scopes=["a", "b"],
)
spec = _to_server_spec(server)
assert spec is not None
assert isinstance(spec.config, TokenExchangeConfig)
assert spec.config.token_exchange_endpoint == "https://idp/exchange"
assert spec.config.client_id == "cid"
assert spec.config.client_secret is not None
assert spec.config.client_secret.get_secret_value() == "csecret"
assert spec.config.scopes == ("a", "b")
assert spec.resource == "https://aud.example" # audience preferred for the binding