From 2c620f10ed43bc16c9f5073db2eed6052314b702 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 20 Aug 2026 20:02:09 +0000 Subject: [PATCH] fix(mcp): honor oauth_token_header for the other v2 oauth grants Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../outbound_credentials/adapter.py | 4 +- .../outbound_credentials/resolver.py | 14 +++--- .../mcp_server/outbound_credentials/types.py | 12 +++++ .../outbound_credentials/test_adapter.py | 49 +++++++++++++++++++ .../outbound_credentials/test_resolver.py | 43 ++++++++++++++++ 5 files changed, 115 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 6688d014739..453a9a278ab 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -122,7 +122,7 @@ def _oauth2_spec(server: MCPServer, resource: str) -> ServerSpec | None: return ServerSpec( server_id=server.server_id, resource=resource, - config=AuthorizationCodeConfig(), + config=AuthorizationCodeConfig(token_header=resolve_oauth_token_header(server.oauth_token_header)), ) return None @@ -182,6 +182,7 @@ def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None: client_secret=SecretStr(server.client_secret), token_endpoint_auth_method=server.token_endpoint_auth_method, scopes=tuple(server.scopes or ()), + token_header=resolve_oauth_token_header(server.oauth_token_header), ), ) @@ -240,6 +241,7 @@ def _id_jag_spec(server: MCPServer, resource: str) -> ServerSpec | None: audience=server.audience, resource=server.id_jag_resource, scopes=tuple(server.scopes or ()), + token_header=resolve_oauth_token_header(server.oauth_token_header), ), ) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 85c48f05f09..3e3b85c8ad8 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -145,8 +145,8 @@ class UpstreamCredentialProvider: return await self._token_exchange(subject, server, config) case IdJagConfig() as config: return await self._id_jag(subject, server, config) - case AuthorizationCodeConfig(): - return await self._authorization_code(subject, server) + case AuthorizationCodeConfig() as config: + return await self._authorization_code(subject, server, config) case AwsSigV4Config(): return _not_implemented(AuthSpecKind.aws_sigv4) assert_never(server.config) @@ -284,15 +284,17 @@ class UpstreamCredentialProvider: match await self._exchanged_tokens.get_or_compute(slot, _exchange, fingerprint=fingerprint): case Ok(access_token): - return Ok(StaticHeaderAuth(f"Bearer {access_token}")) + return Ok(StaticHeaderAuth(f"Bearer {access_token}", header_name=config.token_header)) case Error(err): return Error(err) - async def _authorization_code(self, subject: Subject, server: ServerSpec) -> Result[StaticHeaderAuth, CredError]: + async def _authorization_code( + self, subject: Subject, server: ServerSpec, config: AuthorizationCodeConfig + ) -> Result[StaticHeaderAuth, CredError]: token: Final = await self._authz_token(subject, server) if token is None: return Error(CredError.of_unauthorized("Authorization required: complete the OAuth flow for this server.")) - return Ok(StaticHeaderAuth(f"Bearer {token.access_token}", header_name="Authorization")) + return Ok(StaticHeaderAuth(f"Bearer {token.access_token}", header_name=config.token_header)) async def _client_credentials( self, server_id: str, config: ClientCredentialsConfig @@ -332,7 +334,7 @@ class UpstreamCredentialProvider: inbound.get_secret_value(), server, config, tenant_id=subject.tenant_id ): case Ok(token): - return Ok(StaticHeaderAuth(f"Bearer {token.access_token}", header_name="Authorization")) + return Ok(StaticHeaderAuth(f"Bearer {token.access_token}", header_name=config.token_header)) case Error(err): return Error(err) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index 17a5fa36e7e..7636f1ac743 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -168,6 +168,9 @@ class AuthorizationCodeConfig(BaseModel): (RFC 7591), so the common case carries none of the fields below; they are optional manual overrides for IdPs without discovery / DCR. The per-user token is read from the token store at resolve time, not held here. + + `token_header` is the upstream header the user's token is written to, defaulting to + `Authorization`; see `ClientCredentialsConfig.token_header`. """ model_config = ConfigDict(frozen=True) @@ -177,6 +180,7 @@ class AuthorizationCodeConfig(BaseModel): client_secret: SecretStr | None = None authorization_url: str | None = None token_url: str | None = None + token_header: str = DEFAULT_OAUTH_TOKEN_HEADER class ClientCredentialsConfig(BaseModel): @@ -225,6 +229,9 @@ class TokenExchangeConfig(BaseModel): `audience` (rfc8693 only) is optional and sent only when the operator configured one, since both `audience` and `resource` are optional in RFC 8693 and the authorization server applies its own default when neither is sent (fabricating one risks `invalid_target`). + + `token_header` is the upstream header the exchanged token is written to, defaulting to + `Authorization`; see `ClientCredentialsConfig.token_header`. """ model_config = ConfigDict(frozen=True) @@ -237,6 +244,7 @@ class TokenExchangeConfig(BaseModel): client_secret: SecretStr | None = None token_endpoint_auth_method: Literal["client_secret_basic", "client_secret_post"] | None = None scopes: tuple[str, ...] = () + token_header: str = DEFAULT_OAUTH_TOKEN_HEADER class PrivateKeyJwtAuth(BaseModel): @@ -268,6 +276,9 @@ class IdJagConfig(BaseModel): the upstream resource AS (`resource_token_endpoint`) that swaps the assertion for the access token. The gateway authenticates to both endpoints as `client_id` via `client_auth`. Required fields are enforced at construction so a half-configured server cannot reach the arm. + + `token_header` is the upstream header the minted access token is written to, defaulting to + `Authorization`; see `ClientCredentialsConfig.token_header`. """ model_config = ConfigDict(frozen=True) @@ -280,6 +291,7 @@ class IdJagConfig(BaseModel): audience: str | None = None resource: str | None = None scopes: tuple[str, ...] = () + token_header: str = DEFAULT_OAUTH_TOKEN_HEADER class SharedKey(BaseModel): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index 40f137d2c6e..be9be67562b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -216,6 +216,55 @@ def test_client_credentials_token_header_defaults_to_authorization(): assert spec.config.token_header == "Authorization" +def test_authorization_code_carries_the_configured_token_header(): + spec = to_server_spec(_server(auth_type=MCPAuth.oauth2, oauth_token_header="x-upstream-oauth")) + assert spec is not None + assert isinstance(spec.config, AuthorizationCodeConfig) + assert spec.config.token_header == "x-upstream-oauth" + + +def test_token_exchange_carries_the_configured_token_header(): + spec = to_server_spec( + _server( + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + oauth_token_header="x-upstream-oauth", + ) + ) + assert spec is not None + assert isinstance(spec.config, TokenExchangeConfig) + assert spec.config.token_header == "x-upstream-oauth" + + +def test_id_jag_carries_the_configured_token_header(): + spec = to_server_spec(_id_jag_server(oauth_token_header="x-upstream-oauth")) + assert spec is not None + assert isinstance(spec.config, IdJagConfig) + assert spec.config.token_header == "x-upstream-oauth" + + +@pytest.mark.parametrize( + "server", + [ + _server(auth_type=MCPAuth.oauth2), + _server( + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + ), + _id_jag_server(), + ], +) +def test_oauth_modes_default_their_token_header_to_authorization(server): + spec = to_server_spec(server) + assert spec is not None + assert isinstance(spec.config, (AuthorizationCodeConfig, TokenExchangeConfig, IdJagConfig)) + assert spec.config.token_header == "Authorization" + + def test_client_credentials_with_incomplete_grant_fields_is_owned_for_fail_closed(): # An M2M server missing its grant fields is still owned by v2 (spec, not None) so it fails # closed at the source (misconfigured, 500) rather than deferring to v1, which would connect diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index 500ec87f1c5..e0b3fd98590 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -161,6 +161,19 @@ async def test_authorization_code_emits_bearer_for_a_stored_token(): assert _emitted(result.ok)["Authorization"] == "Bearer at-alice" +@pytest.mark.asyncio +async def test_authorization_code_emits_the_stored_token_on_a_configured_custom_header(): + store = _FakeTokenStore({("alice", "s"): OAuthToken(access_token="at-alice")}) + result = await UpstreamCredentialProvider(oauth_token_store=store).resolve_credentials( + Subject(tenant_id="", subject_id="alice"), + _spec(AuthorizationCodeConfig(token_header="x-upstream-oauth")), + ) + assert isinstance(result, Ok) + emitted = _emitted(result.ok) + assert emitted["x-upstream-oauth"] == "Bearer at-alice" + assert "authorization" not in emitted + + @pytest.mark.asyncio async def test_authorization_code_without_token_is_semantically_unauthorized(): result = await UpstreamCredentialProvider(oauth_token_store=_FakeTokenStore({})).resolve_credentials( @@ -283,6 +296,18 @@ async def test_token_exchange_emits_the_exchanged_bearer(): assert exchanger.calls == [("caller-jwt", "acme", "s")] +@pytest.mark.asyncio +async def test_token_exchange_emits_the_exchanged_bearer_on_a_configured_custom_header(): + exchanger = _FakeExchanger(Ok(OAuthToken(access_token="exchanged-at"))) + subject = Subject(tenant_id="acme", subject_id="alice", inbound_token=SecretStr("caller-jwt")) + config = _OBO.model_copy(update={"token_header": "x-upstream-oauth"}) + result = await UpstreamCredentialProvider(token_exchanger=exchanger).resolve_credentials(subject, _spec(config)) + assert isinstance(result, Ok) + emitted = _emitted(result.ok) + assert emitted["x-upstream-oauth"] == "Bearer exchanged-at" + assert "authorization" not in emitted + + @pytest.mark.asyncio async def test_invalidate_credentials_drops_the_exchanged_token_for_the_subject_and_tenant(): exchanger = _FakeExchanger(Ok(OAuthToken(access_token="exchanged-at"))) @@ -509,6 +534,24 @@ async def test_id_jag_runs_both_legs_and_returns_the_leg2_bearer(): assert leg2_params["assertion"] == "the-id-jag" +@pytest.mark.asyncio +async def test_id_jag_emits_the_leg2_bearer_on_a_configured_custom_header(): + endpoint = _FakeTokenEndpoint( + [ + Ok(ExchangedToken(access_token="the-id-jag", expires_in=300)), + Ok(ExchangedToken(access_token="final-access", expires_in=3600)), + ] + ) + config = _id_jag_config().model_copy(update={"token_header": "x-upstream-oauth"}) + result = await UpstreamCredentialProvider(token_endpoint=endpoint).resolve_credentials( + _with_inbound("user-id-token"), _spec(config) + ) + assert isinstance(result, Ok) + emitted = _emitted(result.ok) + assert emitted["x-upstream-oauth"] == "Bearer final-access" + assert "authorization" not in emitted + + @pytest.mark.asyncio async def test_id_jag_without_inbound_token_or_stored_assertion_is_precondition_required_no_http(): endpoint = _FakeTokenEndpoint([])