From d0ee1109d22d0b5a571c083b61b42950754be514 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 16 Jul 2026 23:09:46 -0700 Subject: [PATCH] fix(mcp): auth scan walks past non-auth responses in the exception tree The consolidation regressed the pre-existing walker semantics: _extract_upstream_auth_failure used to keep scanning until it found a 401/403, while the consolidated helper took the first response of any status and then tested it, so a causal 401 sitting behind an unrelated 5xx (retry attempts, multi-stream task groups) was misclassified as upstream_error and its challenge lost on the listing, tool-call, and probe paths. The traversal is now an iterator in deliberate order and each consumer applies its predicate over the stream: the auth scan takes the first 401/403 even behind non-auth responses, generic classification takes the first response, and classify_list_exception derives its auth arm from the same scan so the carrier choice and the classification can never disagree --- .../mcp_server/faults/list_outcomes.py | 49 ++++++++------ .../mcp_server/faults/test_list_outcomes.py | 64 +++++++++++++++++++ 2 files changed, 94 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py index 96ff9443126..6f27c1c0472 100644 --- a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py +++ b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py @@ -10,6 +10,7 @@ becomes an outcome, never a second failure. from __future__ import annotations +from collections.abc import Iterator from typing import Literal, NamedTuple, NoReturn, TypeAlias import httpx @@ -61,12 +62,14 @@ class AggregateToolListing(NamedTuple): outcomes: dict[str, ServerOutcome] -def _find_upstream_response(exc: BaseException) -> httpx.Response | None: - """Walk the exception tree (``__cause__``/``__context__``/ExceptionGroup members) for an - ``httpx.Response``, mirroring how upstream failures surface through the MCP SDK's task groups. - Explicit links are searched first: each node's ``raise ... from`` cause, then group members in - raise order, then the incidental ``__context__`` chain, so a response raised while handling the - real failure can never shadow the response on the explicit causal chain.""" +def _iter_upstream_responses(exc: BaseException) -> Iterator[httpx.Response]: + """Yield every ``httpx.Response`` in the exception tree (``__cause__``/``__context__``/ + ExceptionGroup members) in deliberate order, mirroring how upstream failures surface through the + MCP SDK's task groups. Explicit links come first: each node's ``raise ... from`` cause, then + group members in raise order, then the incidental ``__context__`` chain, so a response raised + while handling the real failure can never shadow one on the explicit causal chain. Consumers + apply their own predicate over the stream: selecting the first response and THEN testing it + would miss a causal auth response sitting behind an unrelated earlier one.""" seen: set[int] = set() stack = [exc] while stack: @@ -76,7 +79,7 @@ def _find_upstream_response(exc: BaseException) -> httpx.Response | None: seen.add(id(current)) response = getattr(current, "response", None) if isinstance(response, httpx.Response): - return response + yield response if current.__context__ is not None: stack.append(current.__context__) exceptions = getattr(current, "exceptions", None) @@ -84,17 +87,22 @@ def _find_upstream_response(exc: BaseException) -> httpx.Response | None: stack.extend(reversed(exceptions)) if current.__cause__ is not None: stack.append(current.__cause__) - return None + + +def _find_upstream_response(exc: BaseException) -> httpx.Response | None: + return next(_iter_upstream_responses(exc), None) def upstream_auth_challenge(exc: BaseException) -> tuple[int, str | None] | None: - """The upstream 401/403 and its ``WWW-Authenticate`` challenge, both read from the SAME response - the deliberate-order traversal selects, so the status that picks the carrier channel and the - challenge that rides with it can never come from two different responses in the tree.""" - response = _find_upstream_response(exc) - if response is None or response.status_code not in (401, 403): - return None - return response.status_code, response.headers.get("www-authenticate") + """The first upstream 401/403 in deliberate order and its ``WWW-Authenticate`` challenge, both + read from the SAME response, so the status that picks the carrier channel and the challenge that + rides with it can never come from two different responses in the tree. Non-auth responses do not + end the scan: a causal 401 behind an unrelated 5xx must still be found, or the client never + receives the challenge it needs to re-authenticate.""" + for response in _iter_upstream_responses(exc): + if response.status_code in (401, 403): + return response.status_code, response.headers.get("www-authenticate") + return None def raise_classified_list_failure( @@ -131,12 +139,15 @@ def classify_list_exception(exc: BaseException) -> ServerListFault: return ServerListFault(tag="timeout") if isinstance(exc, ConnectionError): return ServerListFault(tag="unreachable") + auth = upstream_auth_challenge(exc) + if auth is not None: + status_code, _ = auth + return ServerListFault( + tag="forbidden" if status_code == 403 else "auth_required", + status_code=status_code, + ) response = _find_upstream_response(exc) if response is not None: - if response.status_code == 401: - return ServerListFault(tag="auth_required", status_code=401) - if response.status_code == 403: - return ServerListFault(tag="forbidden", status_code=403) return ServerListFault(tag="upstream_error", status_code=response.status_code) if isinstance(exc, (httpx.TimeoutException,)): return ServerListFault(tag="timeout") diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py index 1987e42f69f..cb27e992ecb 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py @@ -174,3 +174,67 @@ def test_raise_classified_list_failure_routes_auth_to_upstream_auth_error(): with pytest.raises(MCPServerListError) as fault_info: raise_classified_list_failure(RuntimeError("boom"), "srv") assert fault_info.value.fault.tag == "internal" + + +def test_causal_auth_behind_unrelated_response_is_still_found(): + """The auth scan must not end at the first response of any status: a causal 401 sitting deeper + in the tree than an unrelated 5xx (retry attempts, multi-stream task groups) must still surface + with its challenge, or the client is told upstream_error and never re-authenticates.""" + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import upstream_auth_challenge + + deep_auth = httpx.HTTPStatusError( + "auth", + request=httpx.Request("POST", "https://mcp.example.com/mcp"), + response=httpx.Response( + 401, + headers={"www-authenticate": "Bearer realm=upstream"}, + request=httpx.Request("POST", "https://mcp.example.com/mcp"), + ), + ) + earlier_5xx = httpx.HTTPStatusError( + "flaky attempt", + request=httpx.Request("POST", "https://mcp.example.com/mcp"), + response=httpx.Response(500, request=httpx.Request("POST", "https://mcp.example.com/mcp")), + ) + earlier_5xx.__cause__ = deep_auth + wrapper = RuntimeError("fetch failed") + wrapper.__cause__ = earlier_5xx + + result = upstream_auth_challenge(wrapper) + assert result is not None + assert result == (401, "Bearer realm=upstream") + + +def test_classification_agrees_with_auth_scan_on_nested_auth(): + """classify_list_exception derives its auth arm from the same scan as the carrier choice-point, + so a nested 401 behind a 5xx classifies auth_required, never upstream_error(500).""" + deep_auth = httpx.HTTPStatusError( + "auth", + request=httpx.Request("POST", "https://mcp.example.com/mcp"), + response=httpx.Response(401, request=httpx.Request("POST", "https://mcp.example.com/mcp")), + ) + earlier_5xx = httpx.HTTPStatusError( + "flaky attempt", + request=httpx.Request("POST", "https://mcp.example.com/mcp"), + response=httpx.Response(500, request=httpx.Request("POST", "https://mcp.example.com/mcp")), + ) + earlier_5xx.__cause__ = deep_auth + wrapper = RuntimeError("fetch failed") + wrapper.__cause__ = earlier_5xx + + fault = classify_list_exception(wrapper) + assert fault.tag == "auth_required" + assert fault.status_code == 401 + + +def test_pure_non_auth_response_still_classifies_upstream_error(): + """With no auth response anywhere in the tree, the first response in deliberate order still + drives the generic upstream_error classification.""" + exc = httpx.HTTPStatusError( + "boom", + request=httpx.Request("POST", "https://mcp.example.com/mcp"), + response=httpx.Response(502, request=httpx.Request("POST", "https://mcp.example.com/mcp")), + ) + fault = classify_list_exception(exc) + assert fault.tag == "upstream_error" + assert fault.status_code == 502