refactor(mcp): trim redundant comments and dedupe admission-arm tests

Compress the security rationale in the gateway-session admission path of
user_api_key_auth_mcp.py, keeping the load-bearing "why" and dropping the
restatement, and remove a garbled dead comment in get_allowed_tools_for_server

In the tests, hoist the duplicated _team / _admitted_subject fixtures to
module-level factories and parametrize the four fail-closed session-bearer
variants into one case. No behavior change; the 294 tests in the file still pass
This commit is contained in:
Tin Chi Lo 2026-07-22 15:39:03 -07:00
parent a78130461f
commit ffa0dffbc6
2 changed files with 269 additions and 421 deletions

View file

@ -131,11 +131,9 @@ def _is_mcp_admitted_user_subject(user_api_key_auth: UserAPIKeyAuth | None) -> b
"""True when this auth is a keyless subject admitted by the gateway session / bridge user
path, as opposed to a JWT or other keyless auth that merely lacks a ``team_id``.
Reads the server-only ``UserAPIKeyAuth.mcp_admitted_user_subject`` field, set exclusively by
``_reload_admitted_user`` at admission. It is deliberately NOT a ``metadata`` key: virtual-key
metadata is caller-controlled at key creation, so a metadata marker could be forged on a
personal key to gain the team-inherited grant union or to dodge the caller-Authorization
egress scrub. This field cannot be set from caller input."""
Reads the server-only ``mcp_admitted_user_subject`` field, set only by ``_reload_admitted_user``. It
is deliberately NOT a ``metadata`` key, which is caller-controlled at key creation and so forgeable
on a personal key to gain the team grant union or dodge the egress scrub; this field cannot be."""
return user_api_key_auth is not None and user_api_key_auth.mcp_admitted_user_subject is True
@ -391,11 +389,9 @@ class MCPRequestHandler:
and oauth2_headers
and is_session_bearer_shaped(oauth2_headers["Authorization"])
):
# A gateway DCR session bearer at the aggregate /mcp scope: open the
# identity-only session token and admit under the live litellm user it
# references. A session-shaped bearer that does not open fails closed with
# the aggregate invalid_token challenge; a non-session bearer never reaches
# here (is_session_bearer_shaped is false) and falls through to the oauth2 arm.
# A gateway DCR session bearer at the aggregate /mcp scope: open the identity-only session
# token and admit under the live litellm user. One that does not open fails closed with the
# aggregate invalid_token challenge; a non-session bearer falls through to the oauth2 arm.
validated_user_api_key_auth = await MCPRequestHandler._admit_gateway_session(
authorization_value=oauth2_headers["Authorization"],
request=request,
@ -431,13 +427,10 @@ class MCPRequestHandler:
bearer_presented=False,
)
# Leak-defense (single chokepoint): a gateway admission credential — the session bearer or the
# bridge envelope — is NEVER a valid upstream MCP token. Scrub it from EVERY egress header context
# (top-level Authorization, the deprecated `x-mcp-auth`, and per-server `x-mcp-{alias}-authorization`)
# so no client-forwarded, OBO-subject, or passthrough path can send it upstream, where a hostile
# server could capture and replay it against the aggregate endpoint as this user. Anchored to the
# credential SHAPE, so a legitimate upstream/passthrough token (never session- or envelope-shaped)
# is forwarded unchanged; per-server vaulted credentials (resolved at egress) are unaffected.
# Leak-defense (single chokepoint): a gateway admission credential (session bearer or bridge
# envelope) is NEVER a valid upstream token. Scrub it from EVERY egress context so no
# client-forwarded, OBO, or passthrough path can send it upstream for replay. Anchored to the
# credential SHAPE, so a legitimate upstream/passthrough token is forwarded unchanged.
raw_headers = dict(headers)
(
oauth2_headers,
@ -463,10 +456,9 @@ class MCPRequestHandler:
@staticmethod
def _is_gateway_admission_credential(value: str | None) -> bool:
"""True when a header value is a gateway admission credential — a session bearer (``llm_session_`` /
``llm_srefresh_``) or a bridge envelope. Such a value proves who signed in to the GATEWAY; it is
never a valid credential for an UPSTREAM MCP server, so it must never be forwarded, where a hostile
upstream could capture and replay it against the aggregate ``/mcp`` endpoint as this user."""
"""True when a header value is a gateway admission credential — a session bearer or bridge
envelope. It proves who signed in to the GATEWAY, never a valid UPSTREAM token, so it must never
be forwarded (a hostile upstream could capture and replay it against the aggregate ``/mcp`` scope)."""
return value is not None and (is_session_bearer_shaped(value) or is_bridge_envelope_shaped(value))
@staticmethod
@ -478,13 +470,10 @@ class MCPRequestHandler:
mcp_server_auth_headers: dict[str, dict[str, str]] | None,
) -> tuple[dict[str, str] | None, dict[str, str], str | None, dict[str, dict[str, str]] | None]:
"""Remove any gateway admission credential from EVERY egress header context, keyed on the credential
SHAPE: the top-level ``Authorization`` (``oauth2_headers`` + ``raw_headers``), the deprecated
``x-mcp-auth`` (``mcp_auth_header``), and per-server ``x-mcp-{alias}-authorization``
(``mcp_server_auth_headers``). A legitimate upstream/passthrough token is never session- or
envelope-shaped, so it is forwarded unchanged; the per-server token the bridge arm injects is the
real upstream credential (also not gateway-shaped), so it survives. An admitted subject's top-level
Authorization IS the admission bearer, so it is dropped unconditionally as defense-in-depth even
though it is already gateway-shaped."""
SHAPE: top-level ``Authorization`` (oauth2 + raw), the deprecated ``x-mcp-auth``, and per-server
``x-mcp-{alias}-authorization``. A legitimate upstream/passthrough token is never gateway-shaped so
it survives (including the real upstream token the bridge arm injects per-server); an admitted
subject's top-level Authorization is dropped unconditionally as defense-in-depth."""
cred = MCPRequestHandler._is_gateway_admission_credential
# 1. Top-level Authorization → oauth2_headers.
@ -745,21 +734,12 @@ class MCPRequestHandler:
) -> UserAPIKeyAuth:
"""Open a gateway DCR session bearer and admit the live litellm user it references.
The custody sibling of :meth:`_admit_dcr_bridge_delegate`: the session token seals
no upstream credential (those are vaulted per user and resolved at egress), so this
admits identity only and injects no per-server header. The token's signature proves
the user signed in when it was minted, but authorization is resolved fresh here, the
sealed ``user_id`` reloads the current user record through the SAME
:meth:`_reload_admitted_user` the bridge user-subject path uses, and the admitted
identity runs through the centralized policy gate, so the user's present team, org,
budget, and SCIM state gate the request rather than a snapshot frozen at mint time.
Fails closed with the aggregate ``invalid_token`` challenge on an expired, tampered,
or foreign token, on a refresh token presented at the tool edge, and when the
referenced user is missing, deactivated, or rejected by the policy gate. The
pre-DB gates (size, IP, route allowlist) run first, mirroring the bridge arm and the
standard pipeline, so a caller blocked by IP or route is turned away before any
crypto or DB read."""
Identity-only sibling of :meth:`_admit_dcr_bridge_delegate`: the session token seals no
upstream credential (those are vaulted per user, resolved at egress), so authorization is
resolved fresh via :meth:`_reload_admitted_user` + the centralized policy gate rather than a
mint-time snapshot. Pre-DB gates (size, IP, route allowlist) run first, mirroring the standard
pipeline. Fails closed with the aggregate ``invalid_token`` challenge on an expired, tampered,
foreign, or refresh token, or a missing/deactivated/policy-rejected user."""
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import (
NotSessionBearer,
SessionBearerAdmitted,
@ -842,34 +822,18 @@ class MCPRequestHandler:
@staticmethod
async def _reload_admitted_user(user_id: str) -> UserAPIKeyAuth:
"""Reload the live user an interactively-minted envelope references and admit them as
themselves.
"""Reload the live user an interactively-minted envelope references and admit them as themselves.
The DCR client authenticates via SSO at the bridged authorize, which yields a user
subject rather than a virtual key, so the envelope admits under the user's own
identity: the reloaded ``user_id``, the user's own MCP object permission, and the user's
``org_id`` ride on the returned ``UserAPIKeyAuth``, and the SAME ``get_allowed_mcp_servers``
the key path uses then computes which servers the user may reach, so the user's litellm MCP
grants and access groups gate the request exactly as a key's do. Because the returned auth is
stamped ``mcp_admitted_user_subject`` (below), ``get_allowed_mcp_servers`` unions the servers the
user reaches through ANY of their teams on top of these direct grants a ``UserAPIKeyAuth``
pins one ``team_id`` but a user belongs to many, so the team fan-out happens off the marker, not
the single ``team_id``. Each source is bounded by ITS OWN org: the user's direct grants by the
bound ``org_id`` (their primary org), and each team's grant by that team's owning org inside
``_allowed_mcp_servers_for_single_team`` so a user who spans organizations does not leak one
org's servers past another org's ceiling. The caller's centralized policy gate enforces the
user's live budget and org state, and a SCIM-deactivated owner fails closed.
The user's own object permission and ``org_id`` ride on the returned ``UserAPIKeyAuth``, and the
SAME ``get_allowed_mcp_servers`` the key path uses gates the request. The ``mcp_admitted_user_subject``
marker (set below) makes that resolver union the servers the user reaches through ANY of their teams
on top of these direct grants, each source bounded by ITS OWN org, so a user spanning organizations
cannot leak one org's servers past another's ceiling.
Error handling mirrors the key path's retryable-503 contract, but ``get_user_object`` defeats a
type-based check: where ``get_key_object`` raises a typed ``ProxyException`` for a missing key
and lets a DB outage propagate raw, ``get_user_object`` catches every DB failure and re-raises a
bare ``ValueError``, so a missing user and a real outage look identical and the original error
survives only as ``__context__``. ``_raise_503_if_db_unavailable`` therefore walks the cause
chain: a transient DB outage still surfaces as a retryable 503, while a missing user, or any
other non-outage resolution failure, fails closed as a 401 rather than an opaque 500. The
object-permission load shares this one boundary, so an outage there is classified the same
way (``get_object_permission`` itself swallows a failed load to ``None``, matching how
``get_key_object`` best-effort-loads a key's object permission)."""
Error handling: ``get_user_object`` catches every DB failure and re-raises a bare ``ValueError``, so a
missing user and a real outage look identical (the cause survives only as ``__context__``).
``_raise_503_if_db_unavailable`` walks the cause chain so an outage stays a retryable 503 while any
other failure fails closed as 401, not an opaque 500; the object-permission load shares that boundary."""
from litellm.proxy.auth.auth_checks import get_object_permission, get_user_object
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
@ -907,49 +871,35 @@ class MCPRequestHandler:
org_id=user_object.organization_id,
object_permission=object_permission,
object_permission_id=user_object.object_permission_id,
# Copy the live user's rate limits, exactly as the standard user-subject auth path does
# (user_api_key_auth.py). The parallel limiter reads these off the auth object rather than
# re-fetching, and treats None as sys.maxsize (unlimited), so a keyless admitted user with
# them unset would invoke tools past their configured user RPM/TPM.
#
# Rate-limit model for the keyless admitted subject: bounded by their USER rpm/tpm
# (copied here; those descriptors key off user_id, which is set) AND by the per-server
# mcp_rpm_limit of EVERY team it reaches servers through, stamped below. Per-KEY MCP
# limits genuinely do not apply, because there is no key.
# Copy the live user's rate limits, as the standard user-subject path does: the parallel
# limiter reads these off the auth object and treats None as unlimited, so a keyless subject
# with them unset would outrun its user RPM/TPM. (Per-team mcp_rpm_limit is stamped below;
# per-KEY limits do not apply, there being no key.)
user_tpm_limit=user_object.tpm_limit,
user_rpm_limit=user_object.rpm_limit,
)
# Set the server-only admission marker AFTER construction: the before-validator strips it
# from any validated input, so a post-construction assignment is the only way to set it, and
# caller-supplied data (key metadata, JWT claims) can never forge it.
# Server-only marker, set AFTER construction: the before-validator strips it from any validated
# input, so caller-supplied data (key metadata, JWT claims) can never forge it.
admitted.mcp_admitted_user_subject = True
# Carry each granting team's per-server MCP rpm limit. A key is pinned to one team so the
# limiter reads team_metadata directly; this subject reaches servers through several teams
# under its own identity, so without this the team ceiling silently does not apply to it and
# a cross-team user outruns every team's mcp_rpm_limit. Resolved from the same roster-checked
# sources the grant union uses, so a team can only throttle what it actually granted.
# Carry each granting team's per-server mcp_rpm_limit: this subject reaches servers through
# several teams under its own identity, so without this a cross-team user outruns every team's
# limit. Resolved from the same roster-checked sources as the grant union, so a team throttles
# only what it granted.
admitted.mcp_source_team_rpm_limits = await MCPRequestHandler._admitted_subject_team_rpm_limits(admitted)
return admitted
@staticmethod
async def _admitted_subject_team_rpm_limits(auth: UserAPIKeyAuth) -> dict[str, dict[str, int]] | None:
"""``team_id -> mcp_rpm_limit`` for every team this subject reaches servers through, with each
team's map filtered to the servers THAT team's grant actually reaches.
"""``team_id -> mcp_rpm_limit`` for every team this subject reaches servers through, each map
filtered to the servers THAT team's grant actually reaches.
A limit rides the same scope as the access it bounds: a team's throttle exists to cap usage of
the access the team granted, so a roster team whose grant does not reach a server (not granted,
blocked, org-forbidden, opted out) must not be charged when the user reaches that server
through a DIFFERENT team otherwise this user's calls drain a bucket shared by that team's own
keys for access the team never provided. The grant scope comes from the SAME
``get_allowed_mcp_servers(source)`` call authorization uses, so the throttle scope cannot
diverge from the access scope. Limit maps are keyed by server name/alias (the limiter matches
on the called server's name) while grants are ids, so each key is resolved through
``expand_permission_list`` the one existing name->id owner before the membership check.
Returns None when no team contributes an applicable limit, so the limiter adds no descriptors
rather than empty ones. A lookup failure narrows to None rather than raising: rate limiting
must not be able to deny a request that authorization already allowed, and the user's own
rpm/tpm still bounds them."""
A limit rides the same scope as the access it bounds, so a roster team is charged only for a
server its OWN grant reaches (never one the user reaches through a different team, which would
drain a bucket shared by that team's keys for access it never provided). Grant scope comes from
the SAME ``get_allowed_mcp_servers(source)`` authorization uses; limit-map keys are names/aliases
so each is resolved to an id via ``expand_permission_list`` before the membership check. Returns
None (no descriptors) when nothing applies; a lookup failure narrows to None rather than raising,
since rate limiting must not deny a request authorization already allowed."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
@ -969,12 +919,9 @@ class MCPRequestHandler:
for server_id in global_mcp_server_manager.expand_permission_list([server_name]):
if server_id not in granted_ids:
continue
# Charge ONLY the source the call is attributed to — the same single source
# billing picks, from the same owner. Adding a descriptor for every granting
# team let one cross-team user drain several teams' SHARED buckets at once,
# blocking their other members for access those teams did not provide on
# this call; and when the user's OWN grant reaches the server, no team
# provided it, so no team bucket is charged at all.
# Charge ONLY the source billing attributes the call to (same owner), so one
# cross-team user cannot drain several teams' shared buckets on a single call,
# and a server the user's OWN grant reaches charges no team bucket.
attributed = await MCPRequestHandler.attributing_source_for_server(
auth, server_id, source_grants=source_grants
)
@ -1358,11 +1305,10 @@ class MCPRequestHandler:
from litellm.proxy.proxy_server import general_settings
try:
# A keyless admitted subject is resolved entirely per source, BEFORE any single-source
# rule runs here. Ordering matters: the no_mcp_servers opt-out below reads the caller's
# own object_permission, so leaving it above this branch let a user's own opt-out zero
# their TEAMS' grants too — the sources are independent, and an opt-out on one of them
# must silence only that one (it is applied per source, inside the recursive call).
# A keyless admitted subject resolves per source BEFORE any single-source rule here. Ordering
# matters: the no_mcp_servers opt-out below reads the caller's own object_permission, so above
# this branch a user's own opt-out would wrongly zero their TEAMS' grants too (each source is
# independent; an opt-out silences only its own source, inside the recursive call).
if _is_mcp_admitted_user_subject(user_api_key_auth) and user_api_key_auth is not None:
return await MCPRequestHandler._resolve_admitted_subject_servers(user_api_key_auth)
@ -1484,22 +1430,14 @@ class MCPRequestHandler:
has_lower_level_mcp_restrictions: bool,
keyless_source: bool = False,
) -> list[str]:
"""Cap the resolved server list by this caller's org ceiling. If the org names an explicit MCP
list, lower-level restrictions are intersected with it, else the org list becomes the ceiling.
No org, or an empty org list, leaves the result unchanged.
"""Cap the resolved server list by this caller's org ceiling: an explicit org list intersects
lower-level restrictions (else becomes the ceiling); no org or an empty list leaves it unchanged.
``keyless_source`` marks one grant source of a keyless admitted subject and governs BOTH
org divergences, because they are the same fact about that caller shape.
First, what an UNRESOLVABLE ceiling means. A virtual key keeps the
long-standing fail-open behavior (a DB blip must not lock working keys out mid-incident). A
keyless admitted subject fails CLOSED, because its only org bound is this ceiling: silently
dropping it on a transient fault would widen a cross-org user to servers their team's org
forbids, which is a privilege escalation rather than an availability blip.
Second, whether the org list may SUBSTITUTE for absent lower-level grants. For a key it may
(that is the key ceiling model). For a source it may only ever intersect, because the
admitted model is a union of grants and a ceiling that grants is not a ceiling."""
``keyless_source`` governs both divergences for a keyless admitted source. An UNRESOLVABLE ceiling
fails CLOSED for it (its only org bound is this ceiling, so dropping it on a fault would escalate a
cross-org user) while a key stays fail-open. And an org list may only ever INTERSECT a source (the
admitted model unions grants, so a ceiling must not become one), whereas for a key it may
substitute, that being the key ceiling model."""
if not (user_api_key_auth and user_api_key_auth.org_id):
return allowed_mcp_servers
allowed_mcp_servers_for_org = await MCPRequestHandler._get_allowed_mcp_servers_for_org(user_api_key_auth)
@ -1512,12 +1450,8 @@ class MCPRequestHandler:
if len(allowed_mcp_servers_for_org) == 0:
return allowed_mcp_servers
if has_lower_level_mcp_restrictions or keyless_source:
# Lower-level restrictions exist, so org can only cap them.
#
# A keyless admitted source ALWAYS takes this arm: its model is a union of GRANTS, so an
# org list may only narrow what a source already grants, never become one. Letting it
# substitute would hand every admitted user with an org_id that org's whole server list
# without any direct or team grant — a ceiling silently acting as a grant.
# Org can only cap lower-level restrictions. A keyless admitted source ALWAYS takes this
# arm: its model unions GRANTS, so an org list may only narrow a source, never become one.
capped = [s for s in allowed_mcp_servers if s in allowed_mcp_servers_for_org]
else:
# No lower-level restrictions → org list becomes the ceiling.
@ -1535,15 +1469,12 @@ class MCPRequestHandler:
) -> UserAPIKeyAuth:
"""A plain, UNMARKED auth describing ONE grant source of an admitted subject.
Only the fields the resolver actually consults are carried. Everything else is left at its
default on purpose: ``api_key``/``token`` stay unset (this is not a key), budget, spend and
rate-limit fields stay unset because the admitted subject's own user-level limits are what
the request is metered against and cloning them per source would show the limiter N copies of
the same descriptor, and ``user_role`` stays unset because an admin role would grant every
server if this auth ever reached the server-manager wrapper. The admission marker cannot be
set through the constructor at all (a before-validator pops it), so each source is resolved
as an ordinary caller and cannot re-enter the admitted path.
"""
Only the fields the resolver consults are carried; everything else is left at its default on
purpose: no ``api_key``/``token`` (not a key), no budget/spend/rate-limit (the subject's own
user-level limits meter the request, and per-source copies would double descriptors), no
``user_role`` (an admin role would grant every server at the server-manager wrapper). The
admission marker cannot be set via the constructor (a before-validator pops it), so each source
resolves as an ordinary caller and cannot re-enter the admitted path."""
scoped = UserAPIKeyAuth(
user_id=auth.user_id,
team_id=team_id,
@ -1551,9 +1482,8 @@ class MCPRequestHandler:
parent_otel_span=auth.parent_otel_span,
)
if carry_user_grants:
# The user's OWN grants. A team source deliberately carries none of these: the resolver
# loads that team's object_permission and access groups from team_id itself, and mixing
# the user's in would widen the team source with grants the team never made.
# The user's OWN grants. A team source carries none of these (the resolver loads the team's
# own object_permission from team_id); mixing them in would widen the team with grants it never made.
scoped.object_permission = auth.object_permission
scoped.object_permission_id = auth.object_permission_id
scoped.access_group_ids = auth.access_group_ids
@ -1564,17 +1494,11 @@ class MCPRequestHandler:
"""The independent sources a keyless admitted subject reaches MCP servers through: their own
direct grants, plus every team they are a live roster member of.
Each team source carries that TEAM's org as its ``org_id``, which is what makes the canonical
resolver apply the team's OWN owning-org ceiling to it — a cross-org user's teams are each
bounded by their own org rather than by the caller's home org. A team with no organization
falls back to the user's org so it is bounded rather than unbounded.
Roster membership is checked HERE because it is a property of the source list, not of any one
resolution: a key is structurally pinned to a team it belongs to, while a user's cached
``teams`` array can name a team whose ``members_with_roles`` no longer contains them (SCIM
group sync, or cache lag after a team_member_delete), and JWT auth can rewrite that array
outright. The roster is the source of truth for revocation.
"""
Each team source carries that TEAM's org (falling back to the user's), so the canonical resolver
applies the team's OWN owning-org ceiling — a cross-org user's teams are each bounded by their
own org, not the caller's home org. Roster membership is checked HERE (not per resolution)
because a user's cached ``teams`` array can name a team whose ``members_with_roles`` no longer
lists them; the roster is the source of truth for revocation."""
from litellm.proxy.proxy_server import prisma_client
sources = [
@ -1622,11 +1546,9 @@ class MCPRequestHandler:
proxy_logging_obj=proxy_logging_obj,
)
except Exception as e: # noqa: BLE001 # per-source isolation: one team's blip must not deny the others
# The unit of fault isolation is the SOURCE: a team that cannot be resolved contributes
# nothing this request (fail closed for that team alone — access only ever narrows),
# while the user's own grants and every other resolvable team stand. Raising here
# instead would collapse the whole union to deny-all because one team's row was
# momentarily unreadable, on the servers, tools and throttle axes alike.
# Fault isolation is per SOURCE: an unresolvable team contributes nothing (fail closed for
# it alone, access only narrows) while every other source stands. Raising would collapse the
# whole union to deny-all over one momentarily-unreadable row.
verbose_logger.warning(f"MCP admitted-subject source team {team_id!r} unresolvable, skipping: {str(e)}")
return None
if team_obj is None:
@ -1634,14 +1556,11 @@ class MCPRequestHandler:
member_user_ids = {getattr(m, "user_id", None) for m in (team_obj.members_with_roles or [])} - {None}
if auth.user_id not in member_user_ids:
return None
# A team over its own max budget — or owned by an org over ITS budget — is not a live
# grantor, exactly as it is not for a virtual key pinned to it (common_checks rejects that
# key outright). Enforced with the SAME owners the key path uses (_team_max_budget_check /
# _organization_max_budget_check, cross-pod Redis-first spend), targeted at the TEAM's org
# via the scoped source view, so a cross-org team is judged by its own org's budget. This is
# budget ENFORCEMENT of an already-exceeded state; ATTRIBUTION of new spend stays with the
# user (documented deferral) — the two are different questions. Sitting here, no consumer of
# the source list (servers, tools, throttle stamping) can ever see an over-budget team.
# A team (or its owning org) over budget is not a live grantor, exactly as it is not for a key
# pinned to it. Enforced via the SAME owners the key path uses (_team_max_budget_check /
# _organization_max_budget_check), targeted at the TEAM's org through the scoped source view, so
# no consumer of the source list ever sees an over-budget team. This is ENFORCEMENT of an
# already-exceeded state; ATTRIBUTION of new spend stays with the user (documented deferral).
from litellm.exceptions import BudgetExceededError
from litellm.proxy.auth.auth_checks import (
_organization_max_budget_check,
@ -1696,16 +1615,11 @@ class MCPRequestHandler:
async def billing_auth_for_tool_call(auth: UserAPIKeyAuth, tool_name: str) -> UserAPIKeyAuth:
"""The auth object a tool call's SPEND should be recorded against.
Returns ``auth`` unchanged for every caller that is not a keyless admitted subject, so key
and JWT billing is byte-identical. For an admitted subject whose call is reached through a
team's grant, returns a copy carrying that team's ``team_id`` and its owning ``org_id`` so
the team's budget accumulates and the correct organization is charged.
Inert rather than wrong when the target server cannot be resolved from the tool name (a
display-name override, or a REST caller passing server_id with an unprefixed name): billing
then falls back to today's user-level attribution instead of guessing a team. Resolution
reuses the manager's own tool-name lookup rather than re-deriving prefix rules that live
there."""
``auth`` unchanged for any non-admitted caller (key/JWT billing byte-identical). For an admitted
subject whose call is reached through a team's grant, a copy carrying that team's ``team_id`` and
owning ``org_id`` so the team's budget accumulates and the right org is charged. Falls back to
user-level attribution (rather than guessing a team) when the tool name does not resolve to a
server, reusing the manager's own tool-name lookup."""
if not _is_mcp_admitted_user_subject(auth):
return auth
try:
@ -1733,19 +1647,14 @@ class MCPRequestHandler:
server_id: str,
source_grants: list[tuple[UserAPIKeyAuth, set[str]]] | None = None,
) -> UserAPIKeyAuth | None:
"""The source a billable call to ``server_id`` is attributed to, or None to bill the caller
as themselves (their own grant reaches it, or nothing does).
"""The source a billable call to ``server_id`` is attributed to, or None to bill the caller as
themselves (their own grant reaches it, or nothing does).
A keyless admitted subject carries no ``team_id``, so downstream spend skipped team updates
entirely and charged the user's PRIMARY org — a team-derived call neither accumulated its
team's budget (so that budget could never begin to block) nor charged the org that owns the
granting team. Attribution restores both.
The rule: a user's OWN grant is not "through a team", so it bills the user. Otherwise the
call is billed to a granting team deterministically the lowest ``team_id`` when several
grant the same server, so the choice is stable, reproducible and auditable rather than
dependent on dict ordering. Reads the one grant owner, so the team that gets billed is
always a team that actually granted the server."""
The rule: a user's OWN grant is not "through a team", so it bills the user; otherwise the call
bills a granting team, deterministically the lowest ``team_id`` when several grant the server so
the pick is stable rather than dict-ordering-dependent. Reads the one grant owner, so the billed
team is always one that actually granted the server (restoring the team budget accrual and
owning-org charge that a keyless, team_id-less subject otherwise skipped)."""
source_grants = source_grants or await MCPRequestHandler.admitted_source_grants(auth)
granting = [(source, granted) for source, granted in source_grants if server_id in granted]
if not granting:
@ -1768,14 +1677,10 @@ class MCPRequestHandler:
global_mcp_server_manager,
)
# An OPEN channel (operator-opened allow_all_keys, the user's own BYOM submission) makes the
# server REACHABLE through the user themselves — no grant source names it, so without this
# the union below would return [] and leave it listable but uninvokable. Reachability is ALL
# it confers: it is not a waiver of the ceilings that bound the server. The user's own
# mcp_tool_permissions and their org's tool ceiling still bind, which is what a virtual key
# on the same allow_all server gets (its key_tools and _apply_agent_and_org_tool_ceilings
# both run). Returning None here instead skipped both and let a session holder invoke tools
# their own or their org's policy excludes.
# An OPEN channel (allow_all_keys, the user's own BYOM) makes the server REACHABLE through the
# user, though no grant source names it — without this the union returns [], listable but
# uninvokable. Reachability is ALL it confers, NOT a ceiling waiver: the user's own
# mcp_tool_permissions and org tool ceiling still bind, exactly as a key's do on an allow_all server.
reachable_via_open_channel = server_id in await global_mcp_server_manager.operator_open_server_ids(auth)
allowed: set[str] = set()
@ -1867,12 +1772,9 @@ class MCPRequestHandler:
return None
try:
# FIRST statement, mirroring get_allowed_mcp_servers: a keyless admitted subject is
# resolved per grant source and shares NOTHING with the single-credential prelude below.
# Ordering is the invariant, not a nicety — when this branch sat after the prelude, a
# fault in a lookup the subject never uses (its own mcp_toolsets, its team_obj_perm) hit
# the fail-closed handler and denied tools its teams did grant. Nothing that resolves a
# single credential's scope may run before this line.
# FIRST statement, mirroring get_allowed_mcp_servers: a keyless admitted subject resolves per
# source and shares nothing with the single-credential prelude below. Ordering is the invariant:
# sat after the prelude, a fault in a lookup the subject never uses denied tools its teams grant.
if _is_mcp_admitted_user_subject(user_api_key_auth):
return await MCPRequestHandler._resolve_admitted_subject_tools(server_id, user_api_key_auth)
@ -1917,9 +1819,6 @@ class MCPRequestHandler:
else None
)
# A keyless gateway/bridge-admitted user has no single team_id, so team_obj_perm above is
# None and the single-team lookup yields allow-all — silently dropping every team's
# per-server tool exclusions. Resolve it as the union over the sources that grant the
# Apply same inheritance logic as get_allowed_mcp_servers
if team_tools:
if key_tools:
@ -1938,15 +1837,10 @@ class MCPRequestHandler:
except Exception as e:
verbose_logger.warning(f"Failed to get allowed tools for server: {str(e)}")
# Fail CLOSED for a keyless admitted subject: ANY error resolving the tool allowlist
# (multi-team fan-out, org/agent lookups) must deny the server's tools ([]) for this
# request rather than collapse to allow-all (None), mirroring the fail-closed server
# path. Key/JWT auth keeps its prior allow-all-on-error behavior.
#
# keyless_source matters as much as the marker: each source of an admitted subject is
# resolved through an UNMARKED auth, so without it a fault under a source returned None,
# and None wins the union as allow-all — dropping every team and org tool ceiling on a
# blip. The marker alone only covers a fault raised before the fan-out.
# Fail CLOSED for a keyless admitted subject: ANY error must deny the server's tools ([]),
# not collapse to allow-all (None); key/JWT auth keeps its prior allow-all-on-error. Both
# keyless_source AND the marker are needed: each source resolves through an UNMARKED auth, so
# without keyless_source a fault under a source returns None and wins the union as allow-all.
return [] if (keyless_source or _is_mcp_admitted_user_subject(user_api_key_auth)) else None
@staticmethod
@ -1957,16 +1851,12 @@ class MCPRequestHandler:
keyless_source: bool = False,
) -> list[str] | None:
"""Narrow a key/team tool allowlist by the agent's tool permissions and the caller's org tool
ceiling. Each level only ever intersects, and None at a level means "no restriction from this
level".
ceiling. Each level only intersects; None at a level means no restriction from it.
An UNRESOLVABLE org ceiling (``_get_org_object_permission`` raises: the org names a permission
that cannot be loaded) is decided here, per caller shape, mirroring the servers axis: a
virtual key keeps its long-standing fail-open the org step is skipped and the key/team/agent
restrictions already computed STAND (letting the raise escape would collapse them to
allow-all, which is fail-open WIDER than before the fault). A keyless admitted source
re-raises, and the outer handler denies tools for that one source while the subject's other
sources stand its only org bound is this ceiling, so skipping it would widen access."""
An UNRESOLVABLE org ceiling is decided per caller shape, mirroring the servers axis: a key stays
fail-open (skip the org step, keep the key/team/agent restrictions; letting the raise escape
would collapse them to allow-all, WIDER than before the fault), while a keyless source re-raises
so the outer handler denies that one source (its only org bound is this ceiling)."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
@ -2346,9 +2236,8 @@ class MCPRequestHandler:
verbose_logger.debug("prisma_client is None")
return None
# An ABSENT org is a determinate fact, not a failure: a team's organization_id can point at a
# row that was deleted or has not synced yet, and get_org_object raises for that. It places no
# ceiling, exactly as a key with a dangling org_id is not locked out.
# A team's organization_id can point at a deleted or not-yet-synced row; get_org_object raises
# OrganizationNotFoundError for that. That is a determinate ABSENCE (no ceiling), handled below.
try:
org_obj = await get_org_object(
org_id=user_api_key_auth.org_id,
@ -2358,20 +2247,17 @@ class MCPRequestHandler:
proxy_logging_obj=proxy_logging_obj,
)
except OrganizationNotFoundError as e:
# CONFIRMED absent (deleted org, not-yet-synced organization_id): a determinate fact, so
# it places no ceiling. Every OTHER exception is an operational failure and propagates —
# caught upstream as an unresolvable ceiling, which denies for a keyless source and stays
# fail-open for a key. Catching bare Exception here treated a DB outage as "no org", which
# silently dropped a real org's ceiling for exactly as long as the outage lasted.
# CONFIRMED absent: places no ceiling. Every OTHER exception propagates as an unresolvable
# ceiling (denies for a keyless source, fail-open for a key); catching bare Exception here
# would treat a DB outage as "no org" and silently drop a real ceiling for its duration.
verbose_logger.debug(f"MCP org ceiling: org {user_api_key_auth.org_id!r} does not exist: {e}")
return None
if org_obj is None or not org_obj.object_permission_id:
return None
# From here the org NAMES a permission. Failing to read it is INDETERMINATE, so it must not
# collapse into the same None that means "no ceiling" -- that is what would silently drop a
# real ceiling on a transient fault. Raise and let each caller pick fail-open or fail-closed.
# The org NAMES a permission; failing to read it is INDETERMINATE and must not collapse into the
# None that means "no ceiling". Raise and let each caller pick fail-open or fail-closed.
object_permission = await get_object_permission(
object_permission_id=org_obj.object_permission_id,
prisma_client=prisma_client,
@ -2420,9 +2306,8 @@ class MCPRequestHandler:
all_servers = direct_mcp_servers + access_group_servers + tool_perm_servers
return list(set(all_servers))
except Exception as e:
# None = the org ceiling could NOT be resolved, which is not the same fact as [] = the
# org places no restriction. Collapsing the two is what let a transient DB fault silently
# remove an org's ceiling; the caller picks fail-open or fail-closed from this signal.
# None = ceiling UNRESOLVED, distinct from [] = org places no restriction. Collapsing them
# let a DB fault silently drop a ceiling; the caller picks fail-open/closed from this signal.
verbose_logger.warning(f"Failed to get allowed MCP servers for org: {str(e)}")
return None

View file

@ -6405,28 +6405,33 @@ class TestGatewaySessionAdmission:
# bridge envelope arm); the headers dict is whatever the request carried, here empty.
assert not mcp_server_auth_headers
async def test_expired_session_fails_closed_with_invalid_token_challenge(self):
@pytest.mark.parametrize(
"scenario, expect_challenge",
[("expired", True), ("tampered", False), ("refresh_at_tool_edge", False), ("foreign_key", False)],
)
async def test_bad_session_bearer_fails_closed(self, scenario, expect_challenge):
# Every non-admissible session-shaped bearer fails closed with 401; a valid-but-unusable one
# (expired) additionally carries the invalid_token challenge so the DCR client re-authorizes.
from datetime import datetime, timezone
mint, _refresh, principal, keys = self._session_bearer()
token = mint(principal, keys, datetime(2020, 1, 1, tzinfo=timezone.utc)).token.get_secret_value()
with (
patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY),
):
if scenario == "expired":
mint, _refresh, principal, keys = self._session_bearer()
bearer = mint(principal, keys, datetime(2020, 1, 1, tzinfo=timezone.utc)).token.get_secret_value()
elif scenario == "tampered":
token = self._access_token()
bearer = token[:-3] + ("aaa" if not token.endswith("aaa") else "bbb")
elif scenario == "refresh_at_tool_edge":
_mint, refresh, principal, keys = self._session_bearer()
bearer = refresh(principal, keys, datetime(2030, 1, 1, tzinfo=timezone.utc)).token.get_secret_value()
else: # foreign_key: minted under the real master key, presented while the proxy uses another
bearer = self._access_token()
master_key = "sk-a-totally-different-master-key" if scenario == "foreign_key" else self._MASTER_KEY
with patch("litellm.proxy.proxy_server.master_key", master_key):
with pytest.raises(HTTPException) as exc_info:
await MCPRequestHandler.process_mcp_request(self._scope(token))
assert exc_info.value.status_code == 401
assert 'error="invalid_token"' in (exc_info.value.headers or {})["WWW-Authenticate"]
async def test_tampered_session_fails_closed(self):
token = self._access_token()
tampered = token[:-3] + ("aaa" if not token.endswith("aaa") else "bbb")
with (
patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY),
):
with pytest.raises(HTTPException) as exc_info:
await MCPRequestHandler.process_mcp_request(self._scope(tampered))
await MCPRequestHandler.process_mcp_request(self._scope(bearer))
assert exc_info.value.status_code == 401
if expect_challenge:
assert 'error="invalid_token"' in (exc_info.value.headers or {})["WWW-Authenticate"]
async def test_deactivated_user_fails_with_invalid_token_challenge(self):
"""A cryptographically valid bearer whose referenced user is SCIM-deactivated must fail with
@ -6466,27 +6471,6 @@ class TestGatewaySessionAdmission:
assert oauth2_headers is None
assert not any(k.lower() == "authorization" for k in (raw_headers or {}))
async def test_refresh_token_is_not_admitted_at_the_tool_edge(self):
from datetime import datetime, timezone
_mint, refresh, principal, keys = self._session_bearer()
refresh_token = refresh(principal, keys, datetime(2030, 1, 1, tzinfo=timezone.utc)).token.get_secret_value()
with (
patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY),
):
with pytest.raises(HTTPException) as exc_info:
await MCPRequestHandler.process_mcp_request(self._scope(refresh_token))
assert exc_info.value.status_code == 401
async def test_foreign_key_session_fails_closed(self):
token = self._access_token()
with (
patch("litellm.proxy.proxy_server.master_key", "sk-a-totally-different-master-key"),
):
with pytest.raises(HTTPException) as exc_info:
await MCPRequestHandler.process_mcp_request(self._scope(token))
assert exc_info.value.status_code == 401
async def test_arm_does_not_fire_for_named_server(self):
"""A session-shaped bearer aimed at a named server (path scope) does not enter the
aggregate arm; it is treated as an ordinary bearer on that server."""
@ -6504,24 +6488,41 @@ class TestGatewaySessionAdmission:
mock_auth.assert_called_once()
def _make_team(team_id, mcp_servers, *, org_id=None, tool_perms=None, members=("sso-user",)):
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, Member
return LiteLLM_TeamTable(
team_id=team_id,
organization_id=org_id,
members_with_roles=[Member(user_id=u, role="user") for u in members],
access_group_ids=[],
object_permission=LiteLLM_ObjectPermissionTable(
object_permission_id=f"op-{team_id}", mcp_servers=mcp_servers, mcp_tool_permissions=tool_perms
),
)
def _make_admitted_subject(user_id, *, org_id=None, own_servers=None, own_tool_perms=None):
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
op = None
if own_servers is not None or own_tool_perms is not None:
op = LiteLLM_ObjectPermissionTable(
object_permission_id=f"userop-{user_id}",
mcp_servers=own_servers or [],
mcp_tool_permissions=own_tool_perms,
)
auth = UserAPIKeyAuth(user_id=user_id, api_key=None, org_id=org_id, object_permission=op)
auth.mcp_admitted_user_subject = True
return auth
@pytest.mark.asyncio
class TestUserSubjectTeamUnion:
"""_get_allowed_mcp_servers_for_team unions across ALL a user's teams for a keyless
user-subject caller (the gateway DCR session bearer and bridge user-envelope), while a
key-based caller keeps its single-team behavior byte-identically."""
def _team(self, team_id, mcp_servers, members=("sso-user",), tool_perms=None):
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, Member
return LiteLLM_TeamTable(
team_id=team_id,
members_with_roles=[Member(user_id=u, role="user") for u in members],
access_group_ids=[],
object_permission=LiteLLM_ObjectPermissionTable(
object_permission_id=f"op-{team_id}", mcp_servers=mcp_servers, mcp_tool_permissions=tool_perms
),
)
@contextlib.contextmanager
def _patch(self, *, teams_by_id, user_teams=None, orgs_by_id=None):
async def _get_team_object(team_id, **kw):
@ -6550,15 +6551,9 @@ class TestUserSubjectTeamUnion:
):
yield
@staticmethod
def _admitted_subject(user_id):
auth = UserAPIKeyAuth(user_id=user_id, api_key=None)
auth.mcp_admitted_user_subject = True
return auth
async def test_keyless_user_unions_servers_across_all_their_teams(self):
teams = {"team-a": self._team("team-a", ["srv1", "srv2"]), "team-b": self._team("team-b", ["srv2", "srv3"])}
auth = self._admitted_subject("sso-user")
teams = {"team-a": _make_team("team-a", ["srv1", "srv2"]), "team-b": _make_team("team-b", ["srv2", "srv3"])}
auth = _make_admitted_subject("sso-user")
with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"]):
result = await MCPRequestHandler.get_allowed_mcp_servers(auth)
assert set(result) == {"srv1", "srv2", "srv3"}
@ -6566,7 +6561,7 @@ class TestUserSubjectTeamUnion:
async def test_key_based_caller_uses_single_team_only(self):
"""A key-based caller (api_key set) with a team_id sees ONLY that team, even though the
same user belongs to other teams: key auth must be byte-identical to before."""
teams = {"team-a": self._team("team-a", ["srv1"]), "team-b": self._team("team-b", ["srv2", "srv3"])}
teams = {"team-a": _make_team("team-a", ["srv1"]), "team-b": _make_team("team-b", ["srv2", "srv3"])}
auth = UserAPIKeyAuth(user_id="sso-user", api_key="sk-hash", team_id="team-a")
with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"]):
result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth)
@ -6575,14 +6570,14 @@ class TestUserSubjectTeamUnion:
async def test_keyless_user_with_explicit_team_id_uses_that_team_only(self):
"""A keyless caller that already pins a team_id (not the user-subject fan-out shape)
resolves only that team; the union is strictly for the no-team-id user-subject case."""
teams = {"team-a": self._team("team-a", ["srv1"]), "team-b": self._team("team-b", ["srv2"])}
teams = {"team-a": _make_team("team-a", ["srv1"]), "team-b": _make_team("team-b", ["srv2"])}
auth = UserAPIKeyAuth(user_id="sso-user", api_key=None, team_id="team-a")
with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"]):
result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth)
assert set(result) == {"srv1"}
async def test_keyless_user_with_no_teams_gets_nothing_from_teams(self):
auth = self._admitted_subject("lonely-user")
auth = _make_admitted_subject("lonely-user")
with self._patch(teams_by_id={}, user_teams=[]):
result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth)
assert result == []
@ -6606,7 +6601,7 @@ class TestUserSubjectTeamUnion:
# those pins a team_id, so this helper only ever answers the single-team question. The fan-out
# itself is _admitted_subject_sources' job, asserted below.
with self._patch(teams_by_id={}, user_teams=["t2", "t3"]):
assert await MCPRequestHandler._team_ids_for_mcp_grant(self._admitted_subject("u")) == []
assert await MCPRequestHandler._team_ids_for_mcp_grant(_make_admitted_subject("u")) == []
# keyless, no user_id -> nothing
assert await MCPRequestHandler._team_ids_for_mcp_grant(UserAPIKeyAuth(api_key=None)) == []
# keyless with a user_id but NOT admission-marked (JWT auth) -> nothing (unchanged behavior)
@ -6629,9 +6624,9 @@ class TestUserSubjectTeamUnion:
outage -> the keyless source denies."""
from litellm.proxy.auth.auth_checks import OrganizationNotFoundError
teams = {"t1": self._team("t1", ["srv1"])}
teams = {"t1": _make_team("t1", ["srv1"])}
teams["t1"].organization_id = "org-a"
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
absent = AsyncMock(side_effect=OrganizationNotFoundError("Organization doesn't exist in db."))
with self._patch(teams_by_id=teams, user_teams=["t1"]):
@ -6657,7 +6652,7 @@ class TestUserSubjectTeamUnion:
boom = AsyncMock(side_effect=RuntimeError("org lookup exploded"))
# The subject must actually REACH something, or the assertion passes either way and pins
# nothing (a fail-open mutant survived an earlier version of this test for exactly that).
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
auth.org_id = "org-a"
auth.object_permission = LiteLLM_ObjectPermissionTable(object_permission_id="op-u", mcp_servers=["srv1"])
with self._patch(teams_by_id={}, user_teams=[]):
@ -6667,7 +6662,7 @@ class TestUserSubjectTeamUnion:
assert admitted == [], "admitted subject must fail CLOSED when its org ceiling cannot resolve"
key_auth = UserAPIKeyAuth(user_id="u", api_key="sk-hash", team_id="t1", org_id="org-a")
with self._patch(teams_by_id={"t1": self._team("t1", ["srv1"])}, user_teams=[]):
with self._patch(teams_by_id={"t1": _make_team("t1", ["srv1"])}, user_teams=[]):
with patch.object(MCPRequestHandler, "_get_org_object_permission", boom):
keyed = await MCPRequestHandler.get_allowed_mcp_servers(key_auth)
assert set(keyed) == {"srv1"}, "key auth must keep its long-standing fail-open behavior"
@ -6679,11 +6674,11 @@ class TestUserSubjectTeamUnion:
SAME source billing picks one owner for both, so they cannot disagree."""
from litellm.proxy.hooks.parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3
t1 = self._team("t1", ["srv1"])
t1 = _make_team("t1", ["srv1"])
t1.metadata = {"mcp_rpm_limit": {"srv1": 5}}
t2 = self._team("t2", ["srv1"])
t2 = _make_team("t2", ["srv1"])
t2.metadata = {"mcp_rpm_limit": {"srv1": 9}}
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
with self._patch(teams_by_id={"t1": t1, "t2": t2}, user_teams=["t1", "t2"]):
auth.mcp_source_team_rpm_limits = await MCPRequestHandler._admitted_subject_team_rpm_limits(auth)
billed = await MCPRequestHandler.attributing_source_for_server(auth, "srv1")
@ -6700,9 +6695,9 @@ class TestUserSubjectTeamUnion:
"""When the user's OWN grant reaches the server, no team provided the access, so no team
bucket may be charged the user's own rpm/tpm is what bounds them. Mirrors billing, which
bills the user and their own org for exactly this case."""
t1 = self._team("t1", ["srv1"])
t1 = _make_team("t1", ["srv1"])
t1.metadata = {"mcp_rpm_limit": {"srv1": 5}}
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
auth.object_permission = LiteLLM_ObjectPermissionTable(object_permission_id="op-u", mcp_servers=["srv1"])
with self._patch(teams_by_id={"t1": t1}, user_teams=["t1"]):
limits = await MCPRequestHandler._admitted_subject_team_rpm_limits(auth)
@ -6732,9 +6727,9 @@ class TestUserSubjectTeamUnion:
"""ACCOUNTING half of team budgets. Without attribution the admitted auth kept team_id=None,
so spend skipped team updates (the team's budget never accumulated, so it could never begin
to block) and charged the user's PRIMARY org rather than the org owning the granting team."""
t_grant = self._team("t-grant", ["srv1"])
t_grant = _make_team("t-grant", ["srv1"])
t_grant.organization_id = "org-team"
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
auth.org_id = "org-user-primary"
with self._patch(teams_by_id={"t-grant": t_grant}, user_teams=["t-grant"]):
source = await MCPRequestHandler.attributing_source_for_server(auth, "srv1")
@ -6746,9 +6741,9 @@ class TestUserSubjectTeamUnion:
"""Asserted on billing_auth_for_tool_call itself, not on the source it picks: the source
already carries the team's org by construction, so asserting there leaves the copy step
unpinned (a mutant dropping org_id survived exactly that). This is the object spend reads."""
t_grant = self._team("t-grant", ["srv1"])
t_grant = _make_team("t-grant", ["srv1"])
t_grant.organization_id = "org-team"
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
auth.org_id = "org-user-primary"
server = MagicMock(server_id="srv1")
with self._patch(teams_by_id={"t-grant": t_grant}, user_teams=["t-grant"]):
@ -6766,8 +6761,8 @@ class TestUserSubjectTeamUnion:
would charge that team for access it never provided."""
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
t_other = self._team("t-other", ["srv1"])
auth = self._admitted_subject("sso-user")
t_other = _make_team("t-other", ["srv1"])
auth = _make_admitted_subject("sso-user")
auth.object_permission = LiteLLM_ObjectPermissionTable(object_permission_id="op-u", mcp_servers=["srv1"])
with self._patch(teams_by_id={"t-other": t_other}, user_teams=["t-other"]):
assert await MCPRequestHandler.attributing_source_for_server(auth, "srv1") is None
@ -6775,8 +6770,8 @@ class TestUserSubjectTeamUnion:
async def test_billing_attribution_is_deterministic_across_several_granting_teams(self):
"""When several teams grant the same server the pick must be stable and reproducible rather
than dependent on dict/roster ordering, or the same call bills different teams run to run."""
teams = {"t-b": self._team("t-b", ["srv1"]), "t-a": self._team("t-a", ["srv1"])}
auth = self._admitted_subject("sso-user")
teams = {"t-b": _make_team("t-b", ["srv1"]), "t-a": _make_team("t-a", ["srv1"])}
auth = _make_admitted_subject("sso-user")
with self._patch(teams_by_id=teams, user_teams=["t-b", "t-a"]):
first = await MCPRequestHandler.attributing_source_for_server(auth, "srv1")
with self._patch(teams_by_id=teams, user_teams=["t-a", "t-b"]):
@ -6797,8 +6792,8 @@ class TestUserSubjectTeamUnion:
such a fault hit the fail-closed handler and denied tools its teams did grant."""
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
teams = {"t1": self._team("t1", ["srv1"], tool_perms={"srv1": ["read"]})}
auth = self._admitted_subject("sso-user")
teams = {"t1": _make_team("t1", ["srv1"], tool_perms={"srv1": ["read"]})}
auth = _make_admitted_subject("sso-user")
# The subject must carry a toolset, or the prelude never resolves one and the fault below is
# unreachable — the branch could sit anywhere and the test would still pass (it did).
auth.object_permission = LiteLLM_ObjectPermissionTable(object_permission_id="op-u", mcp_toolsets=["ts-1"])
@ -6825,7 +6820,7 @@ class TestUserSubjectTeamUnion:
manager._get_active_submitted_mcp_server_ids_for_user = AsyncMock(return_value=["srv-byom"])
db_default_perm = LiteLLM_ObjectPermissionTable(object_permission_id="op-u", mcp_servers=[])
admitted = self._admitted_subject("sso-user")
admitted = _make_admitted_subject("sso-user")
admitted.object_permission = db_default_perm
scoped_key = UserAPIKeyAuth(user_id="u", api_key="sk-hash", object_permission=db_default_perm)
@ -6840,7 +6835,7 @@ class TestUserSubjectTeamUnion:
from litellm.proxy._types import LitellmUserRoles
manager = self._manager_with(["srv-granted", "srv-secret"])
admitted = self._admitted_subject("admin-user")
admitted = _make_admitted_subject("admin-user")
admitted.user_role = LitellmUserRoles.PROXY_ADMIN
with patch.object(MCPRequestHandler, "get_allowed_mcp_servers", AsyncMock(return_value=["srv-granted"])):
admitted_view = set(await manager.get_allowed_mcp_servers(admitted))
@ -6863,7 +6858,7 @@ class TestUserSubjectTeamUnion:
opt_out = LiteLLM_ObjectPermissionTable(
object_permission_id="op-u", mcp_servers=[SpecialMCPServerNames.no_mcp_servers.value]
)
admitted = self._admitted_subject("sso-user")
admitted = _make_admitted_subject("sso-user")
admitted.object_permission = opt_out
with patch.object(MCPRequestHandler, "get_allowed_mcp_servers", AsyncMock(return_value=["srv-team"])):
admitted_view = set(await manager.get_allowed_mcp_servers(admitted))
@ -6880,7 +6875,7 @@ class TestUserSubjectTeamUnion:
session holder invoke tools their own policy excludes."""
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
# The user is restricted to `read` on srv-open, and NO grant source names that server —
# it is reachable only through the open channel, which is exactly the bypass path.
auth.object_permission = LiteLLM_ObjectPermissionTable(
@ -6900,7 +6895,7 @@ class TestUserSubjectTeamUnion:
source, so the source union alone returns [] listable but uninvokable. The tools axis asks
the same open-channel owner the server union uses, so the server is default-open for tools
exactly as a virtual key experiences it."""
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
open_ids = AsyncMock(return_value={"srv-open"})
with self._patch(teams_by_id={}, user_teams=[]):
with patch(
@ -6918,13 +6913,13 @@ class TestUserSubjectTeamUnion:
not keep granting servers, tools or throttle scope to a keyless union subject either.
Enforced through the SAME owner the key path uses (_team_max_budget_check). Distinct from
budget ATTRIBUTION of new spend, which stays with the user (documented deferral)."""
t_over = self._team("t-over", ["srv1"])
t_over = _make_team("t-over", ["srv1"])
t_over.max_budget = 10.0
t_over.spend = 11.0
t_ok = self._team("t-ok", ["srv2"])
t_ok = _make_team("t-ok", ["srv2"])
t_ok.max_budget = 10.0
t_ok.spend = 1.0
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
with self._patch(teams_by_id={"t-over": t_over, "t-ok": t_ok}, user_teams=["t-over", "t-ok"]):
servers = set(await MCPRequestHandler.get_allowed_mcp_servers(auth))
limits = await MCPRequestHandler._admitted_subject_team_rpm_limits(auth)
@ -6935,13 +6930,13 @@ class TestUserSubjectTeamUnion:
"""The org axis of the same rule, judged against the TEAM's own org (not the caller's
primary): a team owned by an org over its budget grants nothing, exactly as a key in that
org is rejected by _organization_max_budget_check."""
t_in_broke_org = self._team("t-b", ["srv1"])
t_in_broke_org = _make_team("t-b", ["srv1"])
t_in_broke_org.organization_id = "org-broke"
# object_permission_id=None: the org has NO MCP ceiling, so the source is denied by the
# budget gate alone. A truthy auto-Mock id here made an earlier version of this test pass
# through the org-CEILING fault path with the budget gate deleted — vacuous.
org = MagicMock(object_permission_id=None, litellm_budget_table=MagicMock(max_budget=5.0), spend=9.0)
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
with self._patch(teams_by_id={"t-b": t_in_broke_org}, user_teams=["t-b"], orgs_by_id={"org-broke": org}):
servers = await MCPRequestHandler.get_allowed_mcp_servers(auth)
assert servers == [], "a team in an over-budget org must not grant through the union"
@ -6953,8 +6948,8 @@ class TestUserSubjectTeamUnion:
axis."""
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
t_ok = self._team("t-ok", ["srv1"], tool_perms={"srv1": ["read"]})
auth = self._admitted_subject("sso-user")
t_ok = _make_team("t-ok", ["srv1"], tool_perms={"srv1": ["read"]})
auth = _make_admitted_subject("sso-user")
auth.object_permission = LiteLLM_ObjectPermissionTable(object_permission_id="op-u", mcp_servers=["srv-own"])
teams = {"t-ok": t_ok} # t-boom absent from the map -> our patched get_team_object RAISES for it
@ -6993,17 +6988,17 @@ class TestUserSubjectTeamUnion:
(else this user's calls drain a bucket shared by that team's keys for access the team never
provided), not for map entries beyond its grant, and never when the team is blocked."""
# t-granting grants srv1 and limits it; also names srv9 in its map, which it does NOT grant.
t_granting = self._team("t-granting", ["srv1"])
t_granting = _make_team("t-granting", ["srv1"])
t_granting.metadata = {"mcp_rpm_limit": {"srv1": 5, "srv9": 7}}
# t-other grants only srv2 but retains limit metadata for srv1 -> must not be charged for it.
t_other = self._team("t-other", ["srv2"])
t_other = _make_team("t-other", ["srv2"])
t_other.metadata = {"mcp_rpm_limit": {"srv1": 3}}
# t-blocked grants srv1 and limits it, but is blocked -> grants nothing, charges nothing.
t_blocked = self._team("t-blocked", ["srv1"])
t_blocked = _make_team("t-blocked", ["srv1"])
t_blocked.metadata = {"mcp_rpm_limit": {"srv1": 2}}
t_blocked.blocked = True
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
teams = {"t-granting": t_granting, "t-other": t_other, "t-blocked": t_blocked}
with self._patch(teams_by_id=teams, user_teams=["t-granting", "t-other", "t-blocked"]):
limits = await MCPRequestHandler._admitted_subject_team_rpm_limits(auth)
@ -7015,9 +7010,9 @@ class TestUserSubjectTeamUnion:
async def test_non_roster_team_rpm_limit_does_not_apply(self):
"""The roster gates grants and throttles through one owner, so a team the user was removed
from neither grants servers nor gets charged for their calls."""
stale = self._team("t-stale", ["srv1"], members=("someone-else",))
stale = _make_team("t-stale", ["srv1"], members=("someone-else",))
stale.metadata = {"mcp_rpm_limit": {"srv1": 1}}
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
with self._patch(teams_by_id={"t-stale": stale}, user_teams=["t-stale"]):
limits = await MCPRequestHandler._admitted_subject_team_rpm_limits(auth)
assert limits is None
@ -7029,7 +7024,7 @@ class TestUserSubjectTeamUnion:
org_id their whole org's server list with no direct or team grant behind it."""
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
auth.org_id = "org-a" # org allows srv1+srv2; the user and their teams grant NOTHING
org_perm = AsyncMock(
return_value=LiteLLM_ObjectPermissionTable(object_permission_id="op-org-a", mcp_servers=["srv1", "srv2"])
@ -7043,8 +7038,8 @@ class TestUserSubjectTeamUnion:
"""Each source is resolved through an UNMARKED auth, so a fault under a source must still
deny. Returning None there would win the union as allow-all and drop every team/org tool
ceiling on a DB blip -- the marker alone only covers faults raised before the fan-out."""
auth = self._admitted_subject("sso-user")
teams = {"t1": self._team("t1", ["srv1"])}
auth = _make_admitted_subject("sso-user")
teams = {"t1": _make_team("t1", ["srv1"])}
# Fault INSIDE the tool resolution only. Faulting something the server path also uses would
# make the source grant nothing, so the union would return [] without the tool path ever
# running -- the test would pass while pinning nothing (an earlier version did exactly that).
@ -7061,11 +7056,11 @@ class TestUserSubjectTeamUnion:
virtual KEY still overrides team inheritance -- that is the key ceiling model, unchanged.)"""
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, SpecialMCPServerNames
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
auth.object_permission = LiteLLM_ObjectPermissionTable(
object_permission_id="op-user", mcp_servers=[SpecialMCPServerNames.no_mcp_servers.value]
)
with self._patch(teams_by_id={"t1": self._team("t1", ["srv1"])}, user_teams=["t1"]):
with self._patch(teams_by_id={"t1": _make_team("t1", ["srv1"])}, user_teams=["t1"]):
result = await MCPRequestHandler.get_allowed_mcp_servers(auth)
assert set(result) == {"srv1"}, "the user's own opt-out must not zero their team's grants"
@ -7077,11 +7072,11 @@ class TestUserSubjectTeamUnion:
on. Each team source carries that team's own org, which is what makes the shared resolver apply
the team's owning-org ceiling rather than the caller's home org."""
teams = {
"t-member": self._team("t-member", ["srv1"], members=("sso-user",)),
"t-stale": self._team("t-stale", ["srv2"], members=("someone-else",)),
"t-member": _make_team("t-member", ["srv1"], members=("sso-user",)),
"t-stale": _make_team("t-stale", ["srv2"], members=("someone-else",)),
}
teams["t-member"].organization_id = "org-a"
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
with self._patch(teams_by_id=teams, user_teams=["t-member", "t-stale"]):
sources = await MCPRequestHandler._admitted_subject_sources(auth)
@ -7100,7 +7095,7 @@ class TestUserSubjectTeamUnion:
user_id and (with no team claim) no team_id, but it is NOT admission-marked, so it must
keep its prior behavior of inheriting no team grants rather than silently gaining the
union across every team the user belongs to."""
teams = {"team-a": self._team("team-a", ["srv1"]), "team-b": self._team("team-b", ["srv2"])}
teams = {"team-a": _make_team("team-a", ["srv1"]), "team-b": _make_team("team-b", ["srv2"])}
jwt_auth = UserAPIKeyAuth(user_id="jwt-user", api_key=None) # no admission marker
with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"]):
result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(jwt_auth)
@ -7112,7 +7107,7 @@ class TestUserSubjectTeamUnion:
metadata is caller-controlled at key creation. A user who sets
``mcp_admitted_user_subject: true`` in their own key's metadata (api_key present, no
team_id) must NOT be treated as an admitted subject and must gain no cross-team union."""
teams = {"team-a": self._team("team-a", ["srv1"]), "team-b": self._team("team-b", ["srv2"])}
teams = {"team-a": _make_team("team-a", ["srv1"]), "team-b": _make_team("team-b", ["srv2"])}
forged = UserAPIKeyAuth(
user_id="attacker",
api_key="sk-real-key",
@ -7140,7 +7135,7 @@ class TestUserSubjectTeamUnion:
mcp_tool_permissions={"srv1": ["tool_a"]},
),
)
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
with self._patch(teams_by_id={"team-a": team}, user_teams=["team-a"]):
tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth)
assert tools == ["tool_a"]
@ -7158,8 +7153,8 @@ class TestUserSubjectTeamUnion:
access_group_ids=[],
object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-blk", mcp_servers=["srv-secret"]),
)
teams = {"team-ok": self._team("team-ok", ["srv-ok"]), "team-blocked": blocked}
auth = self._admitted_subject("sso-user")
teams = {"team-ok": _make_team("team-ok", ["srv-ok"]), "team-blocked": blocked}
auth = _make_admitted_subject("sso-user")
with self._patch(teams_by_id=teams, user_teams=["team-ok", "team-blocked"]):
result = await MCPRequestHandler.get_allowed_mcp_servers(auth)
assert set(result) == {"srv-ok"}
@ -7169,8 +7164,8 @@ class TestUserSubjectTeamUnion:
team's roster inherits nothing from it, even when the team id lingers in the user's (stale or
cached) teams array. The team roster is the source of truth, so a removed or foreign
membership revokes access at the union rather than granting it."""
teams = {"team-x": self._team("team-x", ["srv-x"], members=("someone-else",))}
auth = self._admitted_subject("sso-user") # in user.teams for team-x, but NOT on its roster
teams = {"team-x": _make_team("team-x", ["srv-x"], members=("someone-else",))}
auth = _make_admitted_subject("sso-user") # in user.teams for team-x, but NOT on its roster
with self._patch(teams_by_id=teams, user_teams=["team-x"]):
result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth)
assert result == []
@ -7180,7 +7175,7 @@ class TestUserSubjectTeamUnion:
must DENY the server's tools ([]) rather than collapse to allow-all (None). Patches an await
OUTSIDE the multi-team fan-out (the team-object lookup) to prove the whole function fails
closed, not just the one helper mirroring the fail-closed server path."""
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
with patch.object(
MCPRequestHandler,
"_get_team_object_permission",
@ -7209,36 +7204,6 @@ class TestAdmittedSubjectPerTeamOrgCap:
user's own org — never the caller's primary org applied over the whole cross-org union. Guards the
Veria 'team grants bypass their owning policies' finding."""
def _team(self, team_id, mcp_servers, *, org_id=None, tool_perms=None, members=("sso-user",)):
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, Member
return LiteLLM_TeamTable(
team_id=team_id,
organization_id=org_id,
members_with_roles=[Member(user_id=u, role="user") for u in members],
access_group_ids=[],
object_permission=LiteLLM_ObjectPermissionTable(
object_permission_id=f"op-{team_id}",
mcp_servers=mcp_servers,
mcp_tool_permissions=tool_perms,
),
)
@staticmethod
def _admitted_subject(user_id, *, org_id=None, own_servers=None, own_tool_perms=None):
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
op = None
if own_servers is not None or own_tool_perms is not None:
op = LiteLLM_ObjectPermissionTable(
object_permission_id=f"userop-{user_id}",
mcp_servers=own_servers or [],
mcp_tool_permissions=own_tool_perms,
)
auth = UserAPIKeyAuth(user_id=user_id, api_key=None, org_id=org_id, object_permission=op)
auth.mcp_admitted_user_subject = True
return auth
#: sentinel for org_perms: org has an object_permission_id but its load returns None (a swallowed
#: DB error / dangling id), which _object_permission_for_org must treat as fail-closed.
LOAD_FAILS = "__load_fails__"
@ -7315,9 +7280,9 @@ class TestAdmittedSubjectPerTeamOrgCap:
async def test_team_grant_capped_by_its_own_org(self):
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
teams = {"team-a": self._team("team-a", ["srv1", "srv2"], org_id="org-a")}
teams = {"team-a": _make_team("team-a", ["srv1", "srv2"], org_id="org-a")}
org_perms = {"org-a": LiteLLM_ObjectPermissionTable(object_permission_id="orgop-org-a", mcp_servers=["srv1"])}
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
with self._patch(teams_by_id=teams, user_teams=["team-a"], org_perms=org_perms):
result = await MCPRequestHandler.get_allowed_mcp_servers(auth)
assert set(result) == {"srv1"} # srv2 capped out by org-a's ceiling
@ -7326,14 +7291,14 @@ class TestAdmittedSubjectPerTeamOrgCap:
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
teams = {
"team-a": self._team("team-a", ["srv1", "srv2"], org_id="org-a"),
"team-b": self._team("team-b", ["srv3", "srv4"], org_id="org-b"),
"team-a": _make_team("team-a", ["srv1", "srv2"], org_id="org-a"),
"team-b": _make_team("team-b", ["srv3", "srv4"], org_id="org-b"),
}
org_perms = {
"org-a": LiteLLM_ObjectPermissionTable(object_permission_id="orgop-org-a", mcp_servers=["srv1"]),
"org-b": LiteLLM_ObjectPermissionTable(object_permission_id="orgop-org-b", mcp_servers=["srv3"]),
}
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"], org_perms=org_perms):
result = await MCPRequestHandler.get_allowed_mcp_servers(auth)
assert set(result) == {"srv1", "srv3"} # each team clipped by its OWN org, then unioned
@ -7341,9 +7306,9 @@ class TestAdmittedSubjectPerTeamOrgCap:
async def test_all_proxy_grant_capped_by_org(self):
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, SpecialMCPServerName
teams = {"team-a": self._team("team-a", [SpecialMCPServerName.all_proxy_servers.value], org_id="org-a")}
teams = {"team-a": _make_team("team-a", [SpecialMCPServerName.all_proxy_servers.value], org_id="org-a")}
org_perms = {"org-a": LiteLLM_ObjectPermissionTable(object_permission_id="orgop-org-a", mcp_servers=["srv1"])}
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
with self._patch(
teams_by_id=teams, user_teams=["team-a"], org_perms=org_perms, registry=["srv1", "srv2", "srv3"]
):
@ -7353,8 +7318,8 @@ class TestAdmittedSubjectPerTeamOrgCap:
assert set(result) == {"srv1"}
async def test_org_row_without_object_permission_does_not_cap(self):
teams = {"team-a": self._team("team-a", ["srv1", "srv2"], org_id="org-a")}
auth = self._admitted_subject("sso-user")
teams = {"team-a": _make_team("team-a", ["srv1", "srv2"], org_id="org-a")}
auth = _make_admitted_subject("sso-user")
with self._patch(teams_by_id=teams, user_teams=["team-a"], org_perms={"org-a": None}):
result = await MCPRequestHandler.get_allowed_mcp_servers(auth)
assert set(result) == {"srv1", "srv2"} # empty ceiling = no restriction
@ -7362,12 +7327,12 @@ class TestAdmittedSubjectPerTeamOrgCap:
async def test_direct_grants_unioned_with_team_and_capped_by_user_org(self):
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
teams = {"team-a": self._team("team-a", ["srv1"], org_id="org-a")}
teams = {"team-a": _make_team("team-a", ["srv1"], org_id="org-a")}
org_perms = {
"org-a": None, # the team's org imposes no ceiling
"org-u": LiteLLM_ObjectPermissionTable(object_permission_id="orgop-org-u", mcp_servers=["srvD", "srv1"]),
}
auth = self._admitted_subject("sso-user", org_id="org-u", own_servers=["srvD", "srvX"])
auth = _make_admitted_subject("sso-user", org_id="org-u", own_servers=["srvD", "srvX"])
with self._patch(teams_by_id=teams, user_teams=["team-a"], org_perms=org_perms):
result = await MCPRequestHandler.get_allowed_mcp_servers(auth)
# direct {srvD,srvX} ∩ user-org {srvD,srv1} = {srvD}; UNIONed with team {srv1} (not intersected).
@ -7379,7 +7344,7 @@ class TestAdmittedSubjectPerTeamOrgCap:
# A KEY (not admitted): the per-team org cap must NOT fire; the top-level primary-org cap applies,
# byte-identical to before. team-a (org-a) grants {srv1,srv2}; the key's primary org is org-k.
teams = {"team-a": self._team("team-a", ["srv1", "srv2"], org_id="org-a")}
teams = {"team-a": _make_team("team-a", ["srv1", "srv2"], org_id="org-a")}
org_perms = {
"org-a": LiteLLM_ObjectPermissionTable(object_permission_id="orgop-org-a", mcp_servers=["srv2"]),
"org-k": LiteLLM_ObjectPermissionTable(object_permission_id="orgop-org-k", mcp_servers=["srv1"]),
@ -7397,7 +7362,7 @@ class TestAdmittedSubjectPerTeamOrgCap:
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
# team grants srv1 with NO tool restriction; org-a restricts srv1's tools to {tool_a}.
teams = {"team-a": self._team("team-a", ["srv1"], org_id="org-a")}
teams = {"team-a": _make_team("team-a", ["srv1"], org_id="org-a")}
org_perms = {
"org-a": LiteLLM_ObjectPermissionTable(
object_permission_id="orgop-org-a",
@ -7405,7 +7370,7 @@ class TestAdmittedSubjectPerTeamOrgCap:
mcp_tool_permissions={"srv1": ["tool_a"]},
)
}
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
with self._patch(teams_by_id=teams, user_teams=["team-a"], org_perms=org_perms):
tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth)
# Without the per-team org tool ceiling this would be None (all tools) — org-a's tool ceiling
@ -7414,10 +7379,10 @@ class TestAdmittedSubjectPerTeamOrgCap:
async def test_tool_union_across_cross_org_teams(self):
teams = {
"team-a": self._team("team-a", ["srv1"], org_id="org-a", tool_perms={"srv1": ["t1"]}),
"team-b": self._team("team-b", ["srv1"], org_id="org-b", tool_perms={"srv1": ["t2"]}),
"team-a": _make_team("team-a", ["srv1"], org_id="org-a", tool_perms={"srv1": ["t1"]}),
"team-b": _make_team("team-b", ["srv1"], org_id="org-b", tool_perms={"srv1": ["t2"]}),
}
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
with self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"], org_perms={"org-a": None, "org-b": None}):
tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth)
assert set(tools) == {"t1", "t2"}
@ -7425,7 +7390,7 @@ class TestAdmittedSubjectPerTeamOrgCap:
async def test_tool_deny_all_when_team_grant_and_org_tool_ceiling_disjoint(self):
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
teams = {"team-a": self._team("team-a", ["srv1"], org_id="org-a", tool_perms={"srv1": ["t1"]})}
teams = {"team-a": _make_team("team-a", ["srv1"], org_id="org-a", tool_perms={"srv1": ["t1"]})}
org_perms = {
"org-a": LiteLLM_ObjectPermissionTable(
object_permission_id="orgop-org-a",
@ -7433,7 +7398,7 @@ class TestAdmittedSubjectPerTeamOrgCap:
mcp_tool_permissions={"srv1": ["t2"]},
)
}
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
with self._patch(teams_by_id=teams, user_teams=["team-a"], org_perms=org_perms):
tools = await MCPRequestHandler.get_allowed_tools_for_server("srv1", auth)
# team {t1} ∩ org {t2} = {} → deny every tool ([]), NOT allow-all (None).
@ -7446,8 +7411,8 @@ class TestAdmittedSubjectPerTeamOrgCap:
synced). get_org_object RAISES a bare Exception for that; it must be treated as 'no ceiling' and
must NOT lock the admitted subject out of the team's grants (parity with the key path, which
tolerates a deleted org)."""
teams = {"team-a": self._team("team-a", ["srv1", "srv2"], org_id="org-gone")}
auth = self._admitted_subject("sso-user")
teams = {"team-a": _make_team("team-a", ["srv1", "srv2"], org_id="org-gone")}
auth = _make_admitted_subject("sso-user")
with self._patch(teams_by_id=teams, user_teams=["team-a"], org_perms={}): # org-gone absent → raises
result = await MCPRequestHandler.get_allowed_mcp_servers(auth)
assert set(result) == {"srv1", "srv2"}
@ -7456,8 +7421,8 @@ class TestAdmittedSubjectPerTeamOrgCap:
"""The org carries an object_permission_id but the permission load returns None (a swallowed DB
error / dangling id). The ceiling cannot be verified, so the admitted subject must fail CLOSED
for that team NOT skip the ceiling, which would leak org-forbidden servers."""
teams = {"team-a": self._team("team-a", ["srv1", "srv2"], org_id="org-a")}
auth = self._admitted_subject("sso-user")
teams = {"team-a": _make_team("team-a", ["srv1", "srv2"], org_id="org-a")}
auth = _make_admitted_subject("sso-user")
with self._patch(teams_by_id=teams, user_teams=["team-a"], org_perms={"org-a": self.LOAD_FAILS}):
# Asserted through the PUBLIC resolver: the per-source org ceiling is applied there now,
# so calling the single-team helper would return [] for an admitted subject either way
@ -7473,9 +7438,9 @@ class TestAdmittedSubjectPerTeamOrgCap:
grant would bypass every org ceiling and reach servers the user's home org forbids."""
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
teams = {"team-noorg": self._team("team-noorg", ["srv1", "srv2"], org_id=None)}
teams = {"team-noorg": _make_team("team-noorg", ["srv1", "srv2"], org_id=None)}
org_perms = {"org-U": LiteLLM_ObjectPermissionTable(object_permission_id="orgop-org-U", mcp_servers=["srv1"])}
auth = self._admitted_subject("sso-user", org_id="org-U")
auth = _make_admitted_subject("sso-user", org_id="org-U")
with self._patch(teams_by_id=teams, user_teams=["team-noorg"], org_perms=org_perms):
result = await MCPRequestHandler.get_allowed_mcp_servers(auth)
# org-less team falls back to the user's primary org (org-U → {srv1}); srv2 capped out.
@ -7485,8 +7450,8 @@ class TestAdmittedSubjectPerTeamOrgCap:
"""MEDIUM (greptile/cursor): when no source in the tool-resolution view grants the server (a
TOCTOU/cache-lag inconsistency on a server that passed the server gate), the admitted path must
fail CLOSED (deny all tools = []), NOT allow-all (None)."""
teams = {"team-a": self._team("team-a", ["srv1"], org_id="org-a")}
auth = self._admitted_subject("sso-user")
teams = {"team-a": _make_team("team-a", ["srv1"], org_id="org-a")}
auth = _make_admitted_subject("sso-user")
with self._patch(teams_by_id=teams, user_teams=["team-a"], org_perms={"org-a": None}):
# 'srv-nobody' is granted by neither the team nor the user directly → empty contributions.
tools = await MCPRequestHandler.get_allowed_tools_for_server("srv-nobody", auth)
@ -7495,9 +7460,7 @@ class TestAdmittedSubjectPerTeamOrgCap:
async def test_tool_no_db_honors_in_memory_direct_restriction(self):
"""MEDIUM (cursor): with no DB, the tool path must still honor the user's OWN in-memory
object_permission tool restriction (resolvable without a DB) rather than blanket-allow (None)."""
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
auth = self._admitted_subject(
auth = _make_admitted_subject(
"sso-user", own_servers=["srv1"], own_tool_perms={"srv1": ["t1"]}
) # no org_id, direct grant of srv1 restricted to {t1}
with patch("litellm.proxy.proxy_server.prisma_client", None):
@ -7522,11 +7485,11 @@ class TestAdmittedSubjectPerTeamOrgCap:
# team grants the config server BY ALIAS alongside a DB-style bare id; org-a's ceiling lists
# ONLY the config server (also by alias).
teams = {"team-a": self._team("team-a", ["linear_cfg", "srv-db"], org_id="org-a")}
teams = {"team-a": _make_team("team-a", ["linear_cfg", "srv-db"], org_id="org-a")}
org_perms = {
"org-a": LiteLLM_ObjectPermissionTable(object_permission_id="orgop-org-a", mcp_servers=["linear_cfg"])
}
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
with self._patch(
teams_by_id=teams,
user_teams=["team-a"],
@ -7554,12 +7517,12 @@ class TestAdmittedSubjectPerTeamOrgCap:
control.server_id = "control-id"
control.alias = control.server_name = control.name = "control_alias"
teams = {"team-b": self._team("team-b", ["linear_cfg", "control_alias"], org_id="org-b")}
teams = {"team-b": _make_team("team-b", ["linear_cfg", "control_alias"], org_id="org-b")}
# org-b's ceiling allows ONLY the control server, referenced by its RESOLVED server_id.
org_perms = {
"org-b": LiteLLM_ObjectPermissionTable(object_permission_id="orgop-org-b", mcp_servers=["control-id"])
}
auth = self._admitted_subject("sso-user")
auth = _make_admitted_subject("sso-user")
with self._patch(
teams_by_id=teams,
user_teams=["team-b"],