mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
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:
parent
0e6ed18bfc
commit
ea64ef7a2a
2 changed files with 72 additions and 4 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue