diff --git a/litellm/proxy/_experimental/mcp_server/v2_port_bodies.py b/litellm/proxy/_experimental/mcp_server/v2_port_bodies.py index 8216f209ac7..8330a0a4530 100644 --- a/litellm/proxy/_experimental/mcp_server/v2_port_bodies.py +++ b/litellm/proxy/_experimental/mcp_server/v2_port_bodies.py @@ -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: diff --git a/litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py b/litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py index 3641911d9f2..db2b08e1e65 100644 --- a/litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py +++ b/litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py @@ -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 diff --git a/litellm/proxy/gateway/mcp/outbound_credentials/types.py b/litellm/proxy/gateway/mcp/outbound_credentials/types.py index 6e45f7162a8..fc36334bb6e 100644 --- a/litellm/proxy/gateway/mcp/outbound_credentials/types.py +++ b/litellm/proxy/gateway/mcp/outbound_credentials/types.py @@ -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): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_port_bodies.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_port_bodies.py index 4c4ad7af6bf..b28a1c54d66 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_port_bodies.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_port_bodies.py @@ -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" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_resolver_bridge.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_resolver_bridge.py index 0f3d1b542be..657ac116149 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_resolver_bridge.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_resolver_bridge.py @@ -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