mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
Merge remote-tracking branch 'origin/litellm_chat-keys-usage' into litellm_chat-keys-usage
This commit is contained in:
commit
a58930e94e
25 changed files with 1059 additions and 119 deletions
|
|
@ -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:
|
||||
|
|
|
|||
2
.github/pull_request_template.md
vendored
2
.github/pull_request_template.md
vendored
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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 ----
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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 }
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
{
|
||||
"@typescript-eslint/no-explicit-any": 1991,
|
||||
"complexity": 129,
|
||||
"max-depth": 59
|
||||
"max-depth": 59,
|
||||
"no-console": 484
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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: {
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
||||
|
|
|
|||
|
|
@ -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 || "");
|
||||
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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={{
|
||||
|
|
|
|||
|
|
@ -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.";
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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'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) => {
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
||||
|
|
|
|||
|
|
@ -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"],
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue