Merge remote-tracking branch 'origin/litellm_chat-keys-usage' into litellm_chat-keys-usage

This commit is contained in:
Krrish Dholakia 2026-07-03 15:36:49 -07:00
commit a58930e94e
25 changed files with 1059 additions and 119 deletions

View file

@ -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:

View file

@ -4,7 +4,7 @@
## Linear ticket
<!-- if you are an internal contributor (e.g., your username is postfixed with -berri or -berriai), add "Resolves " followed by the Linear ticket e.g., "Resolves LIT-1234" to magically link the Linear ticket to the GitHub PR -->
<!-- if you are an internal contributor, add "Resolves " followed by the Linear ticket e.g., "Resolves LIT-1234" to link the Linear ticket to the GitHub PR. If you don't have one, leave the section blank rather than guessing -->
## Pre-Submission checklist

View file

@ -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

View file

@ -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,

View file

@ -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,
)

View file

@ -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(

View file

@ -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,

View file

@ -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,

View file

@ -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()

View file

@ -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(),
)

View file

@ -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)

View file

@ -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 ----

View file

@ -1,5 +1,6 @@
{
"@typescript-eslint/no-explicit-any": { "max": 2040, "target": 1500 },
"no-console": { "max": 484, "target": 0 },
"complexity": { "max": 140, "target": 80 },
"max-depth": { "max": 70, "target": 30 }
}

View file

@ -1,5 +1,6 @@
{
"@typescript-eslint/no-explicit-any": 1991,
"complexity": 129,
"max-depth": 59
"max-depth": 59,
"no-console": 484
}

View file

@ -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",

View file

@ -7,6 +7,9 @@ const __dirname = path.dirname(__filename);
const nextConfig = {
output: "export",
compiler: {
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
images: {

View file

@ -1479,6 +1479,7 @@ const Teams: React.FC<TeamProps> = ({ 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 || "")}
/>
</Form.Item>

View file

@ -81,8 +81,6 @@ const UsagePage: React.FC<UsagePageProps> = ({ 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 || "");

View file

@ -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(
<MCPServerSelector accessToken="tok" onChange={vi.fn()} value={{ servers: [], accessGroups: [] }} />,
);
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(
<MCPServerSelector
accessToken="tok"
allowAllProxyMcpServers
onChange={onChange}
value={{ servers: ["srv-1"], accessGroups: [] }}
/>,
);
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(
<MCPServerSelector
accessToken="tok"
allowAllProxyMcpServers
onChange={vi.fn()}
value={{ servers: [ALL_PROXY_MCP_SERVERS_SENTINEL], accessGroups: [] }}
/>,
);
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(
<MCPServerSelector
accessToken="tok"
onChange={vi.fn()}
value={{ servers: [ALL_PROXY_MCP_SERVERS_SENTINEL], accessGroups: [] }}
/>,
);
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);
});
});

View file

@ -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<MCPServerSelectorProps> = ({
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<MCPServerSelectorProps> = ({
];
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<MCPServerSelectorProps> = ({
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) && (
<Select.Option
key={ALL_PROXY_MCP_SERVERS_SENTINEL}
value={ALL_PROXY_MCP_SERVERS_SENTINEL}
label="All Proxy MCP Servers"
>
<span style={{ color: "#1890ff", fontWeight: 500 }}>All Proxy MCP Servers</span>
</Select.Option>
)}
{allowNoMcpServers && (
<Select.Option key={NO_MCP_SERVERS_SENTINEL} value={NO_MCP_SERVERS_SENTINEL} label="No MCP Servers">
<div style={{ display: "flex", alignItems: "center", gap: "8px" }}>
@ -126,7 +143,12 @@ const MCPServerSelector: React.FC<MCPServerSelectorProps> = ({
</Select.Option>
)}
{options.map((opt) => (
<Select.Option key={opt.value} value={opt.value} label={opt.label} disabled={hasNoMcpServersSelected}>
<Select.Option
key={opt.value}
value={opt.value}
label={opt.label}
disabled={hasNoMcpServersSelected || hasAllProxyMcpServersSelected}
>
<div style={{ display: "flex", alignItems: "center", gap: "8px" }}>
<span
style={{

View file

@ -1,5 +1,7 @@
// Must match the backend SpecialMCPServerNames.no_mcp_servers enum value.
export const NO_MCP_SERVERS_SENTINEL = "no-mcp-servers";
export const ALL_PROXY_MCP_SERVERS_SENTINEL = "all-proxy-mcpservers";
export const MCP_TOOLS_PREVIEW_FORBIDDEN_MESSAGE =
"Tool preview is not available for submissions. Tools will be verified by an admin during review.";

View file

@ -3,6 +3,7 @@ import { render, screen, waitFor } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import MCPServerPermissions from "./MCPServerPermissions";
import * as networking from "../networking";
import { ALL_PROXY_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants";
vi.mock("../networking");
@ -354,4 +355,21 @@ describe("MCPServerPermissions", () => {
// 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(
<MCPServerPermissions
mcpServers={[ALL_PROXY_MCP_SERVERS_SENTINEL]}
mcpAccessGroups={[]}
mcpToolPermissions={{}}
accessToken={mockAccessToken}
/>,
);
expect(await screen.findByText("All Proxy MCP Servers")).toBeInTheDocument();
expect(screen.getByText("All")).toBeInTheDocument();
expect(screen.queryByText(ALL_PROXY_MCP_SERVERS_SENTINEL)).not.toBeInTheDocument();
});
});

View file

@ -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({
<ServerIcon className="h-4 w-4 text-blue-600" />
<Text className="font-semibold text-gray-900">MCP Servers</Text>
<Badge color={blocksAllMcpServers ? "red" : "blue"} size="xs">
{blocksAllMcpServers ? "Blocked" : totalCount}
{blocksAllMcpServers ? "Blocked" : grantsAllProxyMcpServers ? "All" : totalCount}
</Badge>
</div>
@ -123,6 +124,11 @@ export function MCPServerPermissions({
No MCP servers this key is blocked from all MCP servers, including its team&apos;s servers
</Text>
</div>
) : grantsAllProxyMcpServers ? (
<div className="flex items-center gap-2 px-3 py-2 rounded-lg bg-blue-50 border border-blue-200">
<ServerIcon className="h-4 w-4 text-blue-400" />
<Text className="text-blue-700 text-sm">All Proxy MCP Servers</Text>
</div>
) : totalCount > 0 ? (
<div className="max-h-[400px] overflow-y-auto space-y-2 pr-1">
{mergedItems.map((item, index) => {

View file

@ -1356,6 +1356,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
value={form.getFieldValue("mcp_servers_and_groups")}
accessToken={accessToken || ""}
placeholder="Select MCP servers or access groups (optional)"
allowAllProxyMcpServers={is_proxy_admin}
/>
</Form.Item>

View file

@ -8,6 +8,8 @@ export default defineConfig({
globals: true,
css: true, // lets you import CSS/modules without extra mocks
testTimeout: 30000,
silent: process.env.CI ? "passed-only" : false,
teardownTimeout: 60000,
coverage: {
provider: "v8",
reporter: ["text", "lcov"],