diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 711119b5ab5..08a8b1bc7b3 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -552,13 +552,20 @@ async def get_all_mcp_servers( ) -> list[LiteLLM_MCPServerTable]: """ Returns mcp servers from the db, optionally filtered by approval_status. - Pass approval_status=None to return all servers regardless of approval state. + Pass approval_status=None to return every server except drafts, which back the admin OAuth + session flow, are addressable only by their own server_id, and must never appear in a listing. + NULL approval_status predates the approval workflow, so those rows are kept explicitly rather + than dropped by a bare inequality, which SQL evaluates as NULL and would silently hide them. """ try: - where: Final[prisma_db_types.LiteLLM_MCPServerTableWhereInput] = {} - if approval_status is not None: - where["approval_status"] = approval_status - mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client, where if where else {}) + where: Final[prisma_db_types.LiteLLM_MCPServerTableWhereInput] = ( + {"approval_status": approval_status} + if approval_status is not None + # mutable-ok: prisma where-inputs must be plain dicts, and both `NOT` and `not` drop + # NULL rows (measured), so the OR is the only NULL-preserving way to exclude drafts + else {"OR": [{"approval_status": None}, {"approval_status": {"not": MCPApprovalStatus.draft}}]} + ) + mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client, where) tables: Final = [LiteLLM_MCPServerTable.model_validate(mcp_server.model_dump()) for mcp_server in mcp_servers] for table in tables: @@ -814,6 +821,96 @@ async def create_mcp_server( return new_mcp_server +async def create_draft_mcp_server( + prisma_client: PrismaClient, + data: NewMCPServerRequest, + touched_by: str, + ttl_seconds: int, + server_id: str | None = None, +) -> LiteLLM_MCPServerTable: + """ + Persist a short-lived draft row backing the admin OAuth "Authorize & Fetch Token" flow. + + The draft lives in the database rather than in process memory so that the /register, + /authorize and /token legs resolve it whichever worker or replica accepts each request. + + Writing is strictly create-if-absent. Any existing row for the id is returned untouched, which + covers both a live draft for this same session and a real server the edit form is + re-authorizing against its own id, where writing a draft would collide on the primary key. + Each click of Authorize mints a fresh id, so nothing is lost by never overwriting, and it is + what makes concurrent callers sharing one id safe rather than mutually destructive. + """ + draft_id: Final = server_id or data.server_id or str(uuid.uuid4()) + await _prune_expired_draft_mcp_servers(prisma_client, ttl_seconds) + + existing: Final = await _db_find_mcp_server_row(prisma_client, draft_id) + if existing is not None: + # Already usable by every worker, whether it is a live draft for this same session or a + # real server the edit form is re-authorizing. Either way there is nothing to write, and + # not writing is what keeps concurrent callers for one server_id from racing each other. + return LiteLLM_MCPServerTable.model_validate(existing.model_dump()) + + draft_payload: Final = data.model_copy(update={"server_id": draft_id, "approval_status": MCPApprovalStatus.draft}) + try: + return await create_mcp_server(prisma_client, draft_payload, touched_by) + except Exception: + # Lost the create race: the read above and this create are two statements, not one. The + # winner wrote a draft for this same session, so adopt it rather than failing a caller + # whose session is in fact ready. Anything else still raises. + raced: Final = await _db_find_mcp_server_row(prisma_client, draft_id) + if raced is None or raced.approval_status != MCPApprovalStatus.draft: + raise + return LiteLLM_MCPServerTable.model_validate(raced.model_dump()) + + +async def _prune_expired_draft_mcp_servers(prisma_client: PrismaClient, ttl_seconds: int) -> None: + """Drop drafts already past ``ttl_seconds``, so abandoned OAuth sessions do not accumulate. + + Runs on each draft write rather than on a schedule, mirroring the in-memory cache this + replaces, which pruned on every store. Expired drafts are unreadable by then anyway, so the + only thing at stake is row count, and the work is bounded by how often admins authorize. + """ + cutoff: Final = datetime.now(timezone.utc) - timedelta(seconds=max(1, ttl_seconds)) + # Age is filtered here rather than in the query: the draft set is bounded by how many OAuth + # authorizations are in flight, so it is a handful of rows even on a busy proxy. + drafts: Final = await _db_find_mcp_server_rows( + prisma_client, + where={"approval_status": MCPApprovalStatus.draft}, + ) + for row in drafts: + # A row without a timestamp has no age to judge, so leave it rather than guess it is stale. + # Two workers sweeping the same row is harmless: prisma's delete returns None for a row + # that is already gone rather than raising, so the loser of that race is a no-op. + if row.updated_at is not None and row.updated_at < cutoff: + await delete_mcp_server(prisma_client, row.server_id) + + +async def get_draft_mcp_server( + prisma_client: PrismaClient, server_id: str, ttl_seconds: int +) -> LiteLLM_MCPServerTable | None: + """ + Return the draft row for ``server_id`` if it has not yet aged past ``ttl_seconds``, else None. + + Age is enforced in the query rather than by a sweeper so an expired draft is unreadable the + moment it lapses, regardless of which process last ran a cleanup. + """ + cutoff: Final = datetime.now(timezone.utc) - timedelta(seconds=max(1, ttl_seconds)) + draft_rows: Final = await _db_find_mcp_server_rows( + prisma_client, + where={ + "server_id": server_id, + "approval_status": MCPApprovalStatus.draft, + "updated_at": {"gte": cutoff}, + }, + ) + if not draft_rows: + return None + + table: Final = LiteLLM_MCPServerTable.model_validate(draft_rows[0].model_dump()) + decrypt_global_env_var_values(table.env_vars) + return table + + async def update_mcp_server( prisma_client: PrismaClient, data: UpdateMCPServerRequest, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index bb330d00756..6c01575e271 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1283,6 +1283,9 @@ class MCPApprovalStatus(str, enum.Enum): pending_review = "pending_review" active = "active" rejected = "rejected" + # Short-lived row backing the admin OAuth "Authorize & Fetch Token" flow. Never served: the + # registry loader and every listing exclude it, so it is reachable only by its own server_id. + draft = "draft" from litellm.models.mcp_server import ( # noqa: E402 diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 4d9f0a58bd7..0239368ad0e 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -65,6 +65,7 @@ from litellm.proxy.auth.budget_throttle import ( should_throttle_budget_exceeded, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_headers, @@ -2016,6 +2017,44 @@ async def _cache_team_object( ) +async def delete_cache_team_object( + team_id: str, + team_alias: str | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging | None, +) -> None: + """ + Evict both keys `_cache_team_object` writes. + + `get_team_object` reads the id key and the JWT `team_alias_jwt_field` path reads the alias key, + so leaving either behind keeps a deleted team resolvable for auth until its TTL expires. + + Mirrors `delete_cached_project_object`: evicting locally only reaches the worker handling the + delete, so every key is also broadcast to drop the other workers' in-memory copies. + + Eviction is best-effort, matching `_cache_team_object`. `delete_team` calls this after the team + rows are already gone, so letting an unreachable cache backend raise here would fail a request + whose delete has committed. + """ + keys: Final = (f"team_id:{team_id}", *((f"team_alias:{team_alias}",) if team_alias else ())) + + for key in keys: + try: + user_api_key_cache.delete_cache(key=key) + + ## UPDATE REDIS CACHE ## + if proxy_logging_obj is not None: + await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key) + except Exception as e: # noqa: BLE001 # best-effort invalidation: any cache backend error must not abort the delete + verbose_proxy_logger.warning( + "Failed to invalidate cached team entry %s on delete; " + "a deleted team may be served until its TTL expires: %s", + key, + e, + ) + await publish_auth_cache_invalidation(cache_key=key) + + async def _cache_key_object( hashed_token: str, user_api_key_obj: UserAPIKeyAuth, @@ -2068,6 +2107,44 @@ class TeamNotFoundError(HTTPException): ) +async def delete_cache_key_objects( + hashed_tokens: Sequence[str], + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging | None, +) -> None: + """ + Evict a batch of key objects, for callers that delete keys in bulk rather than through + `/key/delete`. Auth resolves a cached key object without re-reading its team, so a key left + cached after its row is gone keeps buying access until its TTL expires. + + Evicting locally only reaches this worker, so each token is also broadcast: a deleted key left + in a peer worker's in-memory cache still authenticates there until its TTL expires. + + Best-effort per key: the rows are already deleted by the time this runs, so an unreachable + cache backend must not abort the caller partway through its own cascade. + """ + results: Final = await asyncio.gather( + *( + _delete_cache_key_object( + hashed_token=hashed_token, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + for hashed_token in hashed_tokens + ), + return_exceptions=True, + ) + + for hashed_token, result in zip(hashed_tokens, results): + if isinstance(result, BaseException): + verbose_proxy_logger.warning( + "Failed to evict cached key entry for %s; a deleted key may authenticate until its TTL expires: %s", + hashed_token, + result, + ) + await publish_auth_cache_invalidation(cache_key=hashed_token) + + @log_db_metrics async def _get_team_db_check( team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index c4f5ceaba7a..997012dbc65 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -22,7 +22,7 @@ import os from collections.abc import Iterable from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal from fastapi import ( APIRouter, @@ -92,6 +92,9 @@ def does_mcp_server_exist(mcp_server_records: Iterable[Any], mcp_server_id: str) DEFAULT_MCP_REGISTRY_VERSION: Final = "1.0.0" +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + try: importlib.import_module("mcp") except ImportError as e: @@ -115,11 +118,13 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.db import ( approve_mcp_server, + create_draft_mcp_server, create_mcp_server, delete_mcp_server, delete_user_credential, delete_user_env_vars, get_all_mcp_servers_for_user, + get_draft_mcp_server, get_mcp_server, get_mcp_servers, get_mcp_submissions, @@ -467,19 +472,68 @@ if MCP_AVAILABLE: verbose_proxy_logger.debug("Invalid temporary MCP server payload in Redis cache: %s", e) return None + def _get_prisma_client_or_none() -> "PrismaClient | None": + """Non-throwing counterpart to ``get_prisma_client_or_throw`` for paths that degrade + gracefully: a proxy configured without a database keeps the in-memory OAuth session.""" + from litellm.proxy.proxy_server import prisma_client + + return prisma_client + + async def _persist_draft_mcp_server( + payload: NewMCPServerRequest, + server_id: str, + created_by: str, + ) -> None: + """Write the draft row that makes the OAuth session resolvable from any worker. + + A failure here is raised, not swallowed: without the shared row the flow degrades to + the per-process cache and fails intermittently, which is the defect being fixed. + """ + prisma_client: Final = _get_prisma_client_or_none() + if prisma_client is None: + return + await create_draft_mcp_server( + prisma_client, + payload, + created_by, + ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS, + server_id=server_id, + ) + + async def _get_draft_mcp_server_as_mcp_server(server_id: str) -> MCPServer | None: + """Resolve a database-backed draft, which is the only lookup that works across workers.""" + prisma_client: Final = _get_prisma_client_or_none() + if prisma_client is None: + return None + draft: Final = await get_draft_mcp_server( + prisma_client, server_id, ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS + ) + if draft is None: + return None + return await global_mcp_server_manager.build_mcp_server_from_table(draft) + async def get_cached_temporary_mcp_server( server_id: str, ) -> MCPServer | None: _prune_expired_temporary_mcp_servers() entry: Final = _temporary_mcp_servers.get(server_id) - if entry is None: - redis_server: Final = await _get_temporary_mcp_server_from_redis(server_id) - if redis_server is None: - return None - # Intentionally avoid repopulating local cache from Redis to prevent - # extending effective lifetime beyond the remaining Redis TTL. - return redis_server - return entry.server + if entry is not None: + return entry.server + + # A miss here means either an expired session or, on a multi-worker or multi-replica + # proxy, that a different process served /session. The draft row is shared, so it + # resolves the second case; the in-memory hit above still serves single-process + # deployments with no database configured. + draft_server: Final = await _get_draft_mcp_server_as_mcp_server(server_id) + if draft_server is not None: + return draft_server + + redis_server: Final = await _get_temporary_mcp_server_from_redis(server_id) + if redis_server is None: + return None + # Intentionally avoid repopulating local cache from Redis to prevent + # extending effective lifetime beyond the remaining Redis TTL. + return redis_server def _redact_mcp_credentials( mcp_server: LiteLLM_MCPServerTable, @@ -709,12 +763,36 @@ if MCP_AVAILABLE: payload_dict["credentials"] = inherited_credentials return NewMCPServerRequest.model_validate(payload_dict) + async def _resolve_session_server_id(payload: NewMCPServerRequest) -> str: + """Decide the id an OAuth session runs under. + + A caller-supplied id is honoured only when it names a server that really exists, which is + the edit form re-authorizing a saved server against its own id. Anything else gets a fresh + id, so two concurrent sessions can never land on one id and silently adopt each other's + URL or client credentials. Without a database there is nothing shared to collide over, so + the supplied id is kept and behaviour is unchanged. + """ + supplied: Final = payload.server_id + if not supplied: + return str(uuid.uuid4()) + if global_mcp_server_manager.get_mcp_server_by_id(supplied) is not None: + return supplied + prisma_client: Final = _get_prisma_client_or_none() + if prisma_client is None: + return supplied + # A draft is another session's row, not a saved server, so re-supplying an id this + # endpoint previously handed back must not let a later session adopt its configuration. + existing: Final = await get_mcp_server(prisma_client, supplied) + if existing is None or existing.approval_status == MCPApprovalStatus.draft: + return str(uuid.uuid4()) + return supplied + def _build_temporary_mcp_server_record( payload: NewMCPServerRequest, created_by: str | None, + server_id: str, ) -> LiteLLM_MCPServerTable: now: Final = datetime.utcnow() - server_id: Final = payload.server_id or str(uuid.uuid4()) server_name: Final = payload.server_name or payload.alias or server_id return LiteLLM_MCPServerTable( server_id=server_id, @@ -1544,6 +1622,7 @@ if MCP_AVAILABLE: temp_record: Final = _build_temporary_mcp_server_record( payload_with_credentials, created_by, + await _resolve_session_server_id(payload_with_credentials), ) try: @@ -1555,6 +1634,11 @@ if MCP_AVAILABLE: temporary_server, ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS, ) + await _persist_draft_mcp_server( + payload_with_credentials, + temp_record.server_id, + created_by, + ) await _cache_temporary_mcp_server_in_redis( temporary_server, ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 452828edc9e..b72b5d218b2 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -77,6 +77,8 @@ from litellm.proxy.auth.auth_checks import ( _cache_team_object, allowed_route_check_inside_route, can_org_access_model, + delete_cache_key_objects, + delete_cache_team_object, get_org_object, get_team_membership, get_team_object, @@ -313,6 +315,11 @@ class _TeamIdInFilter(TypedDict, total=False): team_id: Mapping[str, Sequence[str]] +_STRIP_DELETED_TEAM_FROM_USERS_SQL: Final = """ +UPDATE "LiteLLM_UserTable" SET teams = array_remove(teams, $1) WHERE $1 = ANY(teams) +""" + + def _team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamTable]": return _typed_table(TeamRepository(prisma_client), LiteLLM_TeamTable) @@ -2657,6 +2664,11 @@ async def _add_team_members_to_team( serialize on the row lock and each appends onto the other's committed result, instead of both rewriting the whole JSON array from a stale snapshot (which silently drops one member on the losing write). + + The same lock serializes this against /team/delete: the delete cannot remove + the row while the reconcile holds it, and a reconcile that finds the row + already gone cleans up after itself rather than leaving the member pointing + at a deleted team id. """ # Process and add new members updated_users, updated_team_memberships = await _process_team_members( @@ -2667,11 +2679,42 @@ async def _add_team_members_to_team( litellm_proxy_admin_name=litellm_proxy_admin_name, ) - async with prisma_client.tx() as tx: - complete_team_data.members_with_roles = await TeamRepository(prisma_client).get_members_with_roles_locked( - tx, data.team_id + updated_team: Final = await _write_members_with_roles_locked( + data=data, + complete_team_data=complete_team_data, + prisma_client=prisma_client, + updated_users=updated_users, + ) + if updated_team is None: + await _sweep_deleted_team_references(team_ids=(data.team_id,), prisma_client=prisma_client) + raise HTTPException( + status_code=404, + detail={"error": f"Team={data.team_id} was deleted while this member add was running"}, ) + return updated_team, updated_users, updated_team_memberships + + +async def _write_members_with_roles_locked( + data: TeamMemberAddRequest, + complete_team_data: LiteLLM_TeamTable, + prisma_client: PrismaClient, + updated_users: list[LiteLLM_UserTable], +) -> LiteLLM_TeamTable | None: + """Reconcile members_with_roles under the team row lock. None when the team row is gone. + + That read is at least as recent as the user and membership writes the caller + already made, so a missing row means /team/delete committed after them. Its + post-delete sweep can have run before those writes landed, which is why the + caller sweeps this team id again rather than only reporting the 404. + """ + async with prisma_client.tx() as tx: + locked_members: Final = await TeamRepository(prisma_client).get_members_with_roles_locked(tx, data.team_id) + if locked_members is None: + return None + + complete_team_data.members_with_roles = locked_members + await _update_team_members_list( data=data, complete_team_data=complete_team_data, @@ -2679,13 +2722,11 @@ async def _add_team_members_to_team( ) _db_team_members: Final = [m.model_dump() for m in complete_team_data.members_with_roles] - updated_team: Final = await tx.litellm_teamtable.update( + return await tx.litellm_teamtable.update( where={"team_id": data.team_id}, data={"members_with_roles": json.dumps(_db_team_members)}, ) - return updated_team, updated_users, updated_team_memberships - def _emit_team_members_metric(team: LiteLLM_TeamTable) -> None: """Update the Prometheus team members gauge after a membership change. @@ -3653,6 +3694,8 @@ async def delete_team( create_audit_log_for_update, litellm_proxy_admin_name, prisma_client, + proxy_logging_obj, + user_api_key_cache, ) if prisma_client is None: @@ -3746,6 +3789,12 @@ async def delete_team( await prisma_client.delete_data(team_id_list=data.team_ids, table_name="key") + await _invalidate_deleted_key_cache( + keys=keys_to_delete, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + ## DELETE ASSOCIATED BYOK MODELS # Runs before the team rows are deleted so a mid-flight failure never leaves # the team gone with its models orphaned. @@ -3779,11 +3828,90 @@ async def delete_team( ) await asyncio.gather(*tasks) + await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client) + ## DELETE TEAMS deleted_teams: Final = await prisma_client.delete_data(team_id_list=data.team_ids, table_name="team") + + # Evict AFTER the rows are gone. Both writers of these keys (`_cache_team_object` and + # `get_team_object_by_alias`) hydrate from the db, so evicting first leaves a window where a + # concurrent auth lookup re-caches the still-present team and the delete looks like it never + # invalidated anything. Nothing fallible runs between the delete and this, or a failure there + # would strand the deleted team in cache. + await _invalidate_deleted_team_cache( + teams=team_rows, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + # Sweep again now the team is gone. A `/team/member_add` that landed between the first sweep + # and the delete would have re-appended the reference; an add still in flight sees the row + # missing under its own row lock and sweeps what it wrote. Both passes are idempotent, and + # keeping the first one means a failure here still leaves a team the admin can retry deleting. + await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client) + return deleted_teams +async def _sweep_deleted_team_references(team_ids: Sequence[str], prisma_client: PrismaClient) -> None: + """ + Strip the deleted team ids from every user row and team-membership row that still references them. + + The per-member `team_member_delete` pass above only reaches users listed in the team's + `members_with_roles`, so a user row that outlived its roster entry is invisible to it and keeps + surfacing the team on `/user/info` after the team is gone. + + #36839 closed the route that created that drift, by resolving member removal off the roster + entry's `user_id` rather than the identifier the caller happened to pass. It does not backfill + rows that already drifted, which is the state this was reported against, so the sweep still has + to run on delete. + + `array_remove` rather than read-filter-write: rewriting the whole array from a snapshot read + outside a transaction drops any team a concurrent `/team/member_add` appended in between. + """ + for team_id in team_ids: + _ = await prisma_client.db.execute_raw(_STRIP_DELETED_TEAM_FROM_USERS_SQL, team_id) + + _ = await _team_membership_db(prisma_client).delete_many(where=_TeamIdInFilter(team_id={"in": tuple(team_ids)})) + + +async def _invalidate_deleted_key_cache( + keys: Sequence[LiteLLM_VerificationToken], + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> None: + """ + Evict the auth cache entry for every key deleted along with the team. + + `/key/delete` evicts as it goes, but the bulk delete above writes straight to the db. Auth + resolves a cached key object without re-reading the team, so a key belonging to a deleted team + keeps buying access until its TTL expires. + """ + await delete_cache_key_objects( + hashed_tokens=tuple(key.token for key in keys), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + +async def _invalidate_deleted_team_cache( + teams: Sequence[LiteLLM_TeamTable], + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> None: + _ = await asyncio.gather( + *( + delete_cache_team_object( + team_id=team.team_id, + team_alias=team.team_alias, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + for team in teams + ) + ) + + def _transform_teams_to_deleted_records( teams: list[LiteLLM_TeamTable], user_api_key_dict: UserAPIKeyAuth, diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py index a4908af561a..7efd32288e4 100644 --- a/litellm/repositories/team_repository.py +++ b/litellm/repositories/team_repository.py @@ -57,9 +57,13 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): return LiteLLM_TeamTable.model_validate(data) - async def get_members_with_roles_locked(self, tx: "Prisma", team_id: str) -> list[Member]: + async def get_members_with_roles_locked(self, tx: "Prisma", team_id: str) -> list[Member] | None: """Return the team's members_with_roles, locking the row FOR UPDATE. + ``None`` when the team row is gone, which a caller holding the lock can + only see if a delete committed under it, as opposed to ``[]`` for a team + that simply has no members. + Must be called inside a transaction so the row lock is held until commit. This serializes concurrent membership writers on the team row so the losing writer appends onto the winner's committed result instead @@ -69,7 +73,9 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): 'SELECT members_with_roles FROM "LiteLLM_TeamTable" WHERE team_id = $1 FOR UPDATE', team_id, ) - raw_value: Final = rows[0]["members_with_roles"] if rows else None + if not rows: + return None + raw_value: Final = rows[0]["members_with_roles"] parsed: Final = json.loads(raw_value) if isinstance(raw_value, str) else raw_value if not parsed: return [] diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index bf119c4fb2f..b6c2fd4c8e3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -1697,6 +1697,362 @@ class TestTemporaryMCPSessionEndpoints: assert result is None assert "expired" not in cache + @pytest.mark.asyncio + async def test_get_cached_temporary_mcp_server_resolves_draft_written_by_another_worker(self): + """Regression: the OAuth session must resolve on a worker that did not serve /session. + + `_temporary_mcp_servers` is per-process, so on a multi-worker or multi-replica proxy the + /authorize and /token legs land on a process whose dict is empty and 404. An empty dict + here IS that other worker. Before the DB-backed draft this returned None. + """ + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + get_cached_temporary_mcp_server, + ) + + draft_row = generate_mock_mcp_server_db_record(server_id="drafted-elsewhere") + rebuilt_server = generate_mock_mcp_server_config_record(server_id="drafted-elsewhere") + mock_manager = MagicMock() + mock_manager.build_mcp_server_from_table = AsyncMock(return_value=rebuilt_server) + get_draft = AsyncMock(return_value=draft_row) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._temporary_mcp_servers", + {}, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_draft_mcp_server", + get_draft, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + ): + result = await get_cached_temporary_mcp_server("drafted-elsewhere") + + assert result is rebuilt_server + # The shared row, not the empty per-process dict, is what answered. + assert get_draft.await_count == 1 + assert get_draft.await_args.args[1] == "drafted-elsewhere" + + @pytest.mark.asyncio + async def test_get_cached_temporary_mcp_server_still_works_without_a_database(self): + """A proxy configured with no database keeps the in-memory session, rather than 404ing. + + Pins the deliberate divergence from a DB-only design: single-process deployments with no + DATABASE_URL must keep working exactly as before. + """ + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _TemporaryMCPServerEntry, + get_cached_temporary_mcp_server, + ) + + server = generate_mock_mcp_server_config_record(server_id="no-db") + entry = _TemporaryMCPServerEntry( + server=server, + expires_at=datetime.utcnow() + timedelta(seconds=300), + ) + get_draft = AsyncMock(return_value=None) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._temporary_mcp_servers", + {"no-db": entry}, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none", + return_value=None, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_draft_mcp_server", + get_draft, + ), + ): + result = await get_cached_temporary_mcp_server("no-db") + + assert result is server + # No database means no draft lookup is even attempted. + assert get_draft.await_count == 0 + + @pytest.mark.asyncio + async def test_create_draft_mcp_server_never_overwrites_a_real_server(self): + """The edit form authorizes against a saved server's own id, so a draft write would + collide on the primary key. That row is already visible to every worker, so it is + returned untouched and no draft is created.""" + from litellm.proxy._experimental.mcp_server.db import create_draft_mcp_server + + real_row = generate_mock_mcp_server_db_record(server_id="already-saved") + real_row.approval_status = "active" + create_call = AsyncMock() + delete_call = AsyncMock() + + with ( + patch( + "litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_row", + AsyncMock(return_value=real_row), + ), + patch( + "litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_rows", + AsyncMock(return_value=[]), + ), + patch("litellm.proxy._experimental.mcp_server.db.create_mcp_server", create_call), + patch("litellm.proxy._experimental.mcp_server.db.delete_mcp_server", delete_call), + ): + result = await create_draft_mcp_server( + MagicMock(), + NewMCPServerRequest(server_id="already-saved", url="https://x.example.com/mcp"), + "tester", + ttl_seconds=300, + ) + + assert result.server_id == "already-saved" + assert create_call.await_count == 0 + assert delete_call.await_count == 0 + + @pytest.mark.asyncio + async def test_create_draft_mcp_server_adopts_the_winner_when_it_loses_a_create_race(self): + """Regression: the read, delete and create are three statements, not one. + + Two concurrent sessions for the same server_id raced and 13 of 20 returned 500 against a + live two-worker proxy. The loser's session is in fact ready, because the winner wrote a + draft for it, so it adopts that row instead of failing the caller. + """ + from litellm.proxy._experimental.mcp_server.db import create_draft_mcp_server + + winner_draft = generate_mock_mcp_server_db_record(server_id="raced") + winner_draft.approval_status = "draft" + # First lookup: nothing yet. After the losing create blows up: the winner's row. + lookups = AsyncMock(side_effect=[None, winner_draft]) + + with ( + patch("litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_row", lookups), + patch( + "litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_rows", + AsyncMock(return_value=[]), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.create_mcp_server", + AsyncMock(side_effect=Exception("duplicate key value violates unique constraint")), + ), + ): + result = await create_draft_mcp_server( + MagicMock(), + NewMCPServerRequest(server_id="raced", url="https://x.example.com/mcp"), + "tester", + ttl_seconds=300, + ) + + assert result.server_id == "raced" + assert lookups.await_count == 2 + + @pytest.mark.asyncio + async def test_create_draft_mcp_server_reraises_when_the_create_failure_was_not_a_race(self): + """A genuine database error must not be swallowed by the race-adoption path.""" + from litellm.proxy._experimental.mcp_server.db import create_draft_mcp_server + + with ( + patch( + "litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_row", + AsyncMock(side_effect=[None, None]), + ), + patch( + "litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_rows", + AsyncMock(return_value=[]), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.create_mcp_server", + AsyncMock(side_effect=Exception("connection refused")), + ), + pytest.raises(Exception, match="connection refused"), + ): + await create_draft_mcp_server( + MagicMock(), + NewMCPServerRequest(server_id="broken", url="https://x.example.com/mcp"), + "tester", + ttl_seconds=300, + ) + + @pytest.mark.asyncio + async def test_create_draft_mcp_server_prunes_drafts_past_their_lifetime(self): + """Regression: abandoned OAuth sessions accumulated forever. Verified against a live + proxy, where 12 drafts aged past the lifetime were still present and a 13th was added.""" + from litellm.proxy._experimental.mcp_server.db import create_draft_mcp_server + + from datetime import timezone + + now = datetime.now(timezone.utc) + stale_one = generate_mock_mcp_server_db_record(server_id="stale-1") + stale_one.updated_at = now - timedelta(hours=1) + stale_two = generate_mock_mcp_server_db_record(server_id="stale-2") + stale_two.updated_at = now - timedelta(hours=1) + # A draft still inside its lifetime must survive the sweep. + fresh_draft = generate_mock_mcp_server_db_record(server_id="still-live") + fresh_draft.updated_at = now + find_rows = AsyncMock(return_value=[stale_one, stale_two, fresh_draft]) + delete_call = AsyncMock() + + with ( + patch("litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_rows", find_rows), + patch( + "litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_row", + AsyncMock(return_value=None), + ), + patch("litellm.proxy._experimental.mcp_server.db.delete_mcp_server", delete_call), + patch( + "litellm.proxy._experimental.mcp_server.db.create_mcp_server", + AsyncMock(return_value=generate_mock_mcp_server_db_record(server_id="fresh")), + ), + ): + await create_draft_mcp_server( + MagicMock(), + NewMCPServerRequest(server_id="fresh", url="https://x.example.com/mcp"), + "tester", + ttl_seconds=300, + ) + + # Only drafts are considered, only the expired ones are removed, and the live one stays. + assert find_rows.await_args.kwargs["where"]["approval_status"] == "draft" + assert sorted(c.args[1] for c in delete_call.await_args_list) == ["stale-1", "stale-2"] + + @pytest.mark.asyncio + async def test_get_all_mcp_servers_hides_drafts_without_hiding_legacy_null_rows(self): + """Drafts are addressable only by their own id and must never appear in a listing, but a + bare inequality would also drop pre-approval-workflow rows, since SQL evaluates + NULL != 'draft' as NULL.""" + from litellm.proxy._experimental.mcp_server.db import get_all_mcp_servers + + find_rows = AsyncMock(return_value=[]) + with patch( + "litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_rows", + find_rows, + ): + await get_all_mcp_servers(MagicMock()) + + where = find_rows.await_args.args[1] + assert where == {"OR": [{"approval_status": None}, {"approval_status": {"not": "draft"}}]} + + @pytest.mark.asyncio + async def test_resolve_session_server_id_refuses_an_unknown_caller_supplied_id(self): + """Regression: two concurrent sessions must never land on one id. + + Honouring an arbitrary caller-supplied id lets a second session adopt the first's draft and + run OAuth against its URL and client credentials, silently. An id naming no real server is + therefore replaced with a fresh one. + """ + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _resolve_session_server_id, + ) + + mock_manager = MagicMock() + mock_manager.get_mcp_server_by_id.return_value = None + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=None), + ), + ): + resolved = await _resolve_session_server_id( + NewMCPServerRequest(server_id="someone-elses-id", url="https://x.example.com/mcp") + ) + + assert resolved != "someone-elses-id" + uuid.UUID(resolved) + + @pytest.mark.asyncio + async def test_resolve_session_server_id_refuses_an_id_that_names_another_sessions_draft(self): + """A draft row is another session's, not a saved server. Replaying an id this endpoint + previously returned must not let a later session inherit the earlier one's config.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _resolve_session_server_id, + ) + + someone_elses_draft = generate_mock_mcp_server_db_record(server_id="earlier-session") + someone_elses_draft.approval_status = "draft" + mock_manager = MagicMock() + mock_manager.get_mcp_server_by_id.return_value = None + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=someone_elses_draft), + ), + ): + resolved = await _resolve_session_server_id( + NewMCPServerRequest(server_id="earlier-session", url="https://x.example.com/mcp") + ) + + assert resolved != "earlier-session" + uuid.UUID(resolved) + + @pytest.mark.asyncio + async def test_resolve_session_server_id_keeps_a_real_servers_id_for_the_edit_flow(self): + """The edit form re-authorizes a saved server against its own id, which must be preserved + or the flow would authorize a throwaway id instead of the server being edited.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _resolve_session_server_id, + ) + + mock_manager = MagicMock() + mock_manager.get_mcp_server_by_id.return_value = generate_mock_mcp_server_config_record(server_id="saved") + + with patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ): + resolved = await _resolve_session_server_id( + NewMCPServerRequest(server_id="saved", url="https://x.example.com/mcp") + ) + + assert resolved == "saved" + + @pytest.mark.asyncio + async def test_resolve_session_server_id_keeps_the_supplied_id_without_a_database(self): + """No database means nothing shared to collide over, so behaviour stays as it is today.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _resolve_session_server_id, + ) + + mock_manager = MagicMock() + mock_manager.get_mcp_server_by_id.return_value = None + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none", + return_value=None, + ), + ): + resolved = await _resolve_session_server_id( + NewMCPServerRequest(server_id="no-db-id", url="https://x.example.com/mcp") + ) + + assert resolved == "no-db-id" + @pytest.mark.asyncio async def test_get_cached_temporary_mcp_server_or_404(self): from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -5705,7 +6061,9 @@ def _edit_endpoint_patches(old_record, update_mock): ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", - AsyncMock(side_effect=old_record) if isinstance(old_record, Exception) else AsyncMock(return_value=old_record), + AsyncMock(side_effect=old_record) + if isinstance(old_record, Exception) + else AsyncMock(return_value=old_record), ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.update_mcp_server", @@ -6114,7 +6472,13 @@ def test_bundled_openapi_registry_parses_and_entries_are_well_formed(): registry_path = os.path.join( os.path.dirname(os.path.abspath(__file__)), - "..", "..", "..", "..", "litellm", "proxy", "openapi_registry.json", + "..", + "..", + "..", + "..", + "litellm", + "proxy", + "openapi_registry.json", ) with open(registry_path) as f: registry = json.load(f) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index fe844db9865..f038af4c7a1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -4,7 +4,7 @@ import os import sys from datetime import datetime, timezone from typing import Optional, cast -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call, patch import pytest from fastapi import HTTPException @@ -39,6 +39,7 @@ from litellm.proxy.management_endpoints.team_endpoints import ( from litellm.proxy.management_endpoints.team_endpoints import ( GetTeamMemberPermissionsResponse, UpdateTeamMemberPermissionsRequest, + _STRIP_DELETED_TEAM_FROM_USERS_SQL, _persist_deleted_team_records, _save_deleted_team_records, _transform_teams_to_deleted_records, @@ -1829,6 +1830,63 @@ async def test_add_team_members_reconciles_against_freshly_locked_row(): assert [m.user_id for m in updated_team.members_with_roles] == ["zed", "alice", "bob"] +@pytest.mark.asyncio +async def test_add_team_members_cleans_up_when_the_team_is_deleted_mid_request(): + """ + Regression pin for the /team/member_add vs /team/delete race. + + The user row and membership writes land before the reconcile takes the team + row lock, so a /team/delete that commits in between has already run its own + reference sweep and cannot see them. The empty locked SELECT is the only + signal that happened, and leaving it at that would strand the member on a + deleted team id, which authorization paths that trust `user.teams` would + treat as membership if the id were ever recreated. So the request must sweep + the references it just wrote and fail, not report success. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + _add_team_members_to_team, + ) + + tx = MagicMock() + tx.query_raw = AsyncMock(return_value=[]) + tx.litellm_teamtable.update = AsyncMock() + + tx_cm = MagicMock() + tx_cm.__aenter__ = AsyncMock(return_value=tx) + tx_cm.__aexit__ = AsyncMock(return_value=None) + + prisma_client = MagicMock() + prisma_client.tx = MagicMock(return_value=tx_cm) + prisma_client.db.execute_raw = AsyncMock() + prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + + with patch( + "litellm.proxy.management_endpoints.team_endpoints._process_team_members", + new=AsyncMock(return_value=([], [])), + ): + with pytest.raises(HTTPException) as exc_info: + await _add_team_members_to_team( + data=TeamMemberAddRequest( + team_id="team-deleted-mid-add", + member=Member(user_id="bob", role="user"), + ), + complete_team_data=LiteLLM_TeamTable(team_id="team-deleted-mid-add", members_with_roles=[]), + prisma_client=cast(object, prisma_client), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + litellm_proxy_admin_name="admin", + ) + + assert exc_info.value.status_code == 404 + tx.litellm_teamtable.update.assert_not_awaited() + + assert prisma_client.db.execute_raw.await_args_list == [ + call(_STRIP_DELETED_TEAM_FROM_USERS_SQL, "team-deleted-mid-add") + ] + prisma_client.db.litellm_teammembership.delete_many.assert_awaited_once_with( + where={"team_id": {"in": ("team-deleted-mid-add",)}} + ) + + def test_add_new_models_to_team_with_existing_models(): """ Test add_new_models_to_team function with existing models @@ -7238,6 +7296,367 @@ async def test_delete_team_persists_deleted_teams(monkeypatch): assert records[0]["litellm_changed_by"] == "admin-user" +@pytest.mark.asyncio +async def test_delete_team_sweeps_references_outside_members_with_roles(monkeypatch): + """ + Regression pin for LIT-5511: a deleted team stayed visible on user records. + + `delete_team` drove all of its cleanup off `team.members_with_roles`, so a user row that + referenced the team by any other route (`/user/update`, SSO sync, a membership row written + without a matching roster entry) kept the dangling team id forever and `/user/info` kept + listing the deleted team. The roster here is deliberately EMPTY, so nothing the per-member + `team_member_delete` path does can make this test pass. + + Both cache keys `_cache_team_object` writes are asserted in the same delete: the id key feeds + `get_team_object` and the alias key feeds the JWT `team_alias_jwt_field` path, so either one + surviving keeps the deleted team resolvable for auth until its TTL expires. + """ + from litellm.proxy._types import DeleteTeamRequest + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + doomed_team = LiteLLM_TeamTable( + team_id="team-doomed", + team_alias="doomed-team", + members_with_roles=[], + metadata={}, + model_max_budget={}, + model_spend={}, + ) + + cache_state_when_rows_deleted = {} + + async def record_cache_state_then_delete(*args, **kwargs): + if kwargs.get("table_name") == "team": + cache_state_when_rows_deleted["doomed_still_cached"] = ( + fresh_cache.get_cache(key="team_id:team-doomed") is not None + ) + return {"deleted_teams": ["team-doomed"]} + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=doomed_team) + mock_prisma_client.delete_data = AsyncMock(side_effect=record_cache_state_then_delete) + mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock() + mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + + mock_execute_raw = AsyncMock() + mock_prisma_client.db.execute_raw = mock_execute_raw + mock_membership_delete_many = AsyncMock() + mock_prisma_client.db.litellm_teammembership.delete_many = mock_membership_delete_many + + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + + fresh_cache = UserApiKeyCache() + for cached_team_id, cached_alias in ( + ("team-doomed", "doomed-team"), + ("team-kept", "kept-team"), + ): + cached_obj = LiteLLM_TeamTableCachedObj( + team_id=cached_team_id, team_alias=cached_alias + ) + fresh_cache.set_cache(key=f"team_id:{cached_team_id}", value=cached_obj) + fresh_cache.set_cache(key=f"team_alias:{cached_alias}", value=cached_obj) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", fresh_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + await delete_team( + data=DeleteTeamRequest(team_ids=["team-doomed"]), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + litellm_changed_by="admin-user", + ) + + # array_remove strips just the deleted id in one statement; a read-filter-write of the whole + # array would drop any team a concurrent /team/member_add appended between read and write + assert "array_remove" in _STRIP_DELETED_TEAM_FROM_USERS_SQL + assert mock_execute_raw.await_args_list == [ + call(_STRIP_DELETED_TEAM_FROM_USERS_SQL, "team-doomed"), + call(_STRIP_DELETED_TEAM_FROM_USERS_SQL, "team-doomed"), + ], "the sweep must run once before the team row is deleted and again after, so a member_add racing the delete cannot leave the reference behind" + + # same two passes: the second one reaps a membership row inserted while the delete was running + assert mock_membership_delete_many.await_args_list == [ + call(where={"team_id": {"in": ("team-doomed",)}}), + call(where={"team_id": {"in": ("team-doomed",)}}), + ] + + assert fresh_cache.get_cache(key="team_id:team-doomed") is None + assert fresh_cache.get_cache(key="team_alias:doomed-team") is None + assert fresh_cache.get_cache(key="team_id:team-kept") is not None + assert fresh_cache.get_cache(key="team_alias:kept-team") is not None + + # Eviction must run AFTER the rows are gone: both writers of these keys hydrate from the db, + # so evicting first lets a concurrent auth lookup re-cache the still-present team. + assert cache_state_when_rows_deleted["doomed_still_cached"] is True + + +@pytest.mark.asyncio +async def test_delete_team_evicts_the_auth_cache_of_the_keys_it_deletes(monkeypatch): + """ + A virtual key scoped to the team is deleted from the db with the team, but auth resolves a + cached key object without re-reading the team, so leaving the cache entry behind lets that key + keep buying access until its TTL expires. Verified live: without this eviction the same key + still returns HTTP 200 on /v1/chat/completions right after /team/delete. + """ + from litellm.proxy._types import DeleteTeamRequest, LiteLLM_VerificationToken + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + team = LiteLLM_TeamTable( + team_id="team-doomed", + team_alias="doomed-team", + members_with_roles=[], + metadata={}, + model_max_budget={}, + model_spend={}, + ) + team_key = LiteLLM_VerificationToken(token="hashed-doomed-key", team_id="team-doomed") + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + mock_prisma_client.delete_data = AsyncMock(return_value={"deleted_teams": ["team-doomed"]}) + mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock() + mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[team_key]) + mock_prisma_client.db.execute_raw = AsyncMock() + mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + + fresh_cache = UserApiKeyCache() + fresh_cache.set_cache(key="hashed-doomed-key", value=UserAPIKeyAuth(token="hashed-doomed-key", team_id="team-doomed")) + fresh_cache.set_cache(key="hashed-unrelated-key", value=UserAPIKeyAuth(token="hashed-unrelated-key")) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", fresh_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + await delete_team( + data=DeleteTeamRequest(team_ids=["team-doomed"]), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + litellm_changed_by="admin-user", + ) + + assert fresh_cache.get_cache(key="hashed-doomed-key") is None + # a key that had nothing to do with the deleted team must survive + assert fresh_cache.get_cache(key="hashed-unrelated-key") is not None + + +@pytest.mark.asyncio +async def test_delete_team_failing_reconcile_sweep_cannot_strand_the_team_in_cache(monkeypatch): + """ + The reconcile sweep runs after the team row is committed deleted. If it ran before cache + eviction, a sweep failure would return an error with the team gone from the db but still + served from cache, which is the exact bug this PR exists to fix. + """ + from litellm.proxy._types import DeleteTeamRequest + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + team = LiteLLM_TeamTable( + team_id="team-doomed", + team_alias="doomed-team", + members_with_roles=[], + metadata={}, + model_max_budget={}, + model_spend={}, + ) + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + mock_prisma_client.delete_data = AsyncMock(return_value={"deleted_teams": ["team-doomed"]}) + mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + # the first sweep succeeds, the post-delete reconcile sweep blows up + mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[None, ConnectionError("db went away")]) + + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + + fresh_cache = UserApiKeyCache() + cached_obj = LiteLLM_TeamTableCachedObj(team_id="team-doomed", team_alias="doomed-team") + fresh_cache.set_cache(key="team_id:team-doomed", value=cached_obj) + fresh_cache.set_cache(key="team_alias:doomed-team", value=cached_obj) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", fresh_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + with pytest.raises(ConnectionError): + await delete_team( + data=DeleteTeamRequest(team_ids=["team-doomed"]), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + litellm_changed_by="admin-user", + ) + + # the delete committed, so the cache must not still be serving the team + assert fresh_cache.get_cache(key="team_id:team-doomed") is None + assert fresh_cache.get_cache(key="team_alias:doomed-team") is None + + +@pytest.mark.asyncio +async def test_delete_team_broadcasts_cache_invalidation_to_other_workers(monkeypatch): + """ + Evicting locally only reaches the worker that handled the delete. Without the broadcast, every + other worker keeps serving the deleted team, and the deleted team's keys, out of its own + in-memory cache until the TTL, so both stay usable for auth cluster-wide. + """ + from litellm.proxy._types import DeleteTeamRequest, LiteLLM_VerificationToken + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + team = LiteLLM_TeamTable( + team_id="team-doomed", + team_alias="doomed-team", + members_with_roles=[], + metadata={}, + model_max_budget={}, + model_spend={}, + ) + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + mock_prisma_client.delete_data = AsyncMock(return_value={"deleted_teams": ["team-doomed"]}) + mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock() + mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[LiteLLM_VerificationToken(token="hashed-doomed-key", team_id="team-doomed")] + ) + mock_prisma_client.db.execute_raw = AsyncMock() + mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + + published = [] + + async def record_publish(cache_key): + published.append(cache_key) + + monkeypatch.setattr("litellm.proxy.auth.auth_checks.publish_auth_cache_invalidation", record_publish) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + await delete_team( + data=DeleteTeamRequest(team_ids=["team-doomed"]), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + litellm_changed_by="admin-user", + ) + + # the deleted key first, then both keys `_cache_team_object` writes: miss the alias one and the + # JWT-by-alias path keeps resolving the team, miss the token and the key still authenticates + assert published == ["hashed-doomed-key", "team_id:team-doomed", "team_alias:doomed-team"] + + +@pytest.mark.asyncio +async def test_delete_team_survives_a_failing_cache_backend(monkeypatch): + """ + Cache eviction runs after the reference sweep has already committed, so a cache backend that + is unreachable must not abort the delete. If it did, `/team/delete` would fail with the team + row still present but its user references and membership rows already gone. + """ + from litellm.proxy._types import DeleteTeamRequest, LiteLLM_VerificationToken + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + team = LiteLLM_TeamTable( + team_id="team-doomed", + team_alias="doomed-team", + members_with_roles=[], + metadata={}, + model_max_budget={}, + model_spend={}, + ) + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + mock_delete_data = AsyncMock(return_value={"deleted_teams": ["team-doomed"]}) + mock_prisma_client.delete_data = mock_delete_data + mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock() + mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock() + # a key to evict: its eviction runs after the key rows are already deleted, so it must not + # raise either + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[LiteLLM_VerificationToken(token="hashed-doomed-key", team_id="team-doomed")] + ) + mock_prisma_client.db.execute_raw = AsyncMock() + mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + + exploding_logging_obj = MagicMock() + exploding_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock( + side_effect=ConnectionError("redis is down") + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", exploding_logging_obj) + monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + result = await delete_team( + data=DeleteTeamRequest(team_ids=["team-doomed"]), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + litellm_changed_by="admin-user", + ) + + assert result == {"deleted_teams": ["team-doomed"]} + mock_delete_data.assert_any_await(team_id_list=["team-doomed"], table_name="team") + assert exploding_logging_obj.internal_usage_cache.dual_cache.async_delete_cache.await_count > 0 + + @pytest.mark.asyncio async def test_team_member_delete_persists_deleted_keys(monkeypatch): from litellm.proxy._types import TeamMemberDeleteRequest diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/test_litellm/repositories/test_repositories.py index 3f567397e1a..38af52f165c 100644 --- a/tests/test_litellm/repositories/test_repositories.py +++ b/tests/test_litellm/repositories/test_repositories.py @@ -546,12 +546,19 @@ class TestTeamRepository: @pytest.mark.asyncio async def test_get_members_with_roles_locked_missing_row(self, repo): + """None, not [], so a caller can tell a deleted team from an empty one. + + /team/member_add reconciles membership under this lock and has to fail, + and clean up the references it already wrote, when a /team/delete + committed underneath it. An empty list would look like a live team with + no members and it would carry on writing. + """ tx = MagicMock() tx.query_raw = AsyncMock(return_value=[]) members = await repo.get_members_with_roles_locked(tx, "missing") - assert members == [] + assert members is None @pytest.mark.asyncio async def test_create_team_all_fields(self, repo): diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx index c704da301ad..237e3b2c842 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx @@ -1,7 +1,8 @@ import { render, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi, beforeEach } from "vitest"; -import UserInfoView, { extractMcpEntitlement } from "./user_info_view"; +import UserInfoView from "./user_info_view"; +import { extractMcpEntitlement } from "@/components/mcp_server_management/mcpEntitlement"; const mockTeamMemberAddCall = vi.fn(); const mockTeamMemberDeleteCall = vi.fn(); @@ -371,7 +372,11 @@ describe("extractMcpEntitlement", () => { const CATALOG = [ { server_id: "srv-1", server_name: "deploy_tracker", alias: "deploy" }, { server_id: "srv-2", server_name: "issue_tracker", alias: null }, - { server_id: "srv-via-group", server_name: "audit_log", alias: null }, + { server_id: "srv-via-group", server_name: "audit_log", alias: null, mcp_access_groups: ["ops_readonly"] }, + ] as any; + + const TOOLSETS = [ + { toolset_id: "ts-1", toolset_name: "audit", tools: [{ server_id: "srv-via-group", tool_name: "read" }] }, ] as any; const form = ( @@ -465,12 +470,43 @@ describe("extractMcpEntitlement", () => { }); it("keeps the allowlist of a deselected server that a retained access group still supplies", () => { - const result = extractMcpEntitlement(form({ accessGroups: ["ops_readonly"] }, { "srv-1": ["read"] }), CATALOG); - expect(result?.mcp_tool_permissions).toEqual({ "srv-1": ["read"] }); + const result = extractMcpEntitlement( + form({ accessGroups: ["ops_readonly"] }, { "srv-via-group": ["read"] }), + CATALOG, + TOOLSETS, + ); + expect(result?.mcp_tool_permissions).toEqual({ "srv-via-group": ["read"] }); }); - it("keeps the allowlist of a deselected server when a toolset is retained", () => { - const result = extractMcpEntitlement(form({ toolsets: ["ts-1"] }, { "srv-1": ["read"] }), CATALOG); + it("drops the allowlist of a deselected server that the retained access group does not contain", () => { + // The gateway grants each allowlist key as its own server, so retaining a group covering only + // srv-via-group must not keep srv-1 callable after the admin removed it. + const result = extractMcpEntitlement( + form({ accessGroups: ["ops_readonly"] }, { "srv-1": ["read"] }), + CATALOG, + TOOLSETS, + ); + expect(result?.mcp_tool_permissions).toEqual({}); + }); + + it("keeps the allowlist of a deselected server that a retained toolset still supplies", () => { + const result = extractMcpEntitlement( + form({ toolsets: ["ts-1"] }, { "srv-via-group": ["read"] }), + CATALOG, + TOOLSETS, + ); + expect(result?.mcp_tool_permissions).toEqual({ "srv-via-group": ["read"] }); + }); + + it("drops the allowlist of a deselected server that the retained toolset does not cover", () => { + const result = extractMcpEntitlement(form({ toolsets: ["ts-1"] }, { "srv-1": ["read"] }), CATALOG, TOOLSETS); + expect(result?.mcp_tool_permissions).toEqual({}); + }); + + it("prunes nothing when a selected toolset is missing from the toolset catalog", () => { + // An unresolvable toolset could supply any server, so pruning against it would be a guess in + // the widening direction. + const result = extractMcpEntitlement(form({ toolsets: ["ts-unknown"] }, { "srv-1": ["read"] }), CATALOG, TOOLSETS); expect(result?.mcp_tool_permissions).toEqual({ "srv-1": ["read"] }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx index 7c7b6b51ef4..34b34066f14 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx @@ -42,77 +42,9 @@ import NotificationsManager from "@/components/molecules/notifications_manager"; import { getBudgetDurationLabel } from "@/components/common_components/budget_duration_dropdown"; import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; import MCPServerPermissions from "@/components/permissions/MCPServerPermissions"; -import { MCPServer } from "@/components/mcp_tools/types"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; - -interface McpEntitlementUpdate { - mcp_servers: string[]; - mcp_access_groups: string[]; - mcp_toolsets: string[]; - mcp_tool_permissions: Record; -} - -const asStringArray = (value: unknown): string[] => - Array.isArray(value) ? value.filter((entry): entry is string => typeof entry === "string") : []; - -const asToolPermissions = (value: unknown): Record => { - if (value === null || typeof value !== "object" || Array.isArray(value)) return {}; - return Object.fromEntries( - Object.entries(value as Record).map(([serverId, tools]) => [serverId, asStringArray(tools)]), - ); -}; - -const mcpServerMatchesIdentifier = (server: MCPServer, identifier: string): boolean => - server.server_id === identifier || server.server_name === identifier || server.alias === identifier; - -/** - * The `object_permission` a save sends, derived from what the editor currently shows. - * - * A tool allowlist is what narrows a grant and an absent one reads as no restriction, so dropping - * an entry is the direction that widens. An entry is kept when an access group or toolset the admin - * retained could still supply its server, and dropped once nothing indirect survives, which is what - * makes removing a grant actually remove it. - * - * A tool-permission key may be a server id, a name or an alias: the gateway normalizes all three - * before looking up the allowlist, so an entry written by the API or by config can use any of them. - * `allServers` is what resolves a key to its servers, plural: names and aliases are not unique, and - * the gateway unions such a key into EVERY server answering to it, so the entry is kept while any - * one of them is still granted. Resolving to the first match instead would make the outcome depend - * on catalog order and could drop a restriction that was also covering a server still granted. A key - * that resolves to nothing is kept too, since a server we cannot identify is one we cannot confirm - * was deselected; that also covers a catalog that has not loaded or failed to load, where every key - * is unresolvable and nothing is pruned. - */ -export const extractMcpEntitlement = ( - formValues: Record, - allServers: MCPServer[], -): McpEntitlementUpdate | null => { - const selection = formValues.mcp_servers_and_groups; - if (selection === null || typeof selection !== "object") return null; - - const { servers, accessGroups, toolsets } = selection as Record; - const mcpServers = asStringArray(servers); - const mcpAccessGroups = asStringArray(accessGroups); - const mcpToolsets = asStringArray(toolsets); - const retainsIndirectGrant = mcpAccessGroups.length > 0 || mcpToolsets.length > 0; - - const grantsServerNamedBy = (permissionKey: string): boolean => { - const named = allServers.filter((candidate) => mcpServerMatchesIdentifier(candidate, permissionKey)); - if (named.length === 0) return true; - return named.some((server) => mcpServers.some((identifier) => mcpServerMatchesIdentifier(server, identifier))); - }; - - return { - mcp_servers: mcpServers, - mcp_access_groups: mcpAccessGroups, - mcp_toolsets: mcpToolsets, - mcp_tool_permissions: Object.fromEntries( - Object.entries(asToolPermissions(formValues.mcp_tool_permissions)).filter( - ([permissionKey]) => retainsIndirectGrant || grantsServerNamedBy(permissionKey), - ), - ), - }; -}; +import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; +import { extractMcpEntitlement } from "@/components/mcp_server_management/mcpEntitlement"; interface UserInfoViewProps { userId: string; @@ -164,6 +96,7 @@ export default function UserInfoView({ const [selectedRole, setSelectedRole] = useState("user"); const [isLoadingTeams, setIsLoadingTeams] = useState(false); const { data: allMcpServers = [] } = useMCPServers(); + const { data: allMcpToolsets = [] } = useMCPToolsets(); React.useEffect(() => { setBaseUrl(getProxyBaseUrl()); @@ -365,7 +298,7 @@ export default function UserInfoView({ try { if (!accessToken || !userData) return; - const mcpEntitlement = extractMcpEntitlement(formValues, allMcpServers); + const mcpEntitlement = extractMcpEntitlement(formValues, allMcpServers, allMcpToolsets); const userFields = Object.fromEntries( Object.entries(formValues).filter( ([field]) => field !== "mcp_servers_and_groups" && field !== "mcp_tool_permissions", diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/mcpEntitlement.ts b/ui/litellm-dashboard/src/components/mcp_server_management/mcpEntitlement.ts new file mode 100644 index 00000000000..66f59f5b71d --- /dev/null +++ b/ui/litellm-dashboard/src/components/mcp_server_management/mcpEntitlement.ts @@ -0,0 +1,91 @@ +import { ALL_PROXY_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; +import { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; + +export interface McpEntitlementUpdate { + mcp_servers: string[]; + mcp_access_groups: string[]; + mcp_toolsets: string[]; + mcp_tool_permissions: Record; +} + +const asStringArray = (value: unknown): string[] => + Array.isArray(value) ? value.filter((entry): entry is string => typeof entry === "string") : []; + +const asToolPermissions = (value: unknown): Record => { + if (value === null || typeof value !== "object" || Array.isArray(value)) return {}; + return Object.fromEntries( + Object.entries(value as Record).map(([serverId, tools]) => [serverId, asStringArray(tools)]), + ); +}; + +const mcpServerMatchesIdentifier = (server: MCPServer, identifier: string): boolean => + server.server_id === identifier || server.server_name === identifier || server.alias === identifier; + +/** + * The `object_permission` a save sends, derived from what the editor currently shows. + * + * A tool allowlist is what narrows a grant and an absent one reads as no restriction, so dropping + * an entry is the direction that widens. An entry is kept while its own server is still reachable, + * directly or through a retained access group or toolset, and dropped once nothing reaches it, which + * is what makes removing a grant actually remove it. Reachability is resolved per server rather than + * per selection: the gateway treats every allowlist key as an independent server grant, so keeping + * every key because some unrelated group survived would leave a deselected server callable. + * + * The catalog carries `mcp_access_groups` on each server and `tools[].server_id` on each toolset, + * which is the same membership the gateway resolves against. A selected toolset missing from the + * catalog is unresolvable, so nothing is pruned in that save. + * + * A tool-permission key may be a server id, a name or an alias: the gateway normalizes all three + * before looking up the allowlist, so an entry written by the API or by config can use any of them. + * `allServers` is what resolves a key to its servers, plural: names and aliases are not unique, and + * the gateway unions such a key into EVERY server answering to it, so the entry is kept while any + * one of them is still granted. Resolving to the first match instead would make the outcome depend + * on catalog order and could drop a restriction that was also covering a server still granted. A key + * that resolves to nothing is kept too, since a server we cannot identify is one we cannot confirm + * was deselected; that also covers a catalog that has not loaded or failed to load, where every key + * is unresolvable and nothing is pruned. + */ +export const extractMcpEntitlement = ( + formValues: Record, + allServers: MCPServer[], + allToolsets: MCPToolset[] = [], +): McpEntitlementUpdate | null => { + const selection = formValues.mcp_servers_and_groups; + if (selection === null || typeof selection !== "object") return null; + + const { servers, accessGroups, toolsets } = selection as Record; + const mcpServers = asStringArray(servers); + const mcpAccessGroups = asStringArray(accessGroups); + const mcpToolsets = asStringArray(toolsets); + const grantsEveryServer = + mcpServers.includes(ALL_PROXY_MCP_SERVERS_SENTINEL) || + mcpToolsets.some((toolsetId) => !allToolsets.some((toolset) => toolset.toolset_id === toolsetId)); + + const toolsetServerIds = new Set( + allToolsets + .filter((toolset) => mcpToolsets.includes(toolset.toolset_id)) + .flatMap((toolset) => toolset.tools.map((tool) => tool.server_id)), + ); + + const grants = (server: MCPServer): boolean => + mcpServers.some((identifier) => mcpServerMatchesIdentifier(server, identifier)) || + (server.mcp_access_groups ?? []).some((group) => mcpAccessGroups.includes(group)) || + toolsetServerIds.has(server.server_id); + + const grantsServerNamedBy = (permissionKey: string): boolean => { + const named = allServers.filter((candidate) => mcpServerMatchesIdentifier(candidate, permissionKey)); + if (named.length === 0) return true; + return named.some(grants); + }; + + return { + mcp_servers: mcpServers, + mcp_access_groups: mcpAccessGroups, + mcp_toolsets: mcpToolsets, + mcp_tool_permissions: Object.fromEntries( + Object.entries(asToolPermissions(formValues.mcp_tool_permissions)).filter( + ([permissionKey]) => grantsEveryServer || grantsServerNamedBy(permissionKey), + ), + ), + }; +}; diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx index 16b1be0d74a..374e36029a0 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -255,6 +255,15 @@ vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ useUISettings: vi.fn().mockReturnValue({ data: { values: {} }, isLoading: false }), })); +// Mock useMCPServers hook (requires QueryClientProvider which is not available in this test) +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ + useMCPServers: vi.fn().mockReturnValue({ data: [] }), +})); + +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPToolsets", () => ({ + useMCPToolsets: vi.fn().mockReturnValue({ data: [] }), +})); + // Mock useResetKeySpend hook (requires QueryClientProvider which is not available in this test) vi.mock("@/app/(dashboard)/hooks/keys/useResetKeySpend", () => ({ useResetKeySpend: vi.fn().mockReturnValue({ diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx index d2390a633d5..1437689958c 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx @@ -32,6 +32,26 @@ vi.mock("@/app/(dashboard)/hooks/projects/useProjects", () => ({ useProjects: vi.fn().mockReturnValue({ data: [], isLoading: false }), })); +const MCP_CATALOG = [ + { server_id: "srv-1", server_name: "deploy_tracker", alias: "deploy" }, + { server_id: "srv-2", server_name: "incident_log", alias: "incidents", mcp_access_groups: ["ops_readonly"] }, +]; + +const MCP_TOOLSETS = [ + { toolset_id: "ts-1", toolset_name: "incidents", tools: [{ server_id: "srv-2", tool_name: "write" }] }, +]; + +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ + useMCPServers: vi.fn(() => ({ data: MCP_CATALOG })), +})); + +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPToolsets", () => ({ + useMCPToolsets: vi.fn(() => ({ data: MCP_TOOLSETS })), +})); + +import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; + vi.mock("../networking", () => ({ keyDeleteCall: vi.fn().mockResolvedValue({}), keyUpdateCall: vi.fn().mockResolvedValue({}), @@ -881,6 +901,182 @@ describe("KeyInfoView", () => { }); }); + describe("MCP tool permissions on save", () => { + const KEY_WITH_TOOL_PERMISSIONS: KeyResponse = { + ...MOCK_KEY_DATA, + user_id: "proxy-admin-user", + object_permission: { + ...MOCK_KEY_DATA.object_permission, + mcp_servers: ["srv-1", "srv-2"], + mcp_access_groups: [], + mcp_tool_permissions: { "srv-1": ["read"], "srv-2": ["write"] }, + }, + } as KeyResponse; + + const enterEditMode = async (keyData: KeyResponse) => { + vi.mocked(useAuthorized).mockReturnValue({ + ...baseUseAuthorizedMock, + userId: "proxy-admin-user", + userRole: "proxy_admin", + }); + renderWithProviders( + {}} keyId="test-key-id" onKeyDataUpdate={() => {}} teams={[]} />, + ); + await userEvent.click(screen.getByRole("tab", { name: /settings/i })); + await userEvent.click(screen.getByRole("button", { name: /edit settings/i })); + await waitFor(() => expect(editViewMocks.onSubmit).toBeDefined()); + }; + + const submittedToolPermissions = () => { + const payload = vi.mocked(keyUpdateCall).mock.calls.at(-1)?.[1] as Record; + return payload.object_permission.mcp_tool_permissions; + }; + + beforeEach(() => { + editViewMocks.onSubmit = undefined; + vi.mocked(keyUpdateCall).mockClear(); + vi.mocked(keyUpdateCall).mockResolvedValue({}); + vi.mocked(useMCPServers).mockReturnValue({ data: MCP_CATALOG } as unknown as ReturnType); + vi.mocked(useMCPToolsets).mockReturnValue({ data: MCP_TOOLSETS } as unknown as ReturnType); + }); + + it("drops the allowlist of every deselected server instead of leaving it entitled", async () => { + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: [] }, + mcp_tool_permissions: { "srv-1": ["read"], "srv-2": ["write"] }, + }); + + expect(submittedToolPermissions()).toEqual({}); + }); + + it("drops only the deselected server and keeps the one still granted", async () => { + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + mcp_servers_and_groups: { servers: ["srv-1"], accessGroups: [], toolsets: [] }, + mcp_tool_permissions: { "srv-1": ["read"], "srv-2": ["write"] }, + }); + + expect(submittedToolPermissions()).toEqual({ "srv-1": ["read"] }); + }); + + it("keeps an allowlist whose server is reachable through a retained access group", async () => { + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + mcp_servers_and_groups: { servers: [], accessGroups: ["ops_readonly"], toolsets: [] }, + mcp_tool_permissions: { "srv-2": ["write"] }, + }); + + expect(submittedToolPermissions()).toEqual({ "srv-2": ["write"] }); + }); + + it("drops an allowlist the retained access group does not reach", async () => { + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + mcp_servers_and_groups: { servers: [], accessGroups: ["ops_readonly"], toolsets: [] }, + mcp_tool_permissions: { "srv-1": ["read"], "srv-2": ["write"] }, + }); + + expect(submittedToolPermissions()).toEqual({ "srv-2": ["write"] }); + }); + + it("drops an allowlist the retained toolset does not cover", async () => { + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: ["ts-1"] }, + mcp_tool_permissions: { "srv-1": ["read"], "srv-2": ["write"] }, + }); + + expect(submittedToolPermissions()).toEqual({ "srv-2": ["write"] }); + }); + + it("refuses to save a permission change while a selected toolset is unresolvable", async () => { + vi.mocked(useMCPToolsets).mockReturnValue({ data: undefined } as unknown as ReturnType); + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: ["ts-1"] }, + mcp_tool_permissions: { "srv-1": ["read"], "srv-2": ["write"] }, + }); + + expect(keyUpdateCall).not.toHaveBeenCalled(); + }); + + it("resolves a name-keyed allowlist against the server catalog", async () => { + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + mcp_servers_and_groups: { servers: ["srv-1"], accessGroups: [], toolsets: [] }, + mcp_tool_permissions: { deploy_tracker: ["read"], incident_log: ["write"] }, + }); + + expect(submittedToolPermissions()).toEqual({ deploy_tracker: ["read"] }); + }); + + it("clears every allowlist when the admin picks the no-MCP-servers sentinel", async () => { + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + mcp_servers_and_groups: { servers: ["no-mcp-servers"], accessGroups: [], toolsets: [] }, + mcp_tool_permissions: { "srv-1": ["read"], "srv-2": ["write"] }, + }); + + expect(submittedToolPermissions()).toEqual({}); + }); + + it("keeps every allowlist when the admin grants all proxy servers", async () => { + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + mcp_servers_and_groups: { servers: ["all-proxy-mcpservers"], accessGroups: [], toolsets: [] }, + mcp_tool_permissions: { "srv-1": ["read"], "srv-2": ["write"] }, + }); + + expect(submittedToolPermissions()).toEqual({ "srv-1": ["read"], "srv-2": ["write"] }); + }); + + it("refuses to save a permission change it cannot compute without the server catalog", async () => { + vi.mocked(useMCPServers).mockReturnValue({ data: undefined } as ReturnType); + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: [] }, + mcp_tool_permissions: { "srv-1": ["read"], "srv-2": ["write"] }, + }); + + expect(keyUpdateCall).not.toHaveBeenCalled(); + }); + + it("preserves a vector-store edit made in the same save", async () => { + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + vector_stores: ["vs-1"], + mcp_servers_and_groups: { servers: ["srv-1"], accessGroups: [], toolsets: [] }, + mcp_tool_permissions: { "srv-1": ["read"] }, + }); + + const payload = vi.mocked(keyUpdateCall).mock.calls.at(-1)?.[1] as Record; + expect(payload.object_permission.vector_stores).toEqual(["vs-1"]); + }); + }); + describe("delete flow", () => { it("invalidates the keys list query after a successful delete so active filters survive (LIT-4080)", async () => { const invalidateSpy = vi.spyOn(QueryClient.prototype, "invalidateQueries"); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index f7720874968..15d14d5abf5 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -24,6 +24,9 @@ import { useResetKeySpend } from "@/app/(dashboard)/hooks/keys/useResetKeySpend" import { useSetKeyBlockedState } from "@/app/(dashboard)/hooks/keys/useSetKeyBlockedState"; import { keyKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import { useQueryClient } from "@tanstack/react-query"; +import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; +import { extractMcpEntitlement } from "../mcp_server_management/mcpEntitlement"; import ObjectPermissionsView from "../object_permissions_view"; import { RegenerateKeyModal } from "../organisms/RegenerateKeyModal"; import { parseErrorMessage } from "../shared/errorUtils"; @@ -74,6 +77,8 @@ export default function KeyInfoView({ const { teams: teamsData } = useTeams(); const { data: projects } = useProjects(); const { data: uiSettingsData } = useUISettings(); + const { data: allMcpServers } = useMCPServers(); + const { data: allMcpToolsets } = useMCPToolsets(); const enableProjectsUI = Boolean(uiSettingsData?.values?.enable_projects_ui); const [isEditing, setIsEditing] = useState(false); const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); @@ -201,33 +206,28 @@ export default function KeyInfoView({ delete formValues.vector_stores; } - if (formValues.mcp_servers_and_groups !== undefined) { - const { servers, accessGroups, toolsets } = formValues.mcp_servers_and_groups || { - servers: [], - accessGroups: [], - toolsets: [], - }; - formValues.object_permission = { - ...currentKeyData.object_permission, - mcp_servers: servers || [], - mcp_access_groups: accessGroups || [], - mcp_toolsets: toolsets || [], - }; - // Remove mcp_servers_and_groups from the top level as it should be in object_permission - delete formValues.mcp_servers_and_groups; - } - - // Handle MCP tool permissions - if (formValues.mcp_tool_permissions !== undefined) { - const mcpToolPermissions = formValues.mcp_tool_permissions || {}; - if (Object.keys(mcpToolPermissions).length > 0) { - formValues.object_permission = { - ...formValues.object_permission, - mcp_tool_permissions: mcpToolPermissions, - }; + const mcpEntitlement = extractMcpEntitlement(formValues, allMcpServers ?? [], allMcpToolsets ?? []); + if (mcpEntitlement) { + // Without a catalog the grants an allowlist key still has are unresolvable, so nothing is + // pruned and a revocation would save as a no-op while reporting success. Refuse instead. + const unresolvableSelection = + allMcpServers === undefined || + mcpEntitlement.mcp_toolsets.some( + (toolsetId) => !(allMcpToolsets ?? []).some((toolset) => toolset.toolset_id === toolsetId), + ); + if (unresolvableSelection && Object.keys(mcpEntitlement.mcp_tool_permissions).length > 0) { + NotificationManager.error( + "MCP server or toolset list is unavailable, so MCP permissions cannot be saved yet. Retry.", + ); + return; } - delete formValues.mcp_tool_permissions; + formValues.object_permission = { + ...(formValues.object_permission ?? currentKeyData.object_permission), + ...mcpEntitlement, + }; } + delete formValues.mcp_servers_and_groups; + delete formValues.mcp_tool_permissions; // Handle agent permissions if (formValues.agents_and_groups !== undefined) {