diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 01e8b1490b2..a27d6b92843 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 5d32a9c740d..b3c0dcd1681 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -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"],