Merge pull request #33183 from BerriAI/litellm_mcp_walker_consolidation

refactor(mcp): consolidate exception-tree walkers into one shared faults traversal

litellm/proxy/_experimental/mcp_server/faults/__init__.py

@@ -15,6 +15,7 @@ from litellm.proxy._experimental.mcp_server.faults.render_oauth import (
     dcr_fault_detail,
     render_token_fault,
 )
+from litellm.proxy._experimental.mcp_server.faults.traversal import iter_exception_tree
 from litellm.proxy._experimental.mcp_server.faults.types import (
     CallerRejected,
     CredentialSource,
@@ -34,5 +35,6 @@ __all__ = [
     "classify_upstream_dcr_rejection",
     "classify_upstream_token_rejection",
     "dcr_fault_detail",
+    "iter_exception_tree",
     "render_token_fault",
 ]

litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py

@@ -22,6 +22,7 @@ from litellm.proxy._experimental.mcp_server.exceptions import (
     MCPServerListError,
     MCPUpstreamAuthError,
 )
+from litellm.proxy._experimental.mcp_server.faults.traversal import iter_exception_tree
 
 ListFaultCategory: TypeAlias = Literal[
     "auth_required",
@@ -63,30 +64,16 @@ class AggregateToolListing(NamedTuple):
 
 
 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:
-        current = stack.pop()
-        if id(current) in seen:
-            continue
-        seen.add(id(current))
+    """Yield every ``httpx.Response`` in the exception tree, in the shared traversal's deliberate
+    order (explicit causes first, ExceptionGroup members in raise order, the incidental
+    ``__context__`` chain last), 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."""
+    for current in iter_exception_tree(exc):
         response = getattr(current, "response", None)
         if isinstance(response, httpx.Response):
             yield response
-        if current.__context__ is not None:
-            stack.append(current.__context__)
-        exceptions = getattr(current, "exceptions", None)
-        if isinstance(exceptions, tuple):
-            stack.extend(reversed(exceptions))
-        if current.__cause__ is not None:
-            stack.append(current.__cause__)
 
 
 def _find_upstream_response(exc: BaseException) -> httpx.Response | None:

litellm/proxy/_experimental/mcp_server/faults/traversal.py

@@ -0,0 +1,35 @@
+"""Shared exception-tree traversal for fault classification.
+
+Failures cross the MCP SDK's anyio task groups wrapped in ``ExceptionGroup``s and chained through
+``raise ... from`` causes, so every classifier that needs an exception buried in the tree (an
+upstream ``httpx.Response``, a context-window overflow) has to walk the same shapes. One traversal
+with one deliberate order keeps blame assignment consistent across classifiers: explicit links are
+searched before incidental ones, so an exception raised while handling the real failure can never
+shadow the failure itself.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+
+
+def iter_exception_tree(exc: BaseException) -> Iterator[BaseException]:
+    """Yield ``exc`` and every exception reachable from it, explicit links first: each node's
+    ``raise ... from`` cause subtree, then ``ExceptionGroup`` members in raise order, then the
+    incidental ``__context__`` chain last. Cycle-safe via identity tracking, and iterative so a
+    deep chain cannot overflow the interpreter stack."""
+    seen: set[int] = set()
+    stack = [exc]
+    while stack:
+        current = stack.pop()
+        if id(current) in seen:
+            continue
+        seen.add(id(current))
+        yield current
+        if current.__context__ is not None:
+            stack.append(current.__context__)
+        exceptions = getattr(current, "exceptions", None)
+        if isinstance(exceptions, tuple):
+            stack.extend(reversed(exceptions))
+        if current.__cause__ is not None:
+            stack.append(current.__cause__)

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -687,7 +687,7 @@ def _extract_upstream_auth_failure(
 ) -> Optional[tuple[int, Optional[str]]]:
     """The upstream 401/403 and its ``WWW-Authenticate`` header from the exception tree, or ``None``.
 
-    Delegates to the shared traversal in ``faults.list_outcomes`` so every consumer (tool listing,
+    Delegates to the shared traversal in ``faults`` so every consumer (tool listing,
     tool calls, the connect-time probe) selects the same response with the same deliberate order:
     explicit ``raise ... from`` causes first, ExceptionGroup members in raise order, the incidental
     ``__context__`` chain last. A response raised while handling the real failure can therefore never

litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py

@@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional
 from litellm._logging import verbose_logger
 from litellm.exceptions import ContextWindowExceededError
 from litellm.litellm_core_utils.exception_mapping_utils import ExceptionCheckers
+from litellm.proxy._experimental.mcp_server.faults import iter_exception_tree
 from litellm.proxy._experimental.mcp_server.utils import MCP_TOOL_PREFIX_SEPARATOR
 
 if TYPE_CHECKING:
@@ -34,18 +35,15 @@ class SemanticToolFilterContextWindowError(Exception):
         )
 
 
-def _is_context_window_error(error: Optional[BaseException], max_depth: int = 5) -> bool:
-    """Detect a context-window overflow anywhere in an exception's cause chain."""
-    current = error
-    for _ in range(max_depth):
-        if current is None:
-            return False
-        if isinstance(current, ContextWindowExceededError):
-            return True
-        if ExceptionCheckers.is_error_str_context_window_exceeded(str(current)):
-            return True
-        current = current.__cause__ or current.__context__
-    return False
+def _is_context_window_error(error: Optional[BaseException]) -> bool:
+    """Detect a context-window overflow anywhere in an exception's tree."""
+    if error is None:
+        return False
+    return any(
+        isinstance(current, ContextWindowExceededError)
+        or ExceptionCheckers.is_error_str_context_window_exceeded(str(current))
+        for current in iter_exception_tree(error)
+    )
 
 
 class SemanticMCPToolFilter:

ruff-strict-budget.json

@@ -60,7 +60,7 @@
     "limit": 4
   },
   "BLE001": {
-    "limit": 2903
+    "limit": 2902
   },
   "C401": {
     "limit": 11
@@ -363,6 +363,6 @@
     "limit": 105
   },
   "UP045": {
-    "limit": 18462
+    "limit": 18461
   }
 }

tests/test_litellm/proxy/_experimental/mcp_server/faults/test_traversal.py

@@ -0,0 +1,48 @@
+"""Traversal contract for the shared exception-tree walk: the root is yielded first, explicit
+links win (the ``raise ... from`` cause subtree, then ExceptionGroup members in raise order,
+then the incidental ``__context__`` chain last), and adversarial shapes terminate."""
+
+from litellm.proxy._experimental.mcp_server.faults import iter_exception_tree
+
+
+def test_yields_the_root_itself_first():
+    exc = ValueError("root")
+    assert list(iter_exception_tree(exc)) == [exc]
+
+
+def test_cause_subtree_is_exhausted_before_context():
+    deep = KeyError("deep")
+    cause = RuntimeError("cause")
+    cause.__cause__ = deep
+    context = OSError("context")
+    root = ValueError("root")
+    root.__cause__ = cause
+    root.__context__ = context
+    assert list(iter_exception_tree(root)) == [root, cause, deep, context]
+
+
+def test_group_members_yield_in_raise_order_between_cause_and_context():
+    first = KeyError("first")
+    second = IndexError("second")
+    group = BaseExceptionGroup("group", [first, second])
+    cause = RuntimeError("cause")
+    context = OSError("context")
+    group.__cause__ = cause
+    group.__context__ = context
+    assert list(iter_exception_tree(group)) == [group, cause, first, second, context]
+
+
+def test_terminates_on_a_cause_cycle():
+    a = ValueError("a")
+    b = RuntimeError("b")
+    a.__cause__ = b
+    b.__cause__ = a
+    assert list(iter_exception_tree(a)) == [a, b]
+
+
+def test_node_reachable_as_both_cause_and_context_yields_once():
+    inner = KeyError("inner")
+    root = ValueError("root")
+    root.__cause__ = inner
+    root.__context__ = inner
+    assert list(iter_exception_tree(root)) == [root, inner]

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py

@@ -50,6 +50,36 @@ def test_extract_upstream_auth_failure_returns_none_for_non_auth():
     assert _extract_upstream_auth_failure(RuntimeError("boom")) is None
 
 
+def _auth_status_error(status_code: int, www_authenticate: str) -> httpx.HTTPStatusError:
+    response = httpx.Response(
+        status_code=status_code,
+        headers={"www-authenticate": www_authenticate},
+        request=httpx.Request("GET", "https://upstream/mcp"),
+    )
+    return httpx.HTTPStatusError(str(status_code), request=response.request, response=response)
+
+
+def test_extract_upstream_auth_failure_finds_401_behind_cause_chain():
+    wrapper = RuntimeError("wrapped")
+    wrapper.__cause__ = _auth_status_error(401, "Bearer")
+    assert _extract_upstream_auth_failure(wrapper) == (401, "Bearer")
+
+
+def test_extract_upstream_auth_failure_finds_401_behind_context_chain():
+    wrapper = RuntimeError("wrapped")
+    wrapper.__context__ = _auth_status_error(401, "Bearer")
+    assert _extract_upstream_auth_failure(wrapper) == (401, "Bearer")
+
+
+def test_extract_upstream_auth_failure_prefers_causal_chain_over_context():
+    """A 403 raised incidentally while handling the real 401 (surviving only as ``__context__``)
+    must not shadow the 401 on the explicit ``raise ... from`` chain."""
+    wrapper = RuntimeError("wrapped")
+    wrapper.__cause__ = _auth_status_error(401, "Bearer realm=real")
+    wrapper.__context__ = _auth_status_error(403, "Bearer realm=incidental")
+    assert _extract_upstream_auth_failure(wrapper) == (401, "Bearer realm=real")
+
+
 @pytest.mark.asyncio
 async def test_fetch_tools_from_passthrough_raises_on_upstream_401():
     manager = MCPServerManager()

tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py

@@ -1913,6 +1913,34 @@ def test_is_context_window_error_detection_variants():
     assert not _is_context_window_error(None)
 
 
+def test_is_context_window_error_sees_through_trees_the_chain_walk_missed():
+    """Overflow shapes the old single-path depth-5 chain walk could not reach: hidden in
+    ``__context__`` behind a non-matching ``__cause__``, buried inside an anyio-style
+    ``ExceptionGroup``, and chained deeper than five links."""
+    import litellm
+    from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
+        _is_context_window_error,
+    )
+
+    def _cwe() -> litellm.ContextWindowExceededError:
+        return litellm.ContextWindowExceededError(message="overflow", model="m", llm_provider="openai")
+
+    shadowed = ValueError("wrapper")
+    shadowed.__cause__ = TypeError("unrelated failure")
+    shadowed.__context__ = _cwe()
+    assert _is_context_window_error(shadowed)
+
+    grouped = BaseExceptionGroup("task group", [RuntimeError("sibling"), _cwe()])
+    assert _is_context_window_error(grouped)
+
+    deep: BaseException = _cwe()
+    for depth in range(6):
+        wrapper = ValueError(f"layer {depth}")
+        wrapper.__cause__ = deep
+        deep = wrapper
+    assert _is_context_window_error(deep)
+
+
 def _make_keyword_embedding_router(recorded_inputs):
     """
     Mock litellm Router whose embeddings are deterministic keyword one-hots:

type-discipline-budget.json

@@ -1,6 +1,6 @@
 {
   "LIT001": {
-    "limit": 23409
+    "limit": 23408
   },
   "LIT002": {
     "limit": 27511