scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

fix(mcp): forward caller's extra headers in true_passthrough auth probe

Browse source 
Motivation:
_probe_upstream_auth predicts whether a true_passthrough MCP server will
accept the caller's request before a stateful session opens, so LiteLLM can
return a 401 up front instead of masking the rejection as an empty tool
list. The probe only ever sent Accept and Authorization upstream. When an
upstream identifies the caller through a different header the admin already
configured in server.extra_headers (a gateway-issued platform header, for
example), the probe always looked anonymous to it, so it answered 401 for
every caller including ones with a valid session, and every connect got the
pre-session challenge regardless of whether the caller was actually
authorized.

Approach:
Added _extra_headers_for_probe, which takes the caller's raw request headers
and the server's configured extra_headers list and returns just the values
for those header names (case-insensitively matched), always excluding
authorization since the probe already sets that separately. _probe_upstream_auth
gained an extra_headers parameter that merges into the outgoing probe
request. The is_true_passthrough branch in
_raise_preemptive_401_for_unauthenticated_servers now passes the caller's
configured extra headers into the probe, using the raw_headers parameter
that function already receives. The is_oauth_passthrough /
_check_passthrough_upstream_auth path is untouched, matching the issue's
stated scope.

Validation:
PYTHONPATH=. pytest tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py -v
passes all 240 tests (235 existing + 5 new). Reverting only the server.py
change (git stash) makes the 5 new/changed tests fail, including one that
raises the exact HTTPException: 401 the issue describes, confirming they are
a real regression guard rather than a structural check. ruff check, ruff
format --check, and basedpyright were run on both changed files: no new
basedpyright errors were introduced in server.py (236 before and after), and
the small increase in the test file (reportPrivateUsage / reportAny noise
from calling internal helpers directly, matching this test file's existing
whitebox pattern) stays well under the repo's basedpyright-code-budget.json
limits for those rules.

Disclosure: dependencies were installed into a local virtualenv rather than
through the maturin/editable install this repo normally uses, because of
disk-space limits in this environment; litellm was imported directly from
source with PYTHONPATH, and the Rust bridge already no-ops gracefully when
its compiled extension is absent. No live MCP server was reachable to
reproduce the issue end to end; the fix and its tests are demonstrated at
the unit level described above.

Report: https://github.com/BerriAI/litellm/issues/40099
Signed-off-by: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com>
Assisted-by: claude-sonnet-5 (via Claude Code)
This commit is contained in:
Pujitha Paladugu 2026-09-18 05:37:56 -07:00
parent 615ed7900f
commit d5c9d17d47
No known key found for this signature in database
2 changed files with 203 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

29
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -1815,7 +1815,11 @@ if MCP_AVAILABLE:
)
},
)
upstream_status, upstream_www_authenticate = await _probe_upstream_auth(server.url or "", "")
upstream_status, upstream_www_authenticate = await _probe_upstream_auth(
server.url or "",
"",
extra_headers=_extra_headers_for_probe(server, raw_headers),
)
if upstream_status == 401 and upstream_www_authenticate:
raise HTTPException(
status_code=401,
@ -1851,10 +1855,27 @@ if MCP_AVAILABLE:
return None
return _get_authorization_header_from_scope(scope)
def _extra_headers_for_probe(server: MCPServer, raw_headers: Mapping[str, str] | None) -> dict[str, str] | None:
"""Caller's values for headers this upstream is configured to read
(``server.extra_headers``), minus ``authorization`` which the probe
sets itself. Lets the pre-session probe re-run the same auth path a
real client request would take instead of an anonymous one.
"""
if not server.extra_headers or not raw_headers:
return None
normalized: Final = {k.lower(): v for k, v in raw_headers.items()}
forwarded: Final = {
header: normalized[header.lower()]
for header in server.extra_headers
if header.lower() != "authorization" and header.lower() in normalized
}
return forwarded or None
async def _probe_upstream_auth(
url: str,
auth_header: str,
timeout: float = 5.0,
extra_headers: dict[str, str] | None = None,
) -> tuple[int, str | None]:
"""JSON-RPC initialize-probe the upstream URL to check whether the token is accepted.
@ -1863,6 +1884,11 @@ if MCP_AVAILABLE:
Fails-open with (200, None) on network errors so a transient hiccup
does not block valid requests.
``extra_headers`` carries any additional headers this upstream is
configured to read (``server.extra_headers``) so the probe predicts the
outcome of the same request a real client would send, not an anonymous
one, when identity travels in a header other than ``Authorization``.
Uses the public ``AsyncHTTPHandler.post()`` interface and catches
``httpx.HTTPStatusError`` separately so the 401/403 we want to surface
is not swallowed by the broad fail-open ``except Exception`` below.
@ -1887,6 +1913,7 @@ if MCP_AVAILABLE:
probe_headers: Final = {
"Accept": "application/json, text/event-stream",
**({"Authorization": auth_header} if auth_header else {}),
**(extra_headers or {}),
}
try:
resp: Final = await client.post(

175
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py
View file

@ -6918,6 +6918,181 @@ async def test_probe_upstream_auth_fails_open_on_network_error():
assert www_auth is None
@pytest.mark.asyncio
async def test_probe_upstream_auth_merges_extra_headers_into_request():
"""extra_headers must reach the outgoing probe request, so an upstream that
identifies the caller via a configured header (not Authorization) sees the
same request a real client would send."""
from litellm.proxy._experimental.mcp_server.server import _probe_upstream_auth
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {}
mock_client = MagicMock()
mock_client.post = AsyncMock(return_value=mock_response)
with patch(
"litellm.proxy._experimental.mcp_server.server.get_async_httpx_client",
return_value=mock_client,
):
await _probe_upstream_auth("http://upstream/mcp", "", extra_headers={"X-Platform-User": "alice"})
_, kwargs = mock_client.post.call_args
assert kwargs["headers"]["X-Platform-User"] == "alice"
assert "Authorization" not in kwargs["headers"]
def test_extra_headers_for_probe_filters_to_configured_non_authorization_headers():
"""Only forwards values for header names the admin listed in
server.extra_headers, matched case-insensitively, and always drops
Authorization (the probe sets that itself)."""
from litellm.proxy._experimental.mcp_server.server import _extra_headers_for_probe
server = MCPServer(
server_id="tp-1",
name="tp_server",
url="https://upstream.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.true_passthrough,
extra_headers=["X-Platform-User", "Authorization"],
)
raw_headers = {
"x-platform-user": "alice",
"authorization": "Bearer caller-key",
"content-type": "application/json",
}
assert _extra_headers_for_probe(server, raw_headers) == {"X-Platform-User": "alice"}
def test_extra_headers_for_probe_returns_none_when_nothing_to_forward():
from litellm.proxy._experimental.mcp_server.server import _extra_headers_for_probe
no_extra_headers_server = MCPServer(
server_id="tp-2",
name="tp_server_2",
url="https://upstream.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.true_passthrough,
extra_headers=None,
)
assert _extra_headers_for_probe(no_extra_headers_server, {"x-platform-user": "alice"}) is None
unmatched_server = MCPServer(
server_id="tp-3",
name="tp_server_3",
url="https://upstream.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.true_passthrough,
extra_headers=["X-Platform-User"],
)
assert _extra_headers_for_probe(unmatched_server, {"content-type": "application/json"}) is None
class TestTruePassthroughProbeForwardsCallerHeaders:
"""Regression test for the true_passthrough pre-session probe sending an
anonymous request to upstreams that identify callers via a header other
than Authorization (the header the admin already listed in
server.extra_headers). Before the fix, _probe_upstream_auth never received
that header, so the probe's answer described a caller that doesn't exist."""
ALIAS: Final = "tp_platform_header"
def _server(self, extra_headers: list[str] | None) -> MCPServer:
return MCPServer(
server_id="tp-platform-header",
name=self.ALIAS,
alias=self.ALIAS,
url="https://upstream.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.true_passthrough,
extra_headers=extra_headers,
)
def _scope(self) -> dict[str, object]:
return {"type": "http", "method": "POST", "path": f"/mcp/{self.ALIAS}", "headers": []}
@pytest.mark.asyncio
async def test_probe_receives_caller_platform_header_and_does_not_challenge(self):
from litellm.proxy._experimental.mcp_server import server as server_module
server = self._server(["X-Platform-User"])
async def fake_upstream_probe(
url: str, auth_header: str, timeout: float = 5.0, extra_headers: dict[str, str] | None = None
) -> tuple[int, str | None]:
# Simulates an upstream that answers 401 for anyone it cannot
# identify by the platform header -- exactly the case the issue
# describes, where Authorization alone is not the caller's identity.
if not extra_headers or extra_headers.get("X-Platform-User") != "alice":
return 401, 'Bearer realm="upstream"'
return 200, None
with (
patch.object(
server_module.global_mcp_server_manager,
"get_mcp_server_by_name",
return_value=server,
),
patch(
"litellm.proxy._experimental.mcp_server.server._probe_upstream_auth",
side_effect=fake_upstream_probe,
),
):
# Must not raise: the caller's platform header reaches the probe,
# so the upstream recognizes them and the pre-session challenge is
# skipped, matching what a real session would get.
await server_module._raise_preemptive_401_for_unauthenticated_servers(
scope=self._scope(),
mcp_servers=[self.ALIAS],
oauth2_headers=None,
mcp_server_auth_headers=None,
user_api_key_auth=UserAPIKeyAuth(),
client_ip=None,
raw_headers={"X-Platform-User": "alice"},
)
@pytest.mark.asyncio
async def test_probe_never_forwards_authorization_via_extra_headers(self):
"""Authorization listed in server.extra_headers must not be replayed a
second time through extra_headers; the probe call already sets its own
(empty) auth_header for true_passthrough."""
from litellm.proxy._experimental.mcp_server import server as server_module
server = self._server(["Authorization", "X-Platform-User"])
captured: dict[str, dict[str, str] | None] = {}
async def fake_upstream_probe(
url: str, auth_header: str, timeout: float = 5.0, extra_headers: dict[str, str] | None = None
) -> tuple[int, str | None]:
captured["extra_headers"] = extra_headers
return 200, None
with (
patch.object(
server_module.global_mcp_server_manager,
"get_mcp_server_by_name",
return_value=server,
),
patch(
"litellm.proxy._experimental.mcp_server.server._probe_upstream_auth",
side_effect=fake_upstream_probe,
),
):
await server_module._raise_preemptive_401_for_unauthenticated_servers(
scope=self._scope(),
mcp_servers=[self.ALIAS],
oauth2_headers=None,
mcp_server_auth_headers=None,
user_api_key_auth=UserAPIKeyAuth(),
client_ip=None,
raw_headers={"Authorization": "Bearer caller-key", "X-Platform-User": "alice"},
)
assert captured["extra_headers"] == {"X-Platform-User": "alice"}
def test_get_forwarded_auth_from_scope_extracts_header():
"""Returns Authorization value when x-litellm-api-key is also present."""
from litellm.proxy._experimental.mcp_server.server import (