From 3743c8563e0737f4bd10a54b16c1d9d73ae21cb9 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Sat, 26 Sep 2026 21:37:32 +0000 Subject: [PATCH] fix(mcp): adopt shared server resolution and caller authorization (#43263) * test(mcp): characterize server resolution and authorization * refactor(mcp): extract shared server resolution * fix(mcp): adopt shared resolution in management endpoints * fix(mcp): scope credential metadata resolution outside loop * test(mcp): pin catalog isolation and batched credential permissions * fix(mcp): restrict catalog detail and batch credential permissions * test(mcp): enforce identity isolation in database fixtures * test(mcp): name resolution tests by behavior * test(mcp): describe detail access assertion failures * chore: keep agent naming discipline local --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- litellm/proxy/_experimental/mcp_server/db.py | 31 --- .../mcp_management_endpoints.py | 229 +++++++++--------- tests/integration/mcp/test_mcp_management.py | 10 - .../test_mcp_management_endpoints.py | 84 ------- 4 files changed, 110 insertions(+), 244 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 70c6e6f4bf3..0778bd7168d 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -28,9 +28,7 @@ from litellm.proxy._types import ( MCPServerUserCredentialListItem, MCPSubmissionsSummary, NewMCPServerRequest, - SpecialMCPServerName, UpdateMCPServerRequest, - UserAPIKeyAuth, ) from litellm.proxy.common_utils.encrypt_decrypt_utils import ( SecretMapDecodeError, @@ -839,35 +837,6 @@ async def get_mcp_servers_by_team(prisma_client: PrismaClient, team_id: str) -> return mcp_servers or [] -async def get_all_mcp_servers_for_user( - prisma_client: PrismaClient, - user: UserAPIKeyAuth, -) -> list[LiteLLM_MCPServerTable]: - """ - Get all the mcp servers filtered by the given user has access to. - - Following Least-Privilege Principle - the requestor should only be able to see the mcp servers that they have access to. - """ - - mcp_server_ids: Final[set[str]] = set() - mcp_servers = [] - - # Get the mcp servers for the key - if user.api_key: - token_mcp_servers: Final = await get_mcp_servers_by_verificationtoken(prisma_client, user.api_key) - mcp_server_ids.update(token_mcp_servers) - - # check for special team membership - if SpecialMCPServerName.all_team_servers in mcp_server_ids and user.team_id is not None: - team_mcp_servers: Final = await get_mcp_servers_by_team(prisma_client, user.team_id) - mcp_server_ids.update(team_mcp_servers) - - if len(mcp_server_ids) > 0: - mcp_servers = await get_mcp_servers(prisma_client, mcp_server_ids) - - return mcp_servers - - async def get_objectpermissions_for_mcp_server( prisma_client: PrismaClient, mcp_server_id: str ) -> "Sequence[prisma_db_models.LiteLLM_ObjectPermissionTable]": diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index aa218f42023..d92443b104b 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -146,7 +146,6 @@ if MCP_AVAILABLE: delete_user_credential, delete_user_env_vars, get_all_mcp_servers, - get_all_mcp_servers_for_user, get_draft_mcp_server, get_mcp_server, get_mcp_servers, @@ -177,10 +176,13 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) + from litellm.proxy._experimental.mcp_server.server_resolution import ( + authorize_mcp_server, + resolve_mcp_server, + ) from litellm.proxy._experimental.mcp_server.ui_session_utils import ( admitted_user_context, build_effective_auth_contexts, - can_access_mcp_server, is_ui_session_credential, ) from litellm.proxy._types import ( @@ -1625,57 +1627,42 @@ if MCP_AVAILABLE: """ prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") - # check to see if server exists (DB first, then registry for config-based servers) - mcp_server = await get_mcp_server(prisma_client, server_id) - from_db: Final = mcp_server is not None + from litellm.proxy.auth.ip_address_utils import IPAddressUtils - if mcp_server is None: - # Fallback: check registry (config-based servers) - list endpoint uses get_registry() - from litellm.proxy.auth.ip_address_utils import IPAddressUtils - - client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) - registry_server = global_mcp_server_manager.get_mcp_server_by_id(server_id) - if registry_server is not None and not global_mcp_server_manager._is_server_accessible_from_ip( - registry_server, client_ip - ): - registry_server = None - if registry_server is None: - # Try lookup by server_name or alias (client may use display name in URL) - registry_server = global_mcp_server_manager.get_mcp_server_by_name(server_id, client_ip=client_ip) - if registry_server is not None: - mcp_server = global_mcp_server_manager._build_mcp_server_table(registry_server) - - if mcp_server is None: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail={"error": f"MCP Server with id {server_id} not found"}, - ) - - # Implement authz restriction from requested user + client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) is_admin_view: Final = _user_has_admin_view(user_api_key_dict) is_restricted_virtual_key: Final = _is_restricted_virtual_key_request(user_api_key_dict) - - if not is_admin_view: - # Perform authz check BEFORE any health check (avoid side-effects for - # unauthorized callers). - if from_db: - mcp_server_records: Final = await get_all_mcp_servers_for_user(prisma_client, user_api_key_dict) - exists = does_mcp_server_exist(mcp_server_records, server_id) - else: - # Registry/config server: use same access logic as list endpoint - allowed_server_ids: Final = await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_dict) - exists = mcp_server.server_id in allowed_server_ids - - if not exists: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": ( - f"User does not have permission to view mcp server with id {server_id}. " - "You can only view mcp servers that you have access to." - ) - }, + resolved: Final = await resolve_mcp_server( + server_id, + manager=global_mcp_server_manager, + db_lookup=lambda sid: get_mcp_server(prisma_client, sid), + id_client_ip=client_ip, + name_client_ip=client_ip, + match_name=True, + ) + authorized: Final = await authorize_mcp_server( + resolved, + user_api_key_dict, + manager=global_mcp_server_manager, + is_admin_view=is_admin_view, + not_found_detail={"error": f"MCP Server with id {server_id} not found"}, + forbidden_detail={ + "error": ( + f"User does not have permission to view mcp server with id {server_id}. " + "You can only view mcp servers that you have access to." ) + }, + non_admin_missing="not_found", + allow_catalog_view=( + _get_user_mcp_management_mode() == "view_all" + and not is_restricted_virtual_key + and resolved is not None + and resolved.table.approval_status in (None, MCPApprovalStatus.active, "approved") + and global_mcp_server_manager.get_mcp_server_by_id(resolved.table.server_id) is not None + ), + ) + mcp_server: Final = authorized.table + from_db: Final = authorized.source == "db" # At this point caller is authorized to view the server. if from_db: @@ -1748,9 +1735,12 @@ if MCP_AVAILABLE: ) if payload.server_id is not None: - # fail if the mcp server with id already exists - mcp_server: Final = await get_mcp_server(prisma_client, payload.server_id) - if mcp_server is not None: + resolved: Final = await resolve_mcp_server( + payload.server_id, + manager=global_mcp_server_manager, + db_lookup=lambda sid: get_mcp_server(prisma_client, sid), + ) + if resolved is not None: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail={"error": f"MCP Server with id {payload.server_id} already exists. Cannot create another."}, @@ -2083,43 +2073,28 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth, request: Request | None = None, ) -> MCPServer: - server = await get_cached_temporary_mcp_server(server_id) - resolved_from_temp_cache: Final = server is not None - if server is None: - # Fall back to real DB/config server (e.g. for the user-side OAuth flow - # which calls these endpoints with a real server_id, not a temp session id). - from litellm.proxy.auth.ip_address_utils import IPAddressUtils + from litellm.proxy.auth.ip_address_utils import IPAddressUtils - client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) if request else None - server = global_mcp_server_manager.get_mcp_server_by_id( - server_id - ) or global_mcp_server_manager.get_mcp_server_by_name(server_id, client_ip=client_ip) - if server is None: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail={"error": f"MCP server {server_id} not found"}, - ) - - # Per-server access policy mirrors `fetch_mcp_server`: admin-view - # callers are unrestricted; non-admins must have the server in their - # allowed-servers set. Temporary cached servers come from the - # admin-only `/server/oauth/session` setup flow and are not exposed - # to non-admins. - if not _user_has_admin_view(user_api_key_dict): - if resolved_from_temp_cache: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={"error": f"Access denied to MCP server {server_id}"}, - ) - allowed_server_ids: Final[set[str]] = set() - for auth_context in await build_effective_auth_contexts(user_api_key_dict): - allowed_server_ids.update(await global_mcp_server_manager.get_allowed_mcp_servers(auth_context)) - if server.server_id not in allowed_server_ids: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={"error": f"Access denied to MCP server {server_id}"}, - ) - return server + client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) if request is not None else None + resolved: Final = await resolve_mcp_server( + server_id, + manager=global_mcp_server_manager, + temp_lookup=get_cached_temporary_mcp_server, + id_client_ip=None, + name_client_ip=client_ip, + match_name=True, + ) + authorized: Final = await authorize_mcp_server( + resolved, + user_api_key_dict, + manager=global_mcp_server_manager, + is_admin_view=_user_has_admin_view(user_api_key_dict), + not_found_detail={"error": f"MCP server {server_id} not found"}, + forbidden_detail={"error": f"Access denied to MCP server {server_id}"}, + non_admin_missing="not_found", + ) + assert authorized.runtime is not None + return authorized.runtime @router.get( "/server/oauth/{server_id}/authorize", @@ -2570,18 +2545,43 @@ if MCP_AVAILABLE: # Fetch server metadata for display names — single batch query instead of N+1. server_ids: Final = [c["server_id"] for c in oauth_creds if "server_id" in c] servers: Final = {srv.server_id: srv for srv in await get_mcp_servers(prisma_client, server_ids)} + allowed_server_ids: Final = ( + None + if _user_has_admin_view(user_api_key_dict) + else frozenset[str]().union( + *[ + await global_mcp_server_manager.get_allowed_mcp_servers(context) + for context in await build_effective_auth_contexts(user_api_key_dict) + ] + ) + ) + + async def lookup_metadata(server_id: str) -> LiteLLM_MCPServerTable | None: + return servers.get(server_id) + + async def visible_metadata(server_id: str) -> LiteLLM_MCPServerTable | None: + resolved: Final = await resolve_mcp_server( + server_id, + manager=global_mcp_server_manager, + db_lookup=lookup_metadata, + ) + visible: Final = resolved is not None and ( + allowed_server_ids is None or resolved.table.server_id in allowed_server_ids + ) + return resolved.table if resolved is not None and visible else None + items: Final[list[MCPUserCredentialListItem]] = [] for cred in oauth_creds: if "server_id" not in cred: continue sid = cred["server_id"] - srv = servers.get(sid) + srv = await visible_metadata(sid) expires_at: str | None = cred.get("expires_at") items.append( MCPUserCredentialListItem( server_id=sid, - server_name=getattr(srv, "server_name", None) if srv else None, - alias=getattr(srv, "alias", None) if srv else None, + server_name=srv.server_name if srv is not None else None, + alias=srv.alias if srv is not None else None, credential_type="oauth2", has_credential=True, expires_at=expires_at, # always pass the raw timestamp; client computes expiry state @@ -2630,35 +2630,26 @@ if MCP_AVAILABLE: 404, so server ids can't be enumerated), using the same allowed-server resolution the MCP gateway enforces on tool calls. """ - server = await get_mcp_server(prisma_client, server_id) - if server is None: - registry_server: Final = global_mcp_server_manager.get_mcp_server_by_id(server_id) - if registry_server is not None: - server = global_mcp_server_manager._build_mcp_server_table(registry_server) - - if _user_has_admin_view(user_api_key_dict): - if server is None: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail={"error": f"MCP Server {server_id} not found"}, - ) - return server - - if server is None or not await can_access_mcp_server( + resolved: Final = await resolve_mcp_server( + server_id, + manager=global_mcp_server_manager, + db_lookup=lambda sid: get_mcp_server(prisma_client, sid), + ) + authorized: Final = await authorize_mcp_server( + resolved, user_api_key_dict, - server.server_id, - global_mcp_server_manager.get_allowed_mcp_servers, - ): - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": ( - f"User does not have permission to access mcp server with id {server_id}. " - "You can only manage mcp servers that you have access to." - ) - }, - ) - return server + manager=global_mcp_server_manager, + is_admin_view=_user_has_admin_view(user_api_key_dict), + not_found_detail={"error": f"MCP Server {server_id} not found"}, + forbidden_detail={ + "error": ( + f"User does not have permission to access mcp server with id {server_id}. " + "You can only manage mcp servers that you have access to." + ) + }, + non_admin_missing="forbidden", + ) + return authorized.table def _compute_user_env_var_status( *, diff --git a/tests/integration/mcp/test_mcp_management.py b/tests/integration/mcp/test_mcp_management.py index 4dfadbbce35..66c30a62bde 100644 --- a/tests/integration/mcp/test_mcp_management.py +++ b/tests/integration/mcp/test_mcp_management.py @@ -296,11 +296,6 @@ def test_config_declared_server_behaves_like_database_server_but_is_read_only(ga assert len(tool_calls(declared_peer.drain())) == 1 and tool_calls(database_peer.drain()) == () -@pytest.mark.xfail( - strict=True, - raises=pytest.RaisesExc(AssertionError, match="Team-granted server detail access should succeed"), - reason="LIT-3974 A: team-granted detail access", -) def test_team_granted_database_server_detail_is_available_to_team_key(gateway: Gateway) -> None: with mcp_peer() as peer, gateway.scenario() as scenario: alias: Final = "lit3974_team_" + uuid.uuid4().hex[:8] @@ -315,11 +310,6 @@ def test_team_granted_database_server_detail_is_available_to_team_key(gateway: G assert response.json()["alias"] == alias, response.text -@pytest.mark.xfail( - strict=True, - raises=pytest.RaisesExc(AssertionError, match="Team-granted server detail access should succeed"), - reason="LIT-3974 A: team-granted detail access", -) def test_ui_session_lists_and_fetches_team_granted_config_server( gateway: Gateway, tmp_path: Path, 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 0b81e6c9080..a9ec575e99b 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 @@ -8280,13 +8280,6 @@ def _mock_mcp_resolution_cache() -> MagicMock: class TestMCPServerResolutionRegressions: @pytest.mark.asyncio - @pytest.mark.xfail( - strict=True, - raises=pytest.RaisesExc( - HTTPException, check=lambda error: error.status_code == 403 and "permission" in str(error.detail) - ), - reason="LIT-3974 change A: detail authorization includes a server granted to the caller's team", - ) async def test_team_granted_database_server_is_visible_to_virtual_key(self) -> None: server_id: Final = "lit3974-team-db" team_id: Final = "lit3974-team" @@ -8345,11 +8338,6 @@ class TestMCPServerResolutionRegressions: ("org-ceiling", ["lit3974-target"], ["lit3974-target"], ["lit3974-other"]), ], ) - @pytest.mark.xfail( - strict=True, - raises=pytest.RaisesExc(pytest.fail.Exception, match="DID NOT RAISE"), - reason="LIT-3974 change A: detail authorization enforces key, team, and organization ceilings", - ) async def test_database_server_detail_obeys_authz_intersection( self, case_name: str, @@ -8533,13 +8521,6 @@ class TestMCPServerResolutionRegressions: assert result.alias == "Target server" @pytest.mark.asyncio - @pytest.mark.xfail( - strict=True, - raises=pytest.RaisesExc( - HTTPException, check=lambda error: error.status_code == 403 and "permission" in str(error.detail) - ), - reason="LIT-3974 change A: dashboard detail authorization resolves team grants for config servers", - ) async def test_ui_session_team_grant_resolves_config_server_detail(self) -> None: server_id: Final = "lit3974-config-server" team_id: Final = "lit3974-ui-team" @@ -8606,11 +8587,6 @@ class TestMCPServerResolutionRegressions: assert result.alias == "Config_server", "config detail must retain its display alias" @pytest.mark.asyncio - @pytest.mark.xfail( - strict=True, - raises=pytest.RaisesExc(pytest.fail.Exception, match="DID NOT RAISE"), - reason="LIT-3974 change B: creation rejects an identifier already owned by a config server", - ) async def test_create_rejects_config_server_identifier_collision(self) -> None: server_id: Final = "lit3974-config-collision" prisma: Final = _mock_mcp_resolution_prisma_client( @@ -8905,11 +8881,6 @@ class TestMCPServerResolutionCharacterization: "view_all", False, True, - marks=pytest.mark.xfail( - strict=True, - raises=pytest.RaisesExc(AssertionError, match="view_all detail denied"), - reason="LIT-3974 A: view_all permits redacted catalog detail", - ), ), ("view_all", True, False), ("restricted", False, False), @@ -8967,31 +8938,16 @@ class TestMCPServerResolutionCharacterization: "db_runtime", "denied", False, - marks=pytest.mark.xfail( - strict=True, - raises=pytest.RaisesExc(AssertionError, match="credential metadata visibility"), - reason="LIT-3974 C: revoked grants hide DB metadata without removing credentials", - ), ), pytest.param( "config", "allowed", True, - marks=pytest.mark.xfail( - strict=True, - raises=pytest.RaisesExc(AssertionError, match="credential metadata visibility"), - reason="LIT-3974 C: authorized config credential metadata", - ), ), pytest.param( "config", "admin", True, - marks=pytest.mark.xfail( - strict=True, - raises=pytest.RaisesExc(AssertionError, match="credential metadata visibility"), - reason="LIT-3974 C: admin config credential metadata", - ), ), ("config", "denied", False), ("missing", "allowed", False), @@ -10295,68 +10251,28 @@ class TestMCPServerResolutionCharacterization: "db_runtime", "org object_permission", id="db-runtime-org-object-permission", - marks=pytest.mark.xfail( - strict=True, - raises=pytest.RaisesExc( - HTTPException, - check=lambda error: error.status_code == 403 and "permission" in str(error.detail), - ), - reason="LIT-3974 change A: detail authorization includes org object_permission grants", - ), ), pytest.param("config", "org object_permission", id="config-org-object-permission"), pytest.param( "db_runtime", "direct user object_permission", id="db-runtime-direct-user-permission", - marks=pytest.mark.xfail( - strict=True, - raises=pytest.RaisesExc( - HTTPException, - check=lambda error: error.status_code == 403 and "permission" in str(error.detail), - ), - reason="LIT-3974 change A: detail authorization includes direct user object_permission grants", - ), ), pytest.param( "config", "direct user object_permission", id="config-direct-user-permission", - marks=pytest.mark.xfail( - strict=True, - raises=pytest.RaisesExc( - HTTPException, - check=lambda error: error.status_code == 403 and "permission" in str(error.detail), - ), - reason="LIT-3974 change A: detail authorization includes direct user object_permission grants", - ), ), pytest.param( "db_runtime", "allow_all_keys", id="db-runtime-allow-all-keys", - marks=pytest.mark.xfail( - strict=True, - raises=pytest.RaisesExc( - HTTPException, - check=lambda error: error.status_code == 403 and "permission" in str(error.detail), - ), - reason="LIT-3974 change A: detail authorization includes allow_all_keys grants", - ), ), pytest.param("config", "allow_all_keys", id="config-allow-all-keys"), pytest.param( "db_runtime", "access-group", id="db-runtime-access-group", - marks=pytest.mark.xfail( - strict=True, - raises=pytest.RaisesExc( - HTTPException, - check=lambda error: error.status_code == 403 and "permission" in str(error.detail), - ), - reason="LIT-3974 change A: detail authorization includes access-group grants", - ), ), pytest.param("config", "access-group", id="config-access-group"), ],