refactor(mcp): consolidate exception-tree walkers into one shared faults traversal - #33183
Conversation
Greptile SummaryIntroduces
Confidence Score: 5/5Safe 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.
|
| 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 Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
|
bugbot run |
There was a problem hiding this comment.
✅ 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.
|
bugbot run |
There was a problem hiding this comment.
✅ 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.
mateo-berri
left a comment
There was a problem hiding this comment.
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
|
bugbot run |
There was a problem hiding this comment.
✅ 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.
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 intoExceptionGroups. Which buried exception won therefore depended on which classifier the failure happened to hit. This PR gives thefaultspackage one shared traversal with one deliberate order and migrates both walkers onto itBehavior changes at a glance
raise ... fromchain, plus an incidental 403 raised while handling it (surviving as__context__)__context__first.response-shaped attribute that is not a realhttpx.Responsehttpx.Responsematches. Plain exceptions carrying a realhttpx.Response(the OBO retry contract) still match__context__behind a non-matching__cause__, buried inside an anyioExceptionGroup, or chained deeper than five linksEverything 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_chainis already a single shared helper inproxy/db/and is deliberately not touched; importing MCPfaultsfrom there would invert the dependency direction)Linear ticket
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito 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-Authenticatechallenge, and an unreachable URL (localhost:9)Before (staging tip b200d66) and after (this branch) produce identical results on all three probes. The healthy server lists its tools
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)The unreachable upstream keeps absorbing to an empty list
Type
🧹 Refactoring
Changes
New module
mcp_server/faults/traversal.pyexportingiter_exception_tree: an iterative, cycle-safe generator that yields the root and every exception reachable from it, explicit links first; each node'sraise ... fromcause subtree, thenExceptionGroupmembers 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_failureshrinks to a loop over the shared traversal plus its own predicate: the first exception bearing a realhttpx.Responsewith a 401/403 wins, and the status andWWW-Authenticateheader are extracted from it. The loosegetattrduck-typing onstatus_code/headersand the try/except around the header read are gone; a realhttpx.Responsemakes those total._is_context_window_errorbecomes anany()over the same traversal with its existing predicate (theContextWindowExceededErrorisinstance 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 openTests pin the axes that were previously unpinned:
faults/test_traversal.pypins 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.pygains cause-only, context-only, and cause-beats-context extraction pins;test_semantic_tool_filter.pygains 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_UpstreamAuthErrorretry tests confirm the plain-exception-with-real-httpx.Responsecontract survives the isinstance tightening. The full walker-adjacent suite (faults, oauth passthrough tools, semantic filter, server manager, MCP auth, discoverable endpoints) passes at 790/790Out 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_chaininproxy/common_request_processing.py(proxy-wide error mapping, not MCP). PR #33153's_find_upstream_responseis the third MCP copy; once both PRs merge it adopts this helper in a one-line follow-upThe ruff-strict and type-discipline budget files ratchet down by the violations this branch removes. The basedpyright budget was left untouched:
make lint-budget-updatewanted 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 commitQA runbook
make bootstrap, then start the proxy:python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --use_v2_migration_resolverWWW-Authenticateheader) on localhost:9401curl -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 exactwww-authenticatevaluepytest 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 needssemantic-routerinstalled)Final Attestation
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_treeinfaults/traversal.pyas the single cycle-safe walk over__cause__,ExceptionGroupmembers, and__context__, with explicit causal links before incidental context.list_outcomes._iter_upstream_responsesdrops its inline walker and uses the helper; upstream auth and listing classification now share that order._is_context_window_errorin the semantic tool filter switches from a depth-5 single-chain walk toany()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.Responseas 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.