mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
fix(mcp): stop 'Team doesn't exist' warnings for UI dashboard sessions (#32348)
UI session tokens carry the virtual team_id litellm-dashboard (UI_TEAM_ID), which is never persisted. The MCP team-permission helpers passed it to get_team_object anyway, so every dashboard MCP listing raised a 404 per lookup that was swallowed into per-server 'Failed to get allowed tools for server' warnings (plus the sibling 'allowed MCP servers for team' and 'MCP access groups for team' warnings) and wasted DB queries. The 404 also escaped past the key-level permission handling in get_allowed_tools_for_server, dropping key tool restrictions for such sessions. Short-circuit the virtual team before the DB lookup in the three helpers, mirroring the existing UI_TEAM_ID handling in agent_permission_handler. Also reject /team/new with the reserved team_id, since a real row would bind its budget and permissions to every UI session
This commit is contained in:
parent
d0c82c308d
commit
7ce573e6e8
4 changed files with 158 additions and 0 deletions
|
|
@ -8,6 +8,7 @@ from starlette.types import Scope
|
|||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.proxy._types import (
|
||||
UI_TEAM_ID,
|
||||
LiteLLM_TeamTable,
|
||||
ProxyException,
|
||||
SpecialHeaders,
|
||||
|
|
@ -738,6 +739,9 @@ class MCPRequestHandler:
|
|||
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,
|
||||
|
|
@ -1033,6 +1037,9 @@ class MCPRequestHandler:
|
|||
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,
|
||||
|
|
@ -1515,6 +1522,9 @@ class MCPRequestHandler:
|
|||
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,
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ from litellm._uuid import uuid
|
|||
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"
|
||||
|
|
|
|||
|
|
@ -3253,6 +3253,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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -9495,3 +9495,43 @@ class TestEmitTeamMembersMetric:
|
|||
# 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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue