diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 387843ee5b23..0ea941a1952b 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -8,6 +8,7 @@ from litellm._logging import verbose_logger from litellm.proxy._types import ( + UI_TEAM_ID, LiteLLM_TeamTable, ProxyException, SpecialHeaders, @@ -726,6 +727,9 @@ async def _get_team_object_permission( if not user_api_key_auth or not user_api_key_auth.team_id or not prisma_client: return None + if user_api_key_auth.team_id == UI_TEAM_ID: + return None + # Get the team object (which has object_permission already loaded) team_obj: Optional[LiteLLM_TeamTable] = await get_team_object( team_id=user_api_key_auth.team_id, @@ -1021,6 +1025,9 @@ async def _get_allowed_mcp_servers_for_team( if user_api_key_auth is None or not user_api_key_auth.team_id or prisma_client is None: return [] + if user_api_key_auth.team_id == UI_TEAM_ID: + return [] + team_obj: Optional[LiteLLM_TeamTable] = await get_team_object( team_id=user_api_key_auth.team_id, prisma_client=prisma_client, @@ -1503,6 +1510,9 @@ async def _get_mcp_access_groups_for_team( verbose_logger.debug("prisma_client is None") return [] + if user_api_key_auth.team_id == UI_TEAM_ID: + return [] + try: team_obj: Optional[LiteLLM_TeamTable] = await get_team_object( team_id=user_api_key_auth.team_id, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 5c41c60fcb1d..8ec7ec707a23 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -26,6 +26,7 @@ from litellm.integrations.prometheus import PrometheusLogger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( + UI_TEAM_ID, BlockTeamRequest, CommonProxyErrors, DeleteTeamRequest, @@ -1046,6 +1047,13 @@ async def new_team( if data.team_id is None: data.team_id = str(uuid.uuid4()) else: + if data.team_id == UI_TEAM_ID: + raise HTTPException( + status_code=400, + detail={ + "error": f"team_id '{UI_TEAM_ID}' is reserved for LiteLLM UI dashboard sessions and cannot be used for a real team. Please use a different team id." + }, + ) # Check if team_id exists already _existing_team_id = await prisma_client.get_data( team_id=data.team_id, table_name="team", query_type="find_unique" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 3607d448aad9..78e6acf794ec 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -3105,6 +3105,106 @@ async def test_get_team_object_permission_with_core_auth_auto_loading(): mock_get_team.assert_called_once() +@pytest.mark.asyncio +async def test_get_team_object_permission_ui_session_team_skips_db_lookup(): + """ + UI session tokens carry the virtual team_id "litellm-dashboard" (UI_TEAM_ID), + which is never persisted. The lookup must short-circuit to None without + calling get_team_object; otherwise every MCP tools listing from the + dashboard logs a "Team doesn't exist in db" warning per server. + """ + from litellm.proxy._types import UI_TEAM_ID + + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id=UI_TEAM_ID, + ) + + mock_prisma = MagicMock() + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + with patch("litellm.proxy.auth.auth_checks.get_team_object") as mock_get_team: + result = await MCPRequestHandler._get_team_object_permission( + mock_user_auth + ) + + assert result is None + mock_get_team.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "helper_name,expected", + [ + ("_get_allowed_mcp_servers_for_team", []), + ("_get_mcp_access_groups_for_team", []), + ], +) +async def test_team_mcp_helpers_ui_session_team_skip_db_lookup(helper_name, expected): + """ + The server-permission and access-group helpers hit get_team_object with the + session's team_id too; for the virtual UI team each used to 404 into its + own swallowed warning per MCP listing. They must short-circuit without a + DB lookup. + """ + from litellm.proxy._types import UI_TEAM_ID + + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id=UI_TEAM_ID, + ) + + mock_prisma = MagicMock() + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + with patch("litellm.proxy.auth.auth_checks.get_team_object") as mock_get_team: + helper = getattr(MCPRequestHandler, helper_name) + result = await helper(mock_user_auth) + + assert result == expected + mock_get_team.assert_not_called() + + +@pytest.mark.asyncio +async def test_get_allowed_tools_for_server_ui_session_team_keeps_key_restrictions(): + """ + Regression: the 404 raised by get_team_object for the virtual UI team used + to escape into get_allowed_tools_for_server's blanket except, dropping + key-level tool restrictions (fail-open) and logging a warning. With the + short-circuit, key restrictions still apply for UI sessions. + """ + from fastapi import HTTPException + + from litellm.proxy._types import UI_TEAM_ID + + user_api_key_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id=UI_TEAM_ID, + ) + key_perm = MagicMock() + key_perm.mcp_tool_permissions = {"server_1": ["tool_a"]} + + mock_prisma = MagicMock() + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + with patch( + "litellm.proxy.auth.auth_checks.get_team_object", + side_effect=HTTPException( + status_code=404, + detail={"error": "Team doesn't exist in db. Team=litellm-dashboard."}, + ), + ): + with patch.object( + MCPRequestHandler, "_get_key_object_permission", return_value=key_perm + ): + result = await MCPRequestHandler.get_allowed_tools_for_server( + server_id="server_1", + user_api_key_auth=user_api_key_auth, + ) + + assert result == ["tool_a"] + + @pytest.mark.asyncio async def test_get_allowed_mcp_servers_for_team_uses_helper(): """ 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 5f3974b46fbf..180fb1d3d8f8 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -9495,3 +9495,43 @@ def test_metric_failure_does_not_break_request(self, restore_callbacks): # A metric failure must be swallowed, not propagated to the handler. _emit_team_members_metric(self._team(1)) fake_logger.set_team_members_metric.assert_called_once() + + +@pytest.mark.asyncio +async def test_new_team_rejects_reserved_ui_session_team_id(): + """ + /team/new must reject team_id "litellm-dashboard" (UI_TEAM_ID): it is the + virtual team stamped on every UI dashboard session token, so a real DB row + with that id would bind its budget and permissions to every UI session. + """ + from fastapi import Request + + from litellm.proxy._types import UI_TEAM_ID, NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + team_request = NewTeamRequest( + team_alias="dashboard-clone", + team_id=UI_TEAM_ID, + ) + dummy_request = MagicMock(spec=Request) + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server._license_check") as mock_license, + ): + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_license.is_team_count_over_limit.return_value = False + mock_prisma.get_data = AsyncMock(return_value=None) + + with pytest.raises(ProxyException) as exc_info: + await new_team( + data=team_request, + http_request=dummy_request, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN + ), + ) + + assert exc_info.value.code == "400" + assert "reserved" in str(exc_info.value.message) + mock_prisma.get_data.assert_not_called()