fix(mcp): surface real status from bridge admission policy gate instead of flattening to 401

Over-budget rendered 401 (should be 429), model-access and other typed
failures collapsed to 401, and a transient DB outage was masked as an auth
error. Mirror UserAPIKeyAuthExceptionHandler: budget maps to 429, a sub-check's
own HTTPException/ProxyException keeps its status, a DB outage is a retryable
503, and only a genuinely unresolvable failure stays the fail-closed 401.
This commit is contained in:
Tin Chi Lo 2026-07-11 12:19:52 -07:00
parent 0e6ed18bfc
commit ea64ef7a2a
2 changed files with 72 additions and 4 deletions

View file

@ -8,6 +8,7 @@ from starlette.requests import Request
from starlette.types import Scope
from typing_extensions import assert_never
import litellm
from litellm._logging import verbose_logger
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import (
BridgeEnvelopeAdmitted,
@ -592,7 +593,10 @@ class MCPRequestHandler:
The standard pipeline enforces this inline in ``_user_api_key_auth_builder`` rather
than in ``common_checks``, so the centralized policy gate does not cover it; without
this mirror, IdP offboarding would leave the user's already-minted envelopes live
until expiry. A failed user lookup skips the gate, matching the builder."""
until expiry. A failed user lookup skips the gate (fail-open), matching the builder:
this is the one deliberately fail-open check in an otherwise fail-closed arm, so a
transient DB outage during this lookup admits the request rather than rejecting it,
keeping parity with how the standard pipeline treats the same lookup failure."""
if key_object.user_id is None:
return
from litellm.proxy.auth.auth_checks import get_user_object
@ -621,8 +625,19 @@ class MCPRequestHandler:
after every builder path, so the envelope identity gets team-block, project-block,
org, and budget enforcement identical to the same key presented directly, and any
policy dimension added to the standard pipeline applies here without this arm
mirroring it. Every failure maps to the arm's uniform 401 so a caller probing with
a stolen envelope learns nothing about why it stopped working."""
mirroring it.
Failures surface with the status the standard pipeline would give them, mirroring
``UserAPIKeyAuthExceptionHandler``: an over-budget identity is a 429, a sub-check that
raised its own ``HTTPException``/``ProxyException`` keeps that status, a transient
database outage is a retryable 503, and only a genuinely unresolvable failure (a
blocked team/project raises a bare ``Exception``, same as the standard pipeline's
fallback) becomes the fail-closed 401. Collapsing every failure to 401 was misleading:
it told an over-budget but validly-authenticated caller their credential was invalid,
which on a DCR client reads as broken auth and can trigger a pointless re-authorize
loop that cannot fix a budget problem, and it masked a DB outage as an auth error."""
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
try:
await _run_centralized_common_checks(
user_api_key_auth_obj=admitted,
@ -630,7 +645,16 @@ class MCPRequestHandler:
request_data=await _read_request_body(request=request),
route=route,
)
except Exception: # noqa: BLE001 # common_checks raises bare Exception for blocked states; narrowing would fail open
except (HTTPException, ProxyException):
raise
except litellm.BudgetExceededError as e:
raise HTTPException(status_code=getattr(e, "status_code", 429), detail=str(e)) from None
except Exception as e: # noqa: BLE001 # untyped gate failure: retryable 503 for a DB outage, else fail closed 401
if PrismaDBExceptionHandler.is_database_service_unavailable_error(e):
raise HTTPException(
status_code=503,
detail="Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.",
) from None
raise HTTPException(status_code=401, detail="Invalid or expired credential") from None
@staticmethod

View file

@ -5243,6 +5243,50 @@ class TestMCPDcrBridgeDelegateAdmission:
assert exc_info.value.status_code == 401
_POLICY_GATE = (
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp._run_centralized_common_checks"
)
async def _enforce_with_gate_error(self, error):
"""Drive _enforce_admitted_live_policy with the centralized gate raising ``error`` and return
the HTTPException the arm maps it to."""
with patch(self._POLICY_GATE, new=AsyncMock(side_effect=error)):
with pytest.raises(HTTPException) as exc_info:
await MCPRequestHandler._enforce_admitted_live_policy(
admitted=UserAPIKeyAuth(user_id="envelope-user-42"),
request=self._mcp_request(),
route="/mcp/bridge_delegate_server",
)
return exc_info.value
async def test_over_budget_admission_surfaces_429_not_401(self):
"""A validly-authenticated but over-budget identity surfaces the standard pipeline's 429, not
a misleading 401. Flattening budget to 401 told the caller their credential was invalid, which
on a DCR client reads as broken auth and triggers a re-authorize that cannot fix a budget
problem. Regression for the status-flattening finding on the live-policy gate."""
import litellm
mapped = await self._enforce_with_gate_error(litellm.BudgetExceededError(current_cost=10.0, max_budget=1.0))
assert mapped.status_code == 429
async def test_db_outage_during_policy_surfaces_503_not_401(self):
"""A transient database outage during the live-policy gate surfaces a retryable 503, not a 401
that masks the outage as an auth failure and tells a valid caller to re-authenticate."""
mapped = await self._enforce_with_gate_error(ConnectionError("could not reach database server"))
assert mapped.status_code == 503
async def test_blocked_state_bare_exception_stays_401(self):
"""A blocked team/project raises a bare Exception (no status) in common_checks, which the
standard pipeline renders as 401; the arm keeps failing those closed as 401, never a 500."""
mapped = await self._enforce_with_gate_error(Exception("Team=team-x is blocked."))
assert mapped.status_code == 401
async def test_subcheck_httpexception_status_preserved(self):
"""A sub-check that raises its own HTTPException (e.g. a 403 model-access denial) keeps that
status through the arm rather than being flattened to 401."""
mapped = await self._enforce_with_gate_error(HTTPException(status_code=403, detail="model not allowed"))
assert mapped.status_code == 403
async def test_alias_only_server_injects_under_alias_egress_can_resolve(self):
"""When server_name is None, the inner token must be keyed under the alias (which egress
resolves), never under server.name (which egress never looks up), so the forwarded token is