From d8b3d6d345d960892921cf924f1bf1631ff61585 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 24 Jun 2026 10:54:29 -0700 Subject: [PATCH] feat(mcp): let CredError.of_unauthorized carry a 401 challenge The unauthorized case becomes a structured Unauthorized (detail + optional WWW-Authenticate header + optional structured body) instead of a bare string, and raise_public emits the header and body when present. This lets a mode reproduce a rich 401 challenge (e.g. BYOK's provisioning prompt) through the generic resolver edge. of_unauthorized's new params are keyword-only and default to None, so existing callers and the summary string are unchanged. --- .../outbound_credentials/adapter.py | 11 ++++++- .../mcp_server/outbound_credentials/types.py | 32 ++++++++++++++++--- .../outbound_credentials/test_adapter.py | 23 +++++++++++++ 3 files changed, 60 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 39db2314aee..91876df762f 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -126,7 +126,16 @@ def raise_public(error: CredError) -> NoReturn: """Map a resolver CredError onto the proxy's public HTTP contract. The one edge that raises.""" match error.tag: case "unauthorized": - raise HTTPException(status_code=401, detail=error.summary) + challenge = error.unauthorized + raise HTTPException( + status_code=401, + detail=challenge.body if challenge.body is not None else error.summary, + headers=( + {"WWW-Authenticate": challenge.www_authenticate} + if challenge.www_authenticate + else None + ), + ) case "misconfigured": raise HTTPException(status_code=500, detail=error.summary) case "upstream_unavailable": diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index 2088dc77252..f0e81272ffd 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -26,7 +26,7 @@ union (see `result.py`), not `expression.Result`. from __future__ import annotations from enum import Enum -from typing import Annotated, Literal +from typing import Annotated, Literal, Mapping, Optional from expression import case, tag, tagged_union from pydantic import BaseModel, ConfigDict, Field, SecretStr @@ -59,6 +59,19 @@ class AuthSpecKind(str, Enum): aws_sigv4 = "aws_sigv4" # AWS SigV4 per-request signing (e.g. Bedrock AgentCore) +class Unauthorized(BaseModel): + """A 401 plus the optional challenge a client needs to recover. + + ``detail`` is the human message; ``www_authenticate`` and ``body`` carry a scheme-specific + challenge (e.g. BYOK's provisioning prompt) so the edge can reproduce it verbatim. + """ + + model_config = ConfigDict(frozen=True) + detail: str + www_authenticate: Optional[str] = None + body: Optional[Mapping[str, str]] = None + + @tagged_union(frozen=True) class CredError: """Why a credential could not be produced. Fail-closed: an arm yields this or an `httpx.Auth`. @@ -76,7 +89,7 @@ class CredError: "not_implemented", ] = tag() - unauthorized: str = ( + unauthorized: Unauthorized = ( case() ) # no usable credential for this (subject, server) -> 401 challenge misconfigured: str = ( @@ -96,8 +109,17 @@ class CredError: ) # the declared mode's resolver arm is not built yet -> 501 (not operator error) @staticmethod - def of_unauthorized(detail: str) -> CredError: - return CredError(unauthorized=detail) + def of_unauthorized( + detail: str, + *, + www_authenticate: Optional[str] = None, + body: Optional[Mapping[str, str]] = None, + ) -> CredError: + return CredError( + unauthorized=Unauthorized( + detail=detail, www_authenticate=www_authenticate, body=body + ) + ) @staticmethod def of_misconfigured(detail: str) -> CredError: @@ -125,7 +147,7 @@ class CredError: # only while that stays true (a `case _` would defeat reportMatchNotExhaustive). match self.tag: case "unauthorized": - return f"unauthorized: {self.unauthorized}" + return f"unauthorized: {self.unauthorized.detail}" case "misconfigured": return f"misconfigured: {self.misconfigured}" case "upstream_unavailable": 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 5481d60a22a..594e9dcc969 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 @@ -139,3 +139,26 @@ def test_raise_public_maps_each_error_to_its_status(error, status): with pytest.raises(HTTPException) as exc_info: raise_public(error) assert exc_info.value.status_code == status + + +def test_raise_public_emits_unauthorized_challenge(): + body = {"error": "byok_auth_required", "server_id": "s1"} + error = CredError.of_unauthorized( + "needs key", www_authenticate='Bearer resource_metadata="/x"', body=body + ) + with pytest.raises(HTTPException) as exc_info: + raise_public(error) + exc = exc_info.value + assert exc.status_code == 401 + assert exc.detail == body + assert exc.headers is not None + assert exc.headers["WWW-Authenticate"] == 'Bearer resource_metadata="/x"' + + +def test_raise_public_plain_unauthorized_has_no_challenge(): + with pytest.raises(HTTPException) as exc_info: + raise_public(CredError.of_unauthorized("nope")) + exc = exc_info.value + assert exc.status_code == 401 + assert exc.detail == "unauthorized: nope" + assert exc.headers is None