Skip to content

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

Merged
tin-berri merged 2 commits into
litellm_internal_stagingfrom
litellm_mcp_walker_consolidation
Jul 20, 2026
Merged

refactor(mcp): consolidate exception-tree walkers into one shared faults traversal#33183
tin-berri merged 2 commits into
litellm_internal_stagingfrom
litellm_mcp_walker_consolidation

Conversation

@tin-berri

@tin-berri tin-berri commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

Relevant issues

Follow-up to the MCP error-handling framework (#33113). Two copies of the same exception-tree walk lived in MCP code, each with different semantics chosen by accident rather than design: _extract_upstream_auth_failure (mcp_server_manager.py) walked the full tree but visited __context__ before __cause__ and duck-typed any .response-shaped attribute, while _is_context_window_error (semantic_tool_filter.py) walked a single chain path bounded at five links and never descended into ExceptionGroups. Which buried exception won therefore depended on which classifier the failure happened to hit. This PR gives the faults package one shared traversal with one deliberate order and migrates both walkers onto it

Behavior changes at a glance

Scenario Before After
Upstream 401 on the explicit raise ... from chain, plus an incidental 403 raised while handling it (surviving as __context__) the incidental 403 could win; traversal popped __context__ first the 401 on the causal chain wins; explicit links are searched first
Exception in the tree carrying a .response-shaped attribute that is not a real httpx.Response duck-typed as the upstream auth failure ignored; only a real httpx.Response matches. Plain exceptions carrying a real httpx.Response (the OBO retry contract) still match
Context-window overflow hidden in __context__ behind a non-matching __cause__, buried inside an anyio ExceptionGroup, or chained deeper than five links missed; the semantic filter treated it as a generic error and passed all tools through detected; the filter fails closed per LIT-4284

Everything else is unchanged: the 401 challenge relay with WWW-Authenticate, the 403 absorb on listing, unreachable upstreams absorbed to an empty list, and DB-outage classification (PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain is already a single shared helper in proxy/db/ and is deliberately not touched; importing MCP faults from there would invert the dependency direction)

Linear ticket

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

  • I have added meaningful tests
  • My PR passes all CI/CD checks (e.g., lint, format, unit tests)
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to re-request a review after pushing changes)

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

Screenshots / Proof of Fix

Live rig: proxy on localhost:4000 backed by a dedicated Postgres database, with three servers created through the management API: a healthy public MCP server (deepwiki), a stub upstream on localhost:9401 answering every request with 401 plus a WWW-Authenticate challenge, and an unreachable URL (localhost:9)

for s in 'walkers_deepwiki|https://mcp.deepwiki.com/mcp' 'walkers_challenge|http://localhost:9401/mcp' 'walkers_unreachable|http://localhost:9/mcp'; do
  n="${s%%|*}"; u="${s##*|}"
  curl -s -X POST http://localhost:4000/v1/mcp/server -H "Authorization: Bearer sk-1234" -H "Content-Type: application/json" \
    -d "{\"server_name\":\"$n\",\"url\":\"$u\",\"transport\":\"http\",\"auth_type\":\"none\"}"
done

Before (staging tip b200d66) and after (this branch) produce identical results on all three probes. The healthy server lists its tools

$ curl -s "http://localhost:4000/mcp-rest/tools/list?server_id=$DEEPWIKI_ID" -H "Authorization: Bearer sk-1234"
tool_count: 3

The 401-challenging upstream's exact challenge is dug out of the MCP SDK's ExceptionGroup wrapping and relayed (this is the consolidated walk running live; the proxy log shows Upstream auth failure from MCP server walkers_challenge: HTTP 401)

$ curl -s -i "http://localhost:4000/mcp-rest/tools/list?server_id=$CHALLENGE_ID" -H "Authorization: Bearer sk-1234"
HTTP/1.1 401 Unauthorized
www-authenticate: Bearer resource_metadata="http://localhost:9401/.well-known/oauth-protected-resource"
{"detail":"Unauthorized"}

The unreachable upstream keeps absorbing to an empty list

$ curl -s "http://localhost:4000/mcp-rest/tools/list?server_id=$UNREACHABLE_ID" -H "Authorization: Bearer sk-1234"
{"tools":[],"error":null,"message":"Successfully retrieved tools"} HTTP:200

Type

🧹 Refactoring

Changes

New module mcp_server/faults/traversal.py exporting iter_exception_tree: an iterative, cycle-safe generator that yields the root 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. Deriving every classifier's search from one traversal makes blame assignment consistent: an exception raised while handling the real failure can never shadow the failure itself, which is the same order PR #33153 pins for its listing classifier

