fix(mcp): probe legacy delegates with only the caller's token, as v1.102 did

The probe is the check on a delegate caller's own token, so sending the server's static headers with it
let a static Authorization stand in for that token. On a delegate server that also sets a static
Authorization, any junk bearer passed the probe and reached the upstream with the operator's token

The probe goes back to sending only the caller's token for every server, as v1.102 and main do. This
drops the static header probing added for the earlier review, so an upstream that needs a companion
static header next to the token gets the same probe it got on v1.102
This commit is contained in:
shrey kharbanda 2026-10-28 21:21:55 -07:00
parent bfeb48bb65
commit 0cade620df
3 changed files with 16 additions and 90 deletions

View file

@ -4176,12 +4176,6 @@ class MCPServerManager:
referenced: Final = collect_env_var_references(strings=static_headers.values())
return bool(referenced & user_var_names)
async def resolved_static_headers(
self, server: MCPServer, user_api_key_auth: UserAPIKeyAuth | None
) -> dict[str, str] | None:
"""``server.static_headers`` interpolated exactly as a tool call sends them."""
return await self._resolve_static_headers_with_env_vars(server, user_api_key_auth)
async def _resolve_static_headers_with_env_vars(
self,
server: MCPServer,

View file

@ -2104,7 +2104,6 @@ if MCP_AVAILABLE:
url: str,
auth_header: str,
timeout: float = 5.0,
headers: Mapping[str, str] | None = None,
) -> tuple[int, str | None]:
"""JSON-RPC initialize-probe the upstream URL to check whether the token is accepted.
@ -2137,7 +2136,6 @@ if MCP_AVAILABLE:
probe_headers: Final = {
"Accept": "application/json, text/event-stream",
**({"Authorization": auth_header} if auth_header else {}),
**(headers or {}),
}
try:
resp: Final = await client.post(
@ -2156,30 +2154,6 @@ if MCP_AVAILABLE:
verbose_logger.debug("_probe_upstream_auth: probe to %s failed (%s), allowing request through", url, exc)
return 200, None
async def _probe_with_outbound_headers(
server: MCPServer,
auth_header: str,
caller_headers: Mapping[str, str],
user_api_key_auth: UserAPIKeyAuth | None,
) -> tuple[int, str | None]:
"""Probe with the headers a real tool call sends: the caller's prepared headers with the server's
resolved static headers on top, so a static ``Authorization`` wins here too. When those cannot be
resolved yet, such as a per-user variable the caller has not set, the probe fails open like a network
error and leaves the error to the call."""
try:
static_headers: Final = await global_mcp_server_manager.resolved_static_headers(server, user_api_key_auth)
except Exception as exc: # noqa: BLE001 # the probe fails open on anything it cannot reproduce
verbose_logger.debug("Skipping upstream auth probe for %s (%s)", server.name, type(exc).__name__)
return 200, None
outbound: Final = merge_mcp_headers(extra_headers=caller_headers, static_headers=static_headers) or {}
authorization: Final = next(
(value for name, value in reversed(outbound.items()) if name.lower() == "authorization"), auth_header
)
companion_headers: Final = {name: value for name, value in outbound.items() if name.lower() != "authorization"}
if not companion_headers:
return await _probe_upstream_auth(server.url or "", authorization)
return await _probe_upstream_auth(server.url or "", authorization, headers=companion_headers)
async def _check_passthrough_upstream_auth(
scope: Scope,
user_api_key_auth: UserAPIKeyAuth | None,
@ -2257,13 +2231,8 @@ if MCP_AVAILABLE:
)
for srv, (auth, extra) in prepared_headers
)
probe_targets: Final[tuple[tuple[MCPServer, str, str, Mapping[str, str]], ...]] = tuple(
(
srv,
auth_header,
srv.name if srv.is_oauth_passthrough else single_requested_name or srv.name,
headers or {},
)
probe_targets: Final[tuple[tuple[MCPServer, str, str], ...]] = tuple(
(srv, auth_header, srv.name if srv.is_oauth_passthrough else single_requested_name or srv.name)
for srv, headers in probe_headers
if (
auth_header := next(
@ -2275,12 +2244,9 @@ if MCP_AVAILABLE:
return
probe_results: Final = await asyncio.gather(
*[
_probe_with_outbound_headers(srv, auth_header, headers, user_api_key_auth)
for srv, auth_header, _, headers in probe_targets
]
*[_probe_upstream_auth(srv.url or "", auth_header) for srv, auth_header, _ in probe_targets]
)
for (_srv, _, challenge_server_name, _), (probe_status, _) in zip(probe_targets, probe_results):
for (_srv, _, challenge_server_name), (probe_status, _) in zip(probe_targets, probe_results):
if probe_status == 401:
# Token is missing or expired: keep pass-through clients on the
# protected-resource discovery flow so they re-authorize against

View file

@ -7437,51 +7437,13 @@ async def test_legacy_delegate_multi_server_scope_is_not_failed_by_one_upstream_
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.parametrize(
"static_headers, expected_authorization",
[
({"X-Tenant": "tenant-1"}, "Bearer expired-upstream-token"),
({"X-Tenant": "tenant-1", "Authorization": "Bearer server-wide-token"}, "Bearer server-wide-token"),
],
ids=["companion-header", "static-authorization-wins"],
)
async def test_legacy_delegate_probe_sends_the_headers_a_tool_call_sends(
static_headers: dict[str, str], expected_authorization: str
) -> None:
"""An upstream that needs a static header next to the token sees it in the probe too, so a token the
upstream accepts on real calls is not rejected before the session opens."""
import respx
server: Final = _delegate_auth_mcp_server().model_copy(update={"static_headers": static_headers})
with (
_patch_delegate_resolver(server, "delegate_test"),
patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server]),
),
respx.mock as upstream,
):
probe: Final = upstream.post(server.url).mock(
side_effect=lambda request: httpx.Response(200 if request.headers.get("X-Tenant") == "tenant-1" else 401)
)
assert await _expired_delegate_token_preflight(["delegate_test"]) is None
assert probe.call_count == 1
assert probe.calls[0].request.headers["Authorization"] == expected_authorization
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport")
async def test_legacy_delegate_probe_steps_aside_when_static_headers_need_unset_user_values() -> None:
"""A static header built from a per-user value the caller has not set cannot be reproduced, so the
probe lets the request through and the tool call reports the missing value."""
async def test_legacy_delegate_probe_checks_the_callers_token_not_a_static_authorization() -> None:
"""The probe is what checks a delegate caller's own token, so a static Authorization the server sends on
tool calls cannot stand in for it: a junk token still gets the re-authorization challenge."""
import respx
server: Final = _delegate_auth_mcp_server().model_copy(
update={
"static_headers": {"X-Tenant-Token": "${TENANT_TOKEN}"},
"env_vars": [{"name": "TENANT_TOKEN", "value": "", "scope": "user"}],
}
update={"static_headers": {"Authorization": "Bearer operator-token"}}
)
with (
_patch_delegate_resolver(server, "delegate_test"),
@ -7491,12 +7453,16 @@ async def test_legacy_delegate_probe_steps_aside_when_static_headers_need_unset_
),
respx.mock as upstream,
):
probe: Final = upstream.post(server.url).respond(
401, headers={"www-authenticate": 'Bearer error="invalid_token"'}
upstream.post(server.url).mock(
side_effect=lambda request: httpx.Response(
200 if request.headers["Authorization"] == "Bearer operator-token" else 401
)
)
outcome: Final = await asyncio.gather(
_expired_delegate_token_preflight(["delegate_test"]), return_exceptions=True
)
assert await _expired_delegate_token_preflight(["delegate_test"]) is None
assert probe.call_count == 0
assert getattr(outcome[0], "status_code", None) == 401
@pytest.mark.asyncio