test(mcp): cover every delegate ownership branch

The master key, bridge and refresh envelopes, sealed credential blobs and dashboard tokens, including one LiteLLM cannot read, now each have a test proving they take normal admission and are never forwarded upstream. The new tests are fully typed

_caused_by walks the cause chain through a small generator so no line is unreachable, and the assert_never guard in _verify_identity follows the repo's pragma convention for exhaustiveness guards
This commit is contained in:
shrey kharbanda 2026-10-28 16:38:50 -07:00
parent 6f8a42d517
commit 65988f3eb8
3 changed files with 140 additions and 57 deletions

View file

@ -2,7 +2,7 @@ import asyncio
import functools
import re
import secrets
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from collections.abc import Set as AbstractSet
from dataclasses import dataclass
from datetime import datetime, timezone
@ -237,15 +237,17 @@ def _is_legacy_delegate_cold_start(mcp_servers: list[str] | None, client_ip: str
return True
def _caused_by(exc: BaseException, error_type: type[BaseException]) -> bool:
def _cause_chain(exc: BaseException) -> Iterator[BaseException]:
link: BaseException | None = exc # rebind-ok: advances one link of the cause chain per iteration
for _ in range(_CAUSE_CHAIN_MAX_DEPTH):
if link is None:
return False
if isinstance(link, error_type):
return True
return
yield link
link = link.__cause__ or link.__context__
return False
def _caused_by(exc: BaseException, error_type: type[BaseException]) -> bool:
return any(isinstance(link, error_type) for link in _cause_chain(exc))
async def _bearer_is_litellm_credential(token: str) -> bool:
@ -310,8 +312,8 @@ async def _verify_identity(check: _IdentityCheck, token: str, request: Request)
await handle_oauth2_proxy_request(request=request)
case "jwt":
await jwt_handler.auth_jwt(token=token) # pyright: ignore[reportUnknownMemberType] # returns a bare dict; only whether it raises matters
case _:
assert_never(check)
case _: # pragma: no cover - exhaustiveness guard, unreachable while the Literal holds
assert_never(check) # pragma: no cover
async def _upstream_bearer_identity(token: str, litellm_api_key: str, request: Request) -> UserAPIKeyAuth | None:

View file

@ -1,7 +1,7 @@
import contextlib
import json
import os
from collections.abc import Sequence
from collections.abc import Callable, Iterator, Sequence
from collections.abc import Set as AbstractSet
from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
@ -1663,7 +1663,7 @@ class TestMCPOAuth2AuthFlow:
as LiteLLM API keys, causing auth failures and empty tool listings.
"""
async def test_oauth2_upstream_token_on_legacy_delegate_is_admitted_anonymously(self):
async def test_oauth2_upstream_token_on_legacy_delegate_is_admitted_anonymously(self) -> None:
"""
A bare upstream OAuth token on the legacy delegated mode is admitted
anonymously without LiteLLM key validation and kept for egress.
@ -2655,7 +2655,7 @@ class TestMCPDelegateAuthToUpstream:
mock_auth.assert_awaited_once()
@staticmethod
def _bearer_scope(token: str) -> dict:
def _bearer_scope(token: str) -> dict[str, object]:
return {
"type": "http",
"method": "POST",
@ -2671,7 +2671,9 @@ class TestMCPDelegateAuthToUpstream:
@staticmethod
@contextlib.contextmanager
def _delegate_admission(admission: AsyncMock, allowed_server_ids: Sequence[str] = ("test-server-id",)):
def _delegate_admission(
admission: AsyncMock, allowed_server_ids: Sequence[str] = ("test-server-id",)
) -> Iterator[None]:
from litellm.types.mcp import MCPAuth
with (
@ -2687,7 +2689,7 @@ class TestMCPDelegateAuthToUpstream:
@staticmethod
@contextlib.contextmanager
def _jwt_auth_enabled(auth_jwt: AsyncMock):
def _jwt_auth_enabled(auth_jwt: AsyncMock) -> Iterator[None]:
from litellm.proxy import proxy_server
from litellm.proxy._types import LiteLLM_JWTAuth
@ -2702,12 +2704,12 @@ class TestMCPDelegateAuthToUpstream:
os.environ.pop("JWT_ISSUER", None)
yield
async def test_delegate_forwards_unknown_sk_token_after_key_table_miss(self):
async def test_delegate_forwards_unknown_sk_token_after_key_table_miss(self) -> None:
"""An sk- token the key table does not hold, such as a provider API key, belongs to the upstream:
it is forwarded anonymously rather than rejected for its prefix."""
from litellm.proxy.auth.resolvers.exceptions import KeyNotFoundError
admission = AsyncMock()
admission: Final = AsyncMock()
with (
TestMCPDelegateAuthToUpstream._delegate_admission(admission),
patch(
@ -2730,10 +2732,12 @@ class TestMCPDelegateAuthToUpstream:
[None, RuntimeError("key table unreachable")],
ids=["key-held", "lookup-fault"],
)
async def test_delegate_litellm_key_refusal_propagates(self, lookup_effect):
async def test_delegate_litellm_key_refusal_propagates(self, lookup_effect: Exception | None) -> None:
"""A key the table holds, or one the lookup could not check, takes normal admission, so its
refusal reaches the caller instead of becoming an anonymous request that forwards the key."""
refusal = ProxyException(message="Budget has been exceeded", type="budget_exceeded", param=None, code=400)
refusal: Final = ProxyException(
message="Budget has been exceeded", type="budget_exceeded", param=None, code=400
)
with (
TestMCPDelegateAuthToUpstream._delegate_admission(AsyncMock(side_effect=refusal)),
patch(
@ -2747,11 +2751,67 @@ class TestMCPDelegateAuthToUpstream:
assert exc_info.value is refusal
async def test_delegate_master_key_is_admitted_and_never_forwarded(self) -> None:
"""The master key is not in the key table, so only recognizing it directly stops it from being
treated as the upstream's token and sent to the upstream server."""
from litellm.proxy import proxy_server
from litellm.proxy.auth.resolvers.exceptions import KeyNotFoundError
admission: Final = AsyncMock(return_value=UserAPIKeyAuth(api_key="sk-master-secret", user_role="proxy_admin"))
with (
patch.object(proxy_server, "master_key", "sk-master-secret"),
TestMCPDelegateAuthToUpstream._delegate_admission(admission),
patch(
"litellm.proxy.auth.resolvers.store.IdentityStore.resolve",
new_callable=AsyncMock,
side_effect=KeyNotFoundError("hashed-master-key"),
),
):
auth_result, _, _, _, oauth2_headers, _ = await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope("sk-master-secret")
)
admission.assert_awaited_once()
assert auth_result.mcp_forwards_caller_authorization is False
assert oauth2_headers is None
@pytest.mark.parametrize(
"token, ui_token_lookup",
[
("llm_env_sealed-bridge-envelope", None),
("llm_refresh_sealed-refresh-envelope", None),
("v2:gcm:sealed-credential-blob", None),
("dashboard-login-token", lambda: MagicMock(return_value=UserAPIKeyAuth(user_id="dashboard-user"))),
("unreadable-dashboard-token", lambda: MagicMock(side_effect=ValueError("Invalid hash key"))),
],
ids=["bridge-envelope", "refresh-envelope", "gcm-blob", "dashboard-token", "unreadable-dashboard-token"],
)
async def test_delegate_sealed_litellm_token_takes_normal_admission(
self, token: str, ui_token_lookup: Callable[[], MagicMock] | None
) -> None:
"""A token only LiteLLM issues is LiteLLM's own, even one it fails to read, so it goes through
normal admission and is never marked for forwarding upstream."""
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
admission: Final = AsyncMock(return_value=UserAPIKeyAuth(user_id="litellm-user"))
ui_lookup: Final = (
patch.object(ExperimentalUIJWTToken, "get_key_object_from_ui_hash_key", ui_token_lookup())
if ui_token_lookup is not None
else contextlib.nullcontext()
)
with TestMCPDelegateAuthToUpstream._delegate_admission(admission), ui_lookup:
auth_result, *_ = await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope(token)
)
admission.assert_awaited_once()
assert auth_result.mcp_forwards_caller_authorization is False
@staticmethod
def _unknown_user_refusal() -> ProxyException:
from litellm.types.proxy.auth.auth_checks import UserNotFoundError
refusal = ProxyException(message="User doesn't exist in db", type="auth_error", param=None, code=401)
refusal: Final = ProxyException(message="User doesn't exist in db", type="auth_error", param=None, code=401)
refusal.__cause__ = UserNotFoundError(user_id="jwt-user")
return refusal
@ -2786,12 +2846,16 @@ class TestMCPDelegateAuthToUpstream:
ids=["granted-user", "user-not-granted", "user-unknown-to-litellm", "signature-rejected"],
)
async def test_delegate_jwt_is_forwarded_as_granted_user_or_anonymously(
self, auth_jwt_effect, admission, allowed_server_ids, expected_user_id
):
self,
auth_jwt_effect: Exception | None,
admission: Callable[[], AsyncMock],
allowed_server_ids: Sequence[str],
expected_user_id: str | None,
) -> None:
"""A JWT LiteLLM also accepts admits that user when granted the server; a user who is not
granted, unknown to LiteLLM, or whose signature fails is admitted anonymously. The JWT is always
forwarded, because on this server it is the upstream's token."""
token = TestMCPDelegateAuthToUpstream._jwt()
token: Final = TestMCPDelegateAuthToUpstream._jwt()
with (
TestMCPDelegateAuthToUpstream._jwt_auth_enabled(AsyncMock(side_effect=auth_jwt_effect)),
TestMCPDelegateAuthToUpstream._delegate_admission(admission(), allowed_server_ids),
@ -2812,7 +2876,7 @@ class TestMCPDelegateAuthToUpstream:
],
ids=["over-budget", "blocked-team"],
)
async def test_delegate_jwt_refusal_after_valid_signature_propagates(self, refusal):
async def test_delegate_jwt_refusal_after_valid_signature_propagates(self, refusal: ProxyException) -> None:
"""Once the signature verifies, LiteLLM knows who the caller is, so a budget or an explicit
disallowment is enforced and the request fails instead of continuing anonymously."""
with (
@ -2828,18 +2892,20 @@ class TestMCPDelegateAuthToUpstream:
@staticmethod
@contextlib.contextmanager
def _idp_auth_enabled(mechanism: str, verify: AsyncMock, premium: bool = True):
def _idp_auth_enabled(
mechanism: Literal["oauth2", "oauth2_proxy"], verify: AsyncMock, premium: bool = True
) -> Iterator[None]:
from litellm.proxy import proxy_server
from litellm.proxy.auth.oauth2_check import Oauth2Handler
settings = {"oauth2": {"enable_oauth2_auth": True}, "oauth2_proxy": {"enable_oauth2_proxy_auth": True}}
verifier = (
settings: Final = {"enable_oauth2_auth": True} if mechanism == "oauth2" else {"enable_oauth2_proxy_auth": True}
verifier: Final = (
patch.object(Oauth2Handler, "check_oauth2_token", verify)
if mechanism == "oauth2"
else patch("litellm.proxy.auth.oauth2_proxy_hook.handle_oauth2_proxy_request", verify)
)
with (
patch.dict(proxy_server.general_settings, settings[mechanism]),
patch.dict(proxy_server.general_settings, settings),
patch.object(proxy_server, "premium_user", premium),
patch.object(proxy_server, "user_custom_auth", None),
verifier,
@ -2857,11 +2923,15 @@ class TestMCPDelegateAuthToUpstream:
ids=["granted-user", "user-not-granted", "identity-rejected"],
)
async def test_delegate_idp_identity_follows_the_jwt_rules(
self, mechanism, verify_effect, allowed_server_ids, expected_user_id
):
self,
mechanism: Literal["oauth2", "oauth2_proxy"],
verify_effect: Exception | None,
allowed_server_ids: Sequence[str],
expected_user_id: str | None,
) -> None:
"""OAuth2 introspection and trusted reverse-proxy headers admit a granted user by name, and a user
who is not granted or not vouched for anonymously. The bearer is forwarded either way."""
admission = AsyncMock(return_value=UserAPIKeyAuth(user_id="idp-user"))
admission: Final = AsyncMock(return_value=UserAPIKeyAuth(user_id="idp-user"))
with (
TestMCPDelegateAuthToUpstream._idp_auth_enabled(mechanism, AsyncMock(side_effect=verify_effect)),
TestMCPDelegateAuthToUpstream._delegate_admission(admission, allowed_server_ids),
@ -2875,9 +2945,11 @@ class TestMCPDelegateAuthToUpstream:
assert oauth2_headers == {"Authorization": "Bearer opaque-idp-token"}
@pytest.mark.parametrize("mechanism", ["oauth2", "oauth2_proxy"])
async def test_delegate_idp_user_refusal_propagates(self, mechanism):
async def test_delegate_idp_user_refusal_propagates(self, mechanism: Literal["oauth2", "oauth2_proxy"]) -> None:
"""Once the provider vouches for the caller, LiteLLM's budget and permission refusals apply."""
refusal = ProxyException(message="Budget has been exceeded", type="budget_exceeded", param=None, code=400)
refusal: Final = ProxyException(
message="Budget has been exceeded", type="budget_exceeded", param=None, code=400
)
with (
TestMCPDelegateAuthToUpstream._idp_auth_enabled(mechanism, AsyncMock()),
TestMCPDelegateAuthToUpstream._delegate_admission(AsyncMock(side_effect=refusal)),
@ -2889,9 +2961,9 @@ class TestMCPDelegateAuthToUpstream:
assert exc_info.value is refusal
async def test_delegate_oauth2_without_premium_never_calls_the_provider(self):
async def test_delegate_oauth2_without_premium_never_calls_the_provider(self) -> None:
"""OAuth2 introspection is enterprise-only, so without a license the bearer is the upstream's alone."""
introspect = AsyncMock()
introspect: Final = AsyncMock()
with (
TestMCPDelegateAuthToUpstream._idp_auth_enabled("oauth2", introspect, premium=False),
TestMCPDelegateAuthToUpstream._delegate_admission(AsyncMock(return_value=UserAPIKeyAuth(user_id="idp-user"))),
@ -2904,10 +2976,10 @@ class TestMCPDelegateAuthToUpstream:
assert auth_result.mcp_forwards_caller_authorization is True
introspect.assert_not_awaited()
async def test_delegate_jwt_from_unconfigured_issuer_is_forwarded_without_jwt_auth(self):
async def test_delegate_jwt_from_unconfigured_issuer_is_forwarded_without_jwt_auth(self) -> None:
"""A JWT from an issuer LiteLLM's JWT auth is not configured for never reaches JWT auth."""
auth_jwt = AsyncMock(return_value={"sub": "jwt-user"})
admission = AsyncMock(return_value=UserAPIKeyAuth(user_id="jwt-user"))
auth_jwt: Final = AsyncMock(return_value={"sub": "jwt-user"})
admission: Final = AsyncMock(return_value=UserAPIKeyAuth(user_id="jwt-user"))
with (
TestMCPDelegateAuthToUpstream._jwt_auth_enabled(auth_jwt),
patch.dict(os.environ, {"JWT_ISSUER": "https://litellm-idp.example"}),
@ -2939,7 +3011,9 @@ class TestMCPDelegateAuthToUpstream:
],
ids=["accepted", "rejected"],
)
async def test_delegate_custom_auth_decides_alone(self, admission, expected_user_id):
async def test_delegate_custom_auth_decides_alone(
self, admission: Callable[[], AsyncMock], expected_user_id: str | None
) -> None:
"""With custom auth configured, its answer alone decides: accepted admits that user, any failure
admits anonymously. The bearer is forwarded either way."""
from litellm.proxy import proxy_server
@ -2956,7 +3030,7 @@ class TestMCPDelegateAuthToUpstream:
assert auth_result.mcp_forwards_caller_authorization is True
assert oauth2_headers == {"Authorization": "Bearer upstream-custom-token"}
def test_forwarding_marker_cannot_be_forged_from_input(self):
def test_forwarding_marker_cannot_be_forged_from_input(self) -> None:
"""Only admission may mark a caller's Authorization for forwarding; caller-shaped input never can."""
assert UserAPIKeyAuth(mcp_forwards_caller_authorization=True).mcp_forwards_caller_authorization is False
assert (
@ -3038,7 +3112,7 @@ class TestMCPDelegateAuthToUpstream:
[[], [(b"authorization", b"Bearer upstream-pkce-token")]],
ids=["no-bearer", "upstream-bearer"],
)
async def test_delegate_mixed_targets_fail_closed(self, bearer_headers):
async def test_delegate_mixed_targets_fail_closed(self, bearer_headers: list[tuple[bytes, bytes]]) -> None:
"""
x-mcp-servers can list multiple targets. If ANY of them does not opt in
to delegate_auth_to_upstream, the bypass must NOT fire — otherwise an
@ -3365,7 +3439,7 @@ class TestMCPDelegateAuthToUpstream:
mock_auth.assert_awaited_once()
assert auth.api_key is None
async def test_get_allowed_servers_includes_interactive_legacy_delegates_for_anonymous(self):
async def test_get_allowed_servers_includes_interactive_legacy_delegates_for_anonymous(self) -> None:
"""Anonymous callers reach interactive legacy delegates, never an M2M server whose stored
client credentials would let them act as LiteLLM's service account."""
from litellm.proxy.mcp_server.mcp_server_manager import (

View file

@ -3,6 +3,7 @@ import contextlib
import contextvars
import json
import os
from collections.abc import Coroutine, Iterator, Sequence
from datetime import datetime, timedelta
from types import SimpleNamespace
from typing import Callable, Final, Literal
@ -674,7 +675,9 @@ def test_prepare_mcp_server_headers_legacy_delegate_strips_repeated_admission_ke
[{}, {"user_id": "jwt-user", "api_key": "upstream-jwt", "jwt_claims": {"sub": "jwt-user"}}],
ids=["anonymous", "admitted-jwt-user"],
)
def test_prepare_mcp_server_headers_legacy_delegate_forwards_marked_upstream_authorization(caller_fields):
def test_prepare_mcp_server_headers_legacy_delegate_forwards_marked_upstream_authorization(
caller_fields: dict[str, object],
) -> None:
"""When admission marked the caller's Authorization as the upstream token, a legacy delegate forwards
it, including a JWT that also admitted the caller and would otherwise be stripped as an admission JWT."""
from litellm.proxy._types import UserAPIKeyAuth
@ -689,7 +692,9 @@ def test_prepare_mcp_server_headers_legacy_delegate_forwards_marked_upstream_aut
auth_type=MCPAuth.oauth2,
delegate_auth_to_upstream=True,
)
caller = UserAPIKeyAuth(**caller_fields).model_copy(update={"mcp_forwards_caller_authorization": True})
caller: Final = UserAPIKeyAuth.model_validate(caller_fields).model_copy(
update={"mcp_forwards_caller_authorization": True}
)
server_auth_header, extra_headers = _prepare_mcp_server_headers(
server=server,
@ -704,7 +709,7 @@ def test_prepare_mcp_server_headers_legacy_delegate_forwards_marked_upstream_aut
assert extra_headers == {"Authorization": "Bearer upstream-jwt"}
def test_prepare_mcp_server_headers_forwarding_marker_is_ignored_off_legacy_delegates():
def test_prepare_mcp_server_headers_forwarding_marker_is_ignored_off_legacy_delegates() -> None:
"""The forwarding marker only applies to legacy delegates, so it can never carry a LiteLLM key to an
oauth_delegate upstream."""
from litellm.proxy._types import UserAPIKeyAuth
@ -718,7 +723,9 @@ def test_prepare_mcp_server_headers_forwarding_marker_is_ignored_off_legacy_dele
transport=MCPTransport.http,
auth_type=MCPAuth.oauth_delegate,
)
caller = UserAPIKeyAuth(api_key="sk-litellm-key").model_copy(update={"mcp_forwards_caller_authorization": True})
caller: Final = UserAPIKeyAuth(api_key="sk-litellm-key").model_copy(
update={"mcp_forwards_caller_authorization": True}
)
_, extra_headers = _prepare_mcp_server_headers(
server=server,
@ -7328,7 +7335,7 @@ async def test_legacy_delegate_bare_token_is_not_probed_upstream(): # test-qual
@pytest.mark.asyncio
async def test_legacy_delegate_dual_credentials_probe_only_the_upstream_token(): # test-quality-ok: the credential sent to the upstream probe is the security contract and has no return value
async def test_legacy_delegate_dual_credentials_probe_only_the_upstream_token() -> None: # test-quality-ok: the credential sent to the upstream probe is the security contract and has no return value
"""With a separate x-litellm-api-key, the preflight probes the upstream with its own token and never
with the LiteLLM key."""
from litellm.proxy._types import UserAPIKeyAuth
@ -7367,7 +7374,7 @@ async def test_legacy_delegate_dual_credentials_probe_only_the_upstream_token():
assert probe.await_args.args == ("http://upstream:9401/mcp", "Bearer upstream-token")
def _expired_delegate_token_preflight(mcp_servers: list[str]):
def _expired_delegate_token_preflight(mcp_servers: Sequence[str]) -> Coroutine[object, object, None]:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.mcp_server.server import (
_check_passthrough_upstream_auth,
@ -7376,7 +7383,7 @@ def _expired_delegate_token_preflight(mcp_servers: list[str]):
return _check_passthrough_upstream_auth(
scope=_delegate_scope([(b"authorization", b"Bearer expired-upstream-token")]),
user_api_key_auth=UserAPIKeyAuth().model_copy(update={"mcp_forwards_caller_authorization": True}),
mcp_servers=mcp_servers,
mcp_servers=list(mcp_servers),
client_ip=None,
oauth2_headers={"Authorization": "Bearer expired-upstream-token"},
mcp_server_auth_headers=None,
@ -7384,8 +7391,9 @@ def _expired_delegate_token_preflight(mcp_servers: list[str]):
)
def _patch_rejecting_delegate_upstream():
return (
@contextlib.contextmanager
def _rejecting_delegate_upstream() -> Iterator[None]:
with (
patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(return_value=[_delegate_auth_mcp_server()]),
@ -7394,28 +7402,27 @@ def _patch_rejecting_delegate_upstream():
"litellm.proxy.mcp_server.server._probe_upstream_auth",
new=AsyncMock(return_value=(401, 'Bearer error="invalid_token"')),
),
)
):
yield
@pytest.mark.asyncio
async def test_legacy_delegate_expired_upstream_token_gets_reauthorization_challenge():
async def test_legacy_delegate_expired_upstream_token_gets_reauthorization_challenge() -> None:
"""An expired upstream token on a single-server route gets the RFC 9728 re-authorization challenge
for the requested name, so the client signs in again instead of seeing an empty tool list."""
allowed_servers, probe = _patch_rejecting_delegate_upstream()
with allowed_servers, probe, pytest.raises(HTTPException) as exc_info:
with _rejecting_delegate_upstream(), pytest.raises(HTTPException) as exc_info:
await _expired_delegate_token_preflight(["delegate_test"])
assert exc_info.value.status_code == 401
challenge = exc_info.value.headers["www-authenticate"]
challenge: Final = exc_info.value.headers["www-authenticate"]
assert 'error="invalid_token"' in challenge
assert 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/delegate_test"' in challenge
@pytest.mark.asyncio
async def test_legacy_delegate_multi_server_scope_is_not_failed_by_one_upstream_token():
async def test_legacy_delegate_multi_server_scope_is_not_failed_by_one_upstream_token() -> None:
"""A multi-server scope is not probed, so one rejected upstream token cannot fail the whole connect."""
allowed_servers, probe = _patch_rejecting_delegate_upstream()
with allowed_servers, probe:
with _rejecting_delegate_upstream():
assert await _expired_delegate_token_preflight(["delegate_test", "other_server"]) is None