_extract_upstream_auth_failure shrinks to a loop over the shared traversal plus its own predicate: the first exception bearing a real httpx.Response with a 401/403 wins, and the status and WWW-Authenticate header are extracted from it. The loose getattr duck-typing on status_code/headers and the try/except around the header read are gone; a real httpx.Response makes those total. _is_context_window_error becomes an any() over the same traversal with its existing predicate (the ContextWindowExceededError isinstance plus the message check), replacing the single-path five-link walk; the overflow shapes the old walk could not reach are exactly the ones the embedding stack produces under task groups, and missing them meant the filter failed open

Tests pin the axes that were previously unpinned: faults/test_traversal.py pins the yield order (cause subtree before group members before context, members in raise order), cycle termination, and single-yield of shared nodes; test_mcp_oauth_passthrough_tools.py gains cause-only, context-only, and cause-beats-context extraction pins; test_semantic_tool_filter.py gains the three newly reachable overflow shapes. A mutation check (reverting the helper to the old context-first order) fails three of the new tests, and the existing _UpstreamAuthError retry tests confirm the plain-exception-with-real-httpx.Response contract survives the isinstance tightening. The full walker-adjacent suite (faults, oauth passthrough tools, semantic filter, server manager, MCP auth, discoverable endpoints) passes at 790/790

Out of scope, surveyed and left alone: the DB-outage chain check in proxy/db/exception_handler.py (already one shared helper; its chain and non-chain variants are used deliberately at different call sites) and _has_attribute_error_in_chain in proxy/common_request_processing.py (proxy-wide error mapping, not MCP). PR #33153's _find_upstream_response is the third MCP copy; once both PRs merge it adopts this helper in a one-line follow-up

The ruff-strict and type-discipline budget files ratchet down by the violations this branch removes. The basedpyright budget was left untouched: make lint-budget-update wanted to ratchet it down by 136 errors across 48 rules, which is stale headroom from earlier merges, not this branch's work, and belongs in its own ratchet commit

QA runbook

  1. make bootstrap, then start the proxy: python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --use_v2_migration_resolver
  2. Start a 401-challenging stub upstream (any HTTP server answering every request with 401 and a WWW-Authenticate header) on localhost:9401
  3. Create the three servers with the curl loop from the proof section
  4. curl -s -i "http://localhost:4000/mcp-rest/tools/list?server_id=<challenge-id>" -H "Authorization: Bearer sk-1234" and confirm HTTP 401 with the stub's exact www-authenticate value
  5. Confirm the deepwiki server still lists tools and the unreachable server returns an empty 200
  6. pytest tests/test_litellm/proxy/_experimental/mcp_server/faults/test_traversal.py tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py -q (the semantic filter file needs semantic-router installed)

Final Attestation

  • The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

Note

Medium Risk
Changes how buried upstream 401/403 and context-window errors are chosen in MCP fault handling and semantic filtering—behavioral fixes with customer-visible auth relay and fail-closed filtering implications, covered by targeted tests.

Overview
Introduces iter_exception_tree in faults/traversal.py as the single cycle-safe walk over __cause__, ExceptionGroup members, and __context__, with explicit causal links before incidental context.

list_outcomes._iter_upstream_responses drops its inline walker and uses the helper; upstream auth and listing classification now share that order. _is_context_window_error in the semantic tool filter switches from a depth-5 single-chain walk to any() over the full tree, so overflows inside groups or deep/shadowed chains are detected and the filter fails closed instead of passing all tools through.

Auth extraction now only treats a real httpx.Response as upstream (no duck-typed .response). Tests cover traversal order, cause-vs-context precedence for 401/403, and the newly reachable context-window shapes; lint budgets ratchet by one.

Reviewed by Cursor Bugbot for commit 8e0d3af. Bugbot is set up for automated code reviews on this repo. Configure here.

@greptile-apps

greptile-apps Bot commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

Introduces iter_exception_tree in faults/traversal.py — a single iterative, cycle-safe DFS that yields the root and every reachable exception in deliberate order (explicit raise ... from cause subtree first, ExceptionGroup members in raise order, incidental __context__ chain last) — and migrates both MCP fault classifiers onto it.

  • list_outcomes._iter_upstream_responses swaps its inline stack loop for iter_exception_tree; the traversal order was already cause-first, so this is a pure deduplication with no behavior change.
  • semantic_tool_filter._is_context_window_error replaces a five-link single-chain walk with a full-tree any() over iter_exception_tree, fixing a fail-open bug where overflows buried inside ExceptionGroups, behind unrelated causes, or deeper than five links were silently missed.
  • New tests pin the traversal contract (yield order, cycle termination, shared-node deduplication), auth-extraction precedence (cause beats context), and the three overflow shapes the old walk could not reach.

Confidence Score: 5/5

Safe to merge. The refactoring is well-scoped, behaviorally equivalent on all existing paths for list_outcomes, and a deliberate improvement for the semantic filter fail-closed contract.

The traversal algorithm is correct and verified by tracing stack operations against the new unit tests. The list_outcomes migration preserves the prior cause-first order exactly. The semantic filter change broadens detection in the intended direction and is guarded by regression tests covering each newly reachable shape. No existing test assertions were weakened and no production call sites were broken.

No files require special attention.

Important Files Changed

Filename Overview
litellm/proxy/_experimental/mcp_server/faults/traversal.py New module exporting iter_exception_tree: iterative, cycle-safe DFS generator with deliberate ordering (cause -> group members -> context); correctly tested and free of issues.
litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py Migrated _iter_upstream_responses from an inline DFS loop to iter_exception_tree; traversal order preserved (was already cause-first), code simplified with no behavior change.
litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py Replaced five-link single-chain walk in _is_context_window_error with full tree traversal; broadens detection to ExceptionGroups and deep chains, fixing a fail-open bug.
litellm/proxy/_experimental/mcp_server/faults/init.py Exports iter_exception_tree from the new traversal module; clean addition.
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py Docstring-only update to _extract_upstream_auth_failure to reflect the correct delegation module name; no logic changes.
tests/test_litellm/proxy/_experimental/mcp_server/faults/test_traversal.py New traversal contract tests: yield order, cause-before-context, group member order, cycle termination, and single-yield of shared nodes.
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py Adds three new extraction tests (cause-only, context-only, cause-beats-context precedence); no existing test assertions were weakened.
tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py Adds regression tests for the three overflow shapes the old five-link walk missed: shadowed by cause, inside ExceptionGroup, and depth > 5; no existing assertions changed.

Reviews (2): Last reviewed commit: "Merge origin/litellm_internal_staging in..." | Re-trigger Greptile

@codecov

codecov Bot commented Jul 14, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 96.29630% with 1 line in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
...y/_experimental/mcp_server/semantic_tool_filter.py 83.33% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

@codspeed-hq

codspeed-hq Bot commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 31 untouched benchmarks


Comparing litellm_mcp_walker_consolidation (8e0d3af) with litellm_internal_staging (214945a)1

Open in CodSpeed

Footnotes

  1. No successful run was found on litellm_internal_staging (8e0d3af) during the generation of this report, so 214945a was used instead as the comparison base. There might be some changes unrelated to this pull request in this report.

@tin-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@cursor cursor Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ Bugbot reviewed your changes and found no new issues!

Comment @cursor review or bugbot run to trigger another review on this PR

Reviewed by Cursor Bugbot for commit 40f02b2. Configure here.

@tin-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@cursor cursor Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ Bugbot reviewed your changes and found no new issues!

Comment @cursor review or bugbot run to trigger another review on this PR

Reviewed by Cursor Bugbot for commit 40f02b2. Configure here.

@tin-berri
tin-berri enabled auto-merge July 20, 2026 18:53

@mateo-berri mateo-berri left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fix the merge conflicts but otherwise impl looks right LGTM; thanks!

…dation

Staging now contains #33153, whose final rounds made _extract_upstream_auth_failure a thin delegate
to upstream_auth_challenge and introduced the response-level iterator this branch predates. The
resolution completes the consolidation both branches were converging on: iter_exception_tree
(faults/traversal.py) is the one tree walk, _iter_upstream_responses is rebuilt on top of it instead
of carrying a second copy of the traversal, the manager keeps the delegate, and the semantic filter
port from this branch stands. Test conflicts were append-append and both sides are kept
@tin-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

@tin-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@cursor cursor Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ Bugbot reviewed your changes and found no new issues!

Comment @cursor review or bugbot run to trigger another review on this PR

Reviewed by Cursor Bugbot for commit 8e0d3af. Configure here.

@tin-berri
tin-berri merged commit b086cd3 into litellm_internal_staging Jul 20, 2026
82 of 83 checks passed
@tin-berri
tin-berri deleted the litellm_mcp_walker_consolidation branch July 20, 2026 23:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants