feat(mcp): union team-inherited MCP grants across all of a user's teams for keyless admission - #33191
Conversation
Greptile SummaryExpands team-based MCP grant resolution so a keyless user-subject caller (DCR session bearer, bridge user-envelope) inherits MCP server access from all of their teams rather than just the single
Confidence Score: 4/5Safe to merge for key-based callers (zero code-path change); keyless user-subject callers gain their intended team grants with a well-contained fail-safe. The refactoring is clean and the behavioral invariant for key-based callers is well-protected. The two observations are about asyncio.gather defensiveness (the broad exception handler in the inner coroutine makes this a non-issue in practice) and a second get_user_object lookup per keyless request that relies on the cache to be cheap. Neither affects correctness today. The new _resolve_user_team_ids path in user_api_key_auth_mcp.py is worth a second read, specifically the interaction between the asyncio.gather fan-out and the exception handling contract of _allowed_mcp_servers_for_single_team.
|
| Filename | Overview |
|---|---|
| litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py | Refactors _get_allowed_mcp_servers_for_team to fan-out across all of a keyless user's teams via asyncio.gather and a new _resolve_user_team_ids DB lookup; key-based callers are unchanged |
| tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py | Adds TestUserSubjectTeamUnion covering union, single-team key path, explicit team_id keyless path, empty-teams case, and UI sentinel; all mocked, no live network calls |
Reviews (1): Last reviewed commit: "feat(mcp): union team-inherited MCP gran..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| per_team = await asyncio.gather( | ||
| *( | ||
| MCPRequestHandler._allowed_mcp_servers_for_single_team(team_id, user_api_key_auth) | ||
| for team_id in team_ids | ||
| ) | ||
| ) |
There was a problem hiding this comment.
asyncio.gather will propagate an uncaught exception from any team coroutine
_allowed_mcp_servers_for_single_team wraps all of its logic in a broad except Exception block, so in practice nothing escapes. However, if that guard ever develops a gap (e.g., a BaseException subclass like asyncio.CancelledError, or an error raised during the generator expression itself), the gather without return_exceptions=True will cancel remaining tasks and surface the exception to the caller, which has no error handling at this level — potentially returning a 500 to the MCP client. Using return_exceptions=True and filtering out Exception instances before the set union would make the fan-out robust regardless of what _allowed_mcp_servers_for_single_team throws.
| async def _resolve_user_team_ids(user_id: str, user_api_key_auth: UserAPIKeyAuth) -> list[str]: | ||
| """The distinct team ids a user belongs to, from the live user record. Returns [] on | ||
| no DB, a missing user, or any resolution failure so a lookup blip narrows access | ||
| rather than raising; the caller's direct grants still apply.""" | ||
| from litellm.proxy.auth.auth_checks import get_user_object | ||
| from litellm.proxy.proxy_server import ( | ||
| prisma_client, | ||
| proxy_logging_obj, | ||
| user_api_key_cache, | ||
| ) | ||
|
|
||
| if prisma_client is None: | ||
| return [] | ||
| try: | ||
| user_object = await get_user_object( | ||
| user_id=user_id, | ||
| prisma_client=prisma_client, | ||
| user_api_key_cache=user_api_key_cache, | ||
| user_id_upsert=False, | ||
| parent_otel_span=user_api_key_auth.parent_otel_span, | ||
| proxy_logging_obj=proxy_logging_obj, | ||
| ) | ||
| except Exception as e: # noqa: BLE001 # a team-resolution blip narrows access, never raises | ||
| verbose_logger.warning(f"Failed to resolve user teams for MCP grant: {str(e)}") | ||
| return [] | ||
| if user_object is None or not user_object.teams: | ||
| return [] | ||
| return list(dict.fromkeys(t for t in user_object.teams if t and t != UI_TEAM_ID)) |
There was a problem hiding this comment.
get_user_object call on every keyless MCP request
For a keyless user-subject caller this path fires get_user_object to look up team membership. The user was already fetched and placed in user_api_key_cache during the upstream admission check, so a warm-cache hit is cheap. On a cold-cache miss, however, this adds a second DB round-trip per request specifically in the MCP auth path, on top of the N parallel get_team_object calls triggered afterward. If UserAPIKeyAuth is ever extended to carry the user's team list (it already carries user_id), this fetch could be eliminated; for now, worth noting the per-request cost for deployments with a large team count per user.
Rule Used: What: Avoid creating new database requests or Rout... (source)
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Greptile SummaryThis PR fixes a gap in MCP server access for keyless user-subject callers (gateway DCR session bearers and bridge user-envelopes): previously
Confidence Score: 4/5Safe to merge; key-based auth paths are byte-identical to before and the new union branch is well-gated. The
|
| Filename | Overview |
|---|---|
| litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py | Refactors team MCP grant resolution: extracts _allowed_mcp_servers_for_single_team, adds _team_ids_for_mcp_grant and _resolve_user_team_ids to union grants across all teams for keyless callers; key-based paths unchanged |
| tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py | Adds TestUserSubjectTeamUnion class covering: keyless union across teams, key-based single-team, explicit-team-id keyless, no-teams, UI_TEAM_ID sentinel, and _team_ids_for_mcp_grant shape gate; no test for the _resolve_user_team_ids exception fallback path |
Reviews (2): Last reviewed commit: "feat(mcp): union team-inherited MCP gran..." | Re-trigger Greptile
| try: | ||
| user_object = await get_user_object( | ||
| user_id=user_id, | ||
| prisma_client=prisma_client, | ||
| user_api_key_cache=user_api_key_cache, | ||
| user_id_upsert=False, | ||
| parent_otel_span=user_api_key_auth.parent_otel_span, | ||
| proxy_logging_obj=proxy_logging_obj, | ||
| ) | ||
| except Exception as e: # noqa: BLE001 # a team-resolution blip narrows access, never raises | ||
| verbose_logger.warning(f"Failed to resolve user teams for MCP grant: {str(e)}") | ||
| return [] | ||
| if user_object is None or not user_object.teams: | ||
| return [] | ||
| return list(dict.fromkeys(t for t in user_object.teams if t and t != UI_TEAM_ID)) |
There was a problem hiding this comment.
The except Exception branch in _resolve_user_team_ids silently returns [] on any DB failure, narrowing the user's team-inherited access to nothing. This is the explicitly documented fail-safe behavior, but there is no test in TestUserSubjectTeamUnion that exercises this path (e.g., get_user_object raising or returning a network error). A test that patches get_user_object to raise and asserts the result is [] would pin this contract so future exception handling changes cannot accidentally widen access or surface unhandled errors.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
| user_id=user_id, | ||
| prisma_client=prisma_client, | ||
| user_api_key_cache=user_api_key_cache, | ||
| user_id_upsert=False, |
There was a problem hiding this comment.
Medium: Team removal does not immediately revoke inherited access
get_user_object reads from user_api_key_cache by default. If a user's team membership is changed without evicting that cache entry, a user holding a valid gateway session can continue executing MCP tools granted by the removed team until the configurable management-object cache TTL expires. Resolve this authorization decision from the database or ensure every membership mutation invalidates the user cache.
| user_id_upsert=False, | |
| user_id_upsert=False, | |
| check_db_only=True, |
PR overviewThis pull request updates MCP keyless admission so a user can inherit MCP server grants from all teams they belong to, rather than relying on a single team context. The touched authentication code resolves user/team-derived MCP access during admission and tool authorization. There are still significant authorization gaps in the current implementation. A user may be able to forge MCP admission through key metadata on a personal key, and inherited team server grants may lose their per-team tool restrictions, allowing broader tool execution than intended. Team membership removals also may not take effect until cached user data expires, delaying revocation. Overall, the PR is moving toward broader inherited access support but still needs tighter credential validation, permission aggregation, and cache invalidation before it is safe. Open issues (3)
Fixed/addressed: 0 · PR risk: 7/10 |
1731965 to
65211a9
Compare
5daf091 to
8e5d0a6
Compare
| return [] | ||
| per_team = await asyncio.gather( | ||
| *( | ||
| MCPRequestHandler._allowed_mcp_servers_for_single_team(team_id, user_api_key_auth) |
There was a problem hiding this comment.
High: Team tool restrictions are not preserved
This grants servers from every team, but get_allowed_tools_for_server() still obtains team permissions through _get_team_object_permission(), which returns no team permission when this admitted auth has no team_id. A user can consequently call any tool on an inherited server even when the granting team limits mcp_tool_permissions; aggregate the corresponding per-team tool permissions alongside these server grants and enforce their union during tool listing and execution.
65211a9 to
001064b
Compare
…ms for keyless admission
…ey absence A JWT-authenticated caller is also keyless with a user_id and, absent a team claim, no team_id, so gating the multi-team union on api_key-is-None silently broadened JWT MCP access to the union of every team the user belongs to. _reload_admitted_user now stamps MCP_ADMITTED_USER_SUBJECT_METADATA and the union fires only for that positive marker, so the gateway session and bridge user paths union while JWT and other keyless auth keep their prior behavior. Regression-tested.
8e5d0a6 to
1010c32
Compare
| return [] | ||
| if user_api_key_auth.team_id: | ||
| return [] if user_api_key_auth.team_id == UI_TEAM_ID else [user_api_key_auth.team_id] | ||
| if not user_api_key_auth.user_id or not _is_mcp_admitted_user_subject(user_api_key_auth): |
There was a problem hiding this comment.
High: Key metadata can forge MCP session admission
UserAPIKeyAuth.metadata contains the metadata supplied when a virtual key is created, and /key/generate permits personal keys with no team_id. A user can therefore create such a key with mcp_admitted_user_subject: true and use it to execute tools on servers inherited from every team, bypassing the key's intended team scope. Keep the positive marker check, but also require that this is actually a keyless credential.
| if not user_api_key_auth.user_id or not _is_mcp_admitted_user_subject(user_api_key_auth): | |
| if ( | |
| user_api_key_auth.api_key is not None | |
| or not user_api_key_auth.user_id | |
| or not _is_mcp_admitted_user_subject(user_api_key_auth) | |
| ): |
d21b5d1
into
litellm_lit3637_session_admission
1 participant
Relevant issues
Stacked on #33190 (session admission arm) -> #33189 -> #33188 -> #33182 -> #33174. The base of this PR is
litellm_lit3637_session_admission; review and merge after thoseLinear ticket
Part of LIT-3637 (PR 4b of the stack: makes the aggregate flow honor team-based server access, which is how servers are meant to be shared)
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito re-request a review after pushing changes)Delays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
Screenshots / Proof of Fix
This is a permission-resolution change with no request surface of its own, so the contract is pinned by unit tests that drive
_get_allowed_mcp_servers_for_teamwith a user in two teams and assert the union, and separately assert a key-based caller with ateam_idstill sees only that one team. The interactive effect (a signed-in user reaching a server granted to their team) is exercised by the aggregate flow on a live proxyType
🐛 Bug Fix
Changes
A
UserAPIKeyAuthcan name at most one team, so the MCP team-grant resolver read grants from that singleteam_id. For a key that is fine, the key names its team. For a user-subject caller admitted without a key (the gateway DCR session bearer from the PR below this, and the already-merged bridge user-envelope) the auth carries auser_idand no team, so team-inherited grants resolved to nothing and a signed-in user saw only servers granted to them directly. Servers are meant to be shared by assigning teams rather than individuals, so in practice that user saw nothing;_reload_admitted_userdocumented this single-team limitation as a follow-up, and this is that follow-upThe per-team computation is factored into
_allowed_mcp_servers_for_single_team, unchanged in what one team grants._get_allowed_mcp_servers_for_teamnow asks_team_ids_for_mcp_grantwhich teams to union: a caller with ateam_id(every key-based caller, and the UI-session sentinel) resolves to that single team exactly as before, so key auth is byte-identical; only a keyless caller with auser_idand noteam_idfans out to the user's full team list, resolved once from the live user record, and the grants are unioned. Theapi_key is Nonegate is what keeps a personal key from silently gaining its owner's every-team access; only an identity-only admission unionsTeam resolution failures (no database, a missing user, a lookup blip) return no teams, so a transient failure narrows access rather than raising or over-granting; the caller's own direct grants still apply. The org ceiling, end-user, and agent intersections downstream are untouched, so unioning teams cannot lift a cap those layers impose
QA runbook
With the aggregate DCR flow running (
mcp_gateway_dcr: true), grant an MCP server to a team, add a user to that team, sign that user in through the flow to get a session bearer, and confirmtools/listat the aggregate/mcpnow includes that server. Confirm a virtual key scoped to one team still lists only that team's servers (unchanged). The unit tests cover the union, the single-team key path, the keyless-with-explicit-team path, the no-teams case, and the UI sentinelFinal Attestation