From d2f9917f02ae255c5af0db251d8ab9d21d63349f Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 25 Jun 2026 17:41:04 -0700 Subject: [PATCH] refactor(mcp): extract the authorization_code arm into a helper Mirror the api_key arm's structure: the inline AuthorizationCodeConfig body moves into _authorization_code(subject, server), keeping resolve_credentials a flat one-line-per-arm dispatch. The helper is annotated with the concrete StaticHeaderAuth it returns rather than the abstract httpx.Auth (which api_key uses) because a new method carrying the unresolved httpx.Auth return would add reportUnknownMemberType; the concrete type is both precise and budget-neutral. --- .../outbound_credentials/resolver.py | 29 +++++++++++-------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index e6a9a90fdd2..87fa66aeab9 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -83,18 +83,7 @@ class UpstreamCredentialProvider: case TokenExchangeConfig(): return _not_implemented(AuthSpecKind.token_exchange) case AuthorizationCodeConfig(): - token = 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 await self._authorization_code(subject, server) case AwsSigV4Config(): return _not_implemented(AuthSpecKind.aws_sigv4) assert_never(server.config) @@ -125,6 +114,22 @@ class UpstreamCredentialProvider: ) assert_never(config.key_source) + async def _authorization_code( + self, subject: Subject, server: ServerSpec + ) -> Result[StaticHeaderAuth, CredError]: + token = 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" + ) + ) + async def _authz_token( self, subject: Subject, server: ServerSpec ) -> OAuthToken | None: