From 57ca48a863b8dfb2bebef996fd723ea4191e3c79 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 3 Jul 2026 13:59:28 -0700 Subject: [PATCH 1/9] feat(mcp): add all-proxy-mcpservers sentinel to grant teams every MCP server (#32012) * feat(mcp): add all-proxy-mcpservers sentinel to grant every MCP server Teams can now be scoped to the all-proxy-mcpservers sentinel so they gain access to every MCP server on the proxy without listing each id. The sentinel expands to the live registry at request time, so a server added later is picked up with no change to the team's stored permission. The team ceiling that validates a key's MCP scope expands the sentinel too, so a key can be scoped to any server (including one registered after the team) and still pass subset validation Expose the option in the team create and edit forms via a new exclusive "All Proxy MCP Servers" choice in MCPServerSelector, mirroring the existing "No MCP Servers" sentinel * Update litellm/proxy/management_helpers/object_permission_utils.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix(mcp): honor all-proxy-mcpservers only on the team path, never per-key The sentinel was expanded inside the shared expand_permission_list, which also feeds the key, org, end_user and agent resolvers. A key whose stored object_permission ever held all-proxy-mcpservers (a stale write, a configured default, or a bug) would silently resolve to every MCP server at runtime, and a teamless key had nothing to cap it, so all servers got injected. Only write-time validation stripping the value stood between that value and a full grant Move the expansion out of expand_permission_list and into _get_allowed_mcp_servers_for_team so the sentinel is honored only where it is settable (a team). Anywhere else it now passes through as an inert literal that matches no registered server and is denied downstream. Reserved-id protection already blocks a real server from taking that id * fix(mcp): require proxy admin to grant a team the all-proxy MCP sentinel Granting a team every MCP server on the proxy is a proxy-wide authorization decision, but team create/update let any caller who can manage a team set object_permission.mcp_servers, with no ceiling check. Org admins reach /team/update by default (org_admin_allowed_routes) and _verify_team_access also admits team admins, so a non-proxy-admin could set all-proxy-mcpservers and self-grant their team access to every MCP server on the proxy, including servers never assigned to that team Gate the grant in new_team and update_team: a non-proxy-admin cannot add the all-proxy-mcpservers sentinel. The check is scoped to newly adding it, so a team a proxy admin already scoped to all-proxy can still be edited by a team admin without being forced to strip the sentinel. The UI only offers the "All Proxy MCP Servers" option to proxy admins in the team create and edit forms * fix(ui): render friendly all-proxy MCP label for non-admins editing an all-proxy team A team scoped to the all-proxy-mcpservers sentinel could be opened in the team edit form by a team admin or org admin (canEditTeam admits them), but the "All Proxy MCP Servers" option in MCPServerSelector was rendered only behind the proxy-admin-gated allowAllProxyMcpServers flag. For a non-proxy-admin the stored sentinel was hydrated into the selected value with no matching Select.Option, so antd showed the raw all-proxy-mcpservers literal as a chip, and adding another server could persist a mixed [all-proxy-mcpservers, ] value. Render the option whenever the sentinel is present in the value, not only when the caller may grant it, and drive the real-option disabling off presence too so the selection stays exclusive. A non-proxy-admin now sees the friendly label read-only and cannot build a mixed state; only a proxy admin can newly add it, which the backend already enforces. Adds regression tests: the selector shows the friendly option (not the raw literal) when the sentinel is stored but the grant flag is off, plus exclusive emit and disabled-real-options coverage, and MCPServerPermissions renders the green "All" state instead of the raw sentinel string. * fix(ui): drop redundant "All servers" hint from the all-proxy MCP chip antd renders a Select option's children inside the selected tag, so the all-proxy option showed both "All Proxy MCP Servers" and the green "All servers" type-hint in the chip, which say the same thing. Collapse the option to a single green "All Proxy MCP Servers" label so the dropdown row and the chip read cleanly without the duplication. * fix(ui): color the all-proxy MCP label blue to match server chips Use the same blue (#1890ff) as regular MCP server entries for the "All Proxy MCP Servers" option/chip instead of green. * fix(ui): make the all-proxy MCP permissions display blue, not green Match the blue used by the selector chip and regular server entries so the "All Proxy MCP Servers" badge and row in MCPServerPermissions are consistent across the team/key/org detail views. The red "Blocked" state for no-mcp-servers is unchanged. --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../mcp_server/auth/user_api_key_auth_mcp.py | 4 + .../management_endpoints/team_endpoints.py | 13 + .../object_permission_utils.py | 60 ++++- .../auth/test_user_api_key_auth_mcp.py | 239 ++++++++++++++++++ .../test_object_permission_utils.py | 174 ++++++++++++- .../src/components/OldTeams.tsx | 1 + .../MCPServerSelector.test.tsx | 85 ++++++- .../MCPServerSelector.tsx | 26 +- .../src/components/mcp_tools/constants.ts | 2 + .../permissions/MCPServerPermissions.test.tsx | 18 ++ .../permissions/MCPServerPermissions.tsx | 12 +- .../src/components/team/TeamInfo.tsx | 1 + 12 files changed, 621 insertions(+), 14 deletions(-) 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 2520c7e82a1..387843ee5b2 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 @@ -11,6 +11,7 @@ from litellm.proxy._types import ( LiteLLM_TeamTable, ProxyException, SpecialHeaders, + SpecialMCPServerName, SpecialMCPServerNames, UserAPIKeyAuth, ) @@ -1041,6 +1042,9 @@ class MCPRequestHandler: if object_permissions is None: return list(set(team_access_group_servers)) + if SpecialMCPServerName.all_proxy_servers.value in (object_permissions.mcp_servers or []): + return list(global_mcp_server_manager.get_registry().keys()) + direct_mcp_servers = global_mcp_server_manager.expand_permission_list(object_permissions.mcp_servers or []) legacy_access_group_servers = await MCPRequestHandler._get_mcp_servers_from_access_groups( diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 0ea2b9e05f9..5c41c60fcb1 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -93,6 +93,7 @@ from litellm.proxy.management_endpoints.tag_management_endpoints import ( ) from litellm.proxy.management_helpers.object_permission_utils import ( _set_object_permission, + enforce_all_proxy_mcp_servers_grant_is_admin_only, handle_update_object_permission_common, ) from litellm.proxy.management_helpers.team_member_permission_checks import ( @@ -1144,6 +1145,12 @@ async def new_team( data_json = data.json() ## Handle Object Permission - MCP, Vector Stores etc. + await enforce_all_proxy_mcp_servers_grant_is_admin_only( + requested_mcp_servers=(data.object_permission.mcp_servers if data.object_permission is not None else None), + existing_object_permission_id=None, + is_proxy_admin=user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN, + prisma_client=prisma_client, + ) data_json = await _set_object_permission( data_json=data_json, prisma_client=prisma_client, @@ -1846,6 +1853,12 @@ async def update_team( # Check object permission if data.object_permission is not None: + await enforce_all_proxy_mcp_servers_grant_is_admin_only( + requested_mcp_servers=data.object_permission.mcp_servers, + existing_object_permission_id=existing_team_row.object_permission_id, + is_proxy_admin=user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN, + prisma_client=prisma_client, + ) updated_kv = await handle_update_object_permission( data_json=updated_kv, existing_team_row=existing_team_row, diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index fe96d9c260a..6bbd41b93ed 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -11,7 +11,7 @@ from fastapi import HTTPException, status from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.proxy._types import ObjectPermissionDict, SpecialMCPServerNames +from litellm.proxy._types import ObjectPermissionDict, SpecialMCPServerName, SpecialMCPServerNames from litellm.proxy.utils import PrismaClient from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.table_repositories import MCPServerRepository @@ -334,6 +334,8 @@ async def _resolve_team_allowed_mcp_servers( ) direct_servers: List[str] = team_object_permission.mcp_servers or [] + if SpecialMCPServerName.all_proxy_servers.value in direct_servers: + return _get_all_mcp_server_ids() access_group_servers: List[str] = await MCPRequestHandler._get_mcp_servers_from_access_groups( team_object_permission.mcp_access_groups or [] ) @@ -359,6 +361,62 @@ def _get_allow_all_keys_server_ids() -> Set[str]: return set(global_mcp_server_manager.get_allow_all_keys_server_ids()) +def _get_all_mcp_server_ids() -> set[str]: + """Return every MCP server id registered on the proxy (config + DB union).""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + return set(global_mcp_server_manager.get_registry().keys()) + + +async def _existing_object_permission_mcp_servers( + object_permission_id: Optional[str], + prisma_client: Optional[PrismaClient], +) -> list[str]: + if not object_permission_id or prisma_client is None: + return [] + existing = await ObjectPermissionRepository(prisma_client).table.find_unique( + where={"object_permission_id": object_permission_id}, + ) + if existing is None: + return [] + return existing.mcp_servers or [] + + +async def enforce_all_proxy_mcp_servers_grant_is_admin_only( + requested_mcp_servers: Optional[list[str]], + existing_object_permission_id: Optional[str], + is_proxy_admin: bool, + prisma_client: Optional[PrismaClient], +) -> None: + """ + Only a proxy admin may newly grant the all-proxy MCP sentinel. + + Scoping a team to every MCP server on the proxy is a proxy-wide authorization + decision, so a caller who is not a proxy admin (e.g. a team admin managing their + own team) cannot add ``all-proxy-mcpservers``. A sentinel a proxy admin already + granted is left untouched, so unrelated edits to such a team still succeed. + + Raises HTTPException(403) when a non-admin tries to add the sentinel. + """ + sentinel = SpecialMCPServerName.all_proxy_servers.value + if is_proxy_admin or sentinel not in (requested_mcp_servers or []): + return + existing_mcp_servers = await _existing_object_permission_mcp_servers( + object_permission_id=existing_object_permission_id, + prisma_client=prisma_client, + ) + if sentinel in existing_mcp_servers: + return + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": "Only a proxy admin can grant a team access to all proxy MCP servers ('all-proxy-mcpservers')." + }, + ) + + async def _get_team_allowed_mcp_servers( team_obj: Optional["LiteLLM_TeamTableCachedObj"], prisma_client: Optional[PrismaClient] = None, 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 b3b0e8adcf6..3607d448aad 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 @@ -4495,3 +4495,242 @@ async def test_get_allowed_mcp_servers_surfaces_ungated_key_access_group_grant_e assert result == ["srv-deepwiki"] finally: _stop_patches(patches) + + +def test_expand_permission_list_does_not_honor_all_proxy_sentinel(): + """The all-proxy sentinel is a team-only grant. The shared expand_permission_list + also feeds the key/org/end_user/agent resolvers, so it must NOT expand the + sentinel to the full registry; it passes through as an inert literal (denied + downstream). Concrete ids still resolve normally. If the sentinel were expanded + here, any stored key/org/end_user permission holding it would silently gain every + server.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import SpecialMCPServerName + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + sentinel = SpecialMCPServerName.all_proxy_servers.value + for sid in ("srv-x", "srv-y"): + global_mcp_server_manager.registry[sid] = MCPServer( + server_id=sid, + name=sid, + server_name=sid, + url=f"https://{sid}.example.com", + transport=MCPTransport.http, + ) + try: + result = global_mcp_server_manager.expand_permission_list([sentinel]) + assert set(result).isdisjoint({"srv-x", "srv-y"}) + assert result == [sentinel] + assert global_mcp_server_manager.expand_permission_list(["srv-x"]) == ["srv-x"] + finally: + for sid in ("srv-x", "srv-y"): + global_mcp_server_manager.registry.pop(sid, None) + + +@pytest.mark.asyncio +async def test_get_allowed_mcp_servers_for_team_expands_all_proxy_sentinel_dynamically(): + """The TEAM resolver expands the all-proxy sentinel to every registered server and + picks up a server registered later, so a team scoped to all-proxy tracks the live + registry without any change to its stored permission. Reverting the team-side + expansion collapses this to the inert literal and the result no longer contains the + real servers.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + LiteLLM_TeamTable, + SpecialMCPServerName, + ) + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + for sid in ("srv-x", "srv-y"): + global_mcp_server_manager.registry[sid] = MCPServer( + server_id=sid, + name=sid, + server_name=sid, + url=f"https://{sid}.example.com", + transport=MCPTransport.http, + ) + try: + team_perm = LiteLLM_ObjectPermissionTable( + object_permission_id="team-perm", + mcp_servers=[SpecialMCPServerName.all_proxy_servers.value], + mcp_access_groups=[], + vector_stores=[], + ) + team_obj = LiteLLM_TeamTable( + team_id="team-1", + access_group_ids=[], + object_permission_id="team-perm", + ) + team_obj.object_permission = team_perm + auth = UserAPIKeyAuth(token="test-token", api_key="sk-test", team_id="team-1") + + patches = _patch_proxy_server_globals_for_mcp() + [ + patch( + "litellm.proxy.auth.auth_checks.get_team_object", + new_callable=AsyncMock, + return_value=team_obj, + ), + patch( + "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", + new_callable=AsyncMock, + return_value=[], + ), + ] + _start_patches(patches) + try: + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth) + assert set(result) == {"srv-x", "srv-y"} + + global_mcp_server_manager.registry["srv-z"] = MCPServer( + server_id="srv-z", + name="srv-z", + server_name="srv-z", + url="https://srv-z.example.com", + transport=MCPTransport.http, + ) + result_after = await MCPRequestHandler._get_allowed_mcp_servers_for_team(auth) + assert "srv-z" in result_after + finally: + _stop_patches(patches) + finally: + for sid in ("srv-x", "srv-y", "srv-z"): + global_mcp_server_manager.registry.pop(sid, None) + + +@pytest.mark.asyncio +async def test_key_with_all_proxy_sentinel_does_not_grant_all_servers(): + """Security regression: the all-proxy sentinel is a team-only grant. A KEY whose + stored object_permission holds the sentinel (via a stale write, a configured + default, or a bug) must NOT be silently widened to every server at runtime. A + teamless key with the sentinel resolves to no real server — never srv-secret or the + full registry. On the pre-hardening code the key path expanded the sentinel and + this key would reach srv-secret.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + SpecialMCPServerName, + ) + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + for sid in ("srv-x", "srv-y", "srv-secret"): + global_mcp_server_manager.registry[sid] = MCPServer( + server_id=sid, + name=sid, + server_name=sid, + url=f"https://{sid}.example.com", + transport=MCPTransport.http, + ) + try: + key_perm = LiteLLM_ObjectPermissionTable( + object_permission_id="key-perm", + mcp_servers=[SpecialMCPServerName.all_proxy_servers.value], + mcp_access_groups=[], + vector_stores=[], + ) + auth = UserAPIKeyAuth(token="test-token", api_key="sk-test", object_permission=key_perm) + + patches = _patch_proxy_server_globals_for_mcp() + _start_patches(patches) + try: + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + finally: + _stop_patches(patches) + + assert "srv-secret" not in result + assert set(result).isdisjoint(global_mcp_server_manager.get_registry().keys()) + finally: + for sid in ("srv-x", "srv-y", "srv-secret"): + global_mcp_server_manager.registry.pop(sid, None) + + +@pytest.mark.asyncio +async def test_get_allowed_mcp_servers_team_all_proxy_key_scoped_to_one_end_to_end(): + """End-to-end: a team scoped to the all-proxy sentinel is a ceiling of every + registered server, so a key scoped to a single server (srv-x) resolves to + exactly that server (key ∩ all-servers == key). If the sentinel branch is + reverted the team ceiling collapses to the literal marker, the intersection + empties, and the result is [] instead of ["srv-x"].""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + LiteLLM_TeamTable, + SpecialMCPServerName, + ) + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + for sid in ("srv-x", "srv-y"): + global_mcp_server_manager.registry[sid] = MCPServer( + server_id=sid, + name=sid, + server_name=sid, + url=f"https://{sid}.example.com", + transport=MCPTransport.http, + ) + try: + key_perm = LiteLLM_ObjectPermissionTable( + object_permission_id="key-perm", + mcp_servers=["srv-x"], + mcp_access_groups=[], + vector_stores=[], + ) + team_perm = LiteLLM_ObjectPermissionTable( + object_permission_id="team-perm", + mcp_servers=[SpecialMCPServerName.all_proxy_servers.value], + mcp_access_groups=[], + vector_stores=[], + ) + team_obj = LiteLLM_TeamTable( + team_id="team-1", + access_group_ids=[], + object_permission_id="team-perm", + ) + team_obj.object_permission = team_perm + + auth = UserAPIKeyAuth( + token="test-token", + api_key="sk-test", + team_id="team-1", + object_permission=key_perm, + ) + + patches = _patch_proxy_server_globals_for_mcp() + [ + patch( + "litellm.proxy.auth.auth_checks.get_team_object", + new_callable=AsyncMock, + return_value=team_obj, + ), + patch( + "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", + new_callable=AsyncMock, + return_value=[], + ), + patch.object( + MCPRequestHandler, + "_get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], + ), + ] + _start_patches(patches) + try: + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + finally: + _stop_patches(patches) + + assert result == ["srv-x"] + finally: + for sid in ("srv-x", "srv-y"): + global_mcp_server_manager.registry.pop(sid, None) diff --git a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py index 0981c4239ee..d797a27aa67 100644 --- a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py @@ -9,13 +9,19 @@ sys.path.insert(0, os.path.abspath("../../../..")) from unittest.mock import AsyncMock, MagicMock, patch -from litellm.proxy._types import LiteLLM_ObjectPermissionBase, LiteLLM_ObjectPermissionTable, ObjectPermissionDict +from litellm.proxy._types import ( + LiteLLM_ObjectPermissionBase, + LiteLLM_ObjectPermissionTable, + ObjectPermissionDict, + SpecialMCPServerName, +) from litellm.proxy.management_helpers.object_permission_utils import ( _extract_requested_mcp_access_groups, _extract_requested_mcp_server_ids, _resolve_team_allowed_mcp_servers, _rewrite_object_permission_mcp_servers, _set_object_permission, + enforce_all_proxy_mcp_servers_grant_is_admin_only, validate_key_mcp_servers_against_team, validate_key_search_tools_against_team, validate_key_vector_stores_against_team, @@ -876,6 +882,172 @@ async def test_resolve_team_allowed_mcp_servers_dict_tool_permissions( assert result == {"server-a"} +# ---- Tests for the all-proxy-mcpservers sentinel (team scoped to every server) ---- + + +@pytest.mark.asyncio +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_resolve_team_all_proxy_sentinel_resolves_dynamically(mock_access_groups): + """A team whose object_permission.mcp_servers holds the all-proxy sentinel + resolves to every registered server id, and picks up a server registered + later without any change to the team's stored permission (this kills the + early-return that maps the sentinel to the live registry).""" + registry = { + "srv-x": _make_mock_mcp_server("srv-x"), + "srv-y": _make_mock_mcp_server("srv-y"), + } + mock_mgr = MagicMock() + mock_mgr.get_registry.return_value = registry + + team_perm = MagicMock(spec=LiteLLM_ObjectPermissionTable) + team_perm.mcp_servers = [SpecialMCPServerName.all_proxy_servers.value] + team_perm.mcp_access_groups = [] + team_perm.mcp_tool_permissions = {} + + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_mgr, + ): + assert await _resolve_team_allowed_mcp_servers(team_perm) == {"srv-x", "srv-y"} + + registry["srv-z"] = _make_mock_mcp_server("srv-z") + assert await _resolve_team_allowed_mcp_servers(team_perm) == { + "srv-x", + "srv-y", + "srv-z", + } + + +@pytest.mark.asyncio +@patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + new=_make_mock_mcp_manager("srv-x", "srv-y", "srv-z"), +) +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_key_scoped_to_server_added_after_team_all_proxy( + mock_access_groups, mock_allow_all +): + """The exact user scenario: a team scoped to the all-proxy sentinel, a server + (srv-z) registered afterwards, and a key scoped to just srv-z. Because the + team ceiling resolves to every registered server, the key passes validation + and keeps srv-z in its normalized permission.""" + team_obj = _make_team_obj(mcp_servers=[SpecialMCPServerName.all_proxy_servers.value]) + object_permission = {"mcp_servers": ["srv-z"]} + result = await validate_key_mcp_servers_against_team( + object_permission=object_permission, + team_obj=team_obj, + ) + assert result is not None + assert result["mcp_servers"] == ["srv-z"] + + +@pytest.mark.asyncio +@patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + new=_make_mock_mcp_manager("srv-x", "srv-z"), +) +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_key_scoped_to_server_rejected_when_team_not_all_proxy( + mock_access_groups, mock_allow_all +): + """Contrast with the sentinel case: a team scoped to a concrete server list + (srv-x, not the sentinel) does NOT unlock srv-z for a key. It is the sentinel + specifically, not a blanket allow, that widens the team ceiling.""" + team_obj = _make_team_obj(mcp_servers=["srv-x"]) + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["srv-z"]}, + team_obj=team_obj, + ) + assert exc_info.value.status_code == 403 + assert "srv-z" in str(exc_info.value.detail) + + +# ---- Tests for the proxy-admin gate on granting a team the all-proxy sentinel ---- + + +@pytest.mark.asyncio +async def test_enforce_all_proxy_mcp_grant_blocks_non_admin_adding_sentinel(): + """A non-proxy-admin (e.g. a team admin) cannot newly grant a team the all-proxy + MCP sentinel. Without this gate a team admin could self-escalate their team to + every MCP server on the proxy via team create/update.""" + with pytest.raises(HTTPException) as exc_info: + await enforce_all_proxy_mcp_servers_grant_is_admin_only( + requested_mcp_servers=[SpecialMCPServerName.all_proxy_servers.value], + existing_object_permission_id=None, + is_proxy_admin=False, + prisma_client=None, + ) + assert exc_info.value.status_code == 403 + assert "all-proxy-mcpservers" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_enforce_all_proxy_mcp_grant_allows_proxy_admin(): + """A proxy admin may grant the sentinel — the intended way to scope a team to all + proxy MCP servers.""" + await enforce_all_proxy_mcp_servers_grant_is_admin_only( + requested_mcp_servers=[SpecialMCPServerName.all_proxy_servers.value], + existing_object_permission_id=None, + is_proxy_admin=True, + prisma_client=None, + ) + + +@pytest.mark.asyncio +async def test_enforce_all_proxy_mcp_grant_allows_non_admin_without_sentinel(): + """A non-admin scoping a team to concrete servers is unaffected by the gate.""" + await enforce_all_proxy_mcp_servers_grant_is_admin_only( + requested_mcp_servers=["srv-x", "srv-y"], + existing_object_permission_id=None, + is_proxy_admin=False, + prisma_client=None, + ) + + +@pytest.mark.asyncio +async def test_enforce_all_proxy_mcp_grant_allows_non_admin_when_sentinel_already_set(): + """The gate blocks only NEW grants: a non-admin editing a team a proxy admin + already scoped to all-proxy is not forced to strip the sentinel, so unrelated + edits still succeed. The existing permission is read from the DB by id.""" + existing_row = MagicMock() + existing_row.mcp_servers = [SpecialMCPServerName.all_proxy_servers.value] + mock_repo = MagicMock() + mock_repo.table.find_unique = AsyncMock(return_value=existing_row) + + with patch( + "litellm.proxy.management_helpers.object_permission_utils.ObjectPermissionRepository", + return_value=mock_repo, + ): + await enforce_all_proxy_mcp_servers_grant_is_admin_only( + requested_mcp_servers=[SpecialMCPServerName.all_proxy_servers.value], + existing_object_permission_id="op-1", + is_proxy_admin=False, + prisma_client=MagicMock(), + ) + mock_repo.table.find_unique.assert_awaited_once() + + # ---- Tests for validate_key_search_tools_against_team ---- diff --git a/ui/litellm-dashboard/src/components/OldTeams.tsx b/ui/litellm-dashboard/src/components/OldTeams.tsx index 9c2106fa2fb..0ecbf39541b 100644 --- a/ui/litellm-dashboard/src/components/OldTeams.tsx +++ b/ui/litellm-dashboard/src/components/OldTeams.tsx @@ -1479,6 +1479,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser value={form.getFieldValue("allowed_mcp_servers_and_groups")} accessToken={accessToken || ""} placeholder="Select MCP servers or access groups (optional)" + allowAllProxyMcpServers={isProxyAdminRole(userRole || "")} /> diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.test.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.test.tsx index 1517b2dfa27..e6e778e49e5 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.test.tsx @@ -3,7 +3,7 @@ import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders } from "../../../tests/test-utils"; import MCPServerSelector from "./MCPServerSelector"; -import { NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants"; +import { ALL_PROXY_MCP_SERVERS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants"; vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ useMCPServers: vi.fn(), @@ -45,15 +45,19 @@ const mockUseMCPServers = vi.mocked(useMCPServers); const mockUseMCPAccessGroups = vi.mocked(useMCPAccessGroups); const mockUseMCPToolsets = vi.mocked(useMCPToolsets); +const setupMcpMocks = () => { + mockUseMCPServers.mockReturnValue({ + data: [{ server_id: "srv-1", server_name: "Server One" }], + isLoading: false, + } as any); + mockUseMCPAccessGroups.mockReturnValue({ data: [], isLoading: false } as any); + mockUseMCPToolsets.mockReturnValue({ data: [], isLoading: false } as any); +}; + describe("MCPServerSelector no-mcp-servers option", () => { beforeEach(() => { vi.clearAllMocks(); - mockUseMCPServers.mockReturnValue({ - data: [{ server_id: "srv-1", server_name: "Server One" }], - isLoading: false, - } as any); - mockUseMCPAccessGroups.mockReturnValue({ data: [], isLoading: false } as any); - mockUseMCPToolsets.mockReturnValue({ data: [], isLoading: false } as any); + setupMcpMocks(); }); const optionByValue = (value: string) => @@ -98,3 +102,70 @@ describe("MCPServerSelector no-mcp-servers option", () => { expect(optionByValue(NO_MCP_SERVERS_SENTINEL)?.disabled).toBe(false); }); }); + +describe("MCPServerSelector all-proxy-mcpservers option", () => { + beforeEach(() => { + vi.clearAllMocks(); + setupMcpMocks(); + }); + + const optionByValue = (value: string) => + Array.from(screen.getByTestId("mcp-select").querySelectorAll("option")).find( + (o) => (o as HTMLOptionElement).value === value, + ) as HTMLOptionElement | undefined; + + it("hides the All Proxy MCP Servers option by default", () => { + renderWithProviders( + , + ); + expect(optionByValue(ALL_PROXY_MCP_SERVERS_SENTINEL)).toBeUndefined(); + }); + + it("emits an exclusive sentinel when All Proxy MCP Servers is selected", async () => { + const onChange = vi.fn(); + renderWithProviders( + , + ); + expect(optionByValue(ALL_PROXY_MCP_SERVERS_SENTINEL)).toBeDefined(); + + await userEvent.selectOptions(screen.getByTestId("mcp-select"), [ALL_PROXY_MCP_SERVERS_SENTINEL]); + + expect(onChange).toHaveBeenCalledWith({ + servers: [ALL_PROXY_MCP_SERVERS_SENTINEL], + accessGroups: [], + toolsets: [], + }); + }); + + it("disables real server options while the sentinel is selected", () => { + renderWithProviders( + , + ); + expect(optionByValue("srv-1")?.disabled).toBe(true); + expect(optionByValue(ALL_PROXY_MCP_SERVERS_SENTINEL)?.disabled).toBe(false); + }); + + it("renders the friendly option, not the raw literal, when the sentinel is already stored but the flag is off", () => { + renderWithProviders( + , + ); + const option = optionByValue(ALL_PROXY_MCP_SERVERS_SENTINEL); + expect(option).toBeDefined(); + expect(option?.textContent).toContain("All Proxy MCP Servers"); + expect(optionByValue("srv-1")?.disabled).toBe(true); + }); +}); diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx index bbda761938e..f31f60d2084 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx @@ -3,7 +3,7 @@ import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers" import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; import { Select } from "antd"; import React from "react"; -import { NO_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; +import { ALL_PROXY_MCP_SERVERS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; interface MCPServerSelectorProps { onChange: (selected: { servers: string[]; accessGroups: string[]; toolsets: string[] }) => void; @@ -18,6 +18,7 @@ interface MCPServerSelectorProps { disabled?: boolean; teamId?: string | null; allowNoMcpServers?: boolean; + allowAllProxyMcpServers?: boolean; } const TOOLSET_PREFIX = "toolset:"; @@ -31,6 +32,7 @@ const MCPServerSelector: React.FC = ({ disabled = false, teamId, allowNoMcpServers = false, + allowAllProxyMcpServers = false, }) => { const { data: mcpServers = [], isLoading: serversLoading } = useMCPServers(teamId); const { data: accessGroups = [], isLoading: groupsLoading } = useMCPAccessGroups(); @@ -81,9 +83,14 @@ const MCPServerSelector: React.FC = ({ ]; const hasNoMcpServersSelected = allowNoMcpServers && selectedValues.includes(NO_MCP_SERVERS_SENTINEL); + const hasAllProxyMcpServersSelected = selectedValues.includes(ALL_PROXY_MCP_SERVERS_SENTINEL); // Handle selection const handleChange = (selected: string[]) => { + if (allowAllProxyMcpServers && selected.includes(ALL_PROXY_MCP_SERVERS_SENTINEL)) { + onChange({ servers: [ALL_PROXY_MCP_SERVERS_SENTINEL], accessGroups: [], toolsets: [] }); + return; + } // "No MCP Servers" is exclusive: picking it clears everything else. if (allowNoMcpServers && selected.includes(NO_MCP_SERVERS_SENTINEL)) { onChange({ servers: [NO_MCP_SERVERS_SENTINEL], accessGroups: [], toolsets: [] }); @@ -113,10 +120,20 @@ const MCPServerSelector: React.FC = ({ disabled={disabled} filterOption={(input, option) => { if (option?.value === NO_MCP_SERVERS_SENTINEL) return true; + if (option?.value === ALL_PROXY_MCP_SERVERS_SENTINEL) return true; const searchText = options.find((opt) => opt.value === option?.value)?.searchText || ""; return searchText.toLowerCase().includes(input.toLowerCase()); }} > + {(allowAllProxyMcpServers || hasAllProxyMcpServersSelected) && ( + + All Proxy MCP Servers + + )} {allowNoMcpServers && (
@@ -126,7 +143,12 @@ const MCPServerSelector: React.FC = ({ )} {options.map((opt) => ( - +
{ // API should not be called without token expect(networking.fetchMCPServers).not.toHaveBeenCalled(); }); + + it("should display the All Proxy MCP Servers state instead of the raw sentinel string", async () => { + vi.mocked(networking.fetchMCPServers).mockResolvedValue([]); + + render( + , + ); + + expect(await screen.findByText("All Proxy MCP Servers")).toBeInTheDocument(); + expect(screen.getByText("All")).toBeInTheDocument(); + expect(screen.queryByText(ALL_PROXY_MCP_SERVERS_SENTINEL)).not.toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx index 9172ff92c69..4dfe41bb515 100644 --- a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx +++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx @@ -4,7 +4,7 @@ import { ServerIcon, ChevronDownIcon, ChevronRightIcon } from "@heroicons/react/ import { Tooltip } from "antd"; import { fetchMCPServers, fetchMCPToolsets } from "../networking"; import { MCPServer, MCPToolset } from "../mcp_tools/types"; -import { NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants"; +import { ALL_PROXY_MCP_SERVERS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants"; interface MCPServerPermissionsProps { mcpServers: string[]; @@ -96,11 +96,12 @@ export function MCPServerPermissions({ }; const blocksAllMcpServers = mcpServers.includes(NO_MCP_SERVERS_SENTINEL); + const grantsAllProxyMcpServers = mcpServers.includes(ALL_PROXY_MCP_SERVERS_SENTINEL); // Merge servers and access groups into one list const mergedItems = [ ...mcpServers - .filter((server) => server !== NO_MCP_SERVERS_SENTINEL) + .filter((server) => server !== NO_MCP_SERVERS_SENTINEL && server !== ALL_PROXY_MCP_SERVERS_SENTINEL) .map((server) => ({ type: "server", value: server })), ...mcpAccessGroups.map((group) => ({ type: "accessGroup", value: group })), ]; @@ -112,7 +113,7 @@ export function MCPServerPermissions({ MCP Servers - {blocksAllMcpServers ? "Blocked" : totalCount} + {blocksAllMcpServers ? "Blocked" : grantsAllProxyMcpServers ? "All" : totalCount}
@@ -123,6 +124,11 @@ export function MCPServerPermissions({ No MCP servers — this key is blocked from all MCP servers, including its team's servers
+ ) : grantsAllProxyMcpServers ? ( +
+ + All Proxy MCP Servers +
) : totalCount > 0 ? (
{mergedItems.map((item, index) => { diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index c5b934d5432..3b80c598d34 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -1356,6 +1356,7 @@ const TeamInfoView: React.FC = ({ value={form.getFieldValue("mcp_servers_and_groups")} accessToken={accessToken || ""} placeholder="Select MCP servers or access groups (optional)" + allowAllProxyMcpServers={is_proxy_admin} /> From fff6a5396c1e80d9be98dc1e7cea4da3425677cb Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 3 Jul 2026 14:00:43 -0700 Subject: [PATCH 2/9] fix(ci): stop ui_unit_tests vitest onTaskUpdate RPC timeout flake The ui_unit_tests job runs vitest with maxForks=8 on an 8-vCPU xlarge container, leaving no headroom for the main vitest process that services worker RPCs. Under full CPU saturation the coordinator misses the onTaskUpdate ack, vitest raises "Timeout calling onTaskUpdate" as an unhandled error, and the job exits 1 even though every test passes. Lower maxForks to 6 so the coordinator, jsdom, and OS keep two cores, and raise teardownTimeout to 60s for extra slack on heavy runs. --- .circleci/config.yml | 2 +- ui/litellm-dashboard/vitest.config.ts | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index f13e9bf66f1..9cbcf440296 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -2629,7 +2629,7 @@ jobs: cd ui/litellm-dashboard CI=true npm run test -- --run \ - --pool forks --poolOptions.forks.maxForks=8 + --pool forks --poolOptions.forks.maxForks=6 e2e_ui_testing: docker: diff --git a/ui/litellm-dashboard/vitest.config.ts b/ui/litellm-dashboard/vitest.config.ts index d2b6e7b43bf..3798f1d0c5a 100644 --- a/ui/litellm-dashboard/vitest.config.ts +++ b/ui/litellm-dashboard/vitest.config.ts @@ -8,6 +8,7 @@ export default defineConfig({ globals: true, css: true, // lets you import CSS/modules without extra mocks testTimeout: 30000, + teardownTimeout: 60000, coverage: { provider: "v8", reporter: ["text", "lcov"], From 2e1d8d2928e37d6e6e9fd3004804224f5e454bd8 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 3 Jul 2026 14:12:32 -0700 Subject: [PATCH 3/9] fix(anthropic): keep context_management working when drop_params is enabled (#32020) * fix(anthropic): keep context_management working when drop_params is enabled drop_params (proxy-wide or per-request) silently disabled the in-gateway context_management polyfill on the /v1/messages -> chat completions adapter path, even though context_management is a LiteLLM-supported param (native on Anthropic, polyfilled elsewhere). Gate the polyfill on an explicit additional_drop_params: ["context_management"] opt-out instead, which also makes that escape hatch actually work on the adapter path. * test(anthropic): cover sync adapter polyfill gate for global drop_params and additional_drop_params --- .../adapters/handler.py | 65 +++--- .../context_management/test_compact.py | 198 +++++++++++++++++- 2 files changed, 231 insertions(+), 32 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index 812e0f62c96..7299fc16897 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -78,7 +78,7 @@ async def _prepare_context_managed_request( system: Optional[Any], context_management_spec: Any, litellm_metadata: Optional[Dict], - drop_params: Optional[bool], + additional_drop_params: Optional[list[str]], llm_router: Any, user_api_key_auth: Any = None, ) -> Optional[PolyfillResult]: @@ -95,7 +95,7 @@ async def _prepare_context_managed_request( # silently drop intermediate turns. polyfill_will_run = _polyfill_will_run( context_management_spec=context_management_spec, - drop_params=drop_params, + additional_drop_params=additional_drop_params, ) if polyfill_will_run: @@ -117,7 +117,7 @@ async def _prepare_context_managed_request( system=working_system, context_management_spec=context_management_spec, litellm_metadata=litellm_metadata, - drop_params=drop_params, + additional_drop_params=additional_drop_params, llm_router=llm_router, user_api_key_auth=user_api_key_auth, ) @@ -143,18 +143,19 @@ async def _prepare_context_managed_request( def _polyfill_will_run( *, context_management_spec: Any, - drop_params: Optional[bool], + additional_drop_params: Optional[list[str]], ) -> bool: """Return True when ``compact_20260112`` will run via the polyfill dispatcher. - Mirrors the gating in ``_run_polyfill_if_enabled``: an empty spec or - effective ``drop_params`` short-circuits the polyfill. The pre-processing - skip only applies when the dispatcher will actually invoke - ``apply_compact_20260112`` (which has its own compaction-block slicing). + Mirrors the gating in ``_run_polyfill_if_enabled``: an empty spec or an + explicit ``context_management`` entry in ``additional_drop_params`` + short-circuits the polyfill. The pre-processing skip only applies when the + dispatcher will actually invoke ``apply_compact_20260112`` (which has its + own compaction-block slicing). """ edits = _normalize_spec_edits( context_management_spec=context_management_spec, - drop_params=drop_params, + additional_drop_params=additional_drop_params, ) if edits is None: return False @@ -169,7 +170,7 @@ def _polyfill_will_run( def _spec_has_non_compact_edits( *, context_management_spec: Any, - drop_params: Optional[bool], + additional_drop_params: Optional[list[str]], ) -> bool: """Return True when the spec includes edits other than ``compact_20260112``. @@ -180,7 +181,7 @@ def _spec_has_non_compact_edits( """ edits = _normalize_spec_edits( context_management_spec=context_management_spec, - drop_params=drop_params, + additional_drop_params=additional_drop_params, ) if edits is None: return False @@ -195,10 +196,22 @@ def _spec_has_non_compact_edits( ) +def _context_management_explicitly_dropped(additional_drop_params: Optional[list[str]]) -> bool: + """True when the caller opted out of context_management via ``additional_drop_params``. + + ``drop_params`` deliberately does NOT gate the polyfill: ``context_management`` + is a LiteLLM-supported param (native on Anthropic, polyfilled elsewhere), and + ``drop_params`` only exists to drop genuinely unsupported params. + """ + if not isinstance(additional_drop_params, list): + return False + return "context_management" in additional_drop_params + + def _normalize_spec_edits( *, context_management_spec: Any, - drop_params: Optional[bool], + additional_drop_params: Optional[list[str]], ) -> Optional[List[Dict[str, Any]]]: """Return the normalized ``edits`` list, or ``None`` if the polyfill won't run. @@ -208,8 +221,7 @@ def _normalize_spec_edits( if not context_management_spec: return None - effective_drop_params = drop_params if drop_params is not None else litellm.drop_params - if effective_drop_params: + if _context_management_explicitly_dropped(additional_drop_params): return None from litellm.llms.anthropic.experimental_pass_through.context_management.dispatcher import ( @@ -230,22 +242,23 @@ async def _run_polyfill_if_enabled( system: Optional[Any], context_management_spec: Any, litellm_metadata: Optional[Dict], - drop_params: Optional[bool], + additional_drop_params: Optional[list[str]], llm_router: Any, user_api_key_auth: Any = None, ) -> Optional[PolyfillResult]: """Run the async context_management polyfill if a spec is present. - Returns ``None`` when the spec is empty or drop_params is on. Raises - ``AnthropicContextManagementError`` so the /v1/messages endpoint can - emit an Anthropic-format 400. All other exceptions are best-effort - swallowed (matches v0 behavior). + Returns ``None`` when the spec is empty or ``context_management`` is + listed in ``additional_drop_params`` (the explicit opt-out; ``drop_params`` + does not disable the polyfill because context_management is a supported + param). Raises ``AnthropicContextManagementError`` so the /v1/messages + endpoint can emit an Anthropic-format 400. All other exceptions are + best-effort swallowed (matches v0 behavior). """ if not context_management_spec: return None - effective_drop_params = drop_params if drop_params is not None else litellm.drop_params - if effective_drop_params: + if _context_management_explicitly_dropped(additional_drop_params): return None try: @@ -274,7 +287,7 @@ async def _run_polyfill_if_enabled( # emits an Anthropic-format error. if _spec_has_non_compact_edits( context_management_spec=context_management_spec, - drop_params=drop_params, + additional_drop_params=additional_drop_params, ): raise AnthropicContextManagementError( status_code=500, @@ -533,7 +546,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: ) -> Union[AnthropicMessagesResponse, AsyncIterator[Any], Iterator[bytes]]: """Handle non-Anthropic models asynchronously using the adapter""" context_management = kwargs.pop("context_management", None) - drop_params: Optional[bool] = kwargs.get("drop_params", None) + additional_drop_params: Optional[list[str]] = kwargs.get("additional_drop_params", None) litellm_router = kwargs.pop("litellm_router", None) if litellm_router is None: try: @@ -555,7 +568,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: system=system, context_management_spec=context_management, litellm_metadata=proxy_litellm_metadata, - drop_params=drop_params, + additional_drop_params=additional_drop_params, llm_router=litellm_router, user_api_key_auth=user_api_key_auth, ) @@ -661,7 +674,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: # ``compact_20260112`` editor can ``await`` the summarization model); # bridge to it via ``run_async_function``. context_management = kwargs.pop("context_management", None) - drop_params: Optional[bool] = kwargs.get("drop_params", None) + additional_drop_params: Optional[list[str]] = kwargs.get("additional_drop_params", None) # Deliberately do NOT auto-attach the proxy ``llm_router`` here: # ``run_async_function`` spawns a new event loop in a worker thread # to bridge to the async dispatcher, but the proxy router's httpx @@ -696,7 +709,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: system=system, context_management_spec=context_management, litellm_metadata=proxy_litellm_metadata, - drop_params=drop_params, + additional_drop_params=additional_drop_params, llm_router=litellm_router, user_api_key_auth=user_api_key_auth, ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py index be430db9eed..9c8df1c79f9 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py @@ -12,11 +12,13 @@ Coverage: - custom instructions → default prompt is not used even when tools present """ +import json from typing import Any, Dict, List from unittest.mock import AsyncMock, MagicMock, patch import pytest +import litellm from litellm.llms.anthropic.experimental_pass_through.context_management import ( AnthropicContextManagementError, apply_context_management, @@ -2042,12 +2044,12 @@ async def test_dispatcher_trigger_below_minimum_raises_through(): # --------------------------------------------------------------------------- -# _run_polyfill_if_enabled: drop_params gate +# _run_polyfill_if_enabled: additional_drop_params gate (drop_params must NOT gate) # --------------------------------------------------------------------------- -async def test_run_polyfill_skipped_when_drop_params_true(): - """When drop_params=True the polyfill must be skipped (returns None).""" +async def test_run_polyfill_skipped_when_context_management_in_additional_drop_params(): + """additional_drop_params=["context_management"] is the explicit opt-out.""" from litellm.llms.anthropic.experimental_pass_through.adapters.handler import ( _run_polyfill_if_enabled, ) @@ -2059,12 +2061,39 @@ async def test_run_polyfill_skipped_when_drop_params_true(): system=None, context_management_spec={"edits": [{"type": "compact_20260112"}]}, litellm_metadata={}, - drop_params=True, + additional_drop_params=["context_management"], llm_router=None, ) assert result is None +async def test_run_polyfill_runs_when_litellm_drop_params_true(monkeypatch): + """drop_params must not disable the polyfill: context_management is a + LiteLLM-supported param (polyfilled where not native), and drop_params only + exists to strip genuinely unsupported params.""" + from litellm.llms.anthropic.experimental_pass_through.adapters.handler import ( + _run_polyfill_if_enabled, + ) + + monkeypatch.setattr(litellm, "drop_params", True) + with patch( + "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting", + return_value=None, + ): + result = await _run_polyfill_if_enabled( + model=MODEL, + messages=_simple_messages(), + tools=None, + system=None, + context_management_spec={"edits": [{"type": "compact_20260112"}]}, + litellm_metadata={}, + additional_drop_params=None, + llm_router=None, + ) + assert result is not None + assert result.applied_edits[0]["type"] == "compact_20260112" + + async def test_run_polyfill_skipped_when_spec_empty(): """Empty context_management_spec must also return None (no polyfill work).""" from litellm.llms.anthropic.experimental_pass_through.adapters.handler import ( @@ -2078,12 +2107,169 @@ async def test_run_polyfill_skipped_when_spec_empty(): system=None, context_management_spec=None, litellm_metadata={}, - drop_params=False, + additional_drop_params=None, llm_router=None, ) assert result is None +# --------------------------------------------------------------------------- +# Adapter handler entry points: polyfill vs drop_params / additional_drop_params +# --------------------------------------------------------------------------- + +_CLEAR_TOOL_USES_SPEC: Dict[str, Any] = { + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "tool_uses", "value": 1}, + "keep": {"type": "tool_uses", "value": 0}, + } + ] +} + +_CLEARED_PLACEHOLDER = "[Cleared by context management]" + + +def _tool_use_messages() -> List[Dict[str, Any]]: + return [ + {"role": "user", "content": "check the weather in two cities"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_01", "name": "get_weather", "input": {"city": "SF"}}], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "sunny in SF"}], + }, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_02", "name": "get_weather", "input": {"city": "NY"}}], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_02", "content": "rainy in NY"}], + }, + {"role": "user", "content": "now summarize both"}, + ] + + +def _openai_chat_response(): + from litellm.types.utils import ModelResponse + + return ModelResponse( + id="chatcmpl-test", + model="gpt-4o", + choices=[{"finish_reason": "stop", "index": 0, "message": {"role": "assistant", "content": "done"}}], + usage={"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12}, + ) + + +async def _call_async_adapter_handler(**handler_kwargs: Any): + from litellm.llms.anthropic.experimental_pass_through.adapters.handler import ( + LiteLLMMessagesToCompletionTransformationHandler, + ) + + captured: Dict[str, Any] = {} + + async def _capture_acompletion(**kwargs): + captured.update(kwargs) + return _openai_chat_response() + + with patch("litellm.acompletion", side_effect=_capture_acompletion): + response = await LiteLLMMessagesToCompletionTransformationHandler.async_anthropic_messages_handler( + max_tokens=128, + messages=_tool_use_messages(), + model=MODEL, + context_management=_CLEAR_TOOL_USES_SPEC, + litellm_router=MagicMock(), + **handler_kwargs, + ) + return response, captured + + +def _assert_polyfill_applied(response: Any, captured: Dict[str, Any]) -> None: + applied_edits = (response.get("context_management") or {}).get("applied_edits") + assert applied_edits, "polyfill must run and report applied_edits" + assert applied_edits[0]["type"] == "clear_tool_uses_20250919" + forwarded = json.dumps(captured["messages"], default=str) + assert _CLEARED_PLACEHOLDER in forwarded + assert "sunny in SF" not in forwarded + assert "rainy in NY" in forwarded + + +async def test_async_handler_runs_polyfill_when_request_drop_params_true(): + """Regression (LIT-3768): per-request drop_params=True silently skipped the + polyfill, so Claude Code requests (where the proxy defaults drop_params on) + lost context editing on non-Anthropic models.""" + response, captured = await _call_async_adapter_handler(drop_params=True) + _assert_polyfill_applied(response, captured) + + +async def test_async_handler_runs_polyfill_when_litellm_drop_params_true(monkeypatch): + """Regression (LIT-3768): proxy-wide litellm.drop_params=True silently + skipped the polyfill too.""" + monkeypatch.setattr(litellm, "drop_params", True) + response, captured = await _call_async_adapter_handler() + _assert_polyfill_applied(response, captured) + + +async def test_async_handler_additional_drop_params_strips_context_management(): + """additional_drop_params=["context_management"] stays the escape hatch: + the polyfill must not run and the request is forwarded untouched.""" + response, captured = await _call_async_adapter_handler(additional_drop_params=["context_management"]) + assert response.get("context_management") is None + forwarded = json.dumps(captured["messages"], default=str) + assert _CLEARED_PLACEHOLDER not in forwarded + assert "sunny in SF" in forwarded + + +def _call_sync_adapter_handler(**handler_kwargs: Any): + from litellm.llms.anthropic.experimental_pass_through.adapters.handler import ( + LiteLLMMessagesToCompletionTransformationHandler, + ) + + captured: Dict[str, Any] = {} + + def _capture_completion(**kwargs): + captured.update(kwargs) + return _openai_chat_response() + + with patch("litellm.completion", side_effect=_capture_completion): + response = LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( + max_tokens=128, + messages=_tool_use_messages(), + model=MODEL, + context_management=_CLEAR_TOOL_USES_SPEC, + litellm_router=None, + **handler_kwargs, + ) + return response, captured + + +def test_sync_handler_runs_polyfill_when_request_drop_params_true(): + """The sync entry point reads its own kwargs; cover its gate separately.""" + response, captured = _call_sync_adapter_handler(drop_params=True) + _assert_polyfill_applied(response, captured) + + +def test_sync_handler_runs_polyfill_when_litellm_drop_params_true(monkeypatch): + """Proxy-wide litellm.drop_params=True must not skip the polyfill on the + sync entry point either.""" + monkeypatch.setattr(litellm, "drop_params", True) + response, captured = _call_sync_adapter_handler() + _assert_polyfill_applied(response, captured) + + +def test_sync_handler_additional_drop_params_strips_context_management(): + """The additional_drop_params=["context_management"] escape hatch is honored + on the sync entry point too: no polyfill, request forwarded untouched.""" + response, captured = _call_sync_adapter_handler(additional_drop_params=["context_management"]) + assert response.get("context_management") is None + forwarded = json.dumps(captured["messages"], default=str) + assert _CLEARED_PLACEHOLDER not in forwarded + assert "sunny in SF" in forwarded + + async def test_prepare_context_managed_request_forwards_proxy_litellm_metadata(): """The handler must hand the polyfill the proxy ``litellm_metadata`` (which carries ``user_api_key`` / ``user_api_key_team_id`` / ...), not the @@ -2120,7 +2306,7 @@ async def test_prepare_context_managed_request_forwards_proxy_litellm_metadata() "user_api_key_user_id": "user-xyz", "litellm_call_id": "call-1", }, - drop_params=False, + additional_drop_params=None, llm_router=_RouterStub(), ) From 5f4b9ad51cfec13c692661cca33623de05968e9f Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 3 Jul 2026 14:14:40 -0700 Subject: [PATCH 4/9] chore: clarify the linear ticket instruction in pr template (#32076) * chore: clarify the linear ticket instruction in pr template * fix: make it more concise * Update CLAUDE.md lol * chore: tell claude not to search for it if it doesn't have it --------- Co-authored-by: ryan-crabbe-berri --- .github/pull_request_template.md | 2 +- CLAUDE.md | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 12ad124fa20..1051459ed44 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -4,7 +4,7 @@ ## Linear ticket - + ## Pre-Submission checklist diff --git a/CLAUDE.md b/CLAUDE.md index 84d2efaab6a..7108dfd4f32 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -21,7 +21,9 @@ End-to-end tests belong in `tests/e2e/` and must follow the harness conventions When creating PRs, don't set base to `main`. `litellm_internal_staging` serves that purpose -Always use @.github/pull_request_template.md as a guide for your PR body +When writing a PR body, treat the comments and imperative instructions inside @.github/pull_request_template.md as rules to follow, not just layout + +If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR From 0115bfa52343ef49d7ed5e5e1c90102577d7b51f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 3 Jul 2026 14:19:30 -0700 Subject: [PATCH 5/9] test(ui): quiet vitest CI logs by silencing passing-test console output The ui_unit_tests CircleCI job logged ~45k lines for a single run, most of it React act() warnings, antd deprecation notices and component stack traces emitted as console output by passing tests, which buried real failures. Set silent: "passed-only" (Vitest 3.2+) gated on process.env.CI so console output from passing tests is suppressed while a failing test still prints its logs and full stack trace. Also drop two stray console.log calls in UsagePageView that dumped the whole currentUser object on every render in production, not just tests. Verified by running the suite the way CI does (CI=true npm run test -- --run --pool forks --poolOptions.forks.maxForks=6): 45,075 lines before, 981 after, all 4075 tests still passing. A throwaway failing test confirms its console.log and assertion diff remain visible. --- .../src/components/UsagePage/components/UsagePageView.tsx | 2 -- ui/litellm-dashboard/vitest.config.ts | 1 + 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx index 432a499a0f6..37c2a25c6a6 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx @@ -81,8 +81,6 @@ const UsagePage: React.FC = ({ teams, organizations }) => { const { data: customers = [] } = useCustomers(); const { data: agentsResponse } = useAgents(); const { data: currentUser } = useCurrentUser(); - console.log(`currentUser: ${JSON.stringify(currentUser)}`); - console.log(`currentUser max budget: ${currentUser?.max_budget}`); const isAdmin = all_admin_roles.includes(userRole || ""); const canViewTagUsage = isAdmin || internalUserRoles.includes(userRole || ""); diff --git a/ui/litellm-dashboard/vitest.config.ts b/ui/litellm-dashboard/vitest.config.ts index d2b6e7b43bf..01c451933b5 100644 --- a/ui/litellm-dashboard/vitest.config.ts +++ b/ui/litellm-dashboard/vitest.config.ts @@ -8,6 +8,7 @@ export default defineConfig({ globals: true, css: true, // lets you import CSS/modules without extra mocks testTimeout: 30000, + silent: process.env.CI ? "passed-only" : false, coverage: { provider: "v8", reporter: ["text", "lcov"], From 76a9f7b5f3be25a2889248be249e1a448c8554de Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 3 Jul 2026 14:35:23 -0700 Subject: [PATCH 6/9] chore(ui): add no-console lint ratchet and strip console from prod builds Introduce a gradual ratchet to remove raw console.* calls from the dashboard, mirroring the existing no-explicit-any budget. The no-console eslint rule is set to warn with allow: [warn, error] so the 486 console.log/debug/info calls are tracked without force-deleting the legitimate console.error/warn error reporting in catch blocks. The count is grandfathered via eslint-budgets.json (max 486, target 0) and eslint-metrics.json, so any newly added console.log fails the budget check and follow-up PRs grind the max down toward zero. Independently, next.config strips console output from production builds via SWC removeConsole (exclude: [error]), gated on NODE_ENV=production so dev keeps full console output. This gives an immediate prod-hygiene net regardless of how long the source cleanup takes. Verified against a real production build: app-code console.log dropped from 675 to 14 in the bundle (remainder is node_modules, which the transform leaves alone), console.warn app calls stripped, console.error preserved 906 to 906. --- ui/litellm-dashboard/eslint-budgets.json | 1 + ui/litellm-dashboard/eslint-metrics.json | 3 ++- ui/litellm-dashboard/eslint.config.mjs | 1 + ui/litellm-dashboard/next.config.mjs | 3 +++ 4 files changed, 7 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/eslint-budgets.json b/ui/litellm-dashboard/eslint-budgets.json index 2139d177512..4ba3581379f 100644 --- a/ui/litellm-dashboard/eslint-budgets.json +++ b/ui/litellm-dashboard/eslint-budgets.json @@ -1,5 +1,6 @@ { "@typescript-eslint/no-explicit-any": { "max": 2040, "target": 1500 }, + "no-console": { "max": 486, "target": 0 }, "complexity": { "max": 140, "target": 80 }, "max-depth": { "max": 70, "target": 30 } } diff --git a/ui/litellm-dashboard/eslint-metrics.json b/ui/litellm-dashboard/eslint-metrics.json index cf754a1bb75..82afa1cbaaa 100644 --- a/ui/litellm-dashboard/eslint-metrics.json +++ b/ui/litellm-dashboard/eslint-metrics.json @@ -1,5 +1,6 @@ { "@typescript-eslint/no-explicit-any": 1991, "complexity": 128, - "max-depth": 59 + "max-depth": 59, + "no-console": 486 } diff --git a/ui/litellm-dashboard/eslint.config.mjs b/ui/litellm-dashboard/eslint.config.mjs index 10caebb5196..acdd0c91309 100644 --- a/ui/litellm-dashboard/eslint.config.mjs +++ b/ui/litellm-dashboard/eslint.config.mjs @@ -17,6 +17,7 @@ const eslintConfig = [ rules: { "unused-imports/no-unused-imports": "error", "@typescript-eslint/no-explicit-any": "warn", + "no-console": ["warn", { allow: ["warn", "error"] }], "@typescript-eslint/no-unused-vars": "off", "@typescript-eslint/no-unused-expressions": "off", "@typescript-eslint/ban-ts-comment": "off", diff --git a/ui/litellm-dashboard/next.config.mjs b/ui/litellm-dashboard/next.config.mjs index 19a2ca298fe..9cfddfd33ca 100644 --- a/ui/litellm-dashboard/next.config.mjs +++ b/ui/litellm-dashboard/next.config.mjs @@ -7,6 +7,9 @@ const __dirname = path.dirname(__filename); const nextConfig = { output: "export", + compiler: { + removeConsole: process.env.NODE_ENV === "production" ? { exclude: ["error"] } : false, + }, // Required with output: "export" — default image optimizer runs only in server mode. // See https://nextjs.org/docs/messages/export-image-api images: { From 5b7c73a5731a4f108b184b9ab4a07e526fa9576e Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 3 Jul 2026 14:37:41 -0700 Subject: [PATCH 7/9] chore(ui): sync no-console budget to 484 after staging merge Merging litellm_internal_staging dropped 2 console.log calls (the currentUser logs removed in #32079), so the no-console budget max and metric move from 486 to 484 to match the current count. --- ui/litellm-dashboard/eslint-budgets.json | 2 +- ui/litellm-dashboard/eslint-metrics.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/eslint-budgets.json b/ui/litellm-dashboard/eslint-budgets.json index 4ba3581379f..8dedb9ac9ca 100644 --- a/ui/litellm-dashboard/eslint-budgets.json +++ b/ui/litellm-dashboard/eslint-budgets.json @@ -1,6 +1,6 @@ { "@typescript-eslint/no-explicit-any": { "max": 2040, "target": 1500 }, - "no-console": { "max": 486, "target": 0 }, + "no-console": { "max": 484, "target": 0 }, "complexity": { "max": 140, "target": 80 }, "max-depth": { "max": 70, "target": 30 } } diff --git a/ui/litellm-dashboard/eslint-metrics.json b/ui/litellm-dashboard/eslint-metrics.json index 82afa1cbaaa..d16be0d2b48 100644 --- a/ui/litellm-dashboard/eslint-metrics.json +++ b/ui/litellm-dashboard/eslint-metrics.json @@ -2,5 +2,5 @@ "@typescript-eslint/no-explicit-any": 1991, "complexity": 128, "max-depth": 59, - "no-console": 486 + "no-console": 484 } From 68d52ac251495ac53e69a75ec694d6e160558df9 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 3 Jul 2026 14:51:14 -0700 Subject: [PATCH 8/9] chore(ui): preserve console.warn in prod builds to match lint allow-list The lint rule allows console.warn (allow: [warn, error]) but removeConsole only excluded error, so approved console.warn calls were silently dropped from production bundles. Add warn to the exclude list so the prod strip and the lint allow-list agree; only console.log/debug/info are stripped now, warn and error both survive (verified: warn 85 to 85, error 906 to 906, log 675 to 14). --- ui/litellm-dashboard/next.config.mjs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/next.config.mjs b/ui/litellm-dashboard/next.config.mjs index 9cfddfd33ca..876df2b49cf 100644 --- a/ui/litellm-dashboard/next.config.mjs +++ b/ui/litellm-dashboard/next.config.mjs @@ -8,7 +8,7 @@ const __dirname = path.dirname(__filename); const nextConfig = { output: "export", compiler: { - removeConsole: process.env.NODE_ENV === "production" ? { exclude: ["error"] } : false, + removeConsole: process.env.NODE_ENV === "production" ? { exclude: ["error", "warn"] } : false, }, // Required with output: "export" — default image optimizer runs only in server mode. // See https://nextjs.org/docs/messages/export-image-api From 099bc973209fa43e56aff37e2cdeff41d4e8c784 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 3 Jul 2026 14:59:34 -0700 Subject: [PATCH 9/9] fix: prevent duplicate budget alert emails on concurrent threshold crossings (#32011) * fix: prevent duplicate budget alert emails on concurrent threshold crossings Budget alert emails were sent more than once for a single threshold crossing. The email dedup guard read the "already sent" marker, awaited the send, then wrote the marker, so concurrent requests crossing the same threshold within the send window all saw no marker and each sent. This affected the multi-threshold path (default_key_max_budget_alert_emails), the legacy single-threshold path (EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE), and the soft budget path, all in EmailBaseCallback.budget_alerts All three branches now claim the send slot atomically before sending via async_increment_cache, which is atomic per event loop for the in-memory cache and across workers via Redis INCR; only the caller that observes a count of 1 sends. On send failure the marker is released with async_delete_cache so a transient failure does not suppress the alert for the full 24h TTL * fix: harden budget alert claim release and skip-path event allocation Addresses review feedback on the claim-before-send change. The claim release in each send-failure handler now logs the send error first and releases the claim best-effort through a shared helper, so a transient cache error during async_delete_cache cannot propagate out of the fire-and-forget budget_alerts task, drop the send-failure log, and leave the claim stuck for the full 24h TTL. In the multi-threshold branch the increment claim now runs before the WebhookEvent is built, so skipped concurrent crossings no longer construct and discard the event, matching the single-threshold and soft budget branches --- .../send_emails/base_email.py | 61 +++--- .../send_emails/test_base_email.py | 199 ++++++++++++++---- 2 files changed, 193 insertions(+), 67 deletions(-) diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py index 9d15f45079f..be80a12c80a 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py @@ -477,9 +477,12 @@ class BaseEmailLogger(CustomLogger): _id = user_info.token or user_info.user_id or "default_id" _cache_key = f"email_budget_alerts:soft_budget_crossed:{_id}" - # Check if we've already sent this alert - result = await _cache.async_get_cache(key=_cache_key) - if result is None: + send_count = await _cache.async_increment_cache( + key=_cache_key, + value=1, + ttl=EMAIL_BUDGET_ALERT_TTL, + ) + if send_count is None or send_count <= 1: # Create WebhookEvent for soft budget alert event_message = f"Soft Budget Crossed - Total Soft Budget: ${user_info.soft_budget}" webhook_event = WebhookEvent( @@ -508,18 +511,12 @@ class BaseEmailLogger(CustomLogger): await self.send_team_soft_budget_alert_email(webhook_event) else: await self.send_soft_budget_alert_email(webhook_event) - - # Cache the alert to prevent duplicate sends - await _cache.async_set_cache( - key=_cache_key, - value="SENT", - ttl=EMAIL_BUDGET_ALERT_TTL, - ) except Exception as e: verbose_proxy_logger.error( f"Error sending soft budget alert email: {e}", exc_info=True, ) + await self._release_budget_alert_claim(_cache, _cache_key) return # For max_budget_alert, check if we've already sent an alert @@ -545,9 +542,12 @@ class BaseEmailLogger(CustomLogger): _id = user_info.token or user_info.user_id or "default_id" _cache_key = f"email_budget_alerts:max_budget_alert:{_id}" - # Check if we've already sent this alert - result = await _cache.async_get_cache(key=_cache_key) - if result is None: + send_count = await _cache.async_increment_cache( + key=_cache_key, + value=1, + ttl=EMAIL_BUDGET_ALERT_TTL, + ) + if send_count is None or send_count <= 1: # Calculate percentage percentage = int( EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100 @@ -576,18 +576,12 @@ class BaseEmailLogger(CustomLogger): try: await self.send_max_budget_alert_email(webhook_event) - - # Cache the alert to prevent duplicate sends - await _cache.async_set_cache( - key=_cache_key, - value="SENT", - ttl=EMAIL_BUDGET_ALERT_TTL, - ) except Exception as e: verbose_proxy_logger.error( f"Error sending max budget alert email: {e}", exc_info=True, ) + await self._release_budget_alert_claim(_cache, _cache_key) return async def _handle_multi_threshold_max_budget_alert( @@ -617,10 +611,6 @@ class BaseEmailLogger(CustomLogger): f"email_budget_alerts:max_budget_alert:{threshold_pct}:{_id}" ) - result = await _cache.async_get_cache(key=_cache_key) - if result is not None: - continue - # Parse emails + auto-include owner emails = _parse_email_list(raw_emails) if user_info.user_email: @@ -634,6 +624,14 @@ class BaseEmailLogger(CustomLogger): continue recipient_emails = list(set(emails)) + send_count = await _cache.async_increment_cache( + key=_cache_key, + value=1, + ttl=EMAIL_BUDGET_ALERT_TTL, + ) + if send_count is not None and send_count > 1: + continue + event_message = f"Max Budget Alert - {threshold_pct}% of Maximum Budget Reached" webhook_event = WebhookEvent( event="max_budget_alert", @@ -660,16 +658,21 @@ class BaseEmailLogger(CustomLogger): threshold_pct=threshold_pct, recipient_emails=recipient_emails, ) - await _cache.async_set_cache( - key=_cache_key, - value="SENT", - ttl=EMAIL_BUDGET_ALERT_TTL, - ) except Exception as e: verbose_proxy_logger.error( f"Error sending multi-threshold max budget alert email for {threshold_pct}%: {e}", exc_info=True, ) + await self._release_budget_alert_claim(_cache, _cache_key) + + async def _release_budget_alert_claim(self, cache: DualCache, cache_key: str) -> None: + try: + await cache.async_delete_cache(key=cache_key) + except Exception: + verbose_proxy_logger.debug( + "Failed to release budget alert claim for %s; it expires with the TTL", + cache_key, + ) async def _get_email_params( self, diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py index db23b712125..c1ccc454305 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py @@ -1,3 +1,4 @@ +import asyncio import json import os import sys @@ -7,6 +8,7 @@ from unittest.mock import patch import pytest from fastapi.testclient import TestClient +from litellm.caching.caching import DualCache from litellm_enterprise.enterprise_callbacks.send_emails.base_email import ( BaseEmailLogger, ) @@ -707,10 +709,9 @@ async def test_budget_alerts_soft_budget_crossed(base_email_logger, mock_send_em event_group=Litellm_EntityType.USER, ) - # Mock the cache to return None (no previous alert sent) + # Mock the cache so the claim is won (increment returns 1) mock_cache = mock.AsyncMock() - mock_cache.async_get_cache = mock.AsyncMock(return_value=None) - mock_cache.async_set_cache = mock.AsyncMock() + mock_cache.async_increment_cache = mock.AsyncMock(return_value=1) base_email_logger.internal_usage_cache = mock_cache with mock.patch.dict( @@ -726,14 +727,14 @@ async def test_budget_alerts_soft_budget_crossed(base_email_logger, mock_send_em call_args = mock_send_email.call_args[1] assert call_args["to_email"] == ["test@example.com"] - # Verify cache was set to prevent duplicate alerts - mock_cache.async_set_cache.assert_called_once() - cache_call_args = mock_cache.async_set_cache.call_args[1] + # Verify the send slot was claimed to prevent duplicate alerts + mock_cache.async_increment_cache.assert_called_once() + cache_call_args = mock_cache.async_increment_cache.call_args[1] assert ( cache_call_args["key"] == "email_budget_alerts:soft_budget_crossed:test_user" ) - assert cache_call_args["value"] == "SENT" + assert cache_call_args["value"] == 1 assert cache_call_args["ttl"] == EMAIL_BUDGET_ALERT_TTL @@ -774,9 +775,9 @@ async def test_budget_alerts_soft_budget_duplicate_prevention( event_group=Litellm_EntityType.USER, ) - # Mock the cache to return "SENT" (previous alert already sent) + # Mock the cache so the slot is already claimed (increment returns > 1) mock_cache = mock.AsyncMock() - mock_cache.async_get_cache = mock.AsyncMock(return_value="SENT") + mock_cache.async_increment_cache = mock.AsyncMock(return_value=2) base_email_logger.internal_usage_cache = mock_cache await base_email_logger.budget_alerts(type="soft_budget", user_info=user_info) @@ -818,10 +819,9 @@ async def test_budget_alerts_uses_token_for_cache_key( event_group=Litellm_EntityType.KEY, ) - # Mock the cache to return None (no previous alert sent) + # Mock the cache so the claim is won (increment returns 1) mock_cache = mock.AsyncMock() - mock_cache.async_get_cache = mock.AsyncMock(return_value=None) - mock_cache.async_set_cache = mock.AsyncMock() + mock_cache.async_increment_cache = mock.AsyncMock(return_value=1) base_email_logger.internal_usage_cache = mock_cache with mock.patch.dict( @@ -833,8 +833,8 @@ async def test_budget_alerts_uses_token_for_cache_key( await base_email_logger.budget_alerts(type="soft_budget", user_info=user_info) # Verify cache key uses token instead of user_id - mock_cache.async_set_cache.assert_called_once() - cache_call_args = mock_cache.async_set_cache.call_args[1] + mock_cache.async_increment_cache.assert_called_once() + cache_call_args = mock_cache.async_increment_cache.call_args[1] assert ( cache_call_args["key"] == "email_budget_alerts:soft_budget_crossed:hashed_token_123" @@ -880,8 +880,7 @@ async def test_budget_alerts_max_budget_alert_crossed( ) mock_cache = mock.AsyncMock() - mock_cache.async_get_cache = mock.AsyncMock(return_value=None) - mock_cache.async_set_cache = mock.AsyncMock() + mock_cache.async_increment_cache = mock.AsyncMock(return_value=1) base_email_logger.internal_usage_cache = mock_cache with mock.patch.dict( @@ -899,12 +898,12 @@ async def test_budget_alerts_max_budget_alert_crossed( assert call_args["to_email"] == ["test@example.com"] assert "Max Budget Alert" in call_args["subject"] - mock_cache.async_set_cache.assert_called_once() - cache_call_args = mock_cache.async_set_cache.call_args[1] + mock_cache.async_increment_cache.assert_called_once() + cache_call_args = mock_cache.async_increment_cache.call_args[1] assert ( cache_call_args["key"] == "email_budget_alerts:max_budget_alert:test_user" ) - assert cache_call_args["value"] == "SENT" + assert cache_call_args["value"] == 1 assert cache_call_args["ttl"] == EMAIL_BUDGET_ALERT_TTL @@ -928,8 +927,7 @@ async def test_multi_threshold_sends_crossed_thresholds( ) mock_cache = mock.AsyncMock() - mock_cache.async_get_cache = mock.AsyncMock(return_value=None) - mock_cache.async_set_cache = mock.AsyncMock() + mock_cache.async_increment_cache = mock.AsyncMock(return_value=1) base_email_logger.internal_usage_cache = mock_cache with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): @@ -941,7 +939,9 @@ async def test_multi_threshold_sends_crossed_thresholds( assert mock_send_email.call_count == 2 # Check cache keys include threshold percentage - cache_keys = [c[1]["key"] for c in mock_cache.async_set_cache.call_args_list] + cache_keys = [ + c[1]["key"] for c in mock_cache.async_increment_cache.call_args_list + ] assert "email_budget_alerts:max_budget_alert:50:hashed_key_1" in cache_keys assert "email_budget_alerts:max_budget_alert:75:hashed_key_1" in cache_keys @@ -964,15 +964,14 @@ async def test_multi_threshold_dedup_cache_prevents_resend( }, ) - # Simulate 50% already sent (cached), 75% not yet sent - async def cache_get(key): + # Simulate 50% already claimed (increment returns >1), 75% first send (returns 1) + async def cache_increment(key, value, ttl=None): if "50:" in key: - return "SENT" - return None + return 2 + return 1 mock_cache = mock.AsyncMock() - mock_cache.async_get_cache = mock.AsyncMock(side_effect=cache_get) - mock_cache.async_set_cache = mock.AsyncMock() + mock_cache.async_increment_cache = mock.AsyncMock(side_effect=cache_increment) base_email_logger.internal_usage_cache = mock_cache with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): @@ -982,7 +981,7 @@ async def test_multi_threshold_dedup_cache_prevents_resend( # Only 75% should fire assert mock_send_email.call_count == 1 - cache_key = mock_cache.async_set_cache.call_args[1]["key"] + cache_key = mock_cache.async_increment_cache.call_args[1]["key"] assert "75:" in cache_key @@ -1004,8 +1003,7 @@ async def test_multi_threshold_owner_email_auto_included( ) mock_cache = mock.AsyncMock() - mock_cache.async_get_cache = mock.AsyncMock(return_value=None) - mock_cache.async_set_cache = mock.AsyncMock() + mock_cache.async_increment_cache = mock.AsyncMock(return_value=1) base_email_logger.internal_usage_cache = mock_cache with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): @@ -1038,8 +1036,7 @@ async def test_multi_threshold_malformed_keys_skipped( ) mock_cache = mock.AsyncMock() - mock_cache.async_get_cache = mock.AsyncMock(return_value=None) - mock_cache.async_set_cache = mock.AsyncMock() + mock_cache.async_increment_cache = mock.AsyncMock(return_value=1) base_email_logger.internal_usage_cache = mock_cache with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): @@ -1069,8 +1066,7 @@ async def test_multi_threshold_empty_emails_only_owner( ) mock_cache = mock.AsyncMock() - mock_cache.async_get_cache = mock.AsyncMock(return_value=None) - mock_cache.async_set_cache = mock.AsyncMock() + mock_cache.async_increment_cache = mock.AsyncMock(return_value=1) base_email_logger.internal_usage_cache = mock_cache with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): @@ -1097,8 +1093,7 @@ async def test_no_map_preserves_old_single_threshold( ) mock_cache = mock.AsyncMock() - mock_cache.async_get_cache = mock.AsyncMock(return_value=None) - mock_cache.async_set_cache = mock.AsyncMock() + mock_cache.async_increment_cache = mock.AsyncMock(return_value=1) base_email_logger.internal_usage_cache = mock_cache with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): @@ -1110,7 +1105,7 @@ async def test_no_map_preserves_old_single_threshold( call_args = mock_send_email.call_args[1] assert call_args["to_email"] == ["test@example.com"] # Old path cache key has no threshold percentage - cache_key = mock_cache.async_set_cache.call_args[1]["key"] + cache_key = mock_cache.async_increment_cache.call_args[1]["key"] assert cache_key == "email_budget_alerts:max_budget_alert:test_user" @@ -1242,3 +1237,131 @@ async def test_send_soft_budget_alert_email_default_footer_when_no_signature( html_body = mock_send_email.call_args[1]["html_body"] assert EMAIL_FOOTER in html_body + + +_BUDGET_ALERT_BRANCHES = [ + ( + "multi_threshold", + "max_budget_alert", + "send_max_budget_alert_email", + dict(max_budget=100.0, spend=80.0, max_budget_alert_emails={"50": ["finance@co.com"]}), + ), + ( + "single_threshold", + "max_budget_alert", + "send_max_budget_alert_email", + dict(max_budget=100.0, spend=85.0), + ), + ( + "soft_budget", + "soft_budget", + "send_soft_budget_alert_email", + dict(soft_budget=50.0, spend=60.0), + ), +] + + +def _budget_alert_user_info(extra: dict) -> CallInfo: + return CallInfo( + token="hashed_key_1", + user_id="test_user", + user_email="owner@co.com", + event_group=Litellm_EntityType.KEY, + **extra, + ) + + +@pytest.mark.parametrize( + "branch, alert_type, send_method, ci_kwargs", + _BUDGET_ALERT_BRANCHES, + ids=[b[0] for b in _BUDGET_ALERT_BRANCHES], +) +@pytest.mark.asyncio +async def test_budget_alert_no_duplicate_on_concurrent_crossing( + base_email_logger, branch, alert_type, send_method, ci_kwargs +): + """Regression for LIT-4172: two requests crossing the same threshold at the + same time must send exactly one email. The old code wrote the dedup marker + only after the send finished awaiting, so both concurrent tasks passed the + 'already sent' check and both sent. Covers all three send branches.""" + base_email_logger.internal_usage_cache = DualCache() + + sends = [] + + async def slow_send(*args, **kwargs): + sends.append(1) + await asyncio.sleep(0.05) + + with mock.patch.object(base_email_logger, send_method, side_effect=slow_send): + with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): + await asyncio.gather( + base_email_logger.budget_alerts( + type=alert_type, user_info=_budget_alert_user_info(ci_kwargs) + ), + base_email_logger.budget_alerts( + type=alert_type, user_info=_budget_alert_user_info(ci_kwargs) + ), + ) + + assert len(sends) == 1 + + +@pytest.mark.parametrize( + "branch, alert_type, send_method, ci_kwargs", + _BUDGET_ALERT_BRANCHES, + ids=[b[0] for b in _BUDGET_ALERT_BRANCHES], +) +@pytest.mark.asyncio +async def test_budget_alert_failed_send_releases_claim_for_retry( + base_email_logger, branch, alert_type, send_method, ci_kwargs +): + """Claiming the send slot before sending must not swallow the alert forever + if the send fails; the claim is released so a later request retries. Covers + all three send branches.""" + base_email_logger.internal_usage_cache = DualCache() + + attempts = [] + + async def flaky_send(*args, **kwargs): + attempts.append(1) + if len(attempts) == 1: + raise ValueError("transient email backend failure") + + with mock.patch.object(base_email_logger, send_method, side_effect=flaky_send): + with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): + await base_email_logger.budget_alerts( + type=alert_type, user_info=_budget_alert_user_info(ci_kwargs) + ) + await base_email_logger.budget_alerts( + type=alert_type, user_info=_budget_alert_user_info(ci_kwargs) + ) + + assert len(attempts) == 2 + + +@pytest.mark.asyncio +async def test_budget_alert_release_failure_does_not_propagate(base_email_logger): + """If the send fails and releasing the claim also fails (transient cache + error), budget_alerts must swallow it and still log the send failure rather + than letting the exception escape the fire-and-forget task.""" + mock_cache = mock.AsyncMock() + mock_cache.async_increment_cache = mock.AsyncMock(return_value=1) + mock_cache.async_delete_cache = mock.AsyncMock( + side_effect=RuntimeError("cache backend unavailable") + ) + base_email_logger.internal_usage_cache = mock_cache + + async def failing_send(*args, **kwargs): + raise ValueError("smtp backend down") + + with mock.patch.object( + base_email_logger, "send_max_budget_alert_email", side_effect=failing_send + ): + with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): + # Must not raise even though both the send and the release fail. + await base_email_logger.budget_alerts( + type="max_budget_alert", + user_info=_budget_alert_user_info(dict(max_budget=100.0, spend=85.0)), + ) + + mock_cache.async_delete_cache.assert_awaited_once()