scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

chore: merge latest litellm_internal_staging

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mateo 2026-08-14 04:51:43 +00:00
commit 6cbfe9013c
143 changed files with 12118 additions and 4643 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

103
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -3,7 +3,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t
"""
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
from typing import TYPE_CHECKING, Any, Dict, Final, List, Optional, Tuple
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
@ -23,6 +23,15 @@ if TYPE_CHECKING:
CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = (
"completed",
"complete",
"failed",
"expired",
"cancelled",
"stale_expired",
)
class CheckBatchCost:
def __init__(
@ -132,11 +141,11 @@ class CheckBatchCost:
in non-terminal states as 'stale_expired'. These will never complete and
should not be polled.
"""
cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
result = await self.prisma_client.db.litellm_managedobjecttable.update_many(
cutoff: Final = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
result: Final = await self.prisma_client.db.litellm_managedobjecttable.update_many(
where={
"file_purpose": "batch",
"status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]},
"status": {"not_in": list(TERMINAL_MANAGED_OBJECT_STATUSES)},
"created_at": {"lt": cutoff},
},
data={"status": "stale_expired"},
@ -147,6 +156,26 @@ class CheckBatchCost:
f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
)
if not self._has_batch_processed_column:
return
# A row already in a terminal status is never rewritten by the sweep above, so
# without this it keeps a poll-page slot forever and starves newer batches.
retired: Final = await self.prisma_client.db.litellm_managedobjecttable.update_many(
where={
"file_purpose": "batch",
"batch_processed": False,
"status": {"in": ["complete", "completed"]},
"created_at": {"lt": cutoff},
},
data={"batch_processed": True},
)
if retired > 0:
verbose_proxy_logger.warning(
f"CheckBatchCost: gave up on {retired} completed managed objects older than "
f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed"
)
async def _fallback_find_jobs(self) -> list:
"""Query batch jobs without the batch_processed filter (for older schemas)."""
return await self.prisma_client.db.litellm_managedobjecttable.find_many(
@ -167,6 +196,68 @@ class CheckBatchCost:
order={"created_at": "asc"},
)
async def _retire_job(self, job: "LiteLLM_ManagedObjectTable", reason: str) -> None:
"""
Take a row that can never be costed out of the poll page. Leaving it selectable
would burn one of the MAX_OBJECTS_PER_POLL_CYCLE slots on every future cycle, and
once enough such rows accumulate no newer batch is ever reached. Older schemas
without batch_processed can only be excluded through the status filter.
"""
data: Final = (
{"batch_processed": True}
if self._has_batch_processed_column
else {"status": "stale_expired"}
)
try:
await self.prisma_client.db.litellm_managedobjecttable.update(
where={"id": job.id},
data=data,
)
except Exception as db_err:
verbose_proxy_logger.error(
f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}"
)
return
verbose_proxy_logger.warning(
f"CheckBatchCost: job {job.id} can never be costed ({reason}), "
"so it will no longer be polled"
)
@staticmethod
def _has_unified_id_without_model(job: "LiteLLM_ManagedObjectTable") -> bool:
"""A unified id that decodes but carries no model_id can never be routed."""
from litellm.proxy.openai_files_endpoints.common_utils import (
convert_b64_uid_to_unified_uid,
get_model_id_from_unified_batch_id,
)
decoded: Final = convert_b64_uid_to_unified_uid(job.unified_object_id)
return (
decoded != job.unified_object_id
and get_model_id_from_unified_batch_id(decoded) is None
)
@staticmethod
def _is_batch_gone_at_provider(error: Exception, batch_id: str) -> bool:
"""
A 404 naming the batch means the provider dropped its record of it, so no later
retrieve can ever succeed. A 404 about anything else, a renamed Azure deployment
or a fallback deployment that never saw this batch, is still fixable in config, so
it keeps retrying.
"""
import openai
from litellm.exceptions import NotFoundError
return isinstance(error, (NotFoundError, openai.NotFoundError)) and batch_id in str(error)
def _batch_deployment_exists(self, model_id: str) -> bool:
"""A 404 only proves the batch is gone when it came from the batch's own
deployment. Once that deployment leaves the router, default fallbacks can
silently send the retrieve to a provider that never saw the batch, so its
404 must not retire the row; the staleness sweep bounds it instead."""
return self.llm_router.get_deployment(model_id=model_id) is not None
@staticmethod
def _record_error(
prom_logger: Optional["PrometheusLogger"], error_type: str
@ -645,6 +736,8 @@ class CheckBatchCost:
for job in jobs:
routing = self._resolve_job_routing(job, prom_logger)
if routing is None:
if self._has_unified_id_without_model(job):
await self._retire_job(job, "unified object id has no model id")
continue
model_id, batch_id = routing
@ -667,6 +760,8 @@ class CheckBatchCost:
)
if prom_logger:
prom_logger.record_check_batch_cost_error("provider_retrieval_error")
if self._is_batch_gone_at_provider(e, batch_id) and self._batch_deployment_exists(model_id):
await self._retire_job(job, f"batch {batch_id} no longer exists at the provider")
continue
## RETRIEVE THE BATCH JOB OUTPUT FILE

6
litellm/constants.py
View file

@ -1491,6 +1491,9 @@ SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int(os.getenv("SPEND_LOG_CLEA
SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS: Final = float(
os.getenv("SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.5)
)
SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS", "300"))
SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS", "30"))
SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP: Final = int(os.getenv("SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP", "100000"))
TOOL_SPEND_TOP_TOOLS: Final = 100
SPEND_LOG_PARTITION_INTERVAL: Final = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day")
SPEND_LOG_PARTITION_PRECREATE_AHEAD: Final = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7))
@ -1742,6 +1745,9 @@ PTU_ROLLUP_LOCK_TTL_SECONDS: Final[int] = 900
# Furthest back the catch-up pass looks for unpriced PTU days when a deployment
# declares no ptu_effective_from, bounding the scan for an open-ended window.
PTU_ROLLUP_MAX_BACKFILL_DAYS: Final[int] = 90
# Deployments named in the lapsed-window alert before it is truncated, so a fleet-wide
# expiry cannot produce an alert too large for the channel delivering it.
PTU_LAPSED_ALERT_LIMIT: Final[int] = 10
# Slack allowed when deciding a sentinel row is stale. The row's updated_at and the
# run's cutoff are stamped by different hosts, so clock skew between them must not let
# one run delete a charge another just wrote. A stale row is hours old and a concurrent

6
litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py
View file

@ -19,9 +19,9 @@ from litellm.llms.bedrock.common_utils import (
convert_bedrock_invoke_output_format_to_inline_schema,
get_anthropic_beta_from_headers,
normalize_bedrock_opus_output_config_effort,
normalize_custom_field_on_tools,
normalize_tool_input_schema_types_for_bedrock_invoke,
pop_bedrock_invoke_output_config_format,
remove_custom_field_from_tools,
)
from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER
from litellm.types.llms.openai import AllMessageValues
@ -243,8 +243,8 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
if "anthropic_version" not in anthropic_request:
anthropic_request["anthropic_version"] = self.anthropic_version
# Remove `custom` field from tools (Bedrock doesn't support it)
remove_custom_field_from_tools(anthropic_request)
# Hoist `custom.defer_loading` then drop `custom` (Bedrock doesn't support it)
normalize_custom_field_on_tools(anthropic_request)
normalize_tool_input_schema_types_for_bedrock_invoke(anthropic_request)
return anthropic_request

21
litellm/llms/bedrock/common_utils.py
View file

@ -176,13 +176,14 @@ def convert_bedrock_invoke_output_format_to_inline_schema(
request_body["messages"] = new_messages
def remove_custom_field_from_tools(request_body: dict) -> None:
def normalize_custom_field_on_tools(request_body: dict) -> None:
"""
Remove ``custom`` field from each tool in the request body.
Drop the ``custom`` field from each tool, first hoisting a boolean
``custom.defer_loading`` onto the top-level ``defer_loading`` flag that
Bedrock and Anthropic actually document, unless the tool already carries one.
Claude Code (v2.1.69+) sends ``custom: {defer_loading: true}`` on tool
definitions, which Anthropic's API accepts but Bedrock rejects with
``"Extra inputs are not permitted"``.
Claude Code (v2.1.69+) is reported to send ``custom: {defer_loading: true}`` on
tool definitions, which Bedrock rejects with ``"Extra inputs are not permitted"``.
Args:
request_body: The request dictionary to modify in-place.
@ -193,8 +194,14 @@ def remove_custom_field_from_tools(request_body: dict) -> None:
if not tools or not isinstance(tools, list):
return
for tool in tools:
if isinstance(tool, dict):
tool.pop("custom", None)
if not isinstance(tool, dict):
continue
custom: dict[str, object] | None = tool.pop("custom", None)
if not isinstance(custom, dict) or "defer_loading" in tool:
continue
deferred: object = custom.get("defer_loading")
if isinstance(deferred, bool):
tool["defer_loading"] = deferred
def normalize_json_schema_custom_types_to_object(schema: dict) -> None:

8
litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
View file

@ -33,9 +33,9 @@ from litellm.llms.bedrock.common_utils import (
get_anthropic_beta_from_headers,
is_claude_4_5_on_bedrock,
normalize_bedrock_opus_output_config_effort,
normalize_custom_field_on_tools,
normalize_tool_input_schema_types_for_bedrock_invoke,
pop_bedrock_invoke_output_config_format,
remove_custom_field_from_tools,
)
from litellm.types.llms.anthropic import (
ANTHROPIC_BETA_HEADER_VALUES,
@ -749,11 +749,9 @@ class AmazonAnthropicClaudeMessagesConfig(
model,
)
# 5b. Remove `custom` field from tools (Bedrock doesn't support it)
# Claude Code sends `custom: {defer_loading: true}` on tool definitions,
# which causes Bedrock to reject the request with "Extra inputs are not permitted"
# 5b. Hoist `custom.defer_loading` then drop `custom` (Bedrock doesn't support it)
# Ref: https://github.com/BerriAI/litellm/issues/22847
remove_custom_field_from_tools(anthropic_messages_request)
normalize_custom_field_on_tools(anthropic_messages_request)
normalize_tool_input_schema_types_for_bedrock_invoke(anthropic_messages_request)
ensure_bedrock_anthropic_messages_tool_names(anthropic_messages_request)

4
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -1190,7 +1190,7 @@ class MCPRequestHandler:
DEPRECATED: This method is deprecated in favor of server-specific auth headers using the format x-mcp-{{server_alias}}-{{header_name}} instead.
"""
mcp_client_side_auth_header_name: Final[str] = MCPRequestHandler._get_mcp_client_side_auth_header_name()
mcp_client_side_auth_header_name: Final[str] = MCPRequestHandler.get_mcp_client_side_auth_header_name()
auth_header: Final = headers.get(mcp_client_side_auth_header_name)
if auth_header:
verbose_logger.warning(
@ -1265,7 +1265,7 @@ class MCPRequestHandler:
return oauth2_headers
@staticmethod
def _get_mcp_client_side_auth_header_name() -> str:
def get_mcp_client_side_auth_header_name() -> str:
"""
Get the header name used to pass the MCP auth header to the MCP server

2
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -118,6 +118,7 @@ from litellm.proxy._experimental.mcp_server.utils import (
is_short_mcp_tool_prefix_enabled,
iter_known_server_prefixes,
iter_known_tool_name_spellings,
logging_safe_mcp_headers,
match_known_server_prefix,
match_known_tool_name,
merge_mcp_headers,
@ -4603,6 +4604,7 @@ class MCPServerManager:
),
"user_api_key_hash": (getattr(user_api_key_auth, "api_key_hash", None) if user_api_key_auth else None),
"incoming_bearer_token": incoming_bearer_token,
"headers": logging_safe_mcp_headers(raw_headers),
}
# Create MCP request object for processing

99
litellm/proxy/_experimental/mcp_server/sampling_handler.py
View file

@ -1042,100 +1042,15 @@ def _build_sampling_request(
raw_headers: dict[str, str] | None = None,
client_ip: str | None = None,
) -> "Request":
"""Build a synthetic FastAPI Request for sampling sub-calls.
"""The synthetic FastAPI Request for sampling sub-calls, carrying the original
MCP connection's headers and client IP."""
from litellm.proxy._experimental.mcp_server.utils import build_synthetic_mcp_request
Converts the original MCP connection's HTTP headers into ASGI
scope format so that ``add_litellm_data_to_request`` can apply
header-dependent guardrails, tag-based routing, trace correlation,
and ``forward_llm_provider_auth_headers``.
Key fields populated:
- **headers**: All original HTTP headers are forwarded (except
hop-by-hop: content-length, transfer-encoding). This ensures
``traceparent``, ``authorization``, ``user-agent``, and
``x-litellm-api-key`` are visible to pre-call utils.
- **client**: The ASGI ``(host, port)`` tuple so that
``request.client.host`` returns the real client IP for
IP-based routing and guardrails.
- **server**: Derived from the running proxy's ``server_host``
/ ``server_port`` when available, avoiding the misleading
``127.0.0.1:0`` placeholder.
- **x-forwarded-for**: Injected from ``client_ip`` if the
original headers don't already carry it, as a fallback for
IP attribution.
"""
from fastapi import Request
# --- Build ASGI headers ---
_scope_headers: Final[list[tuple[bytes, bytes]]] = [(b"content-type", b"application/json")]
# Hop-by-hop headers that must NOT be forwarded into the
# synthetic request (they describe the original HTTP framing,
# not the logical request).
_HOP_BY_HOP: Final = frozenset(
{
"content-length",
"transfer-encoding",
"connection",
"keep-alive",
"upgrade",
"te",
"trailer",
}
return build_synthetic_mcp_request(
path="/mcp/sampling/createMessage",
raw_headers=raw_headers,
client_ip=client_ip,
)
if raw_headers:
for hdr_name, hdr_value in raw_headers.items():
_key = hdr_name.lower()
# Skip content-type (already set), x-forwarded-for (use resolved
# client_ip instead to prevent spoofing), and hop-by-hop headers
if _key in {"content-type", "x-forwarded-for"} or _key in _HOP_BY_HOP:
continue
_scope_headers.append(
(
_key.encode("latin-1", errors="replace"),
hdr_value.encode("utf-8"),
)
)
# Inject x-forwarded-for from captured client_ip if the
# original headers don't already carry it
if client_ip and not any(h[0] == b"x-forwarded-for" for h in _scope_headers):
_scope_headers.append((b"x-forwarded-for", client_ip.encode("utf-8")))
# --- Derive server (host, port) from the running proxy ---
_server_host = "127.0.0.1"
_server_port = 4000 # LiteLLM default
try:
from litellm.proxy import proxy_server
_proxy_host: Final[str | None] = getattr(proxy_server, "server_host", None)
_proxy_port: Final[str | int | None] = getattr(proxy_server, "server_port", None)
if _proxy_host:
_server_host = str(_proxy_host)
if _proxy_port:
_server_port = int(_proxy_port)
except (ImportError, AttributeError, TypeError, ValueError):
pass
# --- Build ASGI client tuple for request.client.host ---
_client_tuple = None
if client_ip:
_client_tuple = (client_ip, 0)
scope: Final[dict[str, object]] = {
"type": "http",
"method": "POST",
"path": "/mcp/sampling/createMessage",
"scheme": "http",
"server": (_server_host, _server_port),
"query_string": b"",
"root_path": "",
"headers": _scope_headers,
}
if _client_tuple is not None:
scope["client"] = _client_tuple
return Request(scope=scope)
async def _build_completion_kwargs(

36
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -58,9 +58,11 @@ from litellm.proxy._experimental.mcp_server.utils import (
LITELLM_MCP_SERVER_VERSION,
MCPMissingUserEnvVarsError,
add_server_prefix_to_name,
build_synthetic_mcp_request,
extract_mcp_tool_result_error_message,
get_server_prefix,
iter_known_server_prefixes,
logging_safe_mcp_headers,
match_known_tool_name,
)
from litellm.proxy._types import (
@ -860,11 +862,11 @@ if MCP_AVAILABLE:
name: str,
arguments: dict[str, object],
user_api_key_auth: UserAPIKeyAuth,
raw_headers: Mapping[str, str] | None = None,
client_ip: str | None = None,
) -> LiteLLMLoggingObj | None:
"""Run the pre-call pipeline (guardrails + logging setup) for a virtual
mcp_tool_call so the SSE path spend-logs like the REST path."""
from fastapi import Request
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
)
@ -874,13 +876,10 @@ if MCP_AVAILABLE:
proxy_logging_obj,
)
request: Final = Request(
scope={
"type": "http",
"method": "POST",
"path": "/mcp/tools/call",
"headers": [(b"content-type", b"application/json")],
}
request: Final = build_synthetic_mcp_request(
path="/mcp/tools/call",
raw_headers=raw_headers,
client_ip=client_ip,
)
_, virtual_logging_obj = await ProxyBaseLLMRequestProcessing(
data={"name": name, "arguments": arguments}
@ -952,7 +951,11 @@ if MCP_AVAILABLE:
assert user_api_key_auth is not None # guaranteed by the flag check above
virtual_logging_obj: Final = await _build_virtual_call_logging_obj(
name=name, arguments=args, user_api_key_auth=user_api_key_auth
name=name,
arguments=args,
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
client_ip=client_ip,
)
return await handle_mcp_tool_call(
tool_name=args.get("tool_name", ""),
@ -979,7 +982,6 @@ if MCP_AVAILABLE:
Raises:
HTTPException: If tool not found or arguments missing
"""
from fastapi import Request
from mcp.server.lowlevel.server import request_ctx
from mcp.types import CallToolResult
@ -1041,13 +1043,10 @@ if MCP_AVAILABLE:
body_data["litellm_trace_id"] = chain_id
body_data["litellm_session_id"] = chain_id
request: Final = Request(
scope={
"type": "http",
"method": "POST",
"path": "/mcp/tools/call",
"headers": [(b"content-type", b"application/json")],
}
request: Final = build_synthetic_mcp_request(
path="/mcp/tools/call",
raw_headers=raw_headers,
client_ip=_client_ip,
)
if user_api_key_auth is not None:
data = await add_litellm_data_to_request(
@ -1905,6 +1904,7 @@ if MCP_AVAILABLE:
"litellm_trace_id": effective_litellm_trace_id,
"metadata": {
"spend_logs_metadata": spend_logs_metadata,
"headers": logging_safe_mcp_headers(raw_headers),
**({"tags": request_tags} if request_tags else {}),
},
# Provide a small input payload for standard logging

147
litellm/proxy/_experimental/mcp_server/utils.py
View file

@ -7,6 +7,7 @@ import importlib
import json
import os
import re
import typing
from collections.abc import Iterable, Iterator, Mapping, MutableMapping, MutableSequence
from collections.abc import Set as AbstractSet
from typing import Any, Final, Protocol
@ -14,6 +15,9 @@ from urllib.parse import quote
from litellm.types.mcp_server.mcp_server_manager import MCPServer
if typing.TYPE_CHECKING:
from fastapi import Request
class _McpServerLike(Protocol):
@property
@ -862,3 +866,146 @@ def set_mcp_tool_result_structured_content(result: object, value: object) -> boo
return True
except (AttributeError, TypeError, ValueError):
return False
_HOP_BY_HOP_HEADERS: Final = frozenset(
{
"content-length",
"transfer-encoding",
"connection",
"keep-alive",
"upgrade",
"te",
"trailer",
}
)
_SYNTHETIC_REQUEST_EXCLUDED_HEADERS: Final = _HOP_BY_HOP_HEADERS | frozenset({"content-type", "x-forwarded-for"})
_SYNTHETIC_REQUEST_SERVER: Final = ("127.0.0.1", 4000)
_MCP_SERVER_AUTH_HEADER_PREFIX: Final = "x-mcp-"
def _custom_litellm_key_header_name() -> str | None:
"""``general_settings.litellm_key_header_name``, the deployment's custom header name for
the proxy virtual key, so it is stripped from observability copies like the standard ones."""
try:
from litellm.proxy.proxy_server import general_settings
except ImportError:
return None
return general_settings.get("litellm_key_header_name") if general_settings else None
def _mcp_client_side_auth_header_name() -> str:
"""The header name the client passes the upstream MCP credential in, falling back to the
default when ``general_settings`` is unavailable (the SDK, outside a running proxy)."""
from .auth.user_api_key_auth_mcp import MCPRequestHandler
try:
return MCPRequestHandler.get_mcp_client_side_auth_header_name()
except ImportError:
return MCPRequestHandler.LITELLM_MCP_AUTH_HEADER_NAME
def _upstream_credential_headers(header_names: Iterable[str]) -> frozenset[str]:
"""Lowercased names of the headers in ``header_names`` that carry an upstream MCP
credential rather than request context: the configured client side auth header and
the per-server ``x-mcp-{alias}-{header}`` family. ``clean_headers`` only knows the
credential headers of the chat completions path, so these are dropped on top of it.
"""
from .auth.user_api_key_auth_mcp import MCPRequestHandler
non_credential: Final = frozenset(
{
MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME.lower(),
MCPRequestHandler.LITELLM_MCP_ACCESS_GROUPS_HEADER_NAME.lower(),
}
)
client_side_auth: Final = _mcp_client_side_auth_header_name().lower()
return frozenset(
name
for name in (raw_name.lower() for raw_name in header_names)
if name == client_side_auth or (name.startswith(_MCP_SERVER_AUTH_HEADER_PREFIX) and name not in non_credential)
)
def build_synthetic_mcp_request(
*,
path: str,
raw_headers: Mapping[str, str] | None = None,
client_ip: str | None = None,
) -> "Request":
"""A synthetic FastAPI ``Request`` carrying the MCP connection's HTTP headers.
The MCP protocol transports do not hand a per-call ``Request`` to the tool
handlers, so one is reconstructed from the connection's ``raw_headers``. That
lets ``add_litellm_data_to_request`` derive ``metadata.headers``,
``proxy_server_request``, header-based tags, guardrails and trace correlation
exactly as on the chat completions path. Hop-by-hop headers describe the
original HTTP framing rather than the logical request, so they are dropped, and
``x-forwarded-for`` comes from the resolved ``client_ip`` to avoid spoofing. Upstream
MCP credentials and the deployment's proxy key header, including a custom
``litellm_key_header_name``, are dropped so they cannot reach a callback or a guardrail
through the derived metadata even when a caller omits ``general_settings``.
"""
from fastapi import Request
custom_key_header: Final = _custom_litellm_key_header_name()
excluded: Final = (
_SYNTHETIC_REQUEST_EXCLUDED_HEADERS
| _upstream_credential_headers(raw_headers.keys() if raw_headers else ())
| (frozenset({custom_key_header.lower()}) if custom_key_header else frozenset())
)
forwarded: Final = tuple(
(
name.lower().encode("latin-1", errors="replace"),
value.encode("utf-8", errors="replace"),
)
for name, value in (raw_headers.items() if raw_headers else ())
if name.lower() not in excluded
)
xff: Final = ((b"x-forwarded-for", client_ip.encode("utf-8")),) if client_ip else ()
return Request(
scope={
"type": "http",
"method": "POST",
"path": path,
"scheme": "http",
"server": _SYNTHETIC_REQUEST_SERVER,
"query_string": b"",
"root_path": "",
"headers": ((b"content-type", b"application/json"), *forwarded, *xff),
**({"client": (client_ip, 0)} if client_ip else {}),
}
)
def logging_safe_mcp_headers(raw_headers: Mapping[str, str] | None) -> Mapping[str, str]:
"""The MCP request's client headers, sanitized the way the chat completions path
sanitizes them before they reach a logging callback or a guardrail: proxy key
headers stripped, including the custom key header name the deployment configured,
upstream MCP credentials dropped, and credential-bearing values masked.
Client-controlled behaviour flags (``litellm-disable-message-redaction``) are dropped
too: these headers are read back out of the metadata to change proxy behaviour, so
leaving one in place would let any MCP client turn off the redaction an admin
configured. This path carries no key or team object to authorize an opt-out with, so
it always strips them."""
from starlette.datastructures import Headers
from litellm.proxy.litellm_pre_call_utils import (
UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS,
clean_headers,
redact_credential_headers,
)
excluded: Final = (
_upstream_credential_headers(raw_headers.keys() if raw_headers else ())
| UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS
)
cleaned: Final = clean_headers(
Headers(raw_headers),
litellm_key_header_name=_custom_litellm_key_header_name(),
)
return redact_credential_headers({name: value for name, value in cleaned.items() if name.lower() not in excluded})

16
litellm/proxy/_types.py
View file

@ -2512,6 +2512,22 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
None,
description="If True and LiteLLM_SpendLogs has been converted to a range-partitioned table (db_scripts/partition_spend_logs.sql), retention cleanup drops expired partitions instead of deleting rows, and pre-creates upcoming partitions. Default is False.",
)
maximum_spend_logs_cleanup_batch_size: int | None = Field(
None,
description="Rows deleted per DELETE statement by the spend log cleanup job. Defaults to 1000.",
)
maximum_spend_logs_cleanup_max_batches: int | None = Field(
None,
description="Maximum DELETE statements the spend log cleanup job issues per table per run. Defaults to 500.",
)
maximum_spend_logs_cleanup_run_budget: str | None = Field(
None,
description="Wall-clock budget for one spend log cleanup run (e.g. '5m'), shared across every table it prunes. A run that hits the budget stops and the next run resumes from where it left off. Defaults to '5m'.",
)
maximum_spend_logs_cleanup_batch_timeout: str | None = Field(
None,
description="Postgres statement_timeout and lock_timeout applied to each spend log cleanup delete batch (e.g. '30s'), so cleanup cannot hold row locks or a connection indefinitely. Defaults to '30s'.",
)
mcp_internal_ip_ranges: list[str] | None = Field(
None,
description="Custom CIDR ranges that define internal/private networks for MCP access control. When set, only these ranges are treated as internal. Defaults to RFC 1918 private ranges (10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16, 127.0.0.0/8).",

29
litellm/proxy/auth/auth_checks.py
View file

@ -14,6 +14,7 @@ import math
import re
import time
from collections.abc import Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast
from fastapi import HTTPException, Request, status
@ -376,6 +377,16 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None
zero_cost_cache[model_name] = False
return False
if _has_ptu_flat_cost(model_name, llm_router):
verbose_proxy_logger.debug(
"Model %s prices reserved PTU capacity as a flat cost, so its zero per-token "
"rate is not a free model (enforce budget)",
safe_name,
)
if zero_cost_cache is not None:
zero_cost_cache[model_name] = False
return False
verbose_proxy_logger.debug(
"Model %s has zero cost explicitly configured (input: %s, output: %s)",
safe_name,
@ -394,6 +405,24 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None
return True
_NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({})
def _has_ptu_flat_cost(model: str, llm_router: "Router") -> bool:
"""Whether any deployment in the model group bills reserved PTU capacity as a flat cost.
Such a deployment carries an explicit zero per-token price so the flat cost is not charged
twice, which otherwise reads here as a free model and waives every budget check for it.
"""
for deployment in llm_router.model_list:
if deployment.get("model_name") != model:
continue
model_info = deployment.get("model_info") or _NO_MODEL_INFO
if model_info.get("ptu_count") is not None and model_info.get("cost_per_ptu_per_hour") is not None:
return True
return False
def _is_cost_explicitly_configured(model: str, llm_router: "Router") -> bool:
"""
Check if any deployment in the model group has cost fields explicitly

83
litellm/proxy/common_utils/callback_utils.py
View file

@ -1,7 +1,10 @@
import copy
import os
from collections.abc import Callable, Iterable
from typing import TYPE_CHECKING, Any, Final, Optional
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias
from typing_extensions import assert_never
import litellm
from litellm import get_secret
@ -50,6 +53,66 @@ if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
@dataclass(frozen=True, slots=True)
class _CallbackResolvedToClass:
entry: str
loaded: type
tag: Literal["resolved_to_class"] = "resolved_to_class"
@dataclass(frozen=True, slots=True)
class _CallbackNotDispatchable:
entry: str
loaded: object
tag: Literal["not_dispatchable"] = "not_dispatchable"
_CallbackLoadError: TypeAlias = _CallbackResolvedToClass | _CallbackNotDispatchable
def _classify_loaded_callback(entry: str, loaded: object) -> CustomLogger | Callable[..., object] | _CallbackLoadError:
"""
Decide whether what a ``litellm_settings.callbacks`` dotted path resolved to can be dispatched.
A dotted path only ever runs as a ``CustomLogger`` instance or as a callback function. Anything
else (most commonly a class instead of an instance) used to load without complaint and then be
skipped on every request, with no log line and no error.
"""
if isinstance(loaded, CustomLogger) or (callable(loaded) and not isinstance(loaded, type)):
return loaded
if isinstance(loaded, type):
return _CallbackResolvedToClass(entry=entry, loaded=loaded)
return _CallbackNotDispatchable(entry=entry, loaded=loaded)
def _raise_callback_load_error(error: _CallbackLoadError) -> NoReturn:
"""The one edge that raises: map a load error onto config load's failure contract."""
match error:
case _CallbackResolvedToClass():
module_path: Final = error.entry.rsplit(".", 1)[0] if "." in error.entry else error.entry
raise ValueError(
f"litellm_settings.callbacks entry '{error.entry}' resolved to the class "
f"{error.loaded.__module__}.{error.loaded.__qualname__}, which is neither a "
"CustomLogger instance nor a callable, so the proxy would never run it."
f" Point it at an instance instead, e.g. add `proxy_handler_instance = {error.loaded.__name__}()` to "
f'{module_path} and set `callbacks: ["{module_path}.proxy_handler_instance"]`.'
)
case _CallbackNotDispatchable():
raise ValueError(
f"litellm_settings.callbacks entry '{error.entry}' resolved to "
f"{type(error.loaded).__name__} {error.loaded!r}, which is neither a "
"CustomLogger instance nor a callable, so the proxy would never run it."
)
assert_never(error)
def _loaded_callback_or_raise(entry: str, loaded: object) -> CustomLogger | Callable[..., object]:
resolved: Final = _classify_loaded_callback(entry=entry, loaded=loaded)
if isinstance(resolved, _CallbackResolvedToClass | _CallbackNotDispatchable):
_raise_callback_load_error(resolved)
return resolved
def initialize_callbacks_on_proxy(
value: Any,
premium_user: bool,
@ -305,9 +368,12 @@ def initialize_callbacks_on_proxy(
"%s attempting to import custom calback=%s %s", blue_color_code, callback, reset_color_code
)
imported_list.append(
get_instance_fn(
value=callback,
config_file_path=config_file_path,
_loaded_callback_or_raise(
entry=callback,
loaded=get_instance_fn(
value=callback,
config_file_path=config_file_path,
),
)
)
if isinstance(litellm.callbacks, list):
@ -321,9 +387,12 @@ def initialize_callbacks_on_proxy(
PrometheusLogger._mount_metrics_endpoint()
else:
litellm.callbacks = [
get_instance_fn(
value=value,
config_file_path=config_file_path,
_loaded_callback_or_raise(
entry=value,
loaded=get_instance_fn(
value=value,
config_file_path=config_file_path,
),
)
]
verbose_proxy_logger.debug("%s Initialized Callbacks - %s %s", blue_color_code, litellm.callbacks, reset_color_code)

453
litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py
View file

@ -1,22 +1,60 @@
import asyncio
import time
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Final
from typing import Final, Literal, TypeAlias
from pydantic import BaseModel, TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.caching import RedisCache
from litellm.constants import (
SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS,
SPEND_LOG_CLEANUP_BATCH_SIZE,
SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS,
SPEND_LOG_CLEANUP_JOB_NAME,
SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES,
SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP,
SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS,
SPEND_LOG_RUN_LOOPS,
)
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import (
RunOutcome,
SpendLogCleanupMetrics,
)
from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import (
RemainingTimeoutMs,
SpendLogsPartitionManager,
)
from litellm.proxy.utils import PrismaClient
StopReason: TypeAlias = Literal["exhausted", "budget_exhausted", "batch_cap_reached", "aborted"]
@dataclass(frozen=True, slots=True)
class TableCleanupResult:
"""Outcome of pruning one table, so the caller can report why a run ended."""
rows_deleted: int
stop_reason: StopReason
class _RemainingRow(BaseModel):
"""One row of the capped outstanding-rows probe, validated out of prisma's untyped result."""
remaining: int
_REMAINING_ROWS: Final = TypeAdapter(list[_RemainingRow])
SPEND_LOG_CLEANUP_BOUND_SETTINGS: Final = (
"maximum_spend_logs_cleanup_batch_size",
"maximum_spend_logs_cleanup_max_batches",
"maximum_spend_logs_cleanup_run_budget",
"maximum_spend_logs_cleanup_batch_timeout",
)
class SpendLogCleanup:
"""
@ -26,6 +64,24 @@ class SpendLogCleanup:
dropping whole partitions (instant, frees disk immediately). Otherwise it
falls back to deleting logs in batches.
Uses PodLockManager to ensure only one pod runs cleanup in multi-pod deployments.
Every run is bounded so it can never monopolise the database: a wall-clock
budget shared across all tables, a per-table batch cap, and a Postgres
statement/lock timeout on every statement the job issues, deletes and the
outstanding-rows probe alike. A run that hits a bound stops cleanly and the
next run resumes from where it left off, because the cutoff is recomputed
and deleted rows are gone.
The budget is a hard wall clock, not an advisory one. Every statement this
job issues, deletes, the outstanding-rows probe and partition DDL alike, is
issued with a timeout clamped to the budget that is still left, so one
started just under the deadline is cancelled by Postgres at the deadline
rather than running a further batch timeout past it. No statement is issued
at all once the budget is spent, which is why the probe is skipped on that
path. Partition DDL additionally carries a lock_timeout, because it takes an
ACCESS EXCLUSIVE lock and would otherwise queue behind a long-running reader
for as long as that reader lives; a partition this run cannot get is left
for the next one.
"""
def __init__(
@ -34,17 +90,88 @@ class SpendLogCleanup:
redis_cache: RedisCache | None = None,
partition_manager: SpendLogsPartitionManager | None = None,
):
self.batch_size = SPEND_LOG_CLEANUP_BATCH_SIZE
self.retention_seconds: int | None = None
self.partition_manager = partition_manager or SpendLogsPartitionManager()
from litellm.proxy.proxy_server import general_settings as default_settings
self.general_settings = general_settings or default_settings
self._refresh_bounds()
from litellm.proxy.proxy_server import proxy_logging_obj
pod_lock_manager: Final = proxy_logging_obj.db_spend_update_writer.pod_lock_manager
self.pod_lock_manager = pod_lock_manager
verbose_proxy_logger.info("SpendLogCleanup initialized with batch size: %s", self.batch_size)
verbose_proxy_logger.info(
"SpendLogCleanup initialized: batch_size=%s max_batches=%s run_budget=%ss batch_timeout=%ss",
self.batch_size,
self.max_batches,
self.run_budget_seconds,
self.batch_timeout_seconds,
)
def _refresh_bounds(self) -> None:
"""
Re-read every bound in SPEND_LOG_CLEANUP_BOUND_SETTINGS from settings.
The scheduler holds one long-lived instance, so a bound captured at
construction would never reflect a dashboard change. general_settings is
the same dict the periodic config reload mutates in place, so reading it
per run is what makes these knobs live. Every bound falls back to its
shipped default, so clearing a field restores that default.
"""
self.batch_size: int = self._positive_int_setting(
"maximum_spend_logs_cleanup_batch_size", SPEND_LOG_CLEANUP_BATCH_SIZE
)
self.max_batches: int = self._positive_int_setting(
"maximum_spend_logs_cleanup_max_batches", SPEND_LOG_RUN_LOOPS
)
self.run_budget_seconds: float = self._duration_setting(
"maximum_spend_logs_cleanup_run_budget", SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS
)
self.batch_timeout_seconds: float = self._duration_setting(
"maximum_spend_logs_cleanup_batch_timeout", SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS
)
def _positive_int_setting(self, setting_name: str, default: int) -> int:
"""
Read a positive-integer knob, falling back to the default when unset or unusable.
"""
raw: Final = self.general_settings.get(setting_name)
if raw is None:
return default
try:
parsed: Final = int(raw)
except (TypeError, ValueError):
verbose_proxy_logger.warning("Invalid %s value: %s, using default %s", setting_name, raw, default)
return default
if parsed <= 0:
verbose_proxy_logger.warning("%s must be positive, got %s, using default %s", setting_name, parsed, default)
return default
return parsed
def _duration_setting(self, setting_name: str, default_seconds: float) -> float:
"""
Read a duration knob (e.g. '5m'), falling back to the default when unset or unusable.
The knob must never be able to remove the bound it exists to enforce, so
anything the parser rejects (including the non-finite spellings 'inf' and
'nan') and anything non-positive falls back rather than being honoured.
"""
raw: Final = self.general_settings.get(setting_name)
if raw is None:
return default_seconds
try:
parsed: Final = float(duration_in_seconds(str(raw)))
except (ValueError, TypeError) as e:
verbose_proxy_logger.warning(
"Invalid %s value: %s (%s), using default %ss", setting_name, raw, e, default_seconds
)
return default_seconds
if parsed <= 0:
verbose_proxy_logger.warning(
"%s must be a positive duration, got %s, using default %ss", setting_name, raw, default_seconds
)
return default_seconds
return parsed
def _retention_seconds_for(self, setting_name: str) -> int | None:
"""
@ -78,6 +205,91 @@ class SpendLogCleanup:
self.retention_seconds = self._retention_seconds_for("maximum_spend_logs_retention_period")
return self.retention_seconds is not None
def _timeout_ms(self, deadline: float) -> int:
"""
The per-statement bound in milliseconds: the batch timeout, or whatever
is left of the run budget, whichever is smaller.
Clamping to the remaining budget is what makes the budget a real
wall-clock bound rather than an advisory one. Postgres offers no "stop
at time T", only a per-statement duration, so a statement issued just
under the deadline would otherwise run a full batch timeout past it, and
with several tables those overruns stack.
Interpolating this into SQL is safe by construction: an int cannot carry
SQL, and SET does not accept a bind parameter.
"""
remaining_ms: Final = int((deadline - time.monotonic()) * 1000)
return max(1, min(int(self.batch_timeout_seconds * 1000), remaining_ms))
def _remaining_timeout_ms(self, deadline: float) -> RemainingTimeoutMs:
"""
The per-statement bound for work this job delegates, as a callable.
Partition maintenance issues one statement per partition, so handing it a
number would bound each statement by the budget that was left before the
FIRST one and never by what remains. Re-evaluating per statement is what
makes the loop itself bounded, and None tells the callee to stop rather
than issue a statement it has no budget for.
"""
def remaining() -> int | None:
return None if time.monotonic() >= deadline else self._timeout_ms(deadline)
return remaining
async def _execute_delete_batch(
self, prisma_client: PrismaClient, delete_sql: str, cutoff_date: datetime, deadline: float
) -> int | None:
"""
Run one delete batch under a Postgres statement and lock timeout.
The timeouts are what actually bound the work: a Prisma transaction
timeout cannot interrupt a statement that is already executing, so
without these a single batch blocked behind a lock would hold its
connection, and the row locks it already took, indefinitely. SET LOCAL
scopes both to this transaction so the pooled connection is unaffected.
Returns the row count, or None when the driver returned something that
is not a row count. That is a contract violation rather than a transient
fault, so the caller stops instead of retrying.
"""
timeout_ms: Final = self._timeout_ms(deadline)
async with prisma_client.db.tx() as tx:
await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}")
await tx.execute_raw(f"SET LOCAL lock_timeout = {timeout_ms}")
deleted_result: Final = await tx.execute_raw(delete_sql, cutoff_date, self.batch_size)
return deleted_result if isinstance(deleted_result, int) else None
async def _count_remaining(
self, prisma_client: PrismaClient, cutoff_date: datetime, table_name: str, time_column: str, deadline: float
) -> int | None:
"""
Count expired rows still outstanding, stopping at a cap.
An uncapped COUNT(*) over an expired backlog would itself be the kind of
long scan this job exists to avoid, so the probe reads at most
SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP index entries. A result equal to
the cap means "at least this many".
"""
count_sql: Final = f"""
SELECT count(*)::int AS remaining FROM (
SELECT 1 FROM "{table_name}"
WHERE "{time_column}" < $1::timestamptz
LIMIT $2
) capped
"""
try:
async with prisma_client.db.tx() as tx:
await tx.execute_raw(f"SET LOCAL statement_timeout = {self._timeout_ms(deadline)}")
rows: Final = _REMAINING_ROWS.validate_python(
await tx.query_raw(count_sql, cutoff_date, SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP)
)
except Exception as e: # noqa: BLE001 - an observability probe must never fail the cleanup run
verbose_proxy_logger.warning("Could not count remaining %s rows: %s", table_name, e)
return None
return rows[0].remaining if rows else None
async def _delete_old_rows_batched(
self,
prisma_client: PrismaClient,
@ -85,10 +297,14 @@ class SpendLogCleanup:
table_name: str,
key_columns: tuple[str, ...],
time_column: str,
) -> int:
deadline: float,
) -> TableCleanupResult:
"""
Helper method to delete a table's rows older than the cutoff in batches.
Returns the total number of rows deleted.
Delete a table's rows older than the cutoff in batches.
Stops at whichever bound is reached first: the backlog running out, the
shared wall-clock deadline, the per-table batch cap, or too many
consecutive batch failures.
"""
key_list: Final = ", ".join(f'"{col}"' for col in key_columns)
delete_sql: Final = f"""
@ -103,23 +319,46 @@ class SpendLogCleanup:
run_count = 0
consecutive_failures = 0
while True:
if run_count > SPEND_LOG_RUN_LOOPS:
if time.monotonic() >= deadline:
verbose_proxy_logger.info(
"Run budget exhausted during %s cleanup after %d rows; the next run resumes from here",
table_name,
total_deleted,
)
return await self._finish_table(
prisma_client, cutoff_date, table_name, time_column, total_deleted, "budget_exhausted", deadline
)
if run_count >= self.max_batches:
verbose_proxy_logger.info(
"Max batches reached for %s cleanup, remaining rows will be deleted in next run", table_name
)
break
# Step 1: Find rows and delete them in one go without fetching to application
# Delete in batches, limited by self.batch_size
try:
deleted_result = await prisma_client.db.execute_raw(
delete_sql,
cutoff_date,
self.batch_size,
return await self._finish_table(
prisma_client, cutoff_date, table_name, time_column, total_deleted, "batch_cap_reached", deadline
)
# Find rows and delete them in one go without fetching to application
batch_started_at = time.monotonic()
try:
batch_result = await self._execute_delete_batch(prisma_client, delete_sql, cutoff_date, deadline)
except Exception as batch_exc:
if time.monotonic() >= deadline:
# The statement timeout was clamped to the budget that was
# left, so this batch was cancelled by the deadline itself.
# That is the bound working, not a database fault, and
# counting it would both inflate the failure metric and push
# every budget-exhausted run toward the abort threshold.
verbose_proxy_logger.info(
"Run budget exhausted mid-batch during %s cleanup after %d rows; "
"the next run resumes from here",
table_name,
total_deleted,
)
return await self._finish_table(
prisma_client, cutoff_date, table_name, time_column, total_deleted, "budget_exhausted", deadline
)
# A single batch failure (e.g. Prisma/DB timeout) must not abort
# the whole run — subsequent batches may still succeed.
consecutive_failures += 1
SpendLogCleanupMetrics.record_batch_failure(table_name)
verbose_proxy_logger.exception(
"%s cleanup batch failed "
"(run_count=%d, consecutive_failures=%d, batch_size=%d, "
@ -140,28 +379,31 @@ class SpendLogCleanup:
consecutive_failures,
total_deleted,
)
break
return await self._finish_table(
prisma_client, cutoff_date, table_name, time_column, total_deleted, "aborted", deadline
)
await asyncio.sleep(SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS)
continue
consecutive_failures = 0
deleted_count = 0
if isinstance(deleted_result, int):
deleted_count = deleted_result
else:
if batch_result is None:
verbose_proxy_logger.error(
"Unexpected execute_raw return type for %s cleanup: %s; aborting cleanup to avoid infinite loop",
"Unexpected execute_raw return type for %s cleanup; aborting cleanup to avoid infinite loop",
table_name,
type(deleted_result),
)
break
return await self._finish_table(
prisma_client, cutoff_date, table_name, time_column, total_deleted, "aborted", deadline
)
consecutive_failures = 0
deleted_count = batch_result
SpendLogCleanupMetrics.record_batch(table_name, deleted_count, time.monotonic() - batch_started_at)
verbose_proxy_logger.info("Deleted %s %s rows in this batch", deleted_count, table_name)
if deleted_count == 0:
verbose_proxy_logger.info("No more %s rows to delete. Total deleted: %s", table_name, total_deleted)
break
return await self._finish_table(
prisma_client, cutoff_date, table_name, time_column, total_deleted, "exhausted", deadline
)
total_deleted += deleted_count
run_count += 1
@ -169,18 +411,49 @@ class SpendLogCleanup:
# Add a small sleep to prevent overwhelming the database
await asyncio.sleep(0.1)
return total_deleted
async def _finish_table(
self,
prisma_client: PrismaClient,
cutoff_date: datetime,
table_name: str,
time_column: str,
rows_deleted: int,
stop_reason: StopReason,
deadline: float,
) -> TableCleanupResult:
"""
Publish how much of this table is still outstanding, then report the run's result.
async def _delete_old_logs(self, prisma_client: PrismaClient, cutoff_date: datetime) -> int:
The probe is skipped once the budget is spent. It is the one piece of
work that would otherwise be ISSUED after the deadline, and every table
exits through here, including the ones a spent run never started, so
keeping it would put one more statement per table past the bound. A run
that ends this way already reports "budget_exhausted", which tells an
operator the backlog was not drained; the gauge simply keeps its value
from the last run that finished inside its budget.
"""
if time.monotonic() >= deadline:
return TableCleanupResult(rows_deleted=rows_deleted, stop_reason=stop_reason)
remaining: Final = await self._count_remaining(prisma_client, cutoff_date, table_name, time_column, deadline)
if remaining is not None:
SpendLogCleanupMetrics.set_rows_remaining(table_name, remaining)
return TableCleanupResult(rows_deleted=rows_deleted, stop_reason=stop_reason)
async def _delete_old_logs(
self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float
) -> TableCleanupResult:
return await self._delete_old_rows_batched(
prisma_client,
cutoff_date,
table_name="LiteLLM_SpendLogs",
key_columns=("request_id", "startTime"),
time_column="startTime",
deadline=deadline,
)
async def _delete_old_tool_index_rows(self, prisma_client: PrismaClient, cutoff_date: datetime) -> int:
async def _delete_old_tool_index_rows(
self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float
) -> TableCleanupResult:
# SpendLogToolIndex rows are derived from spend logs, so they expire on the
# same cutoff; rows older than retention point at already-deleted logs.
return await self._delete_old_rows_batched(
@ -189,17 +462,87 @@ class SpendLogCleanup:
table_name="LiteLLM_SpendLogToolIndex",
key_columns=("request_id", "tool_name"),
time_column="start_time",
deadline=deadline,
)
async def _delete_old_autorouter_session_rows(self, prisma_client: PrismaClient, cutoff_date: datetime) -> int:
async def _delete_old_autorouter_session_rows(
self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float
) -> TableCleanupResult:
return await self._delete_old_rows_batched(
prisma_client,
cutoff_date,
table_name="LiteLLM_AutoRouterSession",
key_columns=("api_key", "session_id", "router_name"),
time_column="last_turn_at",
deadline=deadline,
)
async def _clean_spend_log_tables(
self, prisma_client: PrismaClient, deadline: float
) -> tuple[TableCleanupResult, ...]:
"""
Prune the spend logs and the tool index rows derived from them.
When the table is range-partitioned, whole expired partitions are dropped
first because that reclaims disk immediately. Expired rows can still sit in
the DEFAULT partition (backfill, coverage gaps) or in a partition that spans
the cutoff, so retention still deletes those stragglers row-wise.
"""
cutoff_date: Final = datetime.now(timezone.utc) - timedelta(seconds=float(self.retention_seconds or 0))
verbose_proxy_logger.info("Removing logs older than %s", cutoff_date.isoformat())
# Partition maintenance is DDL taking an ACCESS EXCLUSIVE lock, so it is
# only STARTED while the run still has budget, and each statement carries
# the same timeouts the batches do. Without those, a DROP would queue
# behind any long-running reader for as long as that reader lives, which
# is the one way this job could still outlast its budget without bound.
remaining_timeout_ms: Final = self._remaining_timeout_ms(deadline)
if time.monotonic() >= deadline:
verbose_proxy_logger.info("Run budget already spent, skipping partition maintenance this run")
elif self.general_settings.get(
"use_spend_logs_partitioning", False
) and await self.partition_manager.is_partitioned(prisma_client, remaining_timeout_ms):
await self.partition_manager.ensure_partitions(prisma_client, remaining_timeout_ms)
dropped: Final = await self.partition_manager.drop_partitions_older_than(
prisma_client, cutoff_date, remaining_timeout_ms
)
verbose_proxy_logger.info("Dropped %d expired spend-log partitions: %s", len(dropped), dropped)
logs_result: Final = await self._delete_old_logs(prisma_client, cutoff_date, deadline)
verbose_proxy_logger.info("Deleted %s logs", logs_result.rows_deleted)
index_result: Final = await self._delete_old_tool_index_rows(prisma_client, cutoff_date, deadline)
verbose_proxy_logger.info("Deleted %s expired tool index rows", index_result.rows_deleted)
return (logs_result, index_result)
async def _clean_session_rollup(
self, prisma_client: PrismaClient, retention_seconds: int, deadline: float
) -> tuple[TableCleanupResult, ...]:
"""
Prune auto-router session rollup rows, which carry their own retention horizon.
"""
session_cutoff: Final = datetime.now(timezone.utc) - timedelta(seconds=float(retention_seconds))
sessions_result: Final = await self._delete_old_autorouter_session_rows(prisma_client, session_cutoff, deadline)
verbose_proxy_logger.info("Deleted %s expired auto-router session rollup rows", sessions_result.rows_deleted)
return (sessions_result,)
@staticmethod
def _run_outcome(results: tuple[TableCleanupResult, ...]) -> RunOutcome:
"""
Report the most operationally significant reason the run stopped.
A bound that was hit matters more than a table that simply ran dry, so
those win over "completed", and an abort wins over everything.
"""
reasons: Final = frozenset(result.stop_reason for result in results)
if "aborted" in reasons:
return "aborted"
if "budget_exhausted" in reasons:
return "budget_exhausted"
if "batch_cap_reached" in reasons:
return "batch_cap_reached"
return "completed"
async def cleanup_old_spend_logs(self, prisma_client: PrismaClient) -> None:
"""
Main cleanup function. Deletes old spend logs in batches.
@ -209,16 +552,19 @@ class SpendLogCleanup:
lock_acquired = False
try:
verbose_proxy_logger.info("Cleanup job triggered at %s", datetime.now())
self._refresh_bounds()
delete_spend_logs: Final = self._should_delete_spend_logs()
autorouter_retention_seconds: Final = self._retention_seconds_for(
"maximum_autorouter_session_retention_period"
)
if not delete_spend_logs and autorouter_retention_seconds is None:
SpendLogCleanupMetrics.record_run("skipped_disabled")
return
if delete_spend_logs and self.retention_seconds is None:
verbose_proxy_logger.error("Retention seconds is None, cannot proceed with cleanup")
SpendLogCleanupMetrics.record_run("skipped_disabled")
return
# If we have a pod lock manager, try to acquire the lock
@ -235,43 +581,23 @@ class SpendLogCleanup:
if not lock_acquired:
verbose_proxy_logger.info("Another pod is already running cleanup")
SpendLogCleanupMetrics.record_run("skipped_locked")
return
if delete_spend_logs and self.retention_seconds is not None:
cutoff_date: Final = datetime.now(timezone.utc) - timedelta(seconds=float(self.retention_seconds))
verbose_proxy_logger.info("Removing logs older than %s", cutoff_date.isoformat())
deadline: Final = time.monotonic() + self.run_budget_seconds
if self.general_settings.get(
"use_spend_logs_partitioning", False
) and await self.partition_manager.is_partitioned(prisma_client):
await self.partition_manager.ensure_partitions(prisma_client)
dropped: Final = await self.partition_manager.drop_partitions_older_than(prisma_client, cutoff_date)
verbose_proxy_logger.info(
"Dropped %d expired spend-log partitions: %s",
len(dropped),
dropped,
)
# DROP only reclaims whole expired partitions. Expired rows can
# still sit in the DEFAULT partition (backfill, coverage gaps)
# or in a partition that spans the cutoff, so retention must
# also delete those stragglers row-wise.
total_deleted = await self._delete_old_logs(prisma_client, cutoff_date)
verbose_proxy_logger.info(
"Deleted %s expired logs not covered by dropped partitions", total_deleted
)
else:
total_deleted = await self._delete_old_logs(prisma_client, cutoff_date)
verbose_proxy_logger.info("Deleted %s logs", total_deleted)
spend_log_results: Final = (
await self._clean_spend_log_tables(prisma_client, deadline)
if delete_spend_logs and self.retention_seconds is not None
else ()
)
session_results: Final = (
await self._clean_session_rollup(prisma_client, autorouter_retention_seconds, deadline)
if autorouter_retention_seconds is not None
else ()
)
index_deleted: Final = await self._delete_old_tool_index_rows(prisma_client, cutoff_date)
verbose_proxy_logger.info("Deleted %s expired tool index rows", index_deleted)
if autorouter_retention_seconds is not None:
session_cutoff: Final = datetime.now(timezone.utc) - timedelta(
seconds=float(autorouter_retention_seconds)
)
sessions_deleted: Final = await self._delete_old_autorouter_session_rows(prisma_client, session_cutoff)
verbose_proxy_logger.info("Deleted %s expired auto-router session rollup rows", sessions_deleted)
SpendLogCleanupMetrics.record_run(self._run_outcome(spend_log_results + session_results))
except Exception as e:
# .exception() captures the traceback; str(e) alone on a Prisma/DB
@ -281,6 +607,7 @@ class SpendLogCleanup:
type(e).__name__,
e,
)
SpendLogCleanupMetrics.record_run("aborted")
return # Return after error handling
finally:
# Only release the lock if it was actually acquired

122
litellm/proxy/db/db_transaction_queue/spend_log_cleanup_metrics.py Normal file
View file

@ -0,0 +1,122 @@
"""
Prometheus metrics for the spend-log retention cleanup job.
The job runs in the background on a single elected pod, so its cost is invisible
from request-path metrics. These instruments make a run's database footprint
observable: how much it deleted, how long each batch took, how much work is
still outstanding, and why a run stopped.
``prometheus_client`` is an optional dependency, so every recorder degrades to a
no-op when it is absent.
"""
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
from litellm._logging import verbose_proxy_logger
if TYPE_CHECKING:
# aliased so the annotations below cannot be mistaken for collections.Counter
from prometheus_client import Counter as PrometheusCounter
from prometheus_client import Gauge as PrometheusGauge
from prometheus_client import Histogram as PrometheusHistogram
RunOutcome: TypeAlias = Literal[
"completed",
"budget_exhausted",
"batch_cap_reached",
"skipped_locked",
"skipped_disabled",
"aborted",
]
_BATCH_DURATION_BUCKETS: Final = (0.005, 0.025, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0)
_TABLE_LABEL: Final = ("table",)
_OUTCOME_LABEL: Final = ("outcome",)
class SpendLogCleanupMetrics:
"""
Lazily-registered Prometheus instruments for the retention cleanup job.
Registration is deferred to first use so that importing this module never
touches the Prometheus registry, which keeps it safe to import from the
proxy regardless of whether Prometheus is a configured callback.
"""
_initialized: bool = False
rows_deleted: "PrometheusCounter | None" = None
batch_duration: "PrometheusHistogram | None" = None
rows_remaining: "PrometheusGauge | None" = None
batch_failures: "PrometheusCounter | None" = None
runs: "PrometheusCounter | None" = None
@classmethod
def _ensure_initialized(cls) -> None:
if cls._initialized:
return
cls._initialized = True
try:
# prometheus_client is an optional extra, so it is resolved here rather
# than at module import: this module is reachable from proxy startup
# regardless of whether Prometheus is a configured callback.
from prometheus_client import Counter, Gauge, Histogram
cls.rows_deleted = Counter(
"litellm_spend_log_cleanup_rows_deleted_total",
"Rows deleted by the spend-log retention cleanup job",
labelnames=_TABLE_LABEL,
)
cls.batch_duration = Histogram(
"litellm_spend_log_cleanup_batch_duration_seconds",
"Wall-clock duration of one retention cleanup delete batch",
labelnames=_TABLE_LABEL,
buckets=_BATCH_DURATION_BUCKETS,
)
cls.rows_remaining = Gauge(
"litellm_spend_log_cleanup_rows_remaining",
"Expired rows still awaiting deletion, counted only up to "
"SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP so the probe itself cannot scan a "
"large table; a value equal to that cap means at least that many remain",
labelnames=_TABLE_LABEL,
multiprocess_mode="livemax",
)
cls.batch_failures = Counter(
"litellm_spend_log_cleanup_batch_failures_total",
"Retention cleanup delete batches that raised",
labelnames=_TABLE_LABEL,
)
cls.runs = Counter(
"litellm_spend_log_cleanup_runs_total",
"Retention cleanup runs, labelled by why the run ended",
labelnames=_OUTCOME_LABEL,
)
except Exception as e: # noqa: BLE001 - a metrics problem must never fail the cleanup run
# Covers the extra being absent, a duplicate registration (repeated
# imports under a test runner), and registry misconfiguration alike.
verbose_proxy_logger.warning("Could not register spend-log cleanup metrics: %s", e)
@classmethod
def record_batch(cls, table_name: str, rows_deleted: int, duration_seconds: float) -> None:
cls._ensure_initialized()
if cls.rows_deleted is not None:
cls.rows_deleted.labels(table=table_name).inc(rows_deleted)
if cls.batch_duration is not None:
cls.batch_duration.labels(table=table_name).observe(duration_seconds)
@classmethod
def record_batch_failure(cls, table_name: str) -> None:
cls._ensure_initialized()
if cls.batch_failures is not None:
cls.batch_failures.labels(table=table_name).inc()
@classmethod
def set_rows_remaining(cls, table_name: str, remaining: int) -> None:
cls._ensure_initialized()
if cls.rows_remaining is not None:
cls.rows_remaining.labels(table=table_name).set(remaining)
@classmethod
def record_run(cls, outcome: RunOutcome) -> None:
cls._ensure_initialized()
if cls.runs is not None:
cls.runs.labels(outcome=outcome).inc()

128
litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py
View file

@ -14,8 +14,9 @@ keeps the batched-DELETE path, so existing deployments are untouched.
"""
import re
from collections.abc import Callable
from datetime import date, datetime, timedelta, timezone
from typing import Final
from typing import TYPE_CHECKING, Final, TypeAlias
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
@ -23,8 +24,23 @@ from litellm.constants import (
SPEND_LOG_PARTITION_PRECREATE_AHEAD,
)
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
SPEND_LOGS_TABLE: Final = "LiteLLM_SpendLogs"
RemainingTimeoutMs: TypeAlias = Callable[[], "int | None"]
"""
The per-statement bound in milliseconds, or None once the caller's budget is
spent.
Injected rather than passed as a number so it is re-evaluated before EVERY
statement: a value read once at entry would let a loop issue N statements each
bounded by the budget that was left before the first of them, which is not a
bound on the loop at all. The caller owns the policy; this module only asks how
much time it may still use.
"""
PartitionInterval = str # "day" | "week" | "month"
VALID_PARTITION_INTERVALS: Final = {"day", "week", "month"}
@ -116,21 +132,26 @@ class SpendLogsPartitionManager:
self.interval = interval
self.precreate_ahead = precreate_ahead
async def is_partitioned(self, prisma_client) -> bool:
async def is_partitioned(self, prisma_client: "PrismaClient", remaining_timeout_ms: RemainingTimeoutMs) -> bool:
budget_ms: Final = remaining_timeout_ms()
if budget_ms is None:
return False
try:
rows: Final = await prisma_client.db.query_raw(
"""
SELECT EXISTS (
SELECT 1
FROM pg_partitioned_table pt
JOIN pg_class c ON c.oid = pt.partrelid
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE c.relname = $1
AND n.nspname = current_schema()
) AS partitioned
""",
SPEND_LOGS_TABLE,
)
async with prisma_client.db.tx() as tx:
await tx.execute_raw(f"SET LOCAL statement_timeout = {budget_ms}")
rows: Final = await tx.query_raw(
"""
SELECT EXISTS (
SELECT 1
FROM pg_partitioned_table pt
JOIN pg_class c ON c.oid = pt.partrelid
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE c.relname = $1
AND n.nspname = current_schema()
) AS partitioned
""",
SPEND_LOGS_TABLE,
)
except Exception as e:
verbose_proxy_logger.warning(
"Could not determine if %s is partitioned, assuming it is not: %s",
@ -140,7 +161,25 @@ class SpendLogsPartitionManager:
return False
return bool(rows and rows[0].get("partitioned"))
async def ensure_partitions(self, prisma_client) -> list[str]:
@staticmethod
async def _execute_bounded_ddl(prisma_client: "PrismaClient", statement: str, timeout_ms: int) -> None:
"""
Run one DDL statement under a Postgres statement and lock timeout.
Partition DDL takes an ACCESS EXCLUSIVE lock, so an unbounded statement
queues behind any long-running reader for as long as that reader lives,
and the caller's run budget cannot cut it short. lock_timeout bounds the
wait for the lock and statement_timeout bounds the work itself, so a
partition this run cannot get is simply left for the next one.
"""
async with prisma_client.db.tx() as tx:
await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}")
await tx.execute_raw(f"SET LOCAL lock_timeout = {timeout_ms}")
await tx.execute_raw(statement)
async def ensure_partitions(
self, prisma_client: "PrismaClient", remaining_timeout_ms: RemainingTimeoutMs
) -> list[str]:
"""
Ensure the current and upcoming partitions exist, returning the names
now present. CREATE TABLE IF NOT EXISTS is a no-op for partitions that
@ -150,42 +189,61 @@ class SpendLogsPartitionManager:
for name, lower, upper in upcoming_partitions(
datetime.now(timezone.utc).date(), self.interval, self.precreate_ahead
):
budget_ms = remaining_timeout_ms()
if budget_ms is None:
verbose_proxy_logger.info("Run budget spent, leaving the remaining partitions for the next run")
break
try:
await prisma_client.db.execute_raw(
await self._execute_bounded_ddl(
prisma_client,
f'CREATE TABLE IF NOT EXISTS "{name}" '
f'PARTITION OF "{SPEND_LOGS_TABLE}" '
f"FOR VALUES FROM ('{lower.isoformat()}') TO ('{upper.isoformat()}')"
f"FOR VALUES FROM ('{lower.isoformat()}') TO ('{upper.isoformat()}')",
budget_ms,
)
ensured.append(name)
except Exception as e:
verbose_proxy_logger.warning("Failed to ensure spend-log partition %s: %s", name, e)
return ensured
async def _list_partitions(self, prisma_client) -> list[tuple[str, datetime | None]]:
rows: Final = await prisma_client.db.query_raw(
"""
SELECT c.relname AS name,
pg_get_expr(c.relpartbound, c.oid) AS bound
FROM pg_inherits i
JOIN pg_class c ON c.oid = i.inhrelid
JOIN pg_class p ON p.oid = i.inhparent
JOIN pg_namespace n ON n.oid = p.relnamespace
WHERE p.relname = $1
AND n.nspname = current_schema()
""",
SPEND_LOGS_TABLE,
)
async def _list_partitions(
self, prisma_client: "PrismaClient", timeout_ms: int
) -> list[tuple[str, datetime | None]]:
async with prisma_client.db.tx() as tx:
await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}")
rows: Final = await tx.query_raw(
"""
SELECT c.relname AS name,
pg_get_expr(c.relpartbound, c.oid) AS bound
FROM pg_inherits i
JOIN pg_class c ON c.oid = i.inhrelid
JOIN pg_class p ON p.oid = i.inhparent
JOIN pg_namespace n ON n.oid = p.relnamespace
WHERE p.relname = $1
AND n.nspname = current_schema()
""",
SPEND_LOGS_TABLE,
)
return [(row["name"], parse_partition_upper_bound(row.get("bound") or "")) for row in rows]
async def drop_partitions_older_than(self, prisma_client, cutoff: datetime) -> list[str]:
async def drop_partitions_older_than(
self, prisma_client: "PrismaClient", cutoff: datetime, remaining_timeout_ms: RemainingTimeoutMs
) -> list[str]:
"""DROP every partition whose whole range is older than `cutoff`."""
list_budget_ms: Final = remaining_timeout_ms()
if list_budget_ms is None:
return []
cutoff_naive: Final = cutoff.astimezone(timezone.utc).replace(tzinfo=None)
partitions: Final = await self._list_partitions(prisma_client)
partitions: Final = await self._list_partitions(prisma_client, list_budget_ms)
to_drop: Final = select_partitions_to_drop(partitions, cutoff_naive)
dropped: Final[list[str]] = []
for name in to_drop:
budget_ms = remaining_timeout_ms()
if budget_ms is None:
verbose_proxy_logger.info("Run budget spent, leaving the remaining partitions for the next run")
break
try:
await prisma_client.db.execute_raw(f'DROP TABLE IF EXISTS "{name}"')
await self._execute_bounded_ddl(prisma_client, f'DROP TABLE IF EXISTS "{name}"', budget_ms)
dropped.append(name)
except Exception as e:
verbose_proxy_logger.warning("Failed to drop spend-log partition %s: %s", name, e)

4
litellm/proxy/litellm_pre_call_utils.py
View file

@ -274,7 +274,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = (
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
)
_UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS: Final = frozenset(
UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS: Final = frozenset(
{
"litellm-disable-message-redaction",
}
@ -355,7 +355,7 @@ def _strip_untrusted_request_header_controls(
return
for header_name in list(headers.keys()):
if isinstance(header_name, str) and header_name.lower() in _UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS:
if isinstance(header_name, str) and header_name.lower() in UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS:
if allow_client_message_redaction_opt_out:
continue
headers.pop(header_name, None)

20
litellm/proxy/management_endpoints/access_group_endpoints.py
View file

@ -14,11 +14,11 @@ from litellm.proxy.auth.auth_checks import (
_cache_access_object,
_cache_key_object,
_cache_team_object,
_delete_cache_access_object,
_get_team_object_from_cache,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.management_helpers.access_group_team_sync import invalidate_access_group_cache
from litellm.proxy.utils import get_prisma_client_or_throw
from litellm.repositories.table_repositories import AccessGroupRepository
from litellm.types.access_group import (
@ -146,22 +146,6 @@ async def _cache_access_group_record(record: _AccessGroupRecord) -> None:
)
async def _invalidate_cache_access_group(access_group_id: str) -> None:
"""
Invalidate (delete) an access group entry from both in-memory and Redis caches.
Uses a lazy import of user_api_key_cache and proxy_logging_obj from proxy_server
to avoid circular imports, following the same pattern as key_management_endpoints.
"""
from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache
await _delete_cache_access_object(
access_group_id=access_group_id,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
# ---------------------------------------------------------------------------
# DB sync helpers (called inside a Prisma transaction)
# ---------------------------------------------------------------------------
@ -595,7 +579,7 @@ async def delete_access_group(
from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache
await _invalidate_cache_access_group(access_group_id)
await invalidate_access_group_cache(access_group_id)
await _patch_team_caches_remove_access_group(
affected_team_ids, access_group_id, user_api_key_cache, proxy_logging_obj
)

55
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -88,6 +88,11 @@ from litellm.proxy.management_endpoints.common_utils import (
from litellm.proxy.management_endpoints.model_management_endpoints import (
_add_model_to_db,
)
from litellm.proxy.management_helpers.access_group_key_sync import (
sync_key_access_group_membership,
sync_key_regeneration_access_group_membership,
sync_key_update_access_group_membership,
)
from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at
from litellm.proxy.management_helpers.object_permission_utils import (
_set_object_permission,
@ -2347,6 +2352,17 @@ async def _process_single_key_update(
proxy_logging_obj=proxy_logging_obj,
)
# After the key's own cache entry is dropped, so a failure here cannot leave the key
# authenticating against the access groups it just lost.
await sync_key_update_access_group_membership(
prisma_client=prisma_client,
key_token=_hash_token_if_needed(
_resolve_token_to_update(data=update_key_request, existing_key_row=existing_key_row)
),
data=update_key_request,
existing_key_row=existing_key_row,
)
# Trigger async hook
asyncio.create_task(
KeyManagementEventHooks.async_key_updated_hook(
@ -2828,6 +2844,15 @@ async def update_key_fn(
proxy_logging_obj=proxy_logging_obj,
)
# After the key's own cache entry is dropped, so a failure here cannot leave the key
# authenticating against the access groups it just lost.
await sync_key_update_access_group_membership(
prisma_client=prisma_client,
key_token=_hash_token_if_needed(key),
data=data,
existing_key_row=existing_key_row,
)
if data.spend is not None:
from litellm.proxy.proxy_server import spend_counter_cache
@ -3771,7 +3796,7 @@ async def generate_key_helper_fn(
auto_rotate: bool | None = None,
rotation_interval: str | None = None,
router_settings: dict | None = None,
access_group_ids: list | None = None,
access_group_ids: list[str] | None = None,
budget_limits: list | None = None, # multiple concurrent budget windows
):
from litellm.proxy.proxy_server import premium_user, prisma_client
@ -3979,6 +4004,14 @@ async def generate_key_helper_fn(
create_key_response: Final = await prisma_client.insert_data(data=key_data, table_name="key")
key_data["token_id"] = getattr(create_key_response, "token", None)
created_token_hash: Final = getattr(create_key_response, "token", None)
if isinstance(created_token_hash, str):
await sync_key_access_group_membership(
prisma_client=prisma_client,
key_token=created_token_hash,
previous_access_group_ids=None,
updated_access_group_ids=access_group_ids,
)
key_data["litellm_budget_table"] = getattr(create_key_response, "litellm_budget_table", None)
key_data["created_at"] = getattr(create_key_response, "created_at", None)
key_data["updated_at"] = getattr(create_key_response, "updated_at", None)
@ -4196,6 +4229,7 @@ async def delete_verification_tokens(
deleted_tokens = [key.token for key in authorized_keys]
if len(deleted_tokens) != len(tokens):
failed_tokens = [token for token in tokens if token not in deleted_tokens]
else:
raise Exception("DB not connected. prisma_client is None")
except Exception as e:
@ -4211,6 +4245,16 @@ async def delete_verification_tokens(
hashed_token = hash_token(cast(str, key))
user_api_key_cache.delete_cache(hashed_token)
# After credential invalidation, so a failure here can never keep a deleted key alive.
for deleted_key in authorized_keys:
if deleted_key.token is not None:
await sync_key_access_group_membership(
prisma_client=prisma_client,
key_token=deleted_key.token,
previous_access_group_ids=deleted_key.access_group_ids,
updated_access_group_ids=None,
)
return {
"deleted_keys": deleted_tokens,
"failed_tokens": failed_tokens,
@ -4726,6 +4770,15 @@ async def _execute_virtual_key_regeneration(
proxy_logging_obj=proxy_logging_obj,
)
# After credential invalidation, so a failure here can never keep the old key alive.
await sync_key_regeneration_access_group_membership(
prisma_client=prisma_client,
previous_key_token=hashed_api_key,
new_key_token=new_token_hash,
data=data,
existing_key_row=key_in_db,
)
response: Final = GenerateKeyResponse.model_validate(updated_token_dict)
asyncio.create_task(
KeyManagementEventHooks.async_key_rotated_hook(

171
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -89,6 +89,7 @@ from litellm.types.router import (
ModelInfo,
updateDeployment,
)
from litellm.types.utils import CustomPricingLiteLLMParams
from litellm.utils import get_utc_datetime
router: Final = APIRouter()
@ -242,6 +243,7 @@ def _raise_on_strategy_router_write_violation(
_PTU_MODEL_INFO_FIELDS: Final = ("ptu_count", "cost_per_ptu_per_hour", "ptu_effective_from", "ptu_effective_to")
_PTU_PRICED_PAIR: Final = frozenset({"ptu_count", "cost_per_ptu_per_hour"})
def _explicitly_cleared_ptu_fields(model_info: ModelInfo | None) -> frozenset[str]:
@ -265,9 +267,10 @@ def _merged_ptu_model_info(*, db_model: Deployment, patch_data: updateDeployment
A PTU invariant holds over the deployment as it will exist, not over whichever subset
of fields a caller happened to send.
"""
empty: Final[Mapping[str, object]] = MappingProxyType({})
stored: Final = db_model.model_info.model_dump(exclude_none=True) if db_model.model_info else empty
incoming: Final = patch_data.model_info.model_dump(exclude_none=True) if patch_data.model_info else empty
stored: Final = db_model.model_info.model_dump(exclude_none=True) if db_model.model_info else _EMPTY_MODEL_INFO
incoming: Final = (
patch_data.model_info.model_dump(exclude_none=True) if patch_data.model_info else _EMPTY_MODEL_INFO
)
cleared: Final = _explicitly_cleared_ptu_fields(patch_data.model_info)
return MappingProxyType({k: v for k, v in {**stored, **incoming}.items() if k not in cleared})
@ -339,6 +342,140 @@ def _validate_ptu_model_info(model_info: Mapping[str, object]) -> None:
)
# The six mirrored pricing fields plus the three remaining fields
# Router._inherit_builtin_cache_pricing back-fills from the public cost map. An unset field is
# what that back-fill targets, so a field left out here is one a PTU deployment still bills.
_PTU_ZEROED_PRICING_FIELDS: Final = SPECIAL_MODEL_INFO_PARAMS + (
"cache_creation_input_token_cost_above_1hr",
"cache_creation_input_token_cost_above_200k_tokens",
"cache_read_input_token_cost_above_200k_tokens",
)
_PTU_ZEROED_PRICING: Final[Mapping[str, float]] = MappingProxyType(dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0))
_NO_PRICING_OVERRIDE: Final[Mapping[str, float]] = MappingProxyType({})
_EMPTY_MODEL_INFO: Final[Mapping[str, object]] = _NO_PRICING_OVERRIDE
# Rate fields only. CustomPricingLiteLLMParams also carries settings that are not charges
# (an embedding's output_vector_size, the regional uplift multipliers), and zeroing one of
# those would destroy the deployment's configuration rather than stop a charge.
_CUSTOM_PRICING_FIELDS: Final = frozenset(f for f in CustomPricingLiteLLMParams.model_fields if "cost" in f)
def _is_nonzero_price(value: object) -> bool:
return isinstance(value, (int, float)) and not isinstance(value, bool) and value != 0
def _is_zero_price(value: object) -> bool:
return isinstance(value, (int, float)) and not isinstance(value, bool) and value == 0
def _raise_if_ptu_deployment_is_priced(*, model_info: Mapping[str, object], supplied: Mapping[str, object]) -> None:
"""Refuse a rate the caller supplies for a deployment that bills reserved capacity.
Separate from the zeroing so the team-model path can run it before it touches the team, whose
ACL write autocommits: a refusal raised after it would leave the team changed and the
deployment row never written.
"""
if not is_ptu_cost_attribution_enabled():
return
if model_info.get("ptu_count") is None or model_info.get("cost_per_ptu_per_hour") is None:
return
priced: Final = tuple(sorted(field for field in _CUSTOM_PRICING_FIELDS if _is_nonzero_price(supplied.get(field))))
if not priced:
return
raise HTTPException(
status_code=400,
detail=(
f"A PTU deployment bills by reserved capacity, so {', '.join(priced)} cannot be charged on "
"top of it. Send 0 or no value, or remove ptu_count and cost_per_ptu_per_hour to bill per token."
),
)
def _ptu_zeroed_pricing(
*,
model_info: Mapping[str, object],
litellm_params: Mapping[str, object],
supplied: Mapping[str, object],
) -> Mapping[str, float]:
"""The pricing a PTU deployment must carry, empty unless one is being stored.
Reserved capacity is already billed by the flat cost the rollup writes, so charging the
traffic it serves bills the same tokens twice. Left unset the rate falls back to the public
cost map, which makes the double charge the default rather than an opt-in.
Only a price the caller supplies is refused. A non-zero price already on the row is zeroed
instead, so a deployment priced through a path this rule does not cover heals on its next
save rather than rejecting every later edit of a field that has nothing to do with pricing.
``supplied`` is the caller's litellm_params alone, because that is the blob a price is
authored on. model_info's copy is written by the server, both by the mirror in
``Deployment.__init__`` and by the cost-map defaults /model/info fills in, so a client that
round-trips a model_info blob sends back prices it never chose.
"""
if not is_ptu_cost_attribution_enabled():
return _NO_PRICING_OVERRIDE
if model_info.get("ptu_count") is None or model_info.get("cost_per_ptu_per_hour") is None:
return _NO_PRICING_OVERRIDE
_raise_if_ptu_deployment_is_priced(model_info=model_info, supplied=supplied)
stored: Final = frozenset(
field
for field in _CUSTOM_PRICING_FIELDS
if _is_nonzero_price(model_info.get(field)) or _is_nonzero_price(litellm_params.get(field))
)
if not stored:
return _PTU_ZEROED_PRICING
return MappingProxyType({**_PTU_ZEROED_PRICING, **dict.fromkeys(stored, 0.0)})
def _ptu_pricing_delta(
*,
stored_model_info: Mapping[str, object],
model_info: Mapping[str, object],
litellm_params: Mapping[str, object],
patch: updateDeployment,
) -> tuple[Mapping[str, float], frozenset[str]]:
"""The pricing a patch must write into both blobs, and the pricing it must drop from them.
A patch that takes the deployment off PTU takes the zeroed pricing with it, since the zeros
exist only to stop the double charge. Left behind they would serve the deployment for free.
Reading the stored row rather than the patch alone keeps that release off a deployment that
never carried PTU config, whose zero price is a rate its operator chose. A zero the patch
itself carries is released with the rest, because the dashboard echoes the whole stored
blob on every save, so a supplied zero cannot be told apart from the one this rule wrote.
The release spans every field the zeroing could have written, not just the mirrored ones, or
a rate zeroed on the way in (per-second, per-character tiers) would bill nothing forever.
"""
supplied: Final = patch.litellm_params.model_dump(exclude_none=True) if patch.litellm_params else _EMPTY_MODEL_INFO
zeroed: Final = _ptu_zeroed_pricing(model_info=model_info, litellm_params=litellm_params, supplied=supplied)
if zeroed:
return zeroed, frozenset()
was_ptu: Final = any(stored_model_info.get(field) is not None for field in _PTU_PRICED_PAIR)
if not was_ptu or not _explicitly_cleared_ptu_fields(patch.model_info) & _PTU_PRICED_PAIR:
return _NO_PRICING_OVERRIDE, frozenset()
return _NO_PRICING_OVERRIDE, frozenset(
field
for field in _CUSTOM_PRICING_FIELDS.union(_PTU_ZEROED_PRICING_FIELDS)
if _is_zero_price(model_info.get(field)) or _is_zero_price(litellm_params.get(field))
)
def _ptu_priced_deployment(model_params: Deployment) -> Deployment:
"""``model_params`` with PTU pricing applied, or itself when it configures no PTU."""
model_info: Final = model_params.model_info.model_dump(exclude_none=True)
litellm_params: Final = model_params.litellm_params.model_dump(exclude_none=True)
override: Final = _ptu_zeroed_pricing(model_info=model_info, litellm_params=litellm_params, supplied=litellm_params)
if not override:
return model_params
return model_params.model_copy(
update=MappingProxyType(
{
"litellm_params": model_params.litellm_params.model_copy(update=override),
"model_info": model_params.model_info.model_copy(update=override),
}
)
)
def _parse_ptu_datetime(value: object) -> datetime.datetime | None:
"""``value`` as a datetime, parsing an ISO string, else None."""
if isinstance(value, datetime.datetime):
@ -404,6 +541,19 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr
merged_model_info.pop(field, None)
_validate_ptu_model_info(merged_model_info)
ptu_pricing, ptu_released = _ptu_pricing_delta(
stored_model_info=db_model.model_info.model_dump(exclude_none=True)
if db_model.model_info
else _EMPTY_MODEL_INFO,
model_info=merged_model_info,
litellm_params=merged_litellm_params,
patch=updated_patch,
)
merged_model_info.update(ptu_pricing)
merged_litellm_params.update(ptu_pricing)
for field in ptu_released:
merged_model_info.pop(field, None)
merged_litellm_params.pop(field, None)
# convert to prisma compatible format
@ -863,6 +1013,12 @@ async def _update_team_model_in_db(
if patch_data.model_info is not None:
_raise_if_ptu_cost_attribution_disabled(patch_data.model_info.model_dump(exclude_none=True))
_validate_ptu_model_info(_merged_ptu_model_info(db_model=db_model, patch_data=patch_data))
_raise_if_ptu_deployment_is_priced(
model_info=_merged_ptu_model_info(db_model=db_model, patch_data=patch_data),
supplied=(
patch_data.litellm_params.model_dump(exclude_none=True) if patch_data.litellm_params else _EMPTY_MODEL_INFO
),
)
patch_team_id: Final = patch_data.model_info.team_id if patch_data.model_info else None
@ -1589,6 +1745,7 @@ async def add_new_model(
incoming_model_info: Final = model_params.model_info.model_dump(exclude_none=True)
_raise_if_ptu_cost_attribution_disabled(incoming_model_info)
_validate_ptu_model_info(incoming_model_info)
priced_model_params: Final = _ptu_priced_deployment(model_params)
if store_model_in_db is True:
"""
@ -1602,13 +1759,13 @@ async def add_new_model(
_original_litellm_model_name: Final = model_params.model_name
if model_params.model_info.team_id is None:
model_response = await _add_model_to_db(
model_params=model_params,
model_params=priced_model_params,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
)
else:
model_response = await _add_team_model_to_db(
model_params=model_params,
model_params=priced_model_params,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
)
@ -1620,9 +1777,9 @@ async def add_new_model(
if "slack" in _alerting:
# send notification - new model added
await proxy_logging_obj.slack_alerting_instance.model_added_alert(
model_name=model_params.model_name,
model_name=priced_model_params.model_name,
litellm_model_name=_original_litellm_model_name,
passed_model_info=model_params.model_info,
passed_model_info=priced_model_params.model_info,
)
except Exception as e:
verbose_proxy_logger.exception("Exception in add_new_model: %s", e)

31
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -15,6 +15,7 @@ import math
import traceback
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from types import MappingProxyType
from typing import Annotated, Final, Protocol, TypedDict, TypeVar, cast
import fastapi
@ -106,6 +107,12 @@ from litellm.proxy.management_endpoints.organization_endpoints import (
from litellm.proxy.management_endpoints.tag_management_endpoints import (
get_daily_activity,
)
from litellm.proxy.management_helpers.access_group_team_sync import (
AccessGroupSyncTx,
invalidate_access_group_caches,
reconcile_team_access_group_membership,
sync_team_access_group_membership,
)
from litellm.proxy.management_helpers.object_permission_utils import (
_set_object_permission,
enforce_all_proxy_mcp_servers_grant_is_admin_only,
@ -315,10 +322,17 @@ class _TeamIdInFilter(TypedDict, total=False):
team_id: Mapping[str, Sequence[str]]
class _TeamCreateTx(AccessGroupSyncTx, Protocol):
@property
def litellm_teamtable(self) -> "_PrismaTableActions[LiteLLM_TeamTable]": ...
_STRIP_DELETED_TEAM_FROM_USERS_SQL: Final = """
UPDATE "LiteLLM_UserTable" SET teams = array_remove(teams, $1) WHERE $1 = ANY(teams)
"""
_INCLUDE_MODEL_TABLE: Final = MappingProxyType({"litellm_model_table": True})
def _team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamTable]":
return _typed_table(TeamRepository(prisma_client), LiteLLM_TeamTable)
@ -1511,10 +1525,15 @@ async def new_team(
complete_team_data_dict = prisma_client.jsonify_team_object(db_data=complete_team_data_dict)
team_creation_data: Final[Mapping[str, object]] = complete_team_data_dict
team_row: Final[LiteLLM_TeamTable] = await _team_db(prisma_client).create(
data=team_creation_data,
include={"litellm_model_table": True},
)
tx: _TeamCreateTx
async with prisma_client.db.tx() as tx:
team_row: Final[LiteLLM_TeamTable] = await tx.litellm_teamtable.create(
data=team_creation_data,
include=_INCLUDE_MODEL_TABLE,
)
affected_access_groups: Final = await reconcile_team_access_group_membership(tx, team_row.team_id)
await invalidate_access_group_caches(affected_access_groups)
## ADD TEAM ID TO USER TABLE ##
team_member_add_request: Final = TeamMemberAddRequest(
@ -2217,6 +2236,7 @@ async def update_team(
)
verbose_proxy_logger.info("Successfully updated team - %s, info", team_row.team_id)
await sync_team_access_group_membership(prisma_client=prisma_client, team_id=team_row.team_id)
await _refresh_cached_team(
team_row=team_row,
user_api_key_cache=user_api_key_cache,
@ -3850,6 +3870,9 @@ async def delete_team(
# keeping the first one means a failure here still leaves a team the admin can retry deleting.
await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client)
for deleted_team in team_rows:
await sync_team_access_group_membership(prisma_client=prisma_client, team_id=deleted_team.team_id)
return deleted_teams

173
litellm/proxy/management_helpers/access_group_key_sync.py Normal file
View file

@ -0,0 +1,173 @@
"""
Reverse sync for the key side of the key <-> access group relationship.
`litellm_accessgrouptable.assigned_key_ids` and `litellm_verificationtoken.access_group_ids`
are the two halves of one relationship and BOTH are read: the access group's
attached-keys view reads the former, and so does the grant check in
`auth_checks.get_authorized_resources_from_key_access_groups`, which authorizes a
key only when the group lists the key's token (or the key's team). The access-group
endpoints maintain both halves already; this module is what the key write paths call
so an edit from that side is mirrored back.
Every write is a single guarded statement rather than a read-modify-write. Prisma has no
atomic scalar-list removal (see `TeamRepository.remove_member`), and the read-modify-write
it otherwise forces is not safe here: a lost update would put an already revoked token back
into a group and restore its grants, or drop a grant an admin just made. The guards also
make each statement idempotent, so a retry cannot duplicate an entry. Each statement covers
every group the request touches at once, so the size of the caller's id list does not turn
into a matching number of round trips, and returns the ids it actually moved so only those
groups are dropped from cache.
It deliberately lives outside `access_group_endpoints`, which is a lazily
registered feature router (see `_lazy_features.LAZY_FEATURES`). Importing that
module eagerly from `key_management_endpoints` would put it in `sys.modules`
without its router ever being included, which drops its routes from the OpenAPI
schema.
"""
from collections.abc import Sequence
from typing import Final, Protocol
from pydantic import BaseModel
from litellm.proxy._types import (
LiteLLM_VerificationToken,
RegenerateKeyRequest,
UpdateKeyRequest,
)
from litellm.proxy.auth.auth_checks import (
_delete_cache_access_object, # pyright: ignore[reportPrivateUsage] # the access-group endpoints reach for this same cache primitive
)
from litellm.repositories.table_repositories import AccessGroupRepository
class _MovedGroupRow(BaseModel):
access_group_id: str
class _RawExecutor(Protocol):
async def query_raw(self, query: str, *args: str | Sequence[str]) -> Sequence[object]: ...
_ATTACH_KEY_SQL: Final = (
'UPDATE "LiteLLM_AccessGroupTable" '
'SET "assigned_key_ids" = array_append("assigned_key_ids", $1) '
'WHERE "access_group_id" = ANY($2::text[]) AND NOT ($1 = ANY("assigned_key_ids")) '
'RETURNING "access_group_id"'
)
_DETACH_KEY_SQL: Final = (
'UPDATE "LiteLLM_AccessGroupTable" '
'SET "assigned_key_ids" = array_remove("assigned_key_ids", $1) '
'WHERE "access_group_id" = ANY($2::text[]) AND $1 = ANY("assigned_key_ids") '
'RETURNING "access_group_id"'
)
_REPOINT_KEY_SQL: Final = (
'UPDATE "LiteLLM_AccessGroupTable" '
'SET "assigned_key_ids" = array_append(array_remove(array_remove("assigned_key_ids", $1), $2), $2) '
'WHERE $1 = ANY("assigned_key_ids") '
'RETURNING "access_group_id"'
)
def _raw_executor(prisma_client: object) -> _RawExecutor:
"""Narrow the untyped Prisma client down to the raw-query call this module makes."""
return AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client
async def _invalidate_access_group_cache(access_group_id: str) -> None:
"""
Drop an access group entry from both the in-memory and Redis caches.
Uses a lazy import of user_api_key_cache and proxy_logging_obj from proxy_server
to avoid circular imports, following the same pattern as key_management_endpoints.
"""
from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache
await _delete_cache_access_object(
access_group_id=access_group_id,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
async def _invalidate_moved_groups(moved_rows: Sequence[object]) -> None:
for row in moved_rows:
await _invalidate_access_group_cache(_MovedGroupRow.model_validate(row).access_group_id)
async def _write_membership(prisma_client: object, sql: str, access_group_ids: frozenset[str], key_token: str) -> None:
"""Run one guarded membership statement for every listed group, dropping the cache of those it moved."""
if not access_group_ids:
return
await _invalidate_moved_groups(
await _raw_executor(prisma_client).query_raw(sql, key_token, sorted(access_group_ids))
)
async def sync_key_access_group_membership(
prisma_client: object,
key_token: str,
previous_access_group_ids: Sequence[str] | None,
updated_access_group_ids: Sequence[str] | None,
) -> None:
"""Mirror a key-side change to `access_group_ids` onto each access group's `assigned_key_ids`."""
previous: Final = frozenset(previous_access_group_ids or ())
updated: Final = frozenset(updated_access_group_ids or ())
await _write_membership(prisma_client, _ATTACH_KEY_SQL, updated - previous, key_token)
await _write_membership(prisma_client, _DETACH_KEY_SQL, previous - updated, key_token)
async def sync_key_update_access_group_membership(
prisma_client: object,
key_token: str,
data: UpdateKeyRequest | RegenerateKeyRequest,
existing_key_row: LiteLLM_VerificationToken,
) -> None:
"""
Mirror a key UPDATE onto the group side, honouring `exclude_unset` semantics.
The key row is written from `model_dump(exclude_unset=True)`, so a request that never
mentions `access_group_ids` leaves the key's own list alone and must leave the group's
copy alone too. Reading the attribute instead of `model_fields_set` would see None on
every unrelated edit and withdraw the token from every group it belongs to.
"""
if "access_group_ids" not in data.model_fields_set:
return
await sync_key_access_group_membership(
prisma_client=prisma_client,
key_token=key_token,
previous_access_group_ids=existing_key_row.access_group_ids,
updated_access_group_ids=data.access_group_ids,
)
async def sync_key_regeneration_access_group_membership(
prisma_client: object,
previous_key_token: str,
new_key_token: str,
data: RegenerateKeyRequest | None,
existing_key_row: LiteLLM_VerificationToken,
) -> None:
"""
Re-point every group's copy from the old token to the regenerated one.
Regeneration replaces the token, which is the identity `assigned_key_ids` stores, so
leaving the old hash behind both points the group at a row that no longer exists and
denies the regenerated key the group's grants. The swap is driven by the groups that
hold the old token when the statement runs, not by the key row read earlier, so a group
edited in between is neither resurrected nor skipped. Removing the new token before
appending it keeps a re-run from duplicating it.
"""
await _invalidate_moved_groups(
await _raw_executor(prisma_client).query_raw(_REPOINT_KEY_SQL, previous_key_token, new_key_token)
)
if data is not None:
await sync_key_update_access_group_membership(
prisma_client=prisma_client,
key_token=new_key_token,
data=data,
existing_key_row=existing_key_row,
)

155
litellm/proxy/management_helpers/access_group_team_sync.py Normal file
View file

@ -0,0 +1,155 @@
"""
Reverse sync for the team side of the team <-> access group relationship.
`litellm_accessgrouptable.assigned_team_ids` and `litellm_teamtable.access_group_ids`
are two copies of the same relationship, and both are read: the access group's
attached-teams view reads the former, and so does the key-side grant check in
`auth_checks.get_authorized_resources_from_key_access_groups`. The access-group
endpoints maintain both copies already; this module is what the team write paths
call so an edit from that side is mirrored back.
It deliberately lives outside `access_group_endpoints`, which is a lazily
registered feature router (see `_lazy_features.LAZY_FEATURES`). Importing that
module eagerly from `team_endpoints` would put it in `sys.modules` without its
router ever being included, which drops its routes from the OpenAPI schema.
"""
import asyncio
from collections.abc import Mapping, Sequence
from typing import Final, Protocol
from pydantic import BaseModel, TypeAdapter
from litellm.proxy.auth.auth_checks import _delete_cache_access_object
# hashtext collisions only cost two unrelated teams a little serialization, and the
# lock is never taken by the access-group endpoints, so it cannot join their
# access-group-then-team lock order to form a cycle.
_LOCK_TEAM_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked"
_READ_TEAM_SQL: Final = 'SELECT access_group_ids FROM "LiteLLM_TeamTable" WHERE team_id = $1'
# The groups the team is on either side of the reconcile, so the cache step is driven by
# desired state rather than by which rows this attempt happened to change. A retry after a
# failed invalidation finds the same set even though its statements are already no-ops.
_AFFECTED_SQL: Final = """
SELECT access_group_id FROM "LiteLLM_AccessGroupTable"
WHERE access_group_id = ANY($2::TEXT[])
OR $1 = ANY(COALESCE(assigned_team_ids, ARRAY[]::TEXT[]))
"""
_ATTACH_SQL: Final = """
UPDATE "LiteLLM_AccessGroupTable"
SET assigned_team_ids = array_append(COALESCE(assigned_team_ids, ARRAY[]::TEXT[]), $1)
WHERE access_group_id = ANY($2::TEXT[])
AND NOT ($1 = ANY(COALESCE(assigned_team_ids, ARRAY[]::TEXT[])))
RETURNING access_group_id
"""
_DETACH_SQL: Final = """
UPDATE "LiteLLM_AccessGroupTable"
SET assigned_team_ids = array_remove(assigned_team_ids, $1)
WHERE $1 = ANY(COALESCE(assigned_team_ids, ARRAY[]::TEXT[]))
AND NOT (access_group_id = ANY($2::TEXT[]))
RETURNING access_group_id
"""
class _AffectedGroup(BaseModel):
access_group_id: str
class _TeamGroups(BaseModel):
access_group_ids: tuple[str, ...] | None = None
_AffectedGroups: Final = TypeAdapter(tuple[_AffectedGroup, ...])
_TeamRows: Final = TypeAdapter(tuple[_TeamGroups, ...])
class AccessGroupSyncTx(Protocol):
async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: ...
class _Transaction(Protocol):
async def __aenter__(self) -> AccessGroupSyncTx: ...
async def __aexit__(self, *exc_info: object) -> None: ...
class _PrismaDb(Protocol):
def tx(self) -> _Transaction: ...
class _PrismaClient(Protocol):
@property
def db(self) -> _PrismaDb: ...
async def invalidate_access_group_cache(access_group_id: str) -> None:
"""
Drop an access group entry from both the in-memory and Redis caches.
Uses a lazy import of user_api_key_cache and proxy_logging_obj from proxy_server
to avoid circular imports, following the same pattern as key_management_endpoints.
"""
from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache
await _delete_cache_access_object(
access_group_id=access_group_id,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
async def invalidate_access_group_caches(access_group_ids: Sequence[str]) -> None:
"""
Drop every given access group from the caches, then raise if any drop failed.
Every entry is attempted even when one raises, so a single unreachable cache cannot
leave the rest of the reconciled groups serving a grant the admin revoked.
"""
outcomes: Final = await asyncio.gather(
*(invalidate_access_group_cache(access_group_id) for access_group_id in access_group_ids),
return_exceptions=True,
)
for outcome in outcomes:
if isinstance(outcome, BaseException):
raise outcome
async def reconcile_team_access_group_membership(tx: AccessGroupSyncTx, team_id: str) -> tuple[str, ...]:
"""
Reconcile every access group's `assigned_team_ids` against the team's own
`access_group_ids`, and return the groups whose cache the caller has to drop once the
transaction commits.
Call this inside the transaction that writes the team row, or after that row is
written or deleted: a team with no row reconciles to an empty set, which detaches it
from every group.
The team row is read here rather than passed in, under an advisory lock held for the
rest of the transaction. That is what makes concurrent writes to the same team
converge, since each mirror reconciles against the row as the transaction sees it
instead of against the snapshot its own caller happened to see. It also means a retry
heals a sync that failed partway, where a before/after delta would compute nothing.
Both mirror statements are set-based and mutate the array inside the statement, so a
concurrent write for a different team cannot be lost the way a read-modify-write of
the whole array can, and the pair commits together or not at all.
"""
await tx.query_raw(_LOCK_TEAM_SQL, team_id)
team_rows: Final = _TeamRows.validate_python(await tx.query_raw(_READ_TEAM_SQL, team_id))
desired: Final = (team_rows[0].access_group_ids or ()) if team_rows else ()
affected: Final = _AffectedGroups.validate_python(await tx.query_raw(_AFFECTED_SQL, team_id, desired))
await tx.query_raw(_ATTACH_SQL, team_id, desired)
await tx.query_raw(_DETACH_SQL, team_id, desired)
return tuple(group.access_group_id for group in affected)
async def sync_team_access_group_membership(prisma_client: _PrismaClient, team_id: str) -> None:
"""Reconcile the mirror for an already committed team write, in its own transaction."""
async with prisma_client.db.tx() as tx:
affected: Final = await reconcile_team_access_group_membership(tx, team_id)
await invalidate_access_group_caches(affected)

2
litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py
View file

@ -82,7 +82,7 @@ class CoherePassthroughLoggingHandler(BasePassthroughLoggingHandler):
Handle Cohere passthrough logging with route detection and cost tracking.
"""
# Check if this is an embed endpoint
if "/v1/embed" in url_route:
if "/v1/embed" in url_route and "/v1/embeddings" not in url_route:
model: Final = request_body.get("model", response_body.get("model", ""))
try:
cohere_embed_config: Final = CohereEmbeddingConfig()

40
litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py
View file

@ -31,8 +31,8 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
EndpointType,
PassthroughStandardLoggingPayload,
)
from litellm.types.utils import ImageResponse, LlmProviders, PassthroughCallTypes
from litellm.utils import ModelResponse, TextCompletionResponse
from litellm.types.utils import EmbeddingResponse, ImageResponse, LlmProviders, PassthroughCallTypes
from litellm.utils import ModelResponse, TextCompletionResponse, convert_to_model_response_object
# Hostnames that route to OpenAI-compatible APIs.
#
@ -143,6 +143,14 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
"/v1/responses" in parsed_url.path or "/responses" in parsed_url.path
)
@staticmethod
def is_openai_embeddings_route(url_route: str) -> bool:
"""Check if the URL route is an OpenAI embeddings endpoint."""
if not url_route:
return False
parsed_url: Final = urlparse(url_route)
return _is_openai_compatible_host(parsed_url.hostname) and "/v1/embeddings" in parsed_url.path
def _get_user_from_metadata(
self,
passthrough_logging_payload: PassthroughStandardLoggingPayload,
@ -271,22 +279,21 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
**kwargs,
) -> PassThroughEndpointLoggingTypedDict:
"""
Handle OpenAI passthrough logging with cost tracking for chat completions, image generation, image editing, and responses API.
Handle OpenAI passthrough logging with cost tracking for chat completions,
embeddings, image generation, image editing, and responses API.
"""
# Check if this is a supported endpoint for cost tracking
is_chat_completions: Final = OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(url_route)
is_embeddings: Final = OpenAIPassthroughLoggingHandler.is_openai_embeddings_route(url_route)
is_image_generation: Final = OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(url_route)
is_image_editing: Final = OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(url_route)
is_responses: Final = OpenAIPassthroughLoggingHandler.is_openai_responses_route(url_route)
if not (is_chat_completions or is_image_generation or is_image_editing or is_responses):
# For unsupported endpoints, return None to let the system fall back to generic behavior
if not (is_chat_completions or is_embeddings or is_image_generation or is_image_editing or is_responses):
return {
"result": None,
"kwargs": kwargs,
}
# Extract model from request or response
model: Final = request_body.get("model", response_body.get("model", ""))
if not model:
verbose_proxy_logger.warning("No model found in request or response for OpenAI passthrough cost tracking")
@ -307,7 +314,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
try:
response_cost = 0.0
litellm_model_response: (
ModelResponse | TextCompletionResponse | ImageResponse | ResponsesAPIResponse | None
ModelResponse | TextCompletionResponse | EmbeddingResponse | ImageResponse | ResponsesAPIResponse | None
) = None
handler_instance: Final = OpenAIPassthroughLoggingHandler()
@ -338,6 +345,19 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
model=model,
custom_llm_provider=custom_llm_provider,
)
elif is_embeddings:
litellm_model_response = convert_to_model_response_object(
response_object=response_body,
model_response_object=EmbeddingResponse(),
response_type="embedding",
)
response_cost = litellm.completion_cost(
completion_response=litellm_model_response,
model=model,
custom_llm_provider=custom_llm_provider,
call_type="aembedding",
)
litellm_model_response._hidden_params["response_cost"] = response_cost
elif is_image_generation:
# Handle image generation cost calculation
response_cost = OpenAIPassthroughLoggingHandler._calculate_image_generation_cost(
@ -432,9 +452,13 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
endpoint_type: Final = (
"chat_completions"
if is_chat_completions
else "embeddings"
if is_embeddings
else "image_generation"
if is_image_generation
else "image_editing"
if is_image_editing
else "responses"
)
verbose_proxy_logger.debug(
f"OpenAI passthrough cost tracking - Endpoint: {endpoint_type}, Model: {model}, Cost: ${response_cost:.6f}"

11
litellm/proxy/pass_through_endpoints/success_handler.py
View file

@ -349,10 +349,14 @@ class PassThroughEndpointLogging:
return True
return False
def is_cohere_route(self, url_route: str):
def is_cohere_route(self, url_route: str) -> bool:
for route in self.TRACKED_COHERE_ROUTES:
if route in url_route:
return True
if route not in url_route:
continue
if route == "/v1/embed" and "/v1/embeddings" in url_route:
continue
return True
return False
def is_assemblyai_route(self, url_route: str):
parsed_url: Final = urlparse(url_route)
@ -429,6 +433,7 @@ class PassThroughEndpointLogging:
return (
OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(url_route)
or OpenAIPassthroughLoggingHandler.is_openai_embeddings_route(url_route)
or OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(url_route)
or OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(url_route)
or OpenAIPassthroughLoggingHandler.is_openai_responses_route(url_route)

28
litellm/proxy/proxy_server.py
View file

@ -371,7 +371,10 @@ from litellm.proxy.config_resolvers.alerting import (
)
from litellm.proxy.container_endpoints.endpoints import router as container_router
from litellm.proxy.credential_endpoints.endpoints import router as credential_router
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
SPEND_LOG_CLEANUP_BOUND_SETTINGS,
SpendLogCleanup,
)
from litellm.proxy.db.exception_handler import (
PrismaDBExceptionHandler,
call_with_db_reconnect_retry,
@ -4083,6 +4086,7 @@ class ProxyConfig:
# precedence over stale DB-cached values for these specific keys
# during periodic config reloads (_update_general_settings).
self._yaml_general_settings_keys: set[str] = set() # mutable-ok: populated once at startup, read-only thereafter # fmt: skip
self._yaml_spend_log_cleanup_bounds: dict[str, object] = {} # mutable-ok: snapshot of YAML bounds at load time # fmt: skip
def is_yaml(self, config_file_path: str) -> bool:
if not os.path.isfile(config_file_path):
@ -5019,6 +5023,12 @@ class ProxyConfig:
# These keys take precedence over DB-cached values during periodic
# reloads (see _update_general_settings).
self._yaml_general_settings_keys = set(general_settings.keys()) # mutable-ok: snapshot of YAML keys at load time # fmt: skip
# The VALUES matter for the cleanup bounds, not just which keys were
# set: clearing one from the dashboard has to fall back to what the
# YAML declared, and a set of names cannot answer that.
self._yaml_spend_log_cleanup_bounds = { # mutable-ok: snapshot of YAML bounds at load time # fmt: skip
key: general_settings[key] for key in SPEND_LOG_CLEANUP_BOUND_SETTINGS if key in general_settings
}
### LOAD KEY MANAGEMENT SETTINGS FIRST (needed for custom secret manager) ###
key_management_settings: Final = general_settings.get("key_management_settings", None)
@ -6303,6 +6313,18 @@ class ProxyConfig:
if old_session_value != new_session_value:
await self._reschedule_spend_log_cleanup_job()
## SPEND LOG CLEANUP BOUNDS ##
# The dashboard writes these straight to the DB, so without copying them
# here the running cleanup job never sees them. A key the DB no longer
# carries was cleared from the dashboard, and falls back to whatever
# config.yaml declared, or to None (the shipped default) when it declared
# nothing. Leaving the deleted DB value in memory would keep enforcing the
# bound the operator just removed.
for cleanup_key in SPEND_LOG_CLEANUP_BOUND_SETTINGS:
general_settings[cleanup_key] = _general_settings.get(
cleanup_key, self._yaml_spend_log_cleanup_bounds.get(cleanup_key)
)
for key in (
"user_url_allowed_hosts",
"user_url_validation",
@ -15544,6 +15566,10 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro
"store_model_in_db": "Boolean",
"store_prompts_in_spend_logs": "Boolean",
"maximum_spend_logs_retention_period": "String",
"maximum_spend_logs_cleanup_batch_size": "Integer",
"maximum_spend_logs_cleanup_max_batches": "Integer",
"maximum_spend_logs_cleanup_run_budget": "String",
"maximum_spend_logs_cleanup_batch_timeout": "String",
"mcp_internal_ip_ranges": "List",
"mcp_trusted_proxy_ranges": "List",
"mcp_xff_num_trusted_hops": "Integer",

38
litellm/proxy/public_endpoints/provider_create_fields.json
View file

@ -2062,6 +2062,44 @@
],
"default_model_placeholder": "gpt-3.5-turbo"
},
{
"provider": "NVIDIA_RIVA",
"provider_display_name": "Nvidia Riva",
"litellm_provider": "nvidia_riva",
"credential_fields": [
{
"key": "api_base",
"label": "API Base",
"placeholder": "grpc.nvcf.nvidia.com:443",
"tooltip": "host:port of the Riva gRPC endpoint. Use grpc.nvcf.nvidia.com:443 for NVCF-hosted Riva, or your own host (e.g. localhost:50051) when self-hosting. Riva has no public default, so this is required.",
"required": true,
"field_type": "text",
"options": null,
"default_value": null
},
{
"key": "api_key",
"label": "API Key",
"placeholder": "nvapi-...",
"tooltip": "Sent as gRPC authorization metadata. Required for NVCF-hosted Riva, optional for self-hosted deployments without auth.",
"required": false,
"field_type": "password",
"options": null,
"default_value": null
},
{
"key": "nvcf_function_id",
"label": "NVCF Function ID",
"placeholder": "1598d209-5e27-4d3c-8079-4751568b1081",
"tooltip": "NVCF function id of the hosted Riva model. Setting it turns on TLS and the function-id gRPC metadata. Leave empty for self-hosted Riva.",
"required": false,
"field_type": "text",
"options": null,
"default_value": null
}
],
"default_model_placeholder": "nvidia_riva/nvidia/parakeet-ctc-1_1b-asr"
},
{
"provider": "Ollama",
"provider_display_name": "Ollama",

38
litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py
View file

@ -21,6 +21,7 @@ from typing import TYPE_CHECKING, Final
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
PTU_LAPSED_ALERT_LIMIT,
PTU_PRUNE_SKEW_GRACE_SECONDS,
PTU_ROLLUP_JOB_ID,
PTU_ROLLUP_LOCK_TTL_SECONDS,
@ -45,6 +46,7 @@ class RollupResult:
models_processed: int
rows_written: int
rows_failed: int = 0
lapsed: tuple[str, ...] = ()
@dataclass(frozen=True, slots=True)
@ -387,6 +389,34 @@ async def run_ptu_flat_cost_rollup(
models_processed=len(ptu_models),
rows_written=rows_written,
rows_failed=rows_failed,
lapsed=_lapsed_models(ptu_models, run_started),
)
def _slack_safe(model_name: str) -> str:
"""``model_name`` with the characters Slack reads as markup escaped.
A model name is operator-supplied and this alert is delivered to an operator channel, so an
unescaped name could post a channel-wide mention or a disguised link.
"""
return model_name.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
def _lapsed_models(ptu_models: tuple[PTUModel, ...], now: datetime) -> tuple[str, ...]:
"""PTU deployments whose window has closed, newest bound first.
The provider bills reserved capacity until the deployment is deleted, so a closed window
stops this attribution without stopping the charge. The deployment is left alone: the
window is what the operator asked to be attributed, and per-token pricing would invent a
charge the provider does not make for reserved capacity.
"""
return tuple(
_slack_safe(model.model_name)
for model in sorted(
(m for m in ptu_models if m.effective_to is not None and m.effective_to <= now),
key=lambda m: m.effective_to,
reverse=True,
)
)
@ -585,6 +615,14 @@ async def _run_and_alert(
f"{result.rows_written + result.rows_failed} team charges failed to write. Those teams show no PTU "
f"cost for that date until the rollup is rerun for it.",
)
if result.lapsed:
await _deliver_alert(
alert,
f"PTU flat-cost attribution has stopped for {len(result.lapsed)} deployment(s) whose effective "
f"window has closed: {', '.join(result.lapsed[:PTU_LAPSED_ALERT_LIMIT])}. Reserved capacity is billed "
"until the deployment is deleted, so a deployment still serving traffic is still being charged for "
"by the provider with nothing attributing it here. Extend the window, or retire the deployment.",
)
if target_date is None:
await _backfill_and_alert(prisma_client, alert=alert)
return result

1
litellm/proxy/utils.py
View file

@ -679,6 +679,7 @@ class ProxyLogging:
# (e.g. MCPJWTSigner) to independently verify the caller's identity
# before re-signing an outbound token (FR-5 verify+re-sign).
"incoming_bearer_token": kwargs.get("incoming_bearer_token"),
"metadata": {"headers": kwargs.get("headers") or {}},
}
return synthetic_data

5
litellm/responses/mcp/litellm_proxy_mcp_handler.py
View file

@ -11,6 +11,7 @@ from litellm._logging import verbose_logger
from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._experimental.mcp_server.utils import (
logging_safe_mcp_headers,
split_server_prefix_from_name,
strip_known_server_prefix,
)
@ -653,6 +654,7 @@ class LiteLLM_Proxy_MCP_Handler:
tool_results: Final[list[MCPToolResult]] = []
tool_call_id: str | None = None
rules_obj: Final = Rules()
logging_safe_headers: Final = logging_safe_mcp_headers(raw_headers)
for tool_call in tool_calls:
logging_request_data: dict[str, object] = {}
tool_name: str | None = None
@ -697,6 +699,7 @@ class LiteLLM_Proxy_MCP_Handler:
"tool_call_id": tool_call_id,
"tool_name": sanitized_tool_name,
"server_name": server_name,
"headers": logging_safe_headers,
}
logging_request_data = {
"model": f"MCP: {tool_name}",
@ -708,7 +711,7 @@ class LiteLLM_Proxy_MCP_Handler:
"proxy_server_request": {
"url": "/mcp/tools/call",
"method": "POST",
"headers": {},
"headers": logging_safe_headers,
"body": {
"name": sanitized_tool_name,
"arguments": parsed_arguments,

28
scripts/pre_commit_lint.sh
View file

@ -55,11 +55,13 @@ else
merge_base=$(git merge-base origin/litellm_internal_staging HEAD 2>/dev/null) || {
echo "check: cannot resolve the merge base with origin/litellm_internal_staging." >&2
echo " Fix: git fetch origin litellm_internal_staging" >&2
echo "check: FAIL"
exit 1
}
scope=$(printf '%s\n' "$(git diff --name-only --diff-filter=ACMRD "$merge_base")" "$untracked" | sed '/^$/d' | sort -u)
if [ -z "$scope" ]; then
echo "check: nothing to check (no staged files, no working-tree changes, no branch changes vs origin/litellm_internal_staging)"
echo "check: PASS"
exit 0
fi
echo "check: nothing staged; scoping to the working tree's diff against the merge base with origin/litellm_internal_staging:"
@ -281,4 +283,30 @@ if [ -n "${gen_pid:-}" ]; then
cat "$gen_log"; rm -f "$gen_log"
fi
summary_item() {
local check_name=$1 triggered=$2 skip_reason=$3
if [ -n "$triggered" ]; then
echo " ran: $check_name"
else
echo " skipped: $check_name ($skip_reason)"
fi
}
echo "check: summary"
summary_item "Python lint (make lint)" "$litellm_py_files" "no litellm/ Python files in scope"
summary_item "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_files" "no tests/e2e Python files in scope"
summary_item "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_changed$ui_eslint_changed" "no dashboard files in scope"
summary_item "dashboard API-type sync (npm run gen:api)" "$spec_files" "no litellm/proxy, litellm/types, or generator files in scope"
if [ -z "$litellm_py_files$e2e_py_files$ui_prettier_changed$ui_eslint_changed$spec_files" ]; then
echo "check: NOTE - no gating lint check matches the files in scope, so nothing ran:" >&2
printf '%s\n' "$scope" | sed 's/^/ /' >&2
echo " A pass here is a no-op, not a lint verdict." >&2
fi
if [ "$status" -eq 0 ]; then
echo "check: PASS"
else
echo "check: FAIL"
fi
exit $status

230
tests/proxy_admin_ui_tests/test_access_group_team_sync.py Normal file
View file

@ -0,0 +1,230 @@
"""
Real-Postgres coverage for the team -> access group mirror.
`sync_team_access_group_membership` reconciles `assigned_team_ids` with two raw
statements, and a mocked prisma cannot tell whether that SQL is right: a fake has to
reimplement the array semantics in Python, so it passes no matter what the SQL says.
These tests run the statements against the same Postgres CI seeds for the admin UI
suite, which is the only place a `NOT (... = ANY(...))` guard going missing shows up.
"""
import asyncio
import os
import sys
from contextlib import asynccontextmanager
from datetime import timedelta
from types import SimpleNamespace
from unittest.mock import AsyncMock, patch
import pytest
sys.path.insert(0, os.path.abspath("../.."))
from litellm.proxy.management_helpers.access_group_team_sync import (
reconcile_team_access_group_membership,
sync_team_access_group_membership,
)
TEAM = "ags-team-a"
OTHER_TEAM = "ags-team-b"
GROUPS = ("ags-group-1", "ags-group-2", "ags-group-3")
_DELETE_SEEDED = 'DELETE FROM "LiteLLM_AccessGroupTable" WHERE access_group_id = ANY($1::TEXT[])'
_DELETE_TEAMS = 'DELETE FROM "LiteLLM_TeamTable" WHERE team_id = ANY($1::TEXT[])'
@asynccontextmanager
async def _clean_db():
"""Connects inside the running test's loop. An async fixture would be torn up on a
different loop than the test body, which prisma's engine lock refuses outright."""
from prisma import Prisma
if not os.getenv("DATABASE_URL"):
pytest.fail("DATABASE_URL is required; these tests must not silently skip")
db = Prisma()
await db.connect()
try:
await db.execute_raw(_DELETE_SEEDED, list(GROUPS))
await db.execute_raw(_DELETE_TEAMS, [TEAM, OTHER_TEAM])
yield db
finally:
await db.execute_raw(_DELETE_SEEDED, list(GROUPS))
await db.execute_raw(_DELETE_TEAMS, [TEAM, OTHER_TEAM])
await db.disconnect()
async def _seed(db, assignments):
for group_id, team_ids in assignments.items():
await db.litellm_accessgrouptable.create(
data={
"access_group_id": group_id,
"access_group_name": group_id,
"assigned_team_ids": team_ids,
}
)
async def _read(db):
rows = await db.query_raw(
'SELECT access_group_id, assigned_team_ids FROM "LiteLLM_AccessGroupTable" '
"WHERE access_group_id = ANY($1::TEXT[])",
list(GROUPS),
)
return {row["access_group_id"]: sorted(row["assigned_team_ids"] or []) for row in rows}
async def _set_team_groups(db, team_id, access_group_ids):
"""The mirror reads the committed team row, so the desired state is written there."""
if access_group_ids is None:
await db.execute_raw(_DELETE_TEAMS, [team_id])
return
await db.litellm_teamtable.upsert(
where={"team_id": team_id},
data={
"create": {"team_id": team_id, "access_group_ids": list(access_group_ids)},
"update": {"access_group_ids": list(access_group_ids)},
},
)
async def _sync(db, team_id, access_group_ids):
await _set_team_groups(db, team_id, access_group_ids)
with patch(
"litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache",
new_callable=AsyncMock,
) as invalidate:
await sync_team_access_group_membership(prisma_client=SimpleNamespace(db=db), team_id=team_id)
return {call.args[0] for call in invalidate.call_args_list}
@pytest.mark.asyncio
async def test_reconcile_attaches_and_detaches_without_touching_other_teams():
"""The detach must be scoped to groups the team dropped. Losing that scope would
strip the team from the very groups it just kept, silently revoking live grants."""
async with _clean_db() as db:
await _seed(db, {GROUPS[0]: [TEAM, OTHER_TEAM], GROUPS[1]: [TEAM], GROUPS[2]: [OTHER_TEAM]})
invalidated = await _sync(db, TEAM, [GROUPS[1], GROUPS[2]])
assert await _read(db) == {
GROUPS[0]: [OTHER_TEAM],
GROUPS[1]: [TEAM],
GROUPS[2]: sorted([TEAM, OTHER_TEAM]),
}
assert invalidated == {GROUPS[0], GROUPS[1], GROUPS[2]}
@pytest.mark.asyncio
async def test_reconcile_is_idempotent_so_a_retry_heals_rather_than_duplicates():
"""Reconciling to the same desired state twice must leave the rows alone and still name
the team's groups for the cache step, so a retry after a failed cache drop reaches them.
A delta-based mirror would instead go quiet once the rows match, leaving the caches
serving a grant the admin already revoked."""
async with _clean_db() as db:
await _seed(db, {GROUPS[0]: [], GROUPS[1]: [TEAM], GROUPS[2]: []})
first = await _sync(db, TEAM, [GROUPS[0], GROUPS[1]])
after_first = await _read(db)
second = await _sync(db, TEAM, [GROUPS[0], GROUPS[1]])
assert after_first == {GROUPS[0]: [TEAM], GROUPS[1]: [TEAM], GROUPS[2]: []}
assert await _read(db) == after_first
assert first == {GROUPS[0], GROUPS[1]}
assert second == first
@pytest.mark.asyncio
async def test_reconcile_handles_a_null_array_column():
"""`assigned_team_ids` is nullable in Postgres. Without COALESCE both statements
evaluate their guard to NULL, skip the row, and the grant silently never syncs."""
async with _clean_db() as db:
await _seed(db, {GROUPS[0]: [], GROUPS[1]: []})
await db.execute_raw(
'UPDATE "LiteLLM_AccessGroupTable" SET assigned_team_ids = NULL WHERE access_group_id = $1',
GROUPS[0],
)
await _sync(db, TEAM, [GROUPS[0]])
assert await _read(db) == {GROUPS[0]: [TEAM], GROUPS[1]: []}
@pytest.mark.asyncio
async def test_passing_none_detaches_the_team_from_every_group():
"""Team deletion. A group the deleted row never listed must still let the team go,
otherwise the id dangles under Attached Teams and grants again if it is reused."""
async with _clean_db() as db:
await _seed(db, {GROUPS[0]: [TEAM, OTHER_TEAM], GROUPS[1]: [TEAM], GROUPS[2]: [OTHER_TEAM]})
invalidated = await _sync(db, TEAM, None)
assert await _read(db) == {GROUPS[0]: [OTHER_TEAM], GROUPS[1]: [], GROUPS[2]: [OTHER_TEAM]}
assert invalidated == {GROUPS[0], GROUPS[1]}
@pytest.mark.asyncio
async def test_a_failed_mirror_takes_the_new_team_row_with_it():
"""`/team/new` inserts the team and mirrors it in one transaction. Mirroring in a
transaction of its own instead leaves a committed team whose groups never learned about
it, and the retry with that same team id comes back as a duplicate."""
async with _clean_db() as db:
await _seed(db, {GROUPS[0]: [], GROUPS[1]: [OTHER_TEAM]})
with pytest.raises(RuntimeError):
async with db.tx() as tx:
await tx.litellm_teamtable.create(data={"team_id": TEAM, "access_group_ids": [GROUPS[0]]})
await reconcile_team_access_group_membership(tx, TEAM)
raise RuntimeError("the cache handoff blew up")
assert await _read(db) == {GROUPS[0]: [], GROUPS[1]: [OTHER_TEAM]}
assert await db.litellm_teamtable.find_unique(where={"team_id": TEAM}) is None
@pytest.mark.asyncio
async def test_a_concurrent_writer_cannot_replay_a_stale_team_row_over_a_newer_one():
"""
Two writers edit one team at once. Whichever team row commits last is the admin's
final intent and the mirror must match it, so the mirror has to hold the team's
advisory lock across its read and its writes.
A second connection holds that lock and changes the team underneath, which pins the
interleaving instead of hoping a sleep lands in the gap. With the lock the sync waits
and then reads the new row. Without it the sync reads the old row and writes a group
the admin already moved off, which keeps granting to that team.
"""
from prisma import Prisma
async with _clean_db() as db:
await _seed(db, {GROUPS[0]: [], GROUPS[1]: []})
await _sync(db, TEAM, [GROUPS[0]])
assert await _read(db) == {GROUPS[0]: [TEAM], GROUPS[1]: []}
blocker = Prisma()
await blocker.connect()
sync_started = asyncio.Event()
async def competing_sync():
sync_started.set()
with patch(
"litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache",
new_callable=AsyncMock,
):
await sync_team_access_group_membership(prisma_client=SimpleNamespace(db=db), team_id=TEAM)
try:
async with blocker.tx(timeout=timedelta(seconds=30)) as held:
await held.query_raw("SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked", TEAM)
task = asyncio.create_task(competing_sync())
await sync_started.wait()
await asyncio.sleep(0.2)
assert not task.done(), "the mirror did not wait on the team's advisory lock"
await held.execute_raw(
'UPDATE "LiteLLM_TeamTable" SET access_group_ids = $1 WHERE team_id = $2',
[GROUPS[1]],
TEAM,
)
await asyncio.wait_for(task, timeout=30)
finally:
await blocker.disconnect()
assert await _read(db) == {GROUPS[0]: [], GROUPS[1]: [TEAM]}

238
tests/proxy_unit_tests/test_check_batch_cost.py
View file

@ -1791,3 +1791,241 @@ class TestBatchCostAttribution:
metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
assert metadata["user_api_key_alias"] == "prod-key"
class TestPollPageStarvation:
"""LIT-5462 regression: a row that can never be costed used to keep its slot in the
MAX_OBJECTS_PER_POLL_CYCLE page forever, so once enough of them accumulated no newer
batch was ever polled or costed."""
def _instance(self, prisma, llm_router):
from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
proxy_logging_obj = MagicMock()
proxy_logging_obj.get_proxy_hook.return_value = None
return CheckBatchCost(
proxy_logging_obj=proxy_logging_obj,
prisma_client=prisma,
llm_router=llm_router,
)
def _prisma(self, jobs):
prisma = MagicMock()
prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0)
prisma.db.litellm_managedobjecttable.update = AsyncMock()
prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=jobs)
prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
return prisma
def _job(self, job_id, unified_object_id):
job = MagicMock()
job.id = job_id
job.unified_object_id = unified_object_id
job.created_by = "user-1"
return job
@staticmethod
def _encode(unified_id: str) -> str:
import base64
return base64.urlsafe_b64encode(unified_id.encode()).decode().rstrip("=")
@pytest.mark.asyncio
async def test_unified_id_without_model_id_is_retired(self):
"""A unified id that decodes but carries no model_id is unroutable no matter what
the config says, so it must leave the poll page instead of being retried forever."""
prisma = self._prisma(
[self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]
)
llm_router = MagicMock()
llm_router.aretrieve_batch = AsyncMock()
await self._instance(prisma, llm_router).check_batch_cost()
llm_router.aretrieve_batch.assert_not_awaited()
prisma.db.litellm_managedobjecttable.update.assert_awaited_once()
call = prisma.db.litellm_managedobjecttable.update.call_args[1]
assert call["where"] == {"id": "job-no-model"}
assert call["data"] == {"batch_processed": True}
@pytest.mark.asyncio
async def test_provider_404_retires_job(self):
"""The provider dropping its record of the batch is permanent: no later retrieve
can succeed, so the row must stop occupying a slot."""
import litellm
prisma = self._prisma(
[
self._job(
"job-gone",
self._encode("litellm_proxy;model_id:model-123;llm_batch_id:batch_deadbeef"),
)
]
)
llm_router = MagicMock()
llm_router.aretrieve_batch = AsyncMock(
side_effect=litellm.NotFoundError(
message="No batch found with id 'batch_deadbeef'.",
model="model-123",
llm_provider="openai",
)
)
await self._instance(prisma, llm_router).check_batch_cost()
prisma.db.litellm_managedobjecttable.update.assert_awaited_once()
assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {
"batch_processed": True
}
@pytest.mark.asyncio
async def test_provider_404_with_deployment_gone_keeps_job(self):
"""With the batch's own deployment removed from the router, default fallbacks can
send the retrieve to a provider that never saw the batch. That 404 proves nothing,
so the row must stay unprocessed instead of losing its spend forever."""
import litellm
prisma = self._prisma(
[
self._job(
"job-misrouted",
self._encode("litellm_proxy;model_id:model-gone;llm_batch_id:batch_alive"),
)
]
)
llm_router = MagicMock()
llm_router.get_deployment = MagicMock(return_value=None)
llm_router.aretrieve_batch = AsyncMock(
side_effect=litellm.NotFoundError(
message="No batch found with id 'batch_alive'.",
model="model-gone",
llm_provider="openai",
)
)
await self._instance(prisma, llm_router).check_batch_cost()
prisma.db.litellm_managedobjecttable.update.assert_not_awaited()
@pytest.mark.asyncio
async def test_transient_provider_error_keeps_job_for_retry(self):
"""A failure that may clear up (timeout, 5xx) must still leave the row unprocessed."""
prisma = self._prisma(
[
self._job(
"job-flaky",
self._encode("litellm_proxy;model_id:model-123;llm_batch_id:batch_flaky"),
)
]
)
llm_router = MagicMock()
llm_router.aretrieve_batch = AsyncMock(side_effect=Exception("connection reset"))
await self._instance(prisma, llm_router).check_batch_cost()
prisma.db.litellm_managedobjecttable.update.assert_not_awaited()
@pytest.mark.asyncio
async def test_retirement_falls_back_to_status_without_batch_processed_column(self):
"""Older schemas have no batch_processed column, so the only way to stop selecting
the row is the status filter the poll query already applies."""
prisma = self._prisma(
[self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]
)
instance = self._instance(prisma, MagicMock())
instance._has_batch_processed_column = False
await instance.check_batch_cost()
assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {
"status": "stale_expired"
}
@pytest.mark.asyncio
async def test_stale_cleanup_gives_up_on_never_costed_completed_rows(self):
"""A row already in a terminal status is never rewritten by the staleness sweep, so
it needs its own bound or it starves newer batches indefinitely."""
prisma = self._prisma([])
await self._instance(prisma, MagicMock()).check_batch_cost()
calls = prisma.db.litellm_managedobjecttable.update_many.call_args_list
assert len(calls) == 2, "expected the staleness sweep plus the never-costed sweep"
where = calls[1][1]["where"]
assert where["file_purpose"] == "batch"
assert where["batch_processed"] is False
assert where["status"] == {"in": ["complete", "completed"]}
assert "created_at" in where
assert calls[1][1]["data"] == {"batch_processed": True}
@pytest.mark.asyncio
async def test_newer_batch_is_polled_once_dead_rows_are_retired(self):
"""The end state the customer cares about: dead rows retire on the cycle they are
first seen, and the healthy batch behind them keeps getting polled."""
dead_rows = [
self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model")),
self._job(
"job-gone",
self._encode("litellm_proxy;model_id:model-123;llm_batch_id:batch_deadbeef"),
),
]
live_row = self._job(
"job-live",
self._encode("litellm_proxy;model_id:model-123;llm_batch_id:batch_live"),
)
prisma = self._prisma(dead_rows + [live_row])
import litellm
in_progress = MagicMock()
in_progress.status = "in_progress"
async def _retrieve(model, batch_id, litellm_metadata):
if batch_id == "batch_deadbeef":
raise litellm.NotFoundError(
message=f"No batch found with id '{batch_id}'.",
model=model,
llm_provider="openai",
)
return in_progress
llm_router = MagicMock()
llm_router.aretrieve_batch = AsyncMock(side_effect=_retrieve)
await self._instance(prisma, llm_router).check_batch_cost()
retired = [
call[1]["where"]["id"]
for call in prisma.db.litellm_managedobjecttable.update.call_args_list
]
assert retired == ["job-no-model", "job-gone"]
assert (
llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live"
), "the newer healthy batch must still be polled in the same cycle"
@pytest.mark.asyncio
async def test_404_that_does_not_name_the_batch_keeps_job_for_retry(self):
"""A 404 about something other than the batch, e.g. a renamed Azure deployment, is
fixable in config, so the row must survive to be costed after the fix."""
import litellm
prisma = self._prisma(
[
self._job(
"job-bad-deployment",
self._encode("litellm_proxy;model_id:model-123;llm_batch_id:batch_real"),
)
]
)
llm_router = MagicMock()
llm_router.aretrieve_batch = AsyncMock(
side_effect=litellm.NotFoundError(
message="Error code: 404 - DeploymentNotFound",
model="model-123",
llm_provider="azure",
)
)
await self._instance(prisma, llm_router).check_batch_cost()
prisma.db.litellm_managedobjecttable.update.assert_not_awaited()

84
tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py
View file

@ -16,8 +16,8 @@ sys.path.insert(0, os.path.abspath("../../../../../.."))
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.bedrock.common_utils import (
ensure_bedrock_anthropic_messages_tool_names,
normalize_custom_field_on_tools,
normalize_tool_input_schema_types_for_bedrock_invoke,
remove_custom_field_from_tools,
)
from litellm.constants import (
BEDROCK_MIN_THINKING_BUDGET_TOKENS,
@ -353,12 +353,13 @@ def test_remove_ttl_from_cache_control():
assert request5 == {}
def test_remove_custom_field_from_tools():
def test_normalize_custom_field_on_tools():
"""
Ensure the `custom` field is stripped from every tool definition.
Ensure the `custom` field is stripped from every tool definition, and that a
boolean `custom.defer_loading` is hoisted onto the top-level `defer_loading`
flag Bedrock documents instead of being dropped with the wrapper.
Claude Code v2.1.69+ sends `custom: {defer_loading: true}` on tool
objects. Bedrock does not accept this extra field and returns
Bedrock does not accept a `custom` object on a tool and returns
"Extra inputs are not permitted".
Ref: https://github.com/BerriAI/litellm/issues/22847
@ -381,29 +382,94 @@ def test_remove_custom_field_from_tools():
]
}
remove_custom_field_from_tools(request)
normalize_custom_field_on_tools(request)
for tool in request["tools"]:
assert "custom" not in tool, f"Tool {tool['name']} still has 'custom' field"
# Other fields should be preserved
assert request["tools"][0]["name"] == "Read"
assert request["tools"][1]["name"] == "Write"
# `custom.defer_loading` is hoisted; the tool that never carried it is untouched
assert request["tools"][0]["defer_loading"] is True
assert "defer_loading" not in request["tools"][1]
# Case 2: request without tools key (should not raise error)
request2 = {"messages": [{"role": "user", "content": "hi"}]}
remove_custom_field_from_tools(request2)
normalize_custom_field_on_tools(request2)
assert "tools" not in request2
# Case 3: empty tools list (should not raise error)
request3 = {"tools": []}
remove_custom_field_from_tools(request3)
normalize_custom_field_on_tools(request3)
assert request3["tools"] == []
# Case 4: tools with None value (should not raise error)
request4 = {"tools": None}
remove_custom_field_from_tools(request4)
normalize_custom_field_on_tools(request4)
assert request4["tools"] is None
# Case 5: an explicit top-level flag wins over a conflicting wrapped one
request5 = {
"tools": [
{"name": "Read", "defer_loading": False, "custom": {"defer_loading": True}}
]
}
normalize_custom_field_on_tools(request5)
assert request5["tools"][0] == {"name": "Read", "defer_loading": False}
# Case 6: a non-boolean `custom.defer_loading` is dropped, never forwarded
for junk in ("true", 1, None, {"nested": True}):
request6 = {"tools": [{"name": "Read", "custom": {"defer_loading": junk}}]}
normalize_custom_field_on_tools(request6)
assert request6["tools"][0] == {"name": "Read"}, f"leaked defer_loading={junk!r}"
# Case 7: a `custom` that is not a dict is dropped without raising
request7 = {
"tools": [
{"name": "Read", "custom": "defer_loading"},
{"name": "Write", "custom": None},
]
}
normalize_custom_field_on_tools(request7)
assert request7["tools"] == [{"name": "Read"}, {"name": "Write"}]
@pytest.mark.parametrize(
"deferred_marker", [{"custom": {"defer_loading": True}}, {"defer_loading": True}]
)
def test_bedrock_invoke_messages_transform_emits_top_level_defer_loading(
deferred_marker,
):
"""A deferred tool must reach Bedrock as top-level ``defer_loading``, whether the
client wrapped the flag in ``custom`` or sent it top-level, and the outbound body
must still carry the Bedrock tool-search beta."""
from litellm.types.router import GenericLiteLLMParams
cfg = AmazonAnthropicClaudeMessagesConfig()
result = cfg.transform_anthropic_messages_request(
model="us.anthropic.claude-haiku-4-5-20251001-v1:0",
messages=[{"role": "user", "content": "hi"}],
anthropic_messages_optional_request_params={
"max_tokens": 128,
"stream": False,
"betas": ["advanced-tool-use-2025-11-20"],
"tools": [
{
"name": "Read",
"description": "Read a file",
"input_schema": {"type": "object", "properties": {}},
**deferred_marker,
},
{"type": "tool_search_tool_regex_20251119", "name": "tool_search"},
],
},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert result["tools"][0]["defer_loading"] is True
assert "custom" not in result["tools"][0]
assert result["anthropic_beta"] == ["tool-search-tool-2025-10-19"]
def test_normalize_tool_input_schema_types_for_bedrock_invoke():
"""

6
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
View file

@ -2874,7 +2874,7 @@ class TestMCPCustomHeaderName:
mock_general_settings.get.return_value = general_setting
# Call the method
result = MCPRequestHandler._get_mcp_client_side_auth_header_name()
result = MCPRequestHandler.get_mcp_client_side_auth_header_name()
# Assert the result
assert result == expected_header_name
@ -2938,7 +2938,7 @@ class TestMCPCustomHeaderName:
# Mock the header name method
with patch.object(
MCPRequestHandler,
"_get_mcp_client_side_auth_header_name",
"get_mcp_client_side_auth_header_name",
return_value=custom_header_name,
):
# Create headers from the test data
@ -2963,7 +2963,7 @@ class TestMCPCustomHeaderName:
# Mock the custom header name
with patch.object(
MCPRequestHandler,
"_get_mcp_client_side_auth_header_name",
"get_mcp_client_side_auth_header_name",
return_value="custom-auth-header",
):
# Create ASGI scope with custom header

42
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py
View file

@ -1196,3 +1196,45 @@ class TestOpenApiResolvedUpstreamAuth:
)
assert resolved is None
lookup.assert_not_awaited()
class TestPreCallToolCheckExposesClientHeaders:
"""The pre_mcp_call guardrail payload must carry the caller's sanitized HTTP headers."""
@pytest.mark.asyncio
async def test_sanitized_client_headers_reach_the_guardrail_payload(self):
manager = MCPServerManager()
server = MCPServer(
server_id="test-id",
name="test_server",
server_name="test_server",
url="https://example.com",
transport=MCPTransport.http,
auth_type=MCPAuth.none,
)
captured: Dict[str, Any] = {}
def capture(request_obj, kwargs):
captured.update(kwargs)
return {"model": "fake"}
proxy_logging = MagicMock(spec=ProxyLogging)
proxy_logging._create_mcp_request_object_from_kwargs = MagicMock(return_value=MagicMock())
proxy_logging._convert_mcp_to_llm_format = MagicMock(side_effect=capture)
proxy_logging.pre_call_hook = AsyncMock(return_value=None)
with patch.object(manager, "check_allowed_or_banned_tools", return_value=True):
with patch.object(manager, "check_tool_permission_for_key_team", new_callable=AsyncMock):
with patch.object(manager, "validate_allowed_params"):
await manager.pre_call_tool_check(
name="test_tool",
arguments={"key": "val"},
server_name="test_server",
user_api_key_auth=None,
proxy_logging_obj=proxy_logging,
server=server,
raw_headers={"x-nuid": "nuid-1", "x-litellm-api-key": "sk-proxy"},
)
assert captured["headers"] == {"x-nuid": "nuid-1"}

107
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
View file

@ -77,7 +77,7 @@ async def test_mcp_server_tool_call_body_contains_request_data():
# Mock the add_litellm_data_to_request function to capture the data
captured_data = {}
async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config):
async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config, **kwargs):
captured_data.update(data)
# Simulate the proxy_server_request creation
captured_data["proxy_server_request"] = {
@ -116,6 +116,107 @@ async def test_mcp_server_tool_call_body_contains_request_data():
assert body["arguments"] == tool_arguments
@pytest.mark.asyncio
async def test_mcp_server_tool_call_forwards_client_headers_to_logging():
"""The MCP protocol path must hand the connection's client headers to the pre-call
pipeline, so logging callbacks and guardrails see them the way the REST path does."""
try:
from litellm.proxy._experimental.mcp_server.server import (
mcp_server_tool_call,
set_auth_context,
)
except ImportError:
pytest.skip("MCP server not available")
set_auth_context(
UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
raw_headers={
"x-nuid": "nuid-1",
"x-app-id": "app-1",
"content-length": "42",
"x-forwarded-for": "9.9.9.9",
},
client_ip="1.2.3.4",
)
captured_headers = {}
async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config, **kwargs):
captured_headers.update(request.headers)
return data
async def mock_call_mcp_tool(*args, **kwargs):
return [{"type": "text", "text": "mocked response"}]
with patch(
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request",
mock_add_litellm_data_to_request,
):
with patch(
"litellm.proxy._experimental.mcp_server.server.call_mcp_tool",
mock_call_mcp_tool,
):
with patch("litellm.proxy.proxy_server.proxy_config", MagicMock()):
await mcp_server_tool_call("test_tool", {"param": "value"})
assert captured_headers.get("x-nuid") == "nuid-1"
assert captured_headers.get("x-app-id") == "app-1"
assert "content-length" not in captured_headers
assert captured_headers.get("x-forwarded-for") == "1.2.3.4"
@pytest.mark.asyncio
async def test_mcp_server_tool_call_strips_custom_litellm_key_header():
"""The deployment can rename the proxy key header via general_settings.litellm_key_header_name.
The pre-call pipeline only knows that name if it is passed in, so without it the virtual key
reaches metadata.headers and proxy_server_request.headers in plaintext."""
try:
from litellm.proxy._experimental.mcp_server.server import (
mcp_server_tool_call,
set_auth_context,
)
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
except ImportError:
pytest.skip("MCP server not available")
set_auth_context(
UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
raw_headers={"x-company-key": "sk-proxy-secret", "x-nuid": "nuid-1"},
client_ip="1.2.3.4",
)
captured_data = {}
async def capturing_add_litellm_data_to_request(**kwargs):
data = await add_litellm_data_to_request(**kwargs)
captured_data.update(data)
return data
async def mock_call_mcp_tool(*args, **kwargs):
return [{"type": "text", "text": "mocked response"}]
with patch(
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request",
capturing_add_litellm_data_to_request,
):
with patch(
"litellm.proxy._experimental.mcp_server.server.call_mcp_tool",
mock_call_mcp_tool,
):
with patch("litellm.proxy.proxy_server.proxy_config", MagicMock()):
with patch.dict(
"litellm.proxy.proxy_server.general_settings",
{"litellm_key_header_name": "x-company-key"},
clear=False,
):
await mcp_server_tool_call("test_tool", {"param": "value"})
metadata_headers = captured_data["metadata"]["headers"]
assert metadata_headers.get("x-nuid") == "nuid-1"
assert "x-company-key" not in metadata_headers
assert "x-company-key" not in captured_data["proxy_server_request"]["headers"]
@pytest.mark.asyncio
async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror():
"""The MCP session manager serializes handler exceptions as JSON-RPC errors, so a mid-session
@ -133,7 +234,7 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror():
set_auth_context(UserAPIKeyAuth(api_key="test_key", user_id="test_user"))
async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config):
async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config, **kwargs):
return data
async def mock_call_mcp_tool(*args, **kwargs):
@ -1245,7 +1346,7 @@ async def test_mcp_server_tool_call_body_with_none_arguments():
# Mock the add_litellm_data_to_request function to capture the data
captured_data = {}
async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config):
async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config, **kwargs):
captured_data.update(data)
captured_data["proxy_server_request"] = {
"url": str(request.url),

100
tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py
View file

@ -1,7 +1,11 @@
from unittest.mock import patch
import pytest
from fastapi import HTTPException
from litellm.proxy._experimental.mcp_server.utils import (
build_synthetic_mcp_request,
logging_safe_mcp_headers,
validate_and_normalize_mcp_server_payload,
validate_tool_display_names,
)
@ -47,3 +51,99 @@ class TestValidateAndNormalizeMcpServerPayload:
tool_name_to_display_name={"read_wiki_structure": "browse_repo_docs"},
)
validate_and_normalize_mcp_server_payload(payload)
class TestLoggingSafeMcpHeaders:
def test_returns_empty_for_missing_headers(self):
assert logging_safe_mcp_headers(None) == {}
assert logging_safe_mcp_headers({}) == {}
def test_exposes_custom_headers_and_masks_credentials(self):
safe = logging_safe_mcp_headers(
{
"x-nuid": "nuid-1",
"x-app-id": "app-1",
"x-litellm-api-key": "sk-proxy",
"cookie": "session=secret",
}
)
assert safe == {
"x-nuid": "nuid-1",
"x-app-id": "app-1",
"cookie": "***REDACTED***",
}
def test_strips_custom_litellm_key_header(self):
"""general_settings.litellm_key_header_name carries the proxy virtual key, so it must
never reach a callback or a guardrail even though clean_headers cannot know its name."""
with patch.dict(
"litellm.proxy.proxy_server.general_settings",
{"litellm_key_header_name": "x-company-key"},
clear=False,
):
safe = logging_safe_mcp_headers({"x-company-key": "sk-proxy", "x-nuid": "nuid-1"})
assert safe == {"x-nuid": "nuid-1"}
def test_strips_client_controlled_redaction_opt_out(self):
"""litellm-disable-message-redaction is read back out of the logged metadata to turn off
redaction, so leaving it in place lets any MCP client undo what an admin configured."""
safe = logging_safe_mcp_headers({"litellm-disable-message-redaction": "true", "x-nuid": "nuid-1"})
assert safe == {"x-nuid": "nuid-1"}
def test_strips_upstream_mcp_credentials(self):
safe = logging_safe_mcp_headers(
{
"x-mcp-auth": "Bearer upstream",
"x-mcp-github-authorization": "Bearer gh_token",
"x-mcp-zapier-x-api-key": "zapier-key",
"x-nuid": "nuid-1",
}
)
assert safe == {"x-nuid": "nuid-1"}
def test_strips_custom_mcp_client_side_auth_header(self):
with patch.dict(
"litellm.proxy.proxy_server.general_settings",
{"mcp_client_side_auth_header_name": "x-upstream-token"},
clear=False,
):
safe = logging_safe_mcp_headers({"x-upstream-token": "Bearer upstream", "x-nuid": "nuid-1"})
assert safe == {"x-nuid": "nuid-1"}
class TestBuildSyntheticMcpRequest:
def test_forwards_client_headers_without_upstream_credentials(self):
"""The synthetic request feeds add_litellm_data_to_request, which derives
metadata.headers, so upstream MCP credentials must not ride along."""
request = build_synthetic_mcp_request(
path="/mcp/tools/call",
raw_headers={
"x-nuid": "nuid-1",
"x-mcp-auth": "Bearer upstream",
"x-mcp-github-authorization": "Bearer gh_token",
},
)
assert request.headers.get("x-nuid") == "nuid-1"
assert "x-mcp-auth" not in request.headers
assert "x-mcp-github-authorization" not in request.headers
def test_drops_custom_litellm_key_header(self):
"""Callers such as the sampling flow build metadata off this request, so the
deployment's custom proxy key header must never be forwarded on it."""
with patch.dict(
"litellm.proxy.proxy_server.general_settings",
{"litellm_key_header_name": "x-company-key"},
clear=False,
):
request = build_synthetic_mcp_request(
path="/mcp/sampling/createMessage",
raw_headers={"x-company-key": "sk-proxy-secret", "x-nuid": "nuid-1"},
)
assert request.headers.get("x-nuid") == "nuid-1"
assert "x-company-key" not in request.headers

164
tests/test_litellm/proxy/common_utils/test_callback_utils.py
View file

@ -21,6 +21,10 @@ from litellm.proxy.common_utils.callback_utils import (
strip_callback_config,
)
import litellm
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from unittest.mock import patch
from litellm.proxy.common_utils.callback_utils import process_callback
@ -491,3 +495,163 @@ def test_strip_callback_config_drops_credential_bearing_slots():
@pytest.mark.parametrize("value", [None, "not-a-dict", 42])
def test_strip_callback_config_passes_through_non_dicts(value):
assert strip_callback_config(value) is value
# ---------------------------------------------------------------------------
# initialize_callbacks_on_proxy: dotted-path entries must resolve to something
# the request path can actually dispatch
# ---------------------------------------------------------------------------
_PROBE_MODULE_NAME = "custom_callback_probe"
_PROBE_MODULE_SOURCE = '''
from litellm.integrations.custom_logger import CustomLogger
class FloorMaxTokens(CustomLogger):
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
data["max_tokens"] = 16
return data
class NotALogger:
pass
def log_event_fn(kwargs, response_obj, start_time, end_time):
return None
NOT_A_CALLBACK = "some-plain-string"
proxy_handler_instance = FloorMaxTokens()
'''
@pytest.fixture
def probe_config_path(tmp_path):
"""Write a callback module next to a config.yaml, the layout get_instance_fn's file
branch expects, and restore every global the load + dispatch path touches.
``ProxyLogging._callback_capabilities_cache`` is keyed on the id()s of the
litellm.callbacks members, so an entry left behind here can be read back by an
unrelated test whose (len, ids) signature happens to collide.
"""
(tmp_path / f"{_PROBE_MODULE_NAME}.py").write_text(_PROBE_MODULE_SOURCE)
original_callbacks = (
list(litellm.callbacks) if isinstance(litellm.callbacks, list) else litellm.callbacks
)
litellm.callbacks = []
ProxyLogging._callback_capabilities_cache.clear()
try:
yield str(tmp_path / "config.yaml")
finally:
litellm.callbacks = original_callbacks
ProxyLogging._callback_capabilities_cache.clear()
def _load_callbacks(value, config_file_path):
initialize_callbacks_on_proxy(
value=value,
premium_user=False,
config_file_path=config_file_path,
litellm_settings={},
callback_specific_params={},
)
def test_initialize_callbacks_on_proxy_rejects_class_valued_entry(probe_config_path):
"""A class path loads an object that fails the `isinstance(_callback, CustomLogger)`
dispatch gate in ProxyLogging.pre_call_hook, so the proxy used to boot clean and
silently never run the hook. Config load must fail instead."""
entry = f"{_PROBE_MODULE_NAME}.FloorMaxTokens"
with pytest.raises(ValueError) as exc_info:
_load_callbacks([entry], probe_config_path)
message = str(exc_info.value)
assert entry in message
assert "the class" in message
assert "FloorMaxTokens" in message
assert f"{_PROBE_MODULE_NAME}.proxy_handler_instance" in message
assert litellm.callbacks == []
@pytest.mark.parametrize(
"attribute, expected_fragment",
[
("NotALogger", "the class"),
("NOT_A_CALLBACK", "str 'some-plain-string'"),
],
)
def test_initialize_callbacks_on_proxy_rejects_non_dispatchable_values(
probe_config_path, attribute, expected_fragment
):
entry = f"{_PROBE_MODULE_NAME}.{attribute}"
with pytest.raises(ValueError) as exc_info:
_load_callbacks([entry], probe_config_path)
message = str(exc_info.value)
assert entry in message
assert expected_fragment in message
assert litellm.callbacks == []
def test_initialize_callbacks_on_proxy_rejects_class_valued_non_list_value(probe_config_path):
entry = f"{_PROBE_MODULE_NAME}.FloorMaxTokens"
with pytest.raises(ValueError) as exc_info:
_load_callbacks(entry, probe_config_path)
assert entry in str(exc_info.value)
@pytest.mark.asyncio
async def test_initialize_callbacks_on_proxy_instance_entry_runs_pre_call_hook(probe_config_path):
"""Positive control: the supported shape must still load AND still run. Drives the
real ProxyLogging.pre_call_hook, which is where a class-valued entry goes silent."""
_load_callbacks([f"{_PROBE_MODULE_NAME}.proxy_handler_instance"], probe_config_path)
assert len(litellm.callbacks) == 1
assert isinstance(litellm.callbacks[0], CustomLogger)
ProxyLogging._callback_capabilities_cache.clear()
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
data = await proxy_logging.pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-probe"),
data={
"model": "gpt-4",
"messages": [{"role": "user", "content": "hello"}],
"max_tokens": 1,
"metadata": {},
},
call_type="acompletion",
)
assert data["max_tokens"] == 16
def test_initialize_callbacks_on_proxy_keeps_known_string_callback(probe_config_path):
"""Non-narrowing control: a known callback name never reaches get_instance_fn and
stays a plain string in litellm.callbacks."""
_load_callbacks(["langfuse"], probe_config_path)
assert litellm.callbacks == ["langfuse"]
def test_initialize_callbacks_on_proxy_accepts_plain_function_callback(probe_config_path):
"""Non-narrowing control: litellm.callbacks is typed
`Callable | <known name> | CustomLogger`, so a dotted path resolving to a plain
function is a supported shape and must keep loading."""
_load_callbacks([f"{_PROBE_MODULE_NAME}.log_event_fn"], probe_config_path)
assert [getattr(cb, "__name__", None) for cb in litellm.callbacks] == ["log_event_fn"]
def test_initialize_callbacks_on_proxy_accepts_instance_non_list_value(probe_config_path):
_load_callbacks(f"{_PROBE_MODULE_NAME}.proxy_handler_instance", probe_config_path)
assert len(litellm.callbacks) == 1
assert isinstance(litellm.callbacks[0], CustomLogger)

55
tests/test_litellm/proxy/conftest.py
View file

@ -43,32 +43,51 @@ def disconnected_prisma() -> DisconnectedPrisma:
return DisconnectedPrisma()
@pytest.fixture(autouse=True)
def _isolate_proxy_module_globals():
"""
Snapshot and restore module-level globals on litellm.proxy.proxy_server
that tests sometimes mutate via raw setattr (not monkeypatch).
_MODULE_GLOBAL_MISSING = object()
_proxy_module_globals_snapshot = pytest.StashKey[Dict[str, object]]()
Without this, a leaked value — e.g. master_key set by a sibling test —
@pytest.hookimpl(hookwrapper=True)
def pytest_runtest_setup(item):
"""
Snapshot module-level globals on litellm.proxy.proxy_server before any
fixture runs, and restore them in pytest_runtest_teardown after every
fixture finalizer has run.
Without this, a leaked value (e.g. master_key set by a sibling test)
flips the auth short-circuit in user_api_key_auth and causes unrelated
tests in the same xdist worker to return 401 instead of 200.
This must be a hook pair, not an autouse fixture: an autouse fixture in
the root conftest requests monkeypatch, so monkeypatch's undo stack
unwinds after every other fixture finalizer. A test that monkeypatches a
global while a fixture has it patched records the fixture's mock as the
"original", and monkeypatch.undo re-plants that mock after all restores
have run, poisoning the global for the rest of the xdist worker.
"""
from litellm.proxy import proxy_server
sentinel = object()
snapshot = {
name: getattr(proxy_server, name, sentinel)
item.stash[_proxy_module_globals_snapshot] = {
name: getattr(proxy_server, name, _MODULE_GLOBAL_MISSING)
for name in _PROXY_MODULE_GLOBALS_TO_ISOLATE
}
try:
yield
finally:
for name, value in snapshot.items():
if value is sentinel:
if hasattr(proxy_server, name):
delattr(proxy_server, name)
else:
setattr(proxy_server, name, value)
yield
@pytest.hookimpl(hookwrapper=True)
def pytest_runtest_teardown(item, nextitem):
yield
snapshot = item.stash.get(_proxy_module_globals_snapshot, None)
if snapshot is None:
return
from litellm.proxy import proxy_server
for name, value in snapshot.items():
if value is _MODULE_GLOBAL_MISSING:
if hasattr(proxy_server, name):
delattr(proxy_server, name)
else:
setattr(proxy_server, name, value)
@pytest.fixture(autouse=True)

171
tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py
View file

@ -3,6 +3,7 @@ Tests for SpendLogsPartitionManager: partition naming/bounds math, retention
selection, the non-partitioned no-op safety path, and the drop/ensure SQL flow.
"""
from contextlib import asynccontextmanager
from datetime import date, datetime, timezone
from unittest.mock import AsyncMock, MagicMock
@ -19,6 +20,46 @@ from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import (
)
DDL_TIMEOUT_MS = 30000
def _budget(ms: "int | None" = DDL_TIMEOUT_MS):
"""The injected per-statement bound: a callable re-read before each statement."""
return lambda: ms
def _wire_tx(db) -> list[str]:
"""
Model the prisma seam the partition DDL uses.
Every statement this manager issues, DDL and catalog query alike, runs inside
db.tx() so it can carry SET LOCAL timeouts. Those SET LOCAL statements are
collected in the returned list rather than forwarded, so assertions on
db.execute_raw and db.query_raw still see only the real statements.
"""
session_settings: list[str] = []
@asynccontextmanager
async def _tx():
tx = MagicMock()
async def _execute_raw(sql, *args):
if sql.lstrip().upper().startswith("SET LOCAL"):
session_settings.append(sql.strip())
return 0
return await db.execute_raw(sql, *args)
async def _query_raw(sql, *args):
return await db.query_raw(sql, *args)
tx.execute_raw = _execute_raw
tx.query_raw = _query_raw
yield tx
db.tx = _tx
return session_settings
def test_period_start_per_interval():
d = date(2026, 6, 3) # a Wednesday
assert period_start(d, "day") == date(2026, 6, 3)
@ -78,11 +119,13 @@ async def test_is_partitioned_true_and_false():
client_true = MagicMock()
client_true.db.query_raw = AsyncMock(return_value=[{"partitioned": True}])
assert await mgr.is_partitioned(client_true) is True
_wire_tx(client_true.db)
assert await mgr.is_partitioned(client_true, _budget()) is True
client_false = MagicMock()
client_false.db.query_raw = AsyncMock(return_value=[{"partitioned": False}])
assert await mgr.is_partitioned(client_false) is False
_wire_tx(client_false.db)
assert await mgr.is_partitioned(client_false, _budget()) is False
@pytest.mark.asyncio
@ -94,13 +137,14 @@ async def test_catalog_queries_are_scoped_to_current_schema():
mgr = SpendLogsPartitionManager()
client = MagicMock()
client.db.query_raw = AsyncMock(return_value=[])
_wire_tx(client.db)
await mgr.is_partitioned(client)
await mgr.is_partitioned(client, _budget())
is_partitioned_sql = client.db.query_raw.call_args.args[0]
assert "pg_namespace" in is_partitioned_sql
assert "current_schema()" in is_partitioned_sql
await mgr._list_partitions(client)
await mgr._list_partitions(client, DDL_TIMEOUT_MS)
list_sql = client.db.query_raw.call_args.args[0]
assert "pg_namespace" in list_sql
assert "current_schema()" in list_sql
@ -112,7 +156,10 @@ async def test_is_partitioned_swallows_errors_and_returns_false():
mgr = SpendLogsPartitionManager()
client = MagicMock()
client.db.query_raw = AsyncMock(side_effect=Exception("db down"))
assert await mgr.is_partitioned(client) is False
# Wire the real seam: without it the async with itself raises, and the test
# would pass on the wrong exception.
_wire_tx(client.db)
assert await mgr.is_partitioned(client, _budget()) is False
@pytest.mark.asyncio
@ -133,9 +180,10 @@ async def test_drop_partitions_older_than_drops_expired_only():
]
)
client.db.execute_raw = AsyncMock(return_value=0)
_wire_tx(client.db)
cutoff = datetime(2026, 6, 5, 0, 0, 0, tzinfo=timezone.utc)
dropped = await mgr.drop_partitions_older_than(client, cutoff)
dropped = await mgr.drop_partitions_older_than(client, cutoff, _budget())
assert dropped == ["LiteLLM_SpendLogs_p20260601"]
executed = " ".join(call.args[0] for call in client.db.execute_raw.call_args_list)
@ -149,8 +197,9 @@ async def test_ensure_partitions_issues_create_for_each_period():
mgr = SpendLogsPartitionManager(interval="day", precreate_ahead=2)
client = MagicMock()
client.db.execute_raw = AsyncMock(return_value=0)
_wire_tx(client.db)
created = await mgr.ensure_partitions(client)
created = await mgr.ensure_partitions(client, _budget())
assert len(created) == 3 # current + 2 ahead
assert client.db.execute_raw.await_count == 3
@ -159,6 +208,105 @@ async def test_ensure_partitions_issues_create_for_each_period():
assert "CREATE TABLE IF NOT EXISTS" in first_sql
@pytest.mark.asyncio
async def test_partition_ddl_carries_a_statement_and_lock_timeout():
"""
Partition DDL takes an ACCESS EXCLUSIVE lock, so an unbounded DROP queues
behind any long-running reader for as long as that reader lives. That is the
one path by which cleanup could outlast its run budget without bound, and
lock_timeout is what bounds the wait rather than only the work.
"""
mgr = SpendLogsPartitionManager(interval="day", precreate_ahead=0)
client = MagicMock()
client.db.execute_raw = AsyncMock(return_value=0)
client.db.query_raw = AsyncMock(
return_value=[
{
"name": "LiteLLM_SpendLogs_p20260601",
"bound": "FOR VALUES FROM ('2026-06-01 00:00:00') TO ('2026-06-02 00:00:00')",
}
]
)
session_settings = _wire_tx(client.db)
await mgr.ensure_partitions(client, _budget(7000))
await mgr.drop_partitions_older_than(client, datetime(2026, 6, 5, tzinfo=timezone.utc), _budget(7000))
# Three statements were issued: the CREATE, the catalog list the drop needs,
# and the DROP. All three carry a statement timeout; only the two that take
# a lock also carry a lock timeout, since the catalog read takes none.
assert session_settings.count("SET LOCAL statement_timeout = 7000") == 3
assert session_settings.count("SET LOCAL lock_timeout = 7000") == 2
@pytest.mark.asyncio
async def test_catalog_queries_carry_a_statement_timeout():
"""
Bounding only the DDL leaves the two catalog lookups as statements this job
issues with no bound at all, so a run could still outlast its budget waiting
on one. Every statement the manager issues carries the caller's timeout.
"""
mgr = SpendLogsPartitionManager()
client = MagicMock()
client.db.query_raw = AsyncMock(return_value=[])
session_settings = _wire_tx(client.db)
await mgr.is_partitioned(client, _budget(4000))
assert session_settings == ["SET LOCAL statement_timeout = 4000"], (
f"is_partitioned issued no statement timeout: {session_settings}"
)
session_settings.clear()
await mgr._list_partitions(client, 4000)
assert session_settings == ["SET LOCAL statement_timeout = 4000"], (
f"_list_partitions issued no statement timeout: {session_settings}"
)
@pytest.mark.asyncio
async def test_partition_loops_stop_when_the_budget_runs_out_mid_way():
"""
Each loop issues one statement per partition, so a bound read once at entry
would let N statements each run for the budget that was left before the
first of them. The bound is re-read per statement and the loop stops.
"""
mgr = SpendLogsPartitionManager(interval="day", precreate_ahead=4)
client = MagicMock()
client.db.execute_raw = AsyncMock(return_value=0)
_wire_tx(client.db)
# Budget for two statements, then spent.
calls = {"n": 0}
def budget() -> "int | None":
calls["n"] += 1
return 5000 if calls["n"] <= 2 else None
created = await mgr.ensure_partitions(client, budget)
assert len(created) == 2, f"the loop ran past its budget and created {len(created)}"
assert client.db.execute_raw.await_count == 2
@pytest.mark.asyncio
async def test_partition_maintenance_issues_nothing_when_the_budget_is_already_spent():
"""A run with no budget left must not issue even the catalog lookups."""
mgr = SpendLogsPartitionManager(interval="day", precreate_ahead=2)
client = MagicMock()
client.db.execute_raw = AsyncMock(return_value=0)
client.db.query_raw = AsyncMock(return_value=[])
_wire_tx(client.db)
spent = _budget(None)
assert await mgr.is_partitioned(client, spent) is False
assert await mgr.ensure_partitions(client, spent) == []
assert await mgr.drop_partitions_older_than(client, datetime(2026, 6, 5, tzinfo=timezone.utc), spent) == []
client.db.execute_raw.assert_not_awaited()
client.db.query_raw.assert_not_awaited()
def test_unsupported_interval_raises():
with pytest.raises(ValueError):
period_start(date(2026, 6, 1), "year")
@ -178,8 +326,9 @@ async def test_ensure_partitions_continues_when_one_create_fails():
mgr = SpendLogsPartitionManager(interval="day", precreate_ahead=2)
client = MagicMock()
client.db.execute_raw = AsyncMock(side_effect=[0, Exception("overlap"), 0])
_wire_tx(client.db)
created = await mgr.ensure_partitions(client)
created = await mgr.ensure_partitions(client, _budget())
# the failed partition is skipped, the others still created
assert len(created) == 2
@ -202,8 +351,9 @@ async def test_invalid_interval_does_not_abort_ensure_partitions():
mgr = SpendLogsPartitionManager(interval="fortnight", precreate_ahead=1)
client = MagicMock()
client.db.execute_raw = AsyncMock(return_value=0)
_wire_tx(client.db)
created = await mgr.ensure_partitions(client)
created = await mgr.ensure_partitions(client, _budget())
assert len(created) == 2 # current + 1 ahead, day-based fallback
@ -225,9 +375,10 @@ async def test_drop_partitions_continues_when_one_drop_fails():
]
)
client.db.execute_raw = AsyncMock(side_effect=[Exception("locked"), 0])
_wire_tx(client.db)
cutoff = datetime(2026, 6, 10, 0, 0, 0, tzinfo=timezone.utc)
dropped = await mgr.drop_partitions_older_than(client, cutoff)
dropped = await mgr.drop_partitions_older_than(client, cutoff, _budget())
# both were eligible; the first drop failed so only the second is reported
assert dropped == ["LiteLLM_SpendLogs_p20260602"]

741
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -791,6 +791,7 @@ async def test_generate_key_helper_fn_with_access_group_ids(monkeypatch):
mock_prisma_client.db.litellm_objectpermissiontable.create = AsyncMock(
return_value=MagicMock(object_permission_id=None)
)
mock_prisma_client.db.query_raw = AsyncMock(return_value=[])
captured_key_data = {}
@ -15703,3 +15704,743 @@ async def test_key_generate_omitted_budget_duration_still_filled_by_upperbound(m
assert key_row["budget_duration"] == "30d"
assert key_row["budget_reset_at"] is not None
from litellm.proxy.management_helpers.access_group_key_sync import (
_ATTACH_KEY_SQL,
_DETACH_KEY_SQL,
_REPOINT_KEY_SQL,
)
ACCESS_GROUP_SYNC_TOKEN = "0d62f396c1317066f55a96086517047c737087c61eb2bf016b72e6298927b15b"
def _access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups):
"""
Back the access group table with an in-memory dict so the sync's writes are observable.
The sync writes through guarded set-based SQL statements, so this emulates exactly what
Postgres does with them, including the guards that make each one idempotent and the
`RETURNING` clause that reports which groups actually moved.
"""
def _repoint(previous_token, new_token):
moved = [
group_id
for group_id, stored in access_groups.items()
if previous_token in stored["assigned_key_ids"]
]
for group_id in moved:
current = access_groups[group_id]["assigned_key_ids"]
access_groups[group_id]["assigned_key_ids"] = [
*(t for t in current if t not in (previous_token, new_token)),
new_token,
]
return moved
def _attach(key_token, access_group_ids):
moved = [
group_id
for group_id in access_group_ids
if group_id in access_groups
and key_token not in access_groups[group_id]["assigned_key_ids"]
]
for group_id in moved:
stored = access_groups[group_id]
stored["assigned_key_ids"] = [*stored["assigned_key_ids"], key_token]
return moved
def _detach(key_token, access_group_ids):
moved = [
group_id
for group_id in access_group_ids
if group_id in access_groups
and key_token in access_groups[group_id]["assigned_key_ids"]
]
for group_id in moved:
stored = access_groups[group_id]
stored["assigned_key_ids"] = [
t for t in stored["assigned_key_ids"] if t != key_token
]
return moved
async def _query_raw(query, *args):
if query == _REPOINT_KEY_SQL:
moved = _repoint(*args)
elif query == _ATTACH_KEY_SQL:
moved = _attach(*args)
else:
assert query == _DETACH_KEY_SQL, f"unexpected statement: {query}"
moved = _detach(*args)
return [{"access_group_id": group_id} for group_id in moved]
raw_mock = AsyncMock(side_effect=_query_raw)
mock_prisma_client.db.query_raw = raw_mock
return raw_mock
async def _authorized_models_for_key(access_groups, token, key_access_group_ids):
"""Run the real auth-time reader against the post-sync access group rows."""
from litellm.proxy._types import LiteLLM_AccessGroupTable, LiteLLM_TeamTable
from litellm.proxy.auth.auth_checks import (
get_authorized_resources_from_key_access_groups,
)
async def _get_access_object(*, access_group_id, **_kwargs):
stored = access_groups[access_group_id]
return LiteLLM_AccessGroupTable(
access_group_id=access_group_id,
access_group_name=access_group_id,
access_model_names=list(stored["access_model_names"]),
assigned_team_ids=[],
assigned_key_ids=list(stored["assigned_key_ids"]),
)
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()),
patch(
"litellm.proxy.auth.auth_checks.get_access_object",
new_callable=AsyncMock,
side_effect=_get_access_object,
),
):
return await get_authorized_resources_from_key_access_groups(
valid_token=UserAPIKeyAuth(
token=token,
models=[],
team_id="team-a",
access_group_ids=list(key_access_group_ids),
),
team_object=LiteLLM_TeamTable(team_id="team-a", models=[]),
resource_field="access_model_names",
)
@pytest.mark.asyncio
async def test_update_key_syncs_access_group_assigned_key_ids_in_both_directions(
monkeypatch,
):
"""
A key-side edit of `access_group_ids` must be mirrored onto every affected access
group's `assigned_key_ids`, in one operation, in both directions.
`assigned_key_ids` is not display-only. `get_authorized_resources_from_key_access_groups`
reads it as an authorization input and authorizes only when the group lists the key's
token, so a group the key just added must start granting its resources and a group the
key dropped must stop. A single-direction assertion would pass against a fix that only
ever adds (or only ever removes), so this covers add, remove, untouched, and the
authorization consequence of each.
"""
from litellm.proxy.management_endpoints.key_management_endpoints import (
update_key_fn,
)
key_in_db = LiteLLM_VerificationToken(
token=ACCESS_GROUP_SYNC_TOKEN,
user_id="test-user",
access_group_ids=["ag-drop", "ag-keep"],
)
access_groups = {
"ag-drop": {
"assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN],
"access_model_names": ["dropped-model"],
},
"ag-keep": {
"assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN],
"access_model_names": ["kept-model"],
},
"ag-add": {"assigned_key_ids": [], "access_model_names": ["added-model"]},
}
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=key_in_db
)
mock_prisma_client.db.litellm_verificationtoken.find_first = AsyncMock(
return_value=None
)
mock_prisma_client.update_data = AsyncMock(return_value={"data": {}})
raw_mock = _access_group_table_mocks(
monkeypatch, mock_prisma_client, access_groups
)
_setup_update_key_mocks(monkeypatch, mock_prisma_client)
with (
patch(
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
new_callable=AsyncMock,
),
patch(
"litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache",
new_callable=AsyncMock,
) as invalidate_cache,
):
await update_key_fn(
request=MagicMock(),
data=UpdateKeyRequest(
key=ACCESS_GROUP_SYNC_TOKEN, access_group_ids=["ag-keep", "ag-add"]
),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin-user",
),
litellm_changed_by=None,
)
assert access_groups["ag-drop"]["assigned_key_ids"] == []
assert access_groups["ag-add"]["assigned_key_ids"] == [ACCESS_GROUP_SYNC_TOKEN]
assert access_groups["ag-keep"]["assigned_key_ids"] == [ACCESS_GROUP_SYNC_TOKEN]
# Both halves go out as single guarded statements. A read-modify-write here lets two
# admins editing one group lose each other's change: an attach can vanish, and a detach
# can put an already revoked token back and restore its grants.
assert sorted(call.args for call in raw_mock.call_args_list) == sorted(
[
(_ATTACH_KEY_SQL, ACCESS_GROUP_SYNC_TOKEN, ["ag-add"]),
(_DETACH_KEY_SQL, ACCESS_GROUP_SYNC_TOKEN, ["ag-drop"]),
]
)
assert {call.args[0] for call in invalidate_cache.call_args_list} == {
"ag-drop",
"ag-add",
}
authorized_models = await _authorized_models_for_key(
access_groups,
ACCESS_GROUP_SYNC_TOKEN,
["ag-drop", "ag-keep", "ag-add"],
)
assert sorted(authorized_models) == ["added-model", "kept-model"]
@pytest.mark.asyncio
async def test_update_key_leaves_access_groups_alone_when_field_is_unset(monkeypatch):
"""
An update that never mentions `access_group_ids` must not touch the group rows.
`prepare_key_update_data` writes from `model_dump(exclude_unset=True)`, so an omitted
field leaves the key row's own list intact. Reading the request attribute instead of
its `model_fields_set` would see None and wipe every group's copy of the token on any
unrelated edit, e.g. a max_budget change.
"""
from litellm.proxy.management_endpoints.key_management_endpoints import (
update_key_fn,
)
key_in_db = LiteLLM_VerificationToken(
token=ACCESS_GROUP_SYNC_TOKEN,
user_id="test-user",
access_group_ids=["ag-keep"],
)
access_groups = {
"ag-keep": {
"assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN],
"access_model_names": ["kept-model"],
},
}
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=key_in_db
)
mock_prisma_client.db.litellm_verificationtoken.find_first = AsyncMock(
return_value=None
)
mock_prisma_client.update_data = AsyncMock(return_value={"data": {}})
raw_mock = _access_group_table_mocks(
monkeypatch, mock_prisma_client, access_groups
)
_setup_update_key_mocks(monkeypatch, mock_prisma_client)
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
new_callable=AsyncMock,
):
await update_key_fn(
request=MagicMock(),
data=UpdateKeyRequest(key=ACCESS_GROUP_SYNC_TOKEN, max_budget=50.0),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin-user",
),
litellm_changed_by=None,
)
raw_mock.assert_not_called()
assert access_groups["ag-keep"]["assigned_key_ids"] == [ACCESS_GROUP_SYNC_TOKEN]
assert await _authorized_models_for_key(
access_groups, ACCESS_GROUP_SYNC_TOKEN, ["ag-keep"]
) == ["kept-model"]
@pytest.mark.asyncio
async def test_bulk_update_keys_syncs_access_group_assigned_key_ids(monkeypatch):
"""
/key/bulk_update and /team/keys/bulk_update reach the DB through
`_process_single_key_update`, which is a separate write path from /key/update's own
inline one. Both have to maintain the group's copy or a bulk attach grants nothing.
"""
key_in_db = LiteLLM_VerificationToken(
token=ACCESS_GROUP_SYNC_TOKEN,
user_id="test-user",
access_group_ids=["ag-drop"],
)
access_groups = {
"ag-drop": {
"assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN],
"access_model_names": ["dropped-model"],
},
"ag-add": {"assigned_key_ids": [], "access_model_names": ["added-model"]},
}
mock_prisma_client = AsyncMock()
mock_prisma_client.update_data = AsyncMock(return_value={"data": {}})
_access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups)
_setup_update_key_mocks(monkeypatch, mock_prisma_client)
with (
patch(
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
new_callable=AsyncMock,
),
patch(
"litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache",
new_callable=AsyncMock,
),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
new_callable=AsyncMock,
),
):
await _process_single_key_update(
update_key_request=UpdateKeyRequest(
key=ACCESS_GROUP_SYNC_TOKEN, access_group_ids=["ag-add"]
),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin-user",
),
litellm_changed_by=None,
prisma_client=mock_prisma_client,
user_api_key_cache=AsyncMock(),
proxy_logging_obj=MagicMock(),
llm_router=None,
existing_key_row=key_in_db,
)
assert access_groups["ag-drop"]["assigned_key_ids"] == []
assert access_groups["ag-add"]["assigned_key_ids"] == [ACCESS_GROUP_SYNC_TOKEN]
assert await _authorized_models_for_key(
access_groups, ACCESS_GROUP_SYNC_TOKEN, ["ag-drop", "ag-add"]
) == ["added-model"]
@pytest.mark.asyncio
async def test_delete_key_withdraws_token_from_its_access_groups(monkeypatch):
"""
Deleting a key must withdraw its token from every group that lists it.
Without the withdrawal the group keeps a token that no longer resolves to a row, so
the access group page lists a key that does not exist and the list grows without bound.
"""
key_in_db = LiteLLM_VerificationToken(
token=ACCESS_GROUP_SYNC_TOKEN,
user_id="test-user",
access_group_ids=["ag-keep"],
)
access_groups = {
"ag-keep": {
"assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN, "other-key"],
"access_model_names": ["kept-model"],
},
}
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[key_in_db]
)
mock_prisma_client.delete_data = AsyncMock(return_value={"deleted_keys": 1})
mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock()
_access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups)
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client", mock_prisma_client
)
mock_cache = MagicMock()
mock_cache.delete_cache = MagicMock()
with patch(
"litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache",
new_callable=AsyncMock,
):
await delete_verification_tokens(
tokens=[ACCESS_GROUP_SYNC_TOKEN],
user_api_key_cache=mock_cache,
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin-user",
),
litellm_changed_by="admin-user",
)
assert access_groups["ag-keep"]["assigned_key_ids"] == ["other-key"]
@pytest.mark.asyncio
async def test_generate_key_records_token_in_its_access_groups(monkeypatch):
"""
/key/generate with `access_group_ids` must record the new token on the group side.
The key row's own list alone does not authorize: the group has to list the token back
or `get_authorized_resources_from_key_access_groups` contributes nothing, so a key
created against a group silently gets none of its models.
"""
access_groups = {
"ag-add": {"assigned_key_ids": [], "access_model_names": ["added-model"]},
}
created_key = MagicMock()
created_key.token = ACCESS_GROUP_SYNC_TOKEN
created_key.litellm_budget_table = None
created_key.created_at = None
created_key.updated_at = None
mock_prisma_client = AsyncMock()
mock_prisma_client.insert_data = AsyncMock(return_value=created_key)
_access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups)
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client", mock_prisma_client
)
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
monkeypatch.setattr("litellm.store_audit_logs", False)
with patch(
"litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache",
new_callable=AsyncMock,
):
await generate_key_helper_fn(
request_type="key",
access_group_ids=["ag-add"],
table_name="key",
user_id="test-user",
)
assert access_groups["ag-add"]["assigned_key_ids"] == [ACCESS_GROUP_SYNC_TOKEN]
assert await _authorized_models_for_key(
access_groups, ACCESS_GROUP_SYNC_TOKEN, ["ag-add"]
) == ["added-model"]
@pytest.mark.asyncio
async def test_regenerate_key_repoints_access_group_assigned_key_ids(monkeypatch):
"""
Regeneration replaces the key's token, which is the identity `assigned_key_ids` stores.
Leaving the old hash behind points the group at a token that no longer exists AND
denies the regenerated key the group's grants, so the group's copy has to be
re-pointed from the old hash to the new one in the same operation.
"""
from litellm.proxy._types import RegenerateKeyRequest
from litellm.proxy.management_endpoints.key_management_endpoints import (
_execute_virtual_key_regeneration,
)
from litellm.proxy.utils import hash_token
new_token_hash = hash_token("sk-newtoken1234ab12")
existing_key = LiteLLM_VerificationToken(
token="abc123",
user_id="user-1",
models=["gpt-4"],
access_group_ids=["ag-keep"],
)
access_groups = {
"ag-keep": {
"assigned_key_ids": ["abc123"],
"access_model_names": ["kept-model"],
},
}
mock_prisma_client = _make_regenerate_mock_prisma()
_access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups)
with (
patch(
"litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
new_callable=AsyncMock,
return_value="sk-newtoken1234ab12",
),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
new_callable=AsyncMock,
),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
new_callable=AsyncMock,
),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook",
new_callable=AsyncMock,
),
patch(
"litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache",
new_callable=AsyncMock,
),
):
await _execute_virtual_key_regeneration(
prisma_client=mock_prisma_client,
key_in_db=existing_key,
hashed_api_key="abc123",
key="abc123",
data=RegenerateKeyRequest(),
user_api_key_dict=_make_regenerate_user_api_key_dict(),
litellm_changed_by=None,
user_api_key_cache=MagicMock(),
proxy_logging_obj=MagicMock(),
)
assert access_groups["ag-keep"]["assigned_key_ids"] == [new_token_hash]
assert await _authorized_models_for_key(
access_groups, new_token_hash, ["ag-keep"]
) == ["kept-model"]
assert (
await _authorized_models_for_key(access_groups, "abc123", ["ag-keep"]) == []
)
@pytest.mark.asyncio
async def test_key_write_paths_revoke_the_key_cache_before_syncing_access_groups(
monkeypatch,
):
"""
Credential invalidation must not sit behind the group sync on any key write path.
The cached auth object still carries the key's old `access_group_ids`, so if the sync
raises first, the request fails with the key still authenticating against groups it
just lost, until that entry expires. Ordering it last means a failed sync degrades to
the stale listing this PR fixes rather than to a stale grant.
"""
from litellm.proxy.management_endpoints.key_management_endpoints import (
update_key_fn,
)
order = []
key_in_db = LiteLLM_VerificationToken(
token=ACCESS_GROUP_SYNC_TOKEN,
user_id="test-user",
access_group_ids=["ag-drop"],
)
access_groups = {
"ag-drop": {
"assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN],
"access_model_names": ["dropped-model"],
},
}
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=key_in_db
)
mock_prisma_client.db.litellm_verificationtoken.find_first = AsyncMock(
return_value=None
)
mock_prisma_client.update_data = AsyncMock(return_value={"data": {}})
_access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups)
mock_prisma_client.db.query_raw = AsyncMock(
side_effect=lambda *a, **k: order.append("sync") or []
)
_setup_update_key_mocks(monkeypatch, mock_prisma_client)
with (
patch(
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
new_callable=AsyncMock,
side_effect=lambda **kwargs: order.append("revoke_key_cache"),
),
patch(
"litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache",
new_callable=AsyncMock,
),
):
await update_key_fn(
request=MagicMock(),
data=UpdateKeyRequest(key=ACCESS_GROUP_SYNC_TOKEN, access_group_ids=[]),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin-user",
),
litellm_changed_by=None,
)
assert order == ["revoke_key_cache", "sync"]
@pytest.mark.asyncio
async def test_update_key_syncs_many_access_groups_in_one_statement_per_direction(
monkeypatch,
):
"""
The number of groups on a request must not become a matching number of round trips.
Anyone allowed to assign access groups picks the size of `access_group_ids`, so a
per-group statement lets one /key/update hold a connection for hundreds of sequential
writes. Both halves are set-based, so the cost is two statements no matter the size.
"""
from litellm.proxy.management_endpoints.key_management_endpoints import (
update_key_fn,
)
dropped = [f"ag-drop-{i}" for i in range(60)]
added = [f"ag-add-{i}" for i in range(60)]
key_in_db = LiteLLM_VerificationToken(
token=ACCESS_GROUP_SYNC_TOKEN,
user_id="test-user",
access_group_ids=dropped,
)
access_groups = {
**{
group_id: {
"assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN],
"access_model_names": [f"{group_id}-model"],
}
for group_id in dropped
},
**{
group_id: {"assigned_key_ids": [], "access_model_names": [f"{group_id}-model"]}
for group_id in added
},
}
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=key_in_db
)
mock_prisma_client.db.litellm_verificationtoken.find_first = AsyncMock(
return_value=None
)
mock_prisma_client.update_data = AsyncMock(return_value={"data": {}})
raw_mock = _access_group_table_mocks(
monkeypatch, mock_prisma_client, access_groups
)
_setup_update_key_mocks(monkeypatch, mock_prisma_client)
with (
patch(
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
new_callable=AsyncMock,
),
patch(
"litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache",
new_callable=AsyncMock,
),
):
await update_key_fn(
request=MagicMock(),
data=UpdateKeyRequest(
key=ACCESS_GROUP_SYNC_TOKEN, access_group_ids=added
),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin-user",
),
litellm_changed_by=None,
)
assert [call.args[0] for call in raw_mock.call_args_list] == [
_ATTACH_KEY_SQL,
_DETACH_KEY_SQL,
]
assert sorted(raw_mock.call_args_list[0].args[2]) == sorted(added)
assert sorted(raw_mock.call_args_list[1].args[2]) == sorted(dropped)
assert all(
access_groups[group_id]["assigned_key_ids"] == [ACCESS_GROUP_SYNC_TOKEN]
for group_id in added
)
assert all(access_groups[group_id]["assigned_key_ids"] == [] for group_id in dropped)
@pytest.mark.asyncio
async def test_regenerate_key_repoints_live_membership_not_the_key_row_it_read(
monkeypatch,
):
"""
Regeneration must move whatever the groups hold when it writes, not the key row's list.
That list is read before the new token exists, so replaying it re-adds the key to a
group an admin revoked in between and leaves the dead hash in a group an admin attached
in between, which silently restores one grant and drops another.
"""
from litellm.proxy._types import RegenerateKeyRequest
from litellm.proxy.management_endpoints.key_management_endpoints import (
_execute_virtual_key_regeneration,
)
from litellm.proxy.utils import hash_token
new_token_hash = hash_token("sk-newtoken1234ab12")
existing_key = LiteLLM_VerificationToken(
token="abc123",
user_id="user-1",
models=["gpt-4"],
access_group_ids=["ag-revoked-since"],
)
access_groups = {
"ag-revoked-since": {
"assigned_key_ids": [],
"access_model_names": ["revoked-model"],
},
"ag-attached-since": {
"assigned_key_ids": ["abc123"],
"access_model_names": ["attached-model"],
},
}
mock_prisma_client = _make_regenerate_mock_prisma()
_access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups)
with (
patch(
"litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
new_callable=AsyncMock,
return_value="sk-newtoken1234ab12",
),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
new_callable=AsyncMock,
),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
new_callable=AsyncMock,
),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook",
new_callable=AsyncMock,
),
patch(
"litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache",
new_callable=AsyncMock,
),
):
await _execute_virtual_key_regeneration(
prisma_client=mock_prisma_client,
key_in_db=existing_key,
hashed_api_key="abc123",
key="abc123",
data=RegenerateKeyRequest(),
user_api_key_dict=_make_regenerate_user_api_key_dict(),
litellm_changed_by=None,
user_api_key_cache=MagicMock(),
proxy_logging_obj=MagicMock(),
)
assert access_groups["ag-revoked-since"]["assigned_key_ids"] == []
assert access_groups["ag-attached-since"]["assigned_key_ids"] == [new_token_hash]
assert await _authorized_models_for_key(
access_groups, new_token_hash, ["ag-revoked-since", "ag-attached-since"]
) == ["attached-model"]

403
tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py
View file

@ -4,6 +4,7 @@ import datetime
import json
from contextlib import ExitStack
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import patch as patch_ctx
import pytest
from fastapi import HTTPException
@ -14,15 +15,28 @@ from litellm.proxy._types import (
ReconcileOutcome,
UserAPIKeyAuth,
)
from litellm.proxy.auth.auth_checks import _is_model_cost_zero
from litellm.proxy.management_endpoints.model_management_endpoints import (
_PTU_ZEROED_PRICING_FIELDS,
_merged_ptu_model_info,
_update_team_model_in_db,
_ptu_priced_deployment,
_ptu_zeroed_pricing,
_raise_if_ptu_cost_attribution_disabled,
_validate_ptu_model_info,
add_new_model,
update_db_model,
)
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment
from litellm.router import Router
from litellm.types.router import (
SPECIAL_MODEL_INFO_PARAMS,
Deployment,
LiteLLM_Params,
ModelInfo,
updateDeployment,
updateLiteLLMParams,
)
def test_model_info_accepts_valid_ptu_fields():
@ -717,3 +731,390 @@ class TestAddNewModelPtuGate:
assert result.model_id == "ptu-gate-model"
add_team_model_to_db.assert_called_once()
class TestPtuDeploymentsAreNotBilledPerToken:
"""Reserved capacity is billed by the flat cost the rollup writes, so a PTU deployment must
not also bill the traffic that capacity serves."""
PTU = {"ptu_count": 15, "cost_per_ptu_per_hour": 2.0}
@pytest.fixture(autouse=True)
def _flag_on(self, monkeypatch):
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
# update_db_model encrypts every litellm_params value it is handed, and the salt falls
# back to the master key the proxy sets at boot, which no unit test has.
monkeypatch.setenv("LITELLM_SALT_KEY", "test-salt-key")
@staticmethod
def _zeroed(model_info=None, litellm_params=None, supplied=None):
return _ptu_zeroed_pricing(
model_info=model_info if model_info is not None else {},
litellm_params=litellm_params if litellm_params is not None else {},
supplied=supplied if supplied is not None else {},
)
def test_a_deployment_without_ptu_config_keeps_its_pricing(self):
assert self._zeroed(model_info={"team_id": "t"}, litellm_params={"input_cost_per_token": 5e-07}) == {}
def test_a_half_set_pair_is_not_treated_as_ptu(self):
assert self._zeroed(model_info={"ptu_count": 15}) == {}
def test_every_field_the_cost_map_could_fill_is_zeroed(self):
assert self._zeroed(model_info=self.PTU) == dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0)
def test_nothing_is_zeroed_while_the_feature_is_disabled(self, monkeypatch):
monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
assert self._zeroed(model_info=self.PTU) == {}
@pytest.mark.parametrize("field", ["input_cost_per_token", "cache_read_input_token_cost", "input_cost_per_second"])
def test_a_price_the_caller_supplies_is_refused(self, field):
"""Every custom-pricing field, not only the mirrored ones: per-second pricing bills a
PTU deployment just as surely as per-token pricing does."""
with pytest.raises(HTTPException) as exc:
self._zeroed(model_info=self.PTU, supplied={field: 5e-07})
assert exc.value.status_code == 400
assert field in str(exc.value.detail)
def test_a_price_the_caller_supplies_as_zero_is_accepted(self):
assert self._zeroed(model_info={**self.PTU, "input_cost_per_token": 0}, supplied={"input_cost_per_token": 0})[
"input_cost_per_token"
] == 0
def test_a_price_already_on_the_row_is_zeroed_rather_than_refused(self):
"""A row priced through a path this rule does not cover must heal on its next save. The
alternative refuses every later edit of a field that has nothing to do with pricing."""
zeroed = self._zeroed(model_info={**self.PTU, "input_cost_per_second": 3.0}, litellm_params={})
assert zeroed["input_cost_per_second"] == 0
assert zeroed["input_cost_per_token"] == 0
@pytest.mark.asyncio
async def test_a_refused_price_does_not_leave_the_team_changed(self):
"""The team ACL write autocommits, so the refusal has to run before it. Otherwise a
rejected edit grants the team a model whose settings were never saved."""
db_model = Deployment(
model_name="gpt-4o",
litellm_params=LiteLLM_Params(model="openai/gpt-4o"),
model_info=ModelInfo(
id="dep-0",
team_id="team-1",
ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
**self.PTU,
),
)
patch = updateDeployment(
litellm_params=updateLiteLLMParams(model="openai/gpt-4o", input_cost_per_token=5e-07),
model_info=ModelInfo(id="dep-0", team_id="team-2"),
)
endpoints = "litellm.proxy.management_endpoints.model_management_endpoints"
setup_new = AsyncMock()
update_existing = AsyncMock()
with ExitStack() as stack:
stack.enter_context(
patch_ctx(f"{endpoints}.ModelManagementAuthChecks.allow_team_model_action", AsyncMock(return_value=True))
)
stack.enter_context(patch_ctx(f"{endpoints}._setup_new_team_model_assignment", setup_new))
stack.enter_context(patch_ctx(f"{endpoints}._update_existing_team_model_assignment", update_existing))
stack.enter_context(patch_ctx("litellm.proxy.proxy_server.premium_user", True))
with pytest.raises(HTTPException) as exc:
await _update_team_model_in_db(
db_model=db_model,
patch_data=patch,
user_api_key_dict=UserAPIKeyAuth(user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN),
prisma_client=MagicMock(),
)
assert exc.value.status_code == 400
setup_new.assert_not_called()
update_existing.assert_not_called()
def test_a_setting_that_is_not_a_charge_is_left_alone(self):
"""CustomPricingLiteLLMParams also carries an embedding's output vector size and the
regional uplift multipliers. Zeroing one of those destroys the deployment's config, and
refusing it answers with a message calling a setting a charge."""
priced = _ptu_priced_deployment(
Deployment(
model_name="embeddings",
litellm_params=LiteLLM_Params(
model="azure/text-embedding-3-large",
output_vector_size=1536,
regional_processing_uplift_multiplier_eu=1.15,
),
model_info=ModelInfo(
id="dep-emb",
team_id="team-1",
ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
**self.PTU,
),
)
)
assert priced.litellm_params.get("output_vector_size") == 1536
assert priced.litellm_params.get("regional_processing_uplift_multiplier_eu") == 1.15
assert priced.litellm_params.get("input_cost_per_token") == 0
def test_removing_ptu_config_releases_every_rate_it_zeroed(self):
"""The zeroing covers any stored rate, so a release that only spans the mirrored fields
leaves a per-second deployment billing nothing for that dimension forever."""
on = update_db_model(
db_model=Deployment(
model_name="audio",
litellm_params=LiteLLM_Params(model="azure/whisper", input_cost_per_second=0.006),
model_info=ModelInfo(id="dep-audio", team_id="t"),
),
updated_patch=updateDeployment(
model_info=ModelInfo(
id="dep-audio",
team_id="t",
ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
**self.PTU,
)
),
)
assert json.loads(on["litellm_params"])["input_cost_per_second"] == 0
off = update_db_model(
db_model=Deployment(
model_name="audio",
litellm_params=LiteLLM_Params(**json.loads(on["litellm_params"])),
model_info=ModelInfo(**json.loads(on["model_info"])),
),
updated_patch=updateDeployment(
model_info=ModelInfo(id="dep-audio", ptu_count=None, cost_per_ptu_per_hour=None)
),
)
assert "input_cost_per_second" not in json.loads(off["litellm_params"])
@pytest.mark.parametrize(
"backend", ["azure/gpt-4o", "anthropic/claude-sonnet-4-5", "bedrock/anthropic.claude-sonnet-4-20250514-v1:0"]
)
def test_the_cost_map_contributes_no_price_to_a_priced_ptu_deployment(self, backend):
"""The acceptance criterion, read off the entry the router registers for the deployment.
Zeroing only the per-token pair leaves the cache-tier fields unset, which is exactly what
Router._inherit_builtin_cache_pricing back-fills from the public cost map, so a cached
prompt would still be billed at the public rate."""
priced = _ptu_priced_deployment(
Deployment(
model_name="ptu-deployment",
litellm_params=LiteLLM_Params(model=backend, api_key="fake-key"),
model_info=ModelInfo(
id="dep-ptu",
team_id="team-1",
ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
**self.PTU,
),
)
)
registered = Router._deployment_model_cost_payload(priced)
charged = {k: v for k, v in registered.items() if "cost" in k and k != "cost_per_ptu_per_hour" and v}
assert charged == {}
def test_the_zeroed_pricing_does_not_waive_budget_enforcement(self):
"""A zero price otherwise tells auth the model is free and skips every budget check."""
priced = _ptu_priced_deployment(
Deployment(
model_name="model_name_team-1_dep-ptu",
litellm_params=LiteLLM_Params(model="gemini/gemini-2.5-flash", api_key="fake-key"),
model_info=ModelInfo(
id="dep-ptu",
team_id="team-1",
team_public_model_name="ptu-model",
ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
**self.PTU,
),
)
)
router = Router(model_list=[priced.to_json(exclude_none=True)])
assert _is_model_cost_zero(model="model_name_team-1_dep-ptu", llm_router=router) is False
assert _is_model_cost_zero(model="ptu-model", llm_router=router) is False
def test_an_unrelated_patch_heals_a_deployment_stored_before_this_rule(self):
"""Both blobs, because litellm_params wins over model_info wherever the two are merged."""
written = update_db_model(
db_model=_deployment_with_stored_ptu(),
updated_patch=updateDeployment(model_name="gpt-4o-renamed"),
)
for blob in ("model_info", "litellm_params"):
stored = json.loads(written[blob])
assert all(stored[field] == 0 for field in _PTU_ZEROED_PRICING_FIELDS), blob
def test_an_unrelated_patch_of_a_ptu_row_that_carries_a_price_is_not_refused(self):
"""The pause toggle and the credential-rotation modal send no pricing at all. Refusing
them because the stored row is mispriced blocks flows that cannot fix it."""
priced_ptu = Deployment(
model_name="gpt-4o",
litellm_params=LiteLLM_Params(model="openai/gpt-4o", input_cost_per_token=5e-07),
model_info=ModelInfo(
id="dep-0",
team_id="t",
ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
**self.PTU,
),
)
written = update_db_model(db_model=priced_ptu, updated_patch=updateDeployment(model_name="renamed"))
assert written["model_name"] == "renamed"
assert json.loads(written["litellm_params"])["input_cost_per_token"] == 0
def test_removing_ptu_config_hands_per_token_billing_back(self):
"""Left behind, the zeros this rule wrote would serve the deployment for free forever."""
zeros = dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0)
written = update_db_model(
db_model=Deployment(
model_name="gpt-4o",
litellm_params=LiteLLM_Params(model="openai/gpt-4o", **zeros),
model_info=ModelInfo(
id="dep-0",
team_id="t",
ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
**self.PTU,
**zeros,
),
),
updated_patch=updateDeployment(
model_info=ModelInfo(id="dep-0", ptu_count=None, cost_per_ptu_per_hour=None)
),
)
for blob in ("model_info", "litellm_params"):
stored = json.loads(written[blob])
assert not any(field in stored for field in _PTU_ZEROED_PRICING_FIELDS), blob
def test_the_dashboard_clear_releases_the_zeros_it_echoes_back(self):
"""The edit form re-sends the whole stored model_info on every save, so the clearing
patch carries the zeros this rule wrote. Treating those as a rate the operator chose
left the deployment serving free and reading as a free model to the budget checks."""
zeros = dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0)
written = update_db_model(
db_model=Deployment(
model_name="gpt-4o",
litellm_params=LiteLLM_Params(model="openai/gpt-4o", **zeros),
model_info=ModelInfo(
id="dep-0",
team_id="t",
ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
**self.PTU,
**zeros,
),
),
updated_patch=updateDeployment(
model_info=ModelInfo(id="dep-0", ptu_count=None, cost_per_ptu_per_hour=None, **zeros)
),
)
stored = json.loads(written["model_info"])
assert not any(field in stored for field in _PTU_ZEROED_PRICING_FIELDS)
def test_a_deployment_that_never_had_ptu_keeps_a_price_its_operator_set_to_zero(self):
"""The dashboard sends both PTU keys as null on every save while the feature is on, so a
release keyed on the patch alone would strip a deliberate zero rate from any model."""
free = Deployment(
model_name="free-model",
litellm_params=LiteLLM_Params(model="openai/gpt-4o", input_cost_per_token=0.0),
model_info=ModelInfo(id="dep-free", team_id="t", input_cost_per_token=0.0),
)
written = update_db_model(
db_model=free,
updated_patch=updateDeployment(
model_info=ModelInfo(id="dep-free", ptu_count=None, cost_per_ptu_per_hour=None)
),
)
for blob in ("model_info", "litellm_params"):
assert json.loads(written[blob])["input_cost_per_token"] == 0, blob
def test_a_patch_pricing_a_ptu_deployment_is_refused(self):
with pytest.raises(HTTPException) as exc:
update_db_model(
db_model=_deployment_with_stored_ptu(),
updated_patch=updateDeployment(
litellm_params=updateLiteLLMParams(model="openai/gpt-4o", input_cost_per_token=5e-07)
),
)
assert exc.value.status_code == 400
def test_a_price_the_client_only_echoes_back_is_not_read_as_an_attempt_to_charge(self):
"""/model/info fills missing rates from the public cost map and the edit form re-sends the
whole blob, so a model_info price is one the server wrote. Reading it as the operator's
refused every attempt to put an existing deployment on PTU from the dashboard."""
written = update_db_model(
db_model=_deployment_without_ptu(),
updated_patch=updateDeployment(
model_info=ModelInfo(
id="dep-0",
team_id="t",
input_cost_per_token=3e-07,
output_cost_per_token=2.5e-06,
ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
**self.PTU,
)
),
)
stored = json.loads(written["model_info"])
assert stored["ptu_count"] == 15
assert stored["input_cost_per_token"] == 0
assert stored["output_cost_per_token"] == 0
def test_adding_ptu_config_to_an_already_priced_deployment_is_refused(self):
priced = Deployment(
model_name="gpt-4o",
litellm_params=LiteLLM_Params(model="openai/gpt-4o"),
model_info=ModelInfo(id="dep-0", team_id="t"),
)
with pytest.raises(HTTPException) as exc:
update_db_model(
db_model=priced,
updated_patch=updateDeployment(
litellm_params=updateLiteLLMParams(model="openai/gpt-4o", input_cost_per_token=5e-07),
model_info=ModelInfo(
id="dep-0",
team_id="t",
ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
**self.PTU,
),
),
)
assert exc.value.status_code == 400
def test_a_deployment_without_ptu_config_keeps_its_pricing_through_a_patch(self):
priced = Deployment(
model_name="gpt-4o",
litellm_params=LiteLLM_Params(model="openai/gpt-4o"),
model_info=ModelInfo(id="dep-0", team_id="t", input_cost_per_token=5e-07),
)
stored = json.loads(
update_db_model(db_model=priced, updated_patch=updateDeployment(model_name="renamed"))["model_info"]
)
assert stored["input_cost_per_token"] == 5e-07
@pytest.mark.asyncio
async def test_model_new_stores_zero_pricing_on_both_blobs(self):
(_, add_team_model_to_db), patches = TestAddNewModelPtuGate._patched_proxy("ptu-priced-model")
admin = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN)
with ExitStack() as stack:
for active_patch in patches:
stack.enter_context(active_patch)
await add_new_model(
model_params=TestAddNewModelPtuGate._ptu_deployment("ptu-priced-model"),
user_api_key_dict=admin,
)
written = add_team_model_to_db.call_args.kwargs["model_params"]
assert all(getattr(written.model_info, field, None) == 0 for field in SPECIAL_MODEL_INFO_PARAMS)
assert all(written.litellm_params.get(field) == 0 for field in _PTU_ZEROED_PRICING_FIELDS)
@pytest.mark.asyncio
async def test_model_new_refuses_a_priced_ptu_deployment(self):
(_, add_team_model_to_db), patches = TestAddNewModelPtuGate._patched_proxy("ptu-priced-model")
admin = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN)
base = TestAddNewModelPtuGate._ptu_deployment("ptu-priced-model")
deployment = base.model_copy(
update={"litellm_params": base.litellm_params.model_copy(update={"input_cost_per_token": 5e-07})}
)
with ExitStack() as stack:
for active_patch in patches:
stack.enter_context(active_patch)
with pytest.raises(Exception) as exc:
await add_new_model(model_params=deployment, user_api_key_dict=admin)
assert "input_cost_per_token" in str(exc.value)
add_team_model_to_db.assert_not_called()

374
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
View file

@ -2,7 +2,9 @@ import asyncio
import json
import os
import sys
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from types import SimpleNamespace
from typing import Optional, cast
from unittest.mock import AsyncMock, MagicMock, call, patch
@ -68,6 +70,21 @@ from litellm.types.proxy.management_endpoints.team_endpoints import (
# Setup TestClient
client = TestClient(app)
def _wire_team_create_tx(prisma_client):
"""`/team/new` inserts the team and mirrors it onto the access groups in one transaction,
so a mocked client has to hand its team table back out of `db.tx()`."""
@asynccontextmanager
async def _tx():
yield SimpleNamespace(
litellm_teamtable=prisma_client.db.litellm_teamtable,
query_raw=AsyncMock(return_value=[]),
)
prisma_client.db.tx = lambda *_args, **_kwargs: _tx()
# Mock prisma_client
mock_prisma_client = MagicMock()
# Set up async mock for db operations
@ -400,6 +417,7 @@ async def test_new_team_rejects_a_duration_that_never_advances(
mock_team_create = AsyncMock()
mock_db_client.db.litellm_teamtable = MagicMock()
mock_db_client.db.litellm_teamtable.create = mock_team_create
_wire_team_create_tx(mock_db_client)
with pytest.raises(ProxyException) as exc_info:
await new_team(
@ -481,6 +499,7 @@ async def test_new_team_with_object_permission(mock_db_client, mock_admin_auth):
mock_team_count = AsyncMock(return_value=0)
mock_db_client.db.litellm_teamtable = MagicMock()
mock_db_client.db.litellm_teamtable.create = mock_team_create
_wire_team_create_tx(mock_db_client)
mock_db_client.db.litellm_teamtable.count = mock_team_count
mock_db_client.db.litellm_teamtable.update = AsyncMock(
return_value=team_create_result
@ -570,6 +589,7 @@ async def test_new_team_with_mcp_tool_permissions(mock_db_client, mock_admin_aut
mock_db_client.db.litellm_teamtable.create = AsyncMock(
return_value=team_create_result
)
_wire_team_create_tx(mock_db_client)
mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0)
mock_db_client.db.litellm_teamtable.update = AsyncMock(
return_value=team_create_result
@ -663,6 +683,7 @@ async def test_new_team_disable_auto_add_proxy_admin_flag(
mock_db_client.db.litellm_teamtable.create = AsyncMock(
return_value=team_create_result
)
_wire_team_create_tx(mock_db_client)
mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0)
mock_db_client.db.litellm_usertable = MagicMock()
mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock())
@ -4430,6 +4451,7 @@ async def test_new_team_max_budget_within_user_limit():
mock_prisma.db.litellm_teamtable.create = AsyncMock(
return_value=mock_created_team
)
_wire_team_create_tx(mock_prisma)
mock_prisma.db.litellm_teamtable.update = AsyncMock(
return_value=mock_created_team
)
@ -4573,6 +4595,7 @@ async def test_new_team_org_scoped_budget_bypasses_user_limit():
mock_prisma.db.litellm_teamtable.create = AsyncMock(
return_value=mock_created_team
)
_wire_team_create_tx(mock_prisma)
mock_prisma.db.litellm_teamtable.update = AsyncMock(
return_value=mock_created_team
)
@ -4721,6 +4744,7 @@ async def test_new_team_org_scoped_models_bypasses_user_limit():
mock_prisma.db.litellm_teamtable.create = AsyncMock(
return_value=mock_created_team
)
_wire_team_create_tx(mock_prisma)
mock_prisma.db.litellm_teamtable.update = AsyncMock(
return_value=mock_created_team
)
@ -6567,6 +6591,7 @@ async def test_new_team_org_scoped_tpm_rpm_bypasses_user_limit():
mock_created_team.rpm_limit = 1000
mock_created_team.metadata = None
mock_created_team.members_with_roles = []
mock_created_team.access_group_ids = None
mock_created_team.model_dump.return_value = {
"team_id": "new-bypass-team-id",
"team_alias": "org-bypass-test-team",
@ -6578,6 +6603,7 @@ async def test_new_team_org_scoped_tpm_rpm_bypasses_user_limit():
mock_prisma.db.litellm_teamtable.create = AsyncMock(
return_value=mock_created_team
)
_wire_team_create_tx(mock_prisma)
mock_prisma.db.litellm_teamtable.update = AsyncMock(
return_value=mock_created_team
)
@ -6856,6 +6882,7 @@ async def test_update_team_org_scoped_tpm_rpm_bypasses_user_limit():
mock_updated_team.team_id = "org-team-update-bypass-123"
mock_updated_team.tpm_limit = 10000
mock_updated_team.rpm_limit = 1000
mock_updated_team.access_group_ids = None
mock_updated_team.model_dump.return_value = {
"team_id": "org-team-update-bypass-123",
"tpm_limit": 10000,
@ -7009,6 +7036,7 @@ async def test_update_team_guardrails_with_org_id():
"guardrails": ["aporia-pre-call", "aporia-post-call"]
}
mock_updated_team.litellm_model_table = None
mock_updated_team.access_group_ids = None
mock_updated_team.model_dump.return_value = {
"team_id": "team-guardrails-123",
"organization_id": "test-org-guardrails",
@ -7937,6 +7965,7 @@ async def test_new_team_soft_budget_validation(
mock_prisma.db.litellm_teamtable.create = AsyncMock(
return_value=mock_created_team
)
_wire_team_create_tx(mock_prisma)
mock_prisma.db.litellm_teamtable.update = AsyncMock(
return_value=mock_created_team
)
@ -8236,6 +8265,7 @@ async def test_new_team_with_router_settings(mock_db_client, mock_admin_auth):
mock_team_count = AsyncMock(return_value=0)
mock_db_client.db.litellm_teamtable = MagicMock()
mock_db_client.db.litellm_teamtable.create = mock_team_create
_wire_team_create_tx(mock_db_client)
mock_db_client.db.litellm_teamtable.count = mock_team_count
mock_db_client.db.litellm_teamtable.update = AsyncMock(
return_value=team_create_result
@ -9626,6 +9656,7 @@ async def test_new_team_encrypts_callback_vars(
team_create_result.model_dump.return_value = {"team_id": "team-456"}
mock_team_create = AsyncMock(return_value=team_create_result)
mock_db_client.db.litellm_teamtable.create = mock_team_create
_wire_team_create_tx(mock_db_client)
mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0)
mock_db_client.db.litellm_teamtable.update = AsyncMock(
return_value=team_create_result
@ -10786,6 +10817,7 @@ async def test_new_team_validator_runs_without_metadata_and_rejection_blocks_cre
):
mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0)
mock_prisma.db.litellm_teamtable.create = AsyncMock()
_wire_team_create_tx(mock_prisma)
mock_license.is_team_count_over_limit.return_value = False
with pytest.raises(ProxyException) as exc_info:
@ -10820,6 +10852,7 @@ async def test_new_team_validator_accept_proceeds_to_create(mock_db_client, mock
team_create_result.model_dump.return_value = {"team_id": "team-accept-1"}
mock_db_client.db.litellm_teamtable = MagicMock()
mock_db_client.db.litellm_teamtable.create = AsyncMock(return_value=team_create_result)
_wire_team_create_tx(mock_db_client)
mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0)
mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_create_result)
mock_db_client.db.litellm_usertable = MagicMock()
@ -10859,6 +10892,7 @@ async def test_new_team_rejection_precedes_model_alias_write():
):
mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0)
mock_prisma.db.litellm_teamtable.create = AsyncMock()
_wire_team_create_tx(mock_prisma)
mock_prisma.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model-1"))
mock_license.is_team_count_over_limit.return_value = False
@ -11626,6 +11660,7 @@ def _wire_new_team_prisma(mock_db_client):
mock_db_client.db.litellm_teamtable = MagicMock()
mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0)
mock_db_client.db.litellm_teamtable.create = AsyncMock(return_value=created_team)
_wire_team_create_tx(mock_db_client)
mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=created_team)
mock_db_client.db.litellm_usertable = MagicMock()
mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock())
@ -11714,3 +11749,342 @@ async def test_new_team_explicit_null_max_budget_still_takes_configured_default(
team_data = mock_team_create.call_args.kwargs["data"]
assert team_data.get("max_budget") == 100.0
class _FakeMirrorDb:
"""Stands in for prisma inside the access-group mirror.
Dispatches on the statement so a change to the SQL's shape is visible here, but it
cannot validate the SQL itself: it reimplements the array semantics in Python, so it
passes whatever the statement says. Correctness of the SQL is pinned against a real
Postgres in tests/proxy_admin_ui_tests/test_access_group_team_sync.py.
"""
def __init__(self, access_groups, teams, plain_lists=False):
self._access_groups = access_groups
self._teams = teams
self._plain_lists = plain_lists
self.transactions = []
def _team_ids(self, group_id):
stored = self._access_groups[group_id]
return stored if self._plain_lists else stored["assigned_team_ids"]
async def _query_raw(self, sql, *args):
assert self._open, "mirror statement ran outside a transaction"
if "pg_advisory_xact_lock" in sql:
self.transactions[-1].append("lock")
return [{"locked": False}]
if "LiteLLM_TeamTable" in sql:
self.transactions[-1].append("read")
team_id = args[0]
if team_id not in self._teams:
return []
return [{"access_group_ids": list(self._teams[team_id])}]
team_id, desired = args
if sql.lstrip().startswith("SELECT"):
self.transactions[-1].append("affected")
affected = [g for g in self._access_groups if g in desired or team_id in self._team_ids(g)]
return [{"access_group_id": group_id} for group_id in affected]
if "array_append" in sql:
self.transactions[-1].append("attach")
changed = [
g for g in desired if g in self._access_groups and team_id not in self._team_ids(g)
]
for group_id in changed:
self._team_ids(group_id).append(team_id)
else:
self.transactions[-1].append("detach")
changed = [
g for g in self._access_groups if team_id in self._team_ids(g) and g not in desired
]
for group_id in changed:
self._team_ids(group_id).remove(team_id)
return [{"access_group_id": group_id} for group_id in changed]
async def _create_team(self, data, include=None):
self.transactions[-1].append("create")
team_id = data["team_id"]
self._teams[team_id] = list(data.get("access_group_ids") or ())
return SimpleNamespace(
team_id=team_id,
access_group_ids=list(self._teams[team_id]),
model_dump=lambda: {"team_id": team_id},
)
def tx(self, *_args, **_kwargs):
outer = self
class _Tx:
async def __aenter__(self):
outer.transactions.append([])
outer._open = True
return SimpleNamespace(
query_raw=outer._query_raw,
litellm_teamtable=SimpleNamespace(create=outer._create_team),
)
async def __aexit__(self, *_exc_info):
outer._open = False
return None
return _Tx()
_open = False
@pytest.mark.asyncio
async def test_update_team_syncs_access_group_assigned_team_ids_in_both_directions():
"""
A team-side edit of `access_group_ids` must be mirrored onto every affected access
group's `assigned_team_ids`, in one transaction, in both directions.
`assigned_team_ids` is not display-only. `get_authorized_resources_from_key_access_groups`
reads it as an authorization input, so a group the team dropped must stop granting its
resources to keys on that team, and a group the team added must start granting them.
A single-direction assertion would pass against a fix that only ever removes (or only
ever adds), so this covers add, remove, untouched, and the authorization consequence.
"""
from unittest.mock import Mock
from fastapi import Request
from litellm.proxy._types import LiteLLM_AccessGroupTable
from litellm.proxy.auth.auth_checks import (
get_authorized_resources_from_key_access_groups,
)
access_groups = {
"ag-drop": {"assigned_team_ids": ["team-a"], "access_model_names": ["dropped-model"]},
"ag-keep": {"assigned_team_ids": ["team-a"], "access_model_names": ["kept-model"]},
"ag-add": {"assigned_team_ids": [], "access_model_names": ["added-model"]},
"ag-other-team": {"assigned_team_ids": ["team-b"], "access_model_names": ["other-model"]},
}
committed_team_groups = ["ag-keep", "ag-add"]
fake_db = _FakeMirrorDb(access_groups, {"team-a": committed_team_groups})
existing_team = MagicMock()
existing_team.access_group_ids = ["ag-drop", "ag-keep"]
existing_team.metadata = {}
existing_team.max_budget = None
existing_team.organization_id = None
existing_team.team_alias = "team-a"
existing_team.model_dump.return_value = {"team_id": "team-a", "team_alias": "team-a"}
updated_team = MagicMock()
updated_team.team_id = "team-a"
updated_team.access_group_ids = committed_team_groups
updated_team.model_dump.return_value = {"team_id": "team-a"}
with (
patch("litellm.proxy.proxy_server.prisma_client") as prisma,
patch("litellm.proxy.proxy_server.llm_router"),
patch("litellm.proxy.proxy_server.user_api_key_cache"),
patch("litellm.proxy.proxy_server.proxy_logging_obj"),
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
patch("litellm.proxy.management_endpoints.team_endpoints._refresh_cached_team"),
patch(
"litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache",
new_callable=AsyncMock,
) as invalidate_cache,
):
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team)
prisma.db.litellm_teamtable.update = AsyncMock(return_value=updated_team)
prisma.db.tx = fake_db.tx
prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data)
await update_team(
data=UpdateTeamRequest(team_id="team-a", access_group_ids=committed_team_groups),
http_request=Mock(spec=Request),
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"),
)
assert access_groups["ag-drop"]["assigned_team_ids"] == []
assert access_groups["ag-add"]["assigned_team_ids"] == ["team-a"]
assert access_groups["ag-keep"]["assigned_team_ids"] == ["team-a"]
assert access_groups["ag-other-team"]["assigned_team_ids"] == ["team-b"]
assert fake_db.transactions == [["lock", "read", "affected", "attach", "detach"]]
assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-drop", "ag-keep", "ag-add"}
async def _get_access_object(*, access_group_id, **_kwargs):
stored = access_groups[access_group_id]
return LiteLLM_AccessGroupTable(
access_group_id=access_group_id,
access_group_name=access_group_id,
access_model_names=list(stored["access_model_names"]),
assigned_team_ids=list(stored["assigned_team_ids"]),
assigned_key_ids=[],
)
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()),
patch(
"litellm.proxy.auth.auth_checks.get_access_object",
new_callable=AsyncMock,
side_effect=_get_access_object,
),
):
authorized_models = await get_authorized_resources_from_key_access_groups(
valid_token=UserAPIKeyAuth(
token="sk-hash",
models=[],
team_id="team-a",
access_group_ids=["ag-drop", "ag-keep", "ag-add"],
),
team_object=LiteLLM_TeamTable(team_id="team-a", models=[]),
resource_field="access_model_names",
)
assert sorted(authorized_models) == ["added-model", "kept-model"]
@pytest.mark.asyncio
async def test_sync_reads_the_committed_team_row_rather_than_the_callers_snapshot():
"""
The mirror takes no desired-state argument on purpose. It locks the team and reads
the row as committed, so two concurrent writers for one team converge on the row the
last one committed instead of each replaying its own stale snapshot. Reconciling also
means a retry heals a half-applied sync, where a before/after delta computes nothing.
The same holds for the cache step: the groups to drop come from the reconciled set,
not from the rows this attempt happened to change, so a retry after an unreachable
cache still drops the entries even though its statements are now no-ops.
A team with no row at all is deletion, and must detach from every group.
"""
from litellm.proxy.management_helpers.access_group_team_sync import (
sync_team_access_group_membership,
)
access_groups = {"ag-1": ["team-a", "team-b"], "ag-2": ["team-a"], "ag-3": []}
teams = {"team-a": ["ag-2", "ag-3"]}
fake_db = _FakeMirrorDb(access_groups, teams, plain_lists=True)
prisma_client = SimpleNamespace(db=SimpleNamespace(tx=fake_db.tx))
with patch(
"litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache",
new_callable=AsyncMock,
side_effect=[ConnectionError("redis unreachable"), None, None],
) as invalidate_cache:
with pytest.raises(ConnectionError):
await sync_team_access_group_membership(prisma_client=prisma_client, team_id="team-a")
assert access_groups == {"ag-1": ["team-b"], "ag-2": ["team-a"], "ag-3": ["team-a"]}
assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-1", "ag-2", "ag-3"}
invalidate_cache.reset_mock()
invalidate_cache.side_effect = None
await sync_team_access_group_membership(prisma_client=prisma_client, team_id="team-a")
assert access_groups == {"ag-1": ["team-b"], "ag-2": ["team-a"], "ag-3": ["team-a"]}
assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-2", "ag-3"}
invalidate_cache.reset_mock()
del teams["team-a"]
await sync_team_access_group_membership(prisma_client=prisma_client, team_id="team-a")
assert access_groups == {"ag-1": ["team-b"], "ag-2": [], "ag-3": []}
assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-2", "ag-3"}
assert fake_db.transactions == [["lock", "read", "affected", "attach", "detach"]] * 3
@pytest.mark.asyncio
async def test_new_team_and_delete_team_both_drive_the_mirror():
"""Every writer of `team.access_group_ids` has to reach the mirror, not just update.
These pin the wiring on the other two paths; the mirror's own behavior is covered above.
Creation has to insert the team row and mirror it in one transaction. With the mirror
in a transaction of its own, a sync that fails leaves a committed team whose groups
never learned about it, and the retry is rejected as a duplicate team id."""
from unittest.mock import Mock
from fastapi import Request
from litellm.proxy._types import DeleteTeamRequest, NewTeamRequest
from litellm.proxy.management_endpoints.team_endpoints import delete_team, new_team
access_groups = {"ag-1": [], "ag-2": []}
fake_db = _FakeMirrorDb(access_groups, {}, plain_lists=True)
with (
patch("litellm.proxy.proxy_server.prisma_client") as prisma,
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
patch("litellm.proxy.proxy_server.user_api_key_cache"),
patch("litellm.proxy.proxy_server.proxy_logging_obj"),
patch("litellm.proxy.management_endpoints.team_endpoints._add_team_members_to_team", new_callable=AsyncMock),
patch(
"litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache",
new_callable=AsyncMock,
) as invalidate_cache,
):
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
prisma.db.litellm_teamtable.count = AsyncMock(return_value=0)
prisma.db.tx = fake_db.tx
prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data)
prisma.get_data = AsyncMock(return_value=None)
await new_team(
data=NewTeamRequest(team_id="team-new", team_alias="new", access_group_ids=["ag-1"]),
http_request=Mock(spec=Request),
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"),
)
assert access_groups == {"ag-1": ["team-new"], "ag-2": []}
assert fake_db.transactions == [["create", "lock", "read", "affected", "attach", "detach"]]
assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-1"}
team_row = LiteLLM_TeamTable(team_id="team-gone", models=[], access_group_ids=["ag-1"])
with (
patch("litellm.proxy.proxy_server.prisma_client") as prisma,
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.management_endpoints.team_endpoints._persist_deleted_team_records", new_callable=AsyncMock),
patch("litellm.proxy.management_endpoints.team_endpoints._verify_team_access", new_callable=AsyncMock),
patch(
"litellm.proxy.management_endpoints.team_endpoints.sync_team_access_group_membership",
new_callable=AsyncMock,
) as sync,
):
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
prisma.delete_data = AsyncMock(return_value=[team_row])
prisma.db.execute_raw = AsyncMock(return_value=0)
prisma.db.litellm_teammembership.delete_many = AsyncMock(return_value=0)
await delete_team(
data=DeleteTeamRequest(team_ids=["team-gone"]),
http_request=Mock(spec=Request),
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"),
)
assert sync.await_args_list[0].kwargs["team_id"] == "team-gone"
@pytest.mark.asyncio
async def test_invalidate_access_group_cache_deletes_the_cached_object():
"""The mirror's cache step is what stops a revoked group granting from cache until TTL,
so pin that it actually reaches the delete rather than only being called."""
from litellm.proxy.management_helpers.access_group_team_sync import (
invalidate_access_group_cache,
)
cache, logging_obj = MagicMock(), MagicMock()
with (
patch("litellm.proxy.proxy_server.user_api_key_cache", cache),
patch("litellm.proxy.proxy_server.proxy_logging_obj", logging_obj),
patch(
"litellm.proxy.management_helpers.access_group_team_sync._delete_cache_access_object",
new_callable=AsyncMock,
) as delete_cached,
):
await invalidate_access_group_cache("ag-1")
assert delete_cached.await_args.kwargs == {
"access_group_id": "ag-1",
"user_api_key_cache": cache,
"proxy_logging_obj": logging_obj,
}

18
tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
View file

@ -2,6 +2,7 @@ import asyncio
import json
import os
import sys
from contextlib import asynccontextmanager
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
@ -37,6 +38,20 @@ from litellm.types.proxy.management_endpoints.ui_sso import (
)
def _wire_team_create_tx(prisma_client):
"""`/team/new` inserts the team and mirrors it onto the access groups in one transaction,
so a mocked client has to hand its team table back out of `db.tx()`."""
@asynccontextmanager
async def _tx():
yield SimpleNamespace(
litellm_teamtable=prisma_client.db.litellm_teamtable,
query_raw=AsyncMock(return_value=[]),
)
prisma_client.db.tx = lambda *_args, **_kwargs: _tx()
def test_microsoft_sso_handler_openid_from_response_user_principal_name():
# Arrange
# Create a mock response similar to what Microsoft SSO would return
@ -577,6 +592,7 @@ async def test_default_team_params(team_params):
mock_prisma = MagicMock()
mock_prisma.db.litellm_teamtable.find_first = AsyncMock(return_value=None)
mock_prisma.db.litellm_teamtable.create = AsyncMock()
_wire_team_create_tx(mock_prisma)
mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0)
mock_prisma.get_data = AsyncMock(return_value=None)
mock_prisma.jsonify_team_object = MagicMock(side_effect=mock_jsonify_team_object)
@ -624,6 +640,7 @@ async def test_default_team_params_organization_id_reaches_sso_created_team(team
mock_prisma = MagicMock()
mock_prisma.db.litellm_teamtable.find_first = AsyncMock(return_value=None)
mock_prisma.db.litellm_teamtable.create = AsyncMock()
_wire_team_create_tx(mock_prisma)
mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0)
mock_prisma.get_data = AsyncMock(return_value=None)
mock_prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data)
@ -671,6 +688,7 @@ async def test_create_team_without_default_params():
mock_prisma = MagicMock()
mock_prisma.db.litellm_teamtable.find_first = AsyncMock(return_value=None)
mock_prisma.db.litellm_teamtable.create = AsyncMock()
_wire_team_create_tx(mock_prisma)
mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0)
mock_prisma.get_data = AsyncMock(return_value=None)
mock_prisma.jsonify_team_object = MagicMock(side_effect=mock_jsonify_team_object)

39
tests/test_litellm/proxy/management_helpers/test_access_group_team_sync.py Normal file
View file

@ -0,0 +1,39 @@
import os
import sys
import pytest
sys.path.insert(0, os.path.abspath("../../../.."))
from litellm.proxy.management_helpers.access_group_team_sync import (
invalidate_access_group_caches,
)
@pytest.mark.asyncio
async def test_one_unreachable_cache_does_not_skip_the_other_groups(monkeypatch):
"""
`assigned_team_ids` is an authorization input, so a group whose cache still holds the
revoked grant keeps serving it until the entry is dropped.
A sequential loop would stop at the first failing group and leave the groups behind it
serving stale grants, and swallowing the failure would report success to the admin for
a revoke that never took effect. Every group has to be attempted, and the endpoint has
to fail so the caller can retry.
"""
attempted: list[str] = []
async def _invalidate(access_group_id: str) -> None:
attempted.append(access_group_id)
if access_group_id == "ag-redis-down":
raise ConnectionError("redis unreachable")
monkeypatch.setattr(
"litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache",
_invalidate,
)
with pytest.raises(ConnectionError):
await invalidate_access_group_caches(("ag-redis-down", "ag-2", "ag-3"))
assert attempted == ["ag-redis-down", "ag-2", "ag-3"]

48
tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_cohere_passthrough_logging_handler.py
View file

@ -7,9 +7,7 @@ from unittest.mock import MagicMock, patch
import httpx
import pytest
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.cohere_passthrough_logging_handler import (
@ -69,12 +67,8 @@ class TestCoherePassthroughLoggingHandler:
)
@patch("litellm.completion_cost")
@patch(
"litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload"
)
@patch(
"litellm.llms.cohere.embed.v1_transformation.CohereEmbeddingConfig._transform_response"
)
@patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload")
@patch("litellm.llms.cohere.embed.v1_transformation.CohereEmbeddingConfig._transform_response")
def test_cohere_embed_passthrough_cost_tracking(
self, mock_transform_response, mock_get_standard_logging, mock_completion_cost
):
@ -92,9 +86,7 @@ class TestCoherePassthroughLoggingHandler:
mock_embedding_response.object = "list"
from litellm.types.utils import Usage
mock_embedding_response.usage = Usage(
prompt_tokens=3, completion_tokens=0, total_tokens=3
)
mock_embedding_response.usage = Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3)
mock_transform_response.return_value = mock_embedding_response
mock_completion_cost.return_value = 3.6e-07 # Expected cost for embed-v4.0
@ -151,6 +143,38 @@ class TestCoherePassthroughLoggingHandler:
assert hasattr(result["result"], "model")
assert result["result"].model == "embed-english-v3.0"
@patch(
"litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough_logging_handler.BasePassthroughLoggingHandler.passthrough_chat_handler"
)
@patch("litellm.completion_cost")
def test_openai_embeddings_route_does_not_use_cohere_embed_path(self, mock_completion_cost, mock_chat_handler):
mock_chat_handler.return_value = {"result": None, "kwargs": {}}
response_body = {
"object": "list",
"model": "text-embedding-3-small",
"data": [{"object": "embedding", "index": 0, "embedding": [0.1]}],
"usage": {"prompt_tokens": 6, "total_tokens": 6},
}
result = self.handler.cohere_passthrough_handler(
httpx_response=self._create_mock_httpx_response(response_body),
response_body=response_body,
logging_obj=self._create_mock_logging_obj(),
url_route="https://api.openai.com/v1/embeddings",
result="",
start_time=self.start_time,
end_time=self.end_time,
cache_hit=False,
request_body={"model": "text-embedding-3-small", "input": "PROOF_SENTINEL_TEXT"},
passthrough_logging_payload=PassthroughStandardLoggingPayload(
url="https://api.openai.com/v1/embeddings",
request_body={"model": "text-embedding-3-small", "input": "PROOF_SENTINEL_TEXT"},
request_method="POST",
),
)
mock_completion_cost.assert_not_called()
mock_chat_handler.assert_called_once()
assert result == {"result": None, "kwargs": {}}
if __name__ == "__main__":
pytest.main([__file__])

499
tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py
View file

@ -8,9 +8,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import (
@ -70,9 +68,7 @@ class TestOpenAIPassthroughLoggingHandler:
mock_response.headers = {"content-type": "application/json"}
return mock_response
def _create_passthrough_logging_payload(
self, user: str = "test_user"
) -> PassthroughStandardLoggingPayload:
def _create_passthrough_logging_payload(self, user: str = "test_user") -> PassthroughStandardLoggingPayload:
"""Create a mock passthrough logging payload"""
return PassthroughStandardLoggingPayload(
url="https://api.openai.com/v1/chat/completions",
@ -113,9 +109,7 @@ class TestOpenAIPassthroughLoggingHandler:
# Negative cases
assert (
OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(
"https://api.openai.com/v1/models"
)
OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://api.openai.com/v1/models")
== False
)
assert (
@ -125,15 +119,10 @@ class TestOpenAIPassthroughLoggingHandler:
== False
)
assert (
OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(
"https://api.anthropic.com/v1/messages"
)
== False
)
assert (
OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("")
OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://api.anthropic.com/v1/messages")
== False
)
assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("") == False
def test_is_openai_image_generation_route(self):
"""Test OpenAI image generation route detection"""
@ -159,9 +148,7 @@ class TestOpenAIPassthroughLoggingHandler:
== False
)
assert (
OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(
"https://api.openai.com/v1/images/edits"
)
OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("https://api.openai.com/v1/images/edits")
== False
)
assert (
@ -170,32 +157,23 @@ class TestOpenAIPassthroughLoggingHandler:
)
== False
)
assert (
OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("")
== False
)
assert OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("") == False
def test_is_openai_image_editing_route(self):
"""Test OpenAI image editing route detection"""
# Positive cases
assert (
OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(
"https://api.openai.com/v1/images/edits"
)
OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("https://api.openai.com/v1/images/edits")
== True
)
assert (
OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(
"https://openai.azure.com/v1/images/edits"
)
OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("https://openai.azure.com/v1/images/edits")
== True
)
# Negative cases
assert (
OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(
"https://api.openai.com/v1/chat/completions"
)
OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("https://api.openai.com/v1/chat/completions")
== False
)
assert (
@ -210,118 +188,91 @@ class TestOpenAIPassthroughLoggingHandler:
)
== False
)
assert (
OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("") == False
)
assert OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("") == False
def test_is_openai_responses_route(self):
"""Test OpenAI responses API route detection"""
# Positive cases
assert OpenAIPassthroughLoggingHandler.is_openai_responses_route("https://api.openai.com/v1/responses") == True
assert (
OpenAIPassthroughLoggingHandler.is_openai_responses_route(
"https://api.openai.com/v1/responses"
)
== True
)
assert (
OpenAIPassthroughLoggingHandler.is_openai_responses_route(
"https://openai.azure.com/v1/responses"
)
== True
)
assert (
OpenAIPassthroughLoggingHandler.is_openai_responses_route(
"https://api.openai.com/responses"
)
== True
OpenAIPassthroughLoggingHandler.is_openai_responses_route("https://openai.azure.com/v1/responses") == True
)
assert OpenAIPassthroughLoggingHandler.is_openai_responses_route("https://api.openai.com/responses") == True
# Negative cases
assert (
OpenAIPassthroughLoggingHandler.is_openai_responses_route(
"https://api.openai.com/v1/chat/completions"
)
OpenAIPassthroughLoggingHandler.is_openai_responses_route("https://api.openai.com/v1/chat/completions")
== False
)
assert (
OpenAIPassthroughLoggingHandler.is_openai_responses_route(
"https://api.openai.com/v1/images/generations"
)
OpenAIPassthroughLoggingHandler.is_openai_responses_route("https://api.openai.com/v1/images/generations")
== False
)
assert (
OpenAIPassthroughLoggingHandler.is_openai_responses_route(
"http://localhost:4000/openai/v1/responses"
)
OpenAIPassthroughLoggingHandler.is_openai_responses_route("http://localhost:4000/openai/v1/responses")
== False
)
assert OpenAIPassthroughLoggingHandler.is_openai_responses_route("") == False
def test_is_openai_embeddings_route(self):
assert (
OpenAIPassthroughLoggingHandler.is_openai_embeddings_route("https://api.openai.com/v1/embeddings") is True
)
assert (
OpenAIPassthroughLoggingHandler.is_openai_embeddings_route("https://openai.azure.com/v1/embeddings") is True
)
assert (
OpenAIPassthroughLoggingHandler.is_openai_embeddings_route(
"https://my-resource.cognitiveservices.azure.com/v1/embeddings"
)
is True
)
assert (
OpenAIPassthroughLoggingHandler.is_openai_embeddings_route(
"https://my-resource.openai.azure.com/openai/deployments/text-embedding-3-small/embeddings"
)
is False
)
assert (
OpenAIPassthroughLoggingHandler.is_openai_embeddings_route("https://api.openai.com/v1/chat/completions")
is False
)
assert (
OpenAIPassthroughLoggingHandler.is_openai_embeddings_route(
"http://localhost:4000/openai_passthrough/v1/embeddings"
)
is False
)
assert OpenAIPassthroughLoggingHandler.is_openai_embeddings_route("") is False
def test_is_openai_route_recognizes_cognitiveservices_azure_com(self):
"""Azure OpenAI resources created via the newer "Azure AI Foundry" /
Cognitive Services pathway live on `*.cognitiveservices.azure.com`
subdomains rather than the older `openai.azure.com`. All four
subdomains rather than the older `openai.azure.com`. The
is_openai_*_route methods must recognize both Azure subdomains so
cost tracking applies regardless of which Azure naming the user's
resource happens to be on.
"""
cognitive_chat = (
"https://my-resource.cognitiveservices.azure.com/v1/chat/completions"
)
cognitive_images_gen = (
"https://my-resource.cognitiveservices.azure.com/v1/images/generations"
)
cognitive_images_edit = (
"https://my-resource.cognitiveservices.azure.com/v1/images/edits"
)
cognitive_responses = (
"https://my-resource.cognitiveservices.azure.com/v1/responses"
)
cognitive_chat = "https://my-resource.cognitiveservices.azure.com/v1/chat/completions"
cognitive_images_gen = "https://my-resource.cognitiveservices.azure.com/v1/images/generations"
cognitive_images_edit = "https://my-resource.cognitiveservices.azure.com/v1/images/edits"
cognitive_responses = "https://my-resource.cognitiveservices.azure.com/v1/responses"
cognitive_embeddings = "https://my-resource.cognitiveservices.azure.com/v1/embeddings"
assert (
OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(
cognitive_chat
)
is True
)
assert (
OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(
cognitive_images_gen
)
is True
)
assert (
OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(
cognitive_images_edit
)
is True
)
assert (
OpenAIPassthroughLoggingHandler.is_openai_responses_route(
cognitive_responses
)
is True
)
assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(cognitive_chat) is True
assert OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(cognitive_images_gen) is True
assert OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(cognitive_images_edit) is True
assert OpenAIPassthroughLoggingHandler.is_openai_responses_route(cognitive_responses) is True
assert OpenAIPassthroughLoggingHandler.is_openai_embeddings_route(cognitive_embeddings) is True
# Cross-route negatives still hold for cognitiveservices hosts.
assert (
OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(
cognitive_responses
)
is False
)
assert (
OpenAIPassthroughLoggingHandler.is_openai_responses_route(cognitive_chat)
is False
)
assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(cognitive_responses) is False
assert OpenAIPassthroughLoggingHandler.is_openai_responses_route(cognitive_chat) is False
assert OpenAIPassthroughLoggingHandler.is_openai_embeddings_route(cognitive_chat) is False
@patch("litellm.completion_cost")
@patch(
"litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload"
)
def test_openai_passthrough_handler_success(
self, mock_get_standard_logging, mock_completion_cost
):
@patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload")
def test_openai_passthrough_handler_success(self, mock_get_standard_logging, mock_completion_cost):
"""Test successful cost tracking for OpenAI chat completions"""
# Arrange
mock_completion_cost.return_value = 0.000045
@ -370,9 +321,7 @@ class TestOpenAIPassthroughLoggingHandler:
assert mock_logging_obj.model_call_details["custom_llm_provider"] == "openai"
@patch("litellm.completion_cost")
def test_openai_passthrough_handler_non_chat_completions(
self, mock_completion_cost
):
def test_openai_passthrough_handler_non_chat_completions(self, mock_completion_cost):
"""Test that non-chat-completions routes fall back to base handler"""
# Arrange
mock_httpx_response = self._create_mock_httpx_response()
@ -406,12 +355,8 @@ class TestOpenAIPassthroughLoggingHandler:
# The important thing is that our specific OpenAI handler logic didn't run
@patch("litellm.completion_cost")
@patch(
"litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload"
)
def test_openai_passthrough_handler_with_user_tracking(
self, mock_get_standard_logging, mock_completion_cost
):
@patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload")
def test_openai_passthrough_handler_with_user_tracking(self, mock_get_standard_logging, mock_completion_cost):
"""Test cost tracking with user information"""
# Arrange
mock_completion_cost.return_value = 0.000123
@ -464,15 +409,10 @@ class TestOpenAIPassthroughLoggingHandler:
assert "litellm_params" in result["kwargs"]
assert "proxy_server_request" in result["kwargs"]["litellm_params"]
assert "body" in result["kwargs"]["litellm_params"]["proxy_server_request"]
assert (
result["kwargs"]["litellm_params"]["proxy_server_request"]["body"]["user"]
== "test_user_123"
)
assert result["kwargs"]["litellm_params"]["proxy_server_request"]["body"]["user"] == "test_user_123"
@patch("litellm.completion_cost")
def test_openai_passthrough_handler_cost_calculation_error(
self, mock_completion_cost
):
def test_openai_passthrough_handler_cost_calculation_error(self, mock_completion_cost):
"""Test error handling in cost calculation"""
# Arrange
mock_completion_cost.side_effect = Exception("Cost calculation failed")
@ -521,9 +461,7 @@ class TestOpenAIPassthroughLoggingHandler:
@patch(f"{OpenAIPassthroughLoggingHandler.__module__}.get_standard_logging_object_payload")
@patch("litellm.completion_cost", return_value=3.3e-06)
def test_streaming_responses_cost_uses_completed_response(
self, mock_completion_cost, mock_get_standard_logging
):
def test_streaming_responses_cost_uses_completed_response(self, mock_completion_cost, mock_get_standard_logging):
response_id = "resp_PROOFSENTINEL0123456789abcdef"
completed_event = {
"type": "response.completed",
@ -796,12 +734,8 @@ class TestOpenAIPassthroughLoggingHandler:
mock_completion_cost.assert_not_called()
@patch("litellm.completion_cost")
@patch(
"litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload"
)
def test_different_models_cost_tracking(
self, mock_get_standard_logging, mock_completion_cost
):
@patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload")
def test_different_models_cost_tracking(self, mock_get_standard_logging, mock_completion_cost):
"""Test cost tracking for different OpenAI models"""
# Arrange
mock_get_standard_logging.return_value = {"test": "logging_payload"}
@ -868,12 +802,8 @@ class TestOpenAIPassthroughLoggingHandler:
assert handler.get_provider_config("gpt-4o") is not None
@patch("litellm.completion_cost")
@patch(
"litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload"
)
def test_azure_passthrough_tags_metadata_model_provider(
self, mock_get_standard_logging, mock_completion_cost
):
@patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload")
def test_azure_passthrough_tags_metadata_model_provider(self, mock_get_standard_logging, mock_completion_cost):
"""Test that tags, metadata, model, and custom_llm_provider are preserved for Azure passthrough in UI"""
# Arrange
mock_completion_cost.return_value = 0.000045
@ -929,9 +859,7 @@ class TestOpenAIPassthroughLoggingHandler:
# Verify model and custom_llm_provider are set correctly
assert result["kwargs"]["model"] == "gpt-4o"
assert (
result["kwargs"]["custom_llm_provider"] == "azure"
) # Should preserve Azure, not default to "openai"
assert result["kwargs"]["custom_llm_provider"] == "azure" # Should preserve Azure, not default to "openai"
assert result["kwargs"]["response_cost"] == 0.000045
# Verify metadata tags are preserved in litellm_params
@ -955,12 +883,8 @@ class TestOpenAIPassthroughLoggingHandler:
assert call_args[1]["custom_llm_provider"] == "azure"
@patch("litellm.completion_cost")
@patch(
"litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload"
)
@patch(
"litellm.llms.openai.responses.transformation.OpenAIResponsesAPIConfig.transform_response_api_response"
)
@patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload")
@patch("litellm.llms.openai.responses.transformation.OpenAIResponsesAPIConfig.transform_response_api_response")
def test_responses_api_cost_tracking(
self,
mock_transform_responses,
@ -1052,9 +976,7 @@ class TestOpenAIPassthroughLoggingHandler:
assert mock_logging_obj.model_call_details["custom_llm_provider"] == "openai"
@patch("litellm.completion_cost")
@patch(
"litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload"
)
@patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload")
def test_responses_api_uses_responses_transformer_not_chat_completions(
self, mock_get_standard_logging, mock_completion_cost
):
@ -1185,9 +1107,7 @@ class TestOpenAIPassthroughIntegration:
mock_response.headers = {"content-type": "application/json"}
return mock_response
def _create_passthrough_logging_payload(
self, user: str = "test_user"
) -> PassthroughStandardLoggingPayload:
def _create_passthrough_logging_payload(self, user: str = "test_user") -> PassthroughStandardLoggingPayload:
"""Create a mock passthrough logging payload"""
return PassthroughStandardLoggingPayload(
url="https://api.openai.com/v1/chat/completions",
@ -1201,59 +1121,32 @@ class TestOpenAIPassthroughIntegration:
def test_is_openai_route_detection(self):
"""Test OpenAI route detection in the main success handler"""
# Positive cases
assert (
self.handler.is_openai_route("https://api.openai.com/v1/chat/completions")
== True
)
assert (
self.handler.is_openai_route("https://openai.azure.com/v1/chat/completions")
== True
)
assert self.handler.is_openai_route("https://api.openai.com/v1/chat/completions") == True
assert self.handler.is_openai_route("https://openai.azure.com/v1/chat/completions") == True
assert self.handler.is_openai_route("https://api.openai.com/v1/models") == True
# Azure OpenAI on the shared Cognitive Services domain, identified by an
# OpenAI-style path segment.
assert (
self.handler.is_openai_route(
"https://my-resource.cognitiveservices.azure.com/v1/chat/completions"
)
== True
self.handler.is_openai_route("https://my-resource.cognitiveservices.azure.com/v1/chat/completions") == True
)
# Negative cases
assert (
self.handler.is_openai_route(
"http://localhost:4000/openai/v1/chat/completions"
)
== False
)
assert (
self.handler.is_openai_route("https://api.anthropic.com/v1/messages")
== False
)
assert (
self.handler.is_openai_route("https://api.assemblyai.com/v2/transcript")
== False
)
assert self.handler.is_openai_route("http://localhost:4000/openai/v1/chat/completions") == False
assert self.handler.is_openai_route("https://api.anthropic.com/v1/messages") == False
assert self.handler.is_openai_route("https://api.assemblyai.com/v2/transcript") == False
# Non-OpenAI Azure Cognitive Services share the `cognitiveservices.azure.com`
# domain but must NOT be classified as OpenAI routes (no OpenAI path segment).
assert (
self.handler.is_openai_route(
"https://my-resource.cognitiveservices.azure.com/speechtotext/v3.1/recognize"
)
self.handler.is_openai_route("https://my-resource.cognitiveservices.azure.com/speechtotext/v3.1/recognize")
== False
)
assert (
self.handler.is_openai_route(
"https://my-resource.cognitiveservices.azure.com/vision/v3.2/analyze"
)
== False
self.handler.is_openai_route("https://my-resource.cognitiveservices.azure.com/vision/v3.2/analyze") == False
)
# A look-alike domain that merely contains an OpenAI host as a substring
# must be rejected by the suffix-based hostname match.
assert (
self.handler.is_openai_route(
"https://cognitiveservices.azure.com.attacker.example/v1/chat/completions"
)
self.handler.is_openai_route("https://cognitiveservices.azure.com.attacker.example/v1/chat/completions")
== False
)
assert self.handler.is_openai_route("") == False
@ -1274,52 +1167,188 @@ class TestOpenAIPassthroughIntegration:
remove Responses from the OR-chain without a test failure.
"""
# Responses must be supported on api.openai.com and openai.azure.com.
assert (
self.handler._is_supported_openai_endpoint(
"https://api.openai.com/v1/responses"
)
is True
)
assert (
self.handler._is_supported_openai_endpoint(
"https://openai.azure.com/v1/responses"
)
is True
)
assert self.handler._is_supported_openai_endpoint("https://api.openai.com/v1/responses") is True
assert self.handler._is_supported_openai_endpoint("https://openai.azure.com/v1/responses") is True
# The other supported endpoints stay supported (no regression).
assert (
self.handler._is_supported_openai_endpoint(
"https://api.openai.com/v1/chat/completions"
)
is True
)
assert (
self.handler._is_supported_openai_endpoint(
"https://api.openai.com/v1/images/generations"
)
is True
)
assert (
self.handler._is_supported_openai_endpoint(
"https://api.openai.com/v1/images/edits"
)
is True
)
assert self.handler._is_supported_openai_endpoint("https://api.openai.com/v1/chat/completions") is True
assert self.handler._is_supported_openai_endpoint("https://api.openai.com/v1/images/generations") is True
assert self.handler._is_supported_openai_endpoint("https://api.openai.com/v1/images/edits") is True
# Unsupported OpenAI endpoints (e.g. /v1/models) still return False.
assert self.handler._is_supported_openai_endpoint("https://api.openai.com/v1/models") is False
assert (
self.handler._is_supported_openai_endpoint(
"https://api.openai.com/v1/models"
"https://my-resource.openai.azure.com/openai/deployments/text-embedding-3-small/embeddings"
)
is False
)
def test_is_supported_openai_endpoint_includes_embeddings(self):
assert self.handler._is_supported_openai_endpoint("https://api.openai.com/v1/embeddings") is True
assert self.handler._is_supported_openai_endpoint("https://openai.azure.com/v1/embeddings") is True
def test_is_cohere_route_does_not_match_openai_embeddings(self):
assert self.handler.is_cohere_route("https://api.cohere.com/v1/embed") is True
assert self.handler.is_cohere_route("https://api.cohere.com/v2/chat") is True
assert self.handler.is_cohere_route("https://api.openai.com/v1/embeddings") is False
assert self.handler.is_cohere_route("https://api.cohere.com/v1/rerank") is False
assert self.handler.is_cohere_route("http://localhost:4000/openai_passthrough/v1/embeddings") is False
@patch("litellm.completion_cost")
@patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload")
def test_openai_passthrough_handler_embeddings_sets_response_cost(
self, mock_get_standard_logging, mock_completion_cost
):
mock_completion_cost.return_value = 2.8e-07
mock_get_standard_logging.return_value = {"test": "logging_payload"}
response_body = {
"object": "list",
"model": "text-embedding-3-small",
"data": [
{
"object": "embedding",
"index": 0,
"embedding": [0.1, 0.2],
}
],
"usage": {"prompt_tokens": 14, "total_tokens": 14},
}
mock_httpx_response = self._create_mock_httpx_response(response_body)
mock_logging_obj = self._create_mock_logging_obj()
passthrough_payload = PassthroughStandardLoggingPayload(
url="https://api.openai.com/v1/embeddings",
request_body={
"model": "text-embedding-3-small",
"input": "PROOF_SENTINEL_TEXT",
},
request_method="POST",
)
kwargs = {
"passthrough_logging_payload": passthrough_payload,
"litellm_params": {},
}
result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler(
httpx_response=mock_httpx_response,
response_body=response_body,
logging_obj=mock_logging_obj,
url_route="https://api.openai.com/v1/embeddings",
result="",
start_time=self.start_time,
end_time=self.end_time,
cache_hit=False,
request_body={
"model": "text-embedding-3-small",
"input": "PROOF_SENTINEL_TEXT",
},
**kwargs,
)
assert result["result"] is not None
assert result["kwargs"]["response_cost"] == 2.8e-07
assert result["kwargs"]["model"] == "text-embedding-3-small"
assert result["kwargs"]["custom_llm_provider"] == "openai"
assert result["result"]._hidden_params["response_cost"] == 2.8e-07
mock_completion_cost.assert_called_once()
assert mock_completion_cost.call_args.kwargs["call_type"] == "aembedding"
assert mock_logging_obj.model_call_details["response_cost"] == 2.8e-07
@patch(
"litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.passthrough_chat_handler"
)
@patch("litellm.completion_cost")
def test_openai_passthrough_handler_embeddings_without_model_falls_back(
self, mock_completion_cost, mock_chat_handler
):
mock_chat_handler.return_value = {"result": None, "kwargs": {}}
response_body = {
"object": "list",
"data": [{"object": "embedding", "index": 0, "embedding": [0.1]}],
"usage": {"prompt_tokens": 1, "total_tokens": 1},
}
result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler(
httpx_response=self._create_mock_httpx_response(response_body),
response_body=response_body,
logging_obj=self._create_mock_logging_obj(),
url_route="https://api.openai.com/v1/embeddings",
result="",
start_time=self.start_time,
end_time=self.end_time,
cache_hit=False,
request_body={"input": "PROOF_SENTINEL_TEXT"},
passthrough_logging_payload=PassthroughStandardLoggingPayload(
url="https://api.openai.com/v1/embeddings",
request_body={"input": "PROOF_SENTINEL_TEXT"},
request_method="POST",
),
)
mock_completion_cost.assert_not_called()
mock_chat_handler.assert_called_once()
assert result == {"result": None, "kwargs": {}}
@patch(
"litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.openai_passthrough_handler"
)
@pytest.mark.asyncio
async def test_success_handler_dispatches_responses_api_to_openai_handler(
self, mock_openai_handler
):
async def test_success_handler_dispatches_embeddings_to_openai_handler(self, mock_openai_handler):
mock_openai_handler.return_value = {
"result": {"object": "list"},
"kwargs": {
"response_cost": 2.8e-07,
"model": "text-embedding-3-small",
"custom_llm_provider": "openai",
},
}
mock_httpx_response = MagicMock(spec=httpx.Response)
mock_httpx_response.text = (
'{"object":"list","model":"text-embedding-3-small",'
'"data":[{"object":"embedding","index":0,"embedding":[0.1]}],'
'"usage":{"prompt_tokens":14,"total_tokens":14}}'
)
mock_logging_obj = AsyncMock()
mock_logging_obj.model_call_details = {}
mock_logging_obj.async_success_handler = AsyncMock()
passthrough_payload = PassthroughStandardLoggingPayload(
url="https://api.openai.com/v1/embeddings",
request_body={
"model": "text-embedding-3-small",
"input": "PROOF_SENTINEL_TEXT",
},
request_method="POST",
)
await self.handler.pass_through_async_success_handler(
httpx_response=mock_httpx_response,
response_body={
"object": "list",
"model": "text-embedding-3-small",
"data": [{"object": "embedding", "index": 0, "embedding": [0.1]}],
"usage": {"prompt_tokens": 14, "total_tokens": 14},
},
logging_obj=mock_logging_obj,
url_route="https://api.openai.com/v1/embeddings",
result="",
start_time=datetime.now(),
end_time=datetime.now(),
cache_hit=False,
request_body={
"model": "text-embedding-3-small",
"input": "PROOF_SENTINEL_TEXT",
},
passthrough_logging_payload=passthrough_payload,
)
mock_openai_handler.assert_called_once()
assert mock_openai_handler.call_args.kwargs["url_route"] == "https://api.openai.com/v1/embeddings"
@patch(
"litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.openai_passthrough_handler"
)
@pytest.mark.asyncio
async def test_success_handler_dispatches_responses_api_to_openai_handler(self, mock_openai_handler):
"""End-to-end dispatch test for the Responses API path.
Pre-fix: `_is_supported_openai_endpoint` returned False for
@ -1395,9 +1424,7 @@ class TestOpenAIPassthroughIntegration:
}
mock_httpx_response = MagicMock(spec=httpx.Response)
mock_httpx_response.text = (
'{"id": "chatcmpl-123", "choices": [{"message": {"content": "Hello"}}]}'
)
mock_httpx_response.text = '{"id": "chatcmpl-123", "choices": [{"message": {"content": "Hello"}}]}'
mock_logging_obj = AsyncMock()
mock_logging_obj.model_call_details = {}
@ -1590,14 +1617,10 @@ class TestOpenAIPassthroughIntegration:
# Test the _response_cost_calculator method
calculated_cost = logging_obj._response_cost_calculator(result=image_response)
assert (
calculated_cost == test_cost
), f"Expected {test_cost}, got {calculated_cost}"
assert calculated_cost == test_cost, f"Expected {test_cost}, got {calculated_cost}"
@patch("litellm.cost_calculator.default_image_cost_calculator")
def test_openai_passthrough_handler_image_generation(
self, mock_image_cost_calculator
):
def test_openai_passthrough_handler_image_generation(self, mock_image_cost_calculator):
"""Test successful cost tracking for OpenAI image generation"""
# Arrange
mock_image_cost_calculator.return_value = 0.040

28
tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
View file

@ -243,6 +243,34 @@ def test_bedrock_mantle_provider_fields():
assert fields_by_key["api_base"]["field_type"] == "text"
def test_nvidia_riva_provider_fields():
app_instance = FastAPI()
app_instance.include_router(router)
test_client = TestClient(app_instance)
response = test_client.get("/public/providers/fields")
assert response.status_code == 200
providers = response.json()
riva = next((p for p in providers if p["provider"] == "NVIDIA_RIVA"), None)
assert riva is not None, "NVIDIA Riva provider entry not found"
assert riva["provider_display_name"] == "Nvidia Riva"
assert riva["litellm_provider"] == LlmProviders.NVIDIA_RIVA.value
assert riva["default_model_placeholder"].startswith("nvidia_riva/")
fields_by_key = {f["key"]: f for f in riva["credential_fields"]}
assert fields_by_key["api_base"]["required"] is True
assert fields_by_key["api_base"]["field_type"] == "text"
assert fields_by_key["api_key"]["required"] is False
assert fields_by_key["api_key"]["field_type"] == "password"
assert "nvcf_function_id" in fields_by_key
assert fields_by_key["nvcf_function_id"]["required"] is False
def test_google_ai_studio_provider_fields_expose_api_base():
"""The Google AI Studio (gemini) credential form must let admins set a custom
api_base so they can point at a Gemini-compatible gateway (e.g. a self-hosted

81
tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py
View file

@ -657,6 +657,87 @@ async def test_scheduled_rollup_stays_quiet_when_every_charge_landed():
alert.assert_not_awaited()
@pytest.mark.asyncio
async def test_scheduled_rollup_alerts_once_a_ptu_window_has_closed():
"""Reserved capacity is billed until the deployment is deleted, so a closed window stops
the attribution without stopping the charge. Nobody notices unless it is escalated."""
ptu = {
"ptu_count": 5,
"cost_per_ptu_per_hour": 2.0,
"team_id": "t",
"ptu_effective_from": "2020-01-01T00:00:00Z",
"ptu_effective_to": "2020-02-01T00:00:00Z",
}
prisma, _ = _prisma_with_models([_model_row(model_id="dep-lapsed", model_info=ptu)])
alert = AsyncMock()
result = await run_scheduled_ptu_rollup(prisma, target_date=DAY, alert=alert)
assert result.lapsed == ("gpt-4o-mini-ptu",)
alert.assert_awaited_once()
message = alert.await_args.args[0]
assert "window has closed" in message
assert "gpt-4o-mini-ptu" in message
@pytest.mark.asyncio
async def test_a_model_name_cannot_smuggle_slack_markup_into_the_alert():
"""The alert lands in an operator channel and a model name is operator-supplied, so an
unescaped name could post a channel-wide mention."""
ptu = {
"ptu_count": 5,
"cost_per_ptu_per_hour": 2.0,
"team_id": "t",
"ptu_effective_from": "2020-01-01T00:00:00Z",
"ptu_effective_to": "2020-02-01T00:00:00Z",
}
row = _model_row(model_id="dep-x", model_name="<!channel> & <https://evil.example|click>", model_info=ptu)
prisma, _ = _prisma_with_models([row])
alert = AsyncMock()
await run_scheduled_ptu_rollup(prisma, target_date=DAY, alert=alert)
message = alert.await_args.args[0]
assert "<!channel>" not in message
assert "&lt;!channel&gt;" in message
@pytest.mark.asyncio
async def test_an_open_ptu_window_raises_no_lapsed_alert():
ptu = {
"ptu_count": 5,
"cost_per_ptu_per_hour": 2.0,
"team_id": "t",
"ptu_effective_from": "2020-01-01T00:00:00Z",
"ptu_effective_to": "2999-01-01T00:00:00Z",
}
prisma, _ = _prisma_with_models([_model_row(model_id="dep-open", model_info=ptu)])
alert = AsyncMock()
result = await run_scheduled_ptu_rollup(prisma, target_date=DAY, alert=alert)
assert result.lapsed == ()
alert.assert_not_awaited()
@pytest.mark.asyncio
async def test_an_open_ended_ptu_window_raises_no_lapsed_alert():
"""No end bound means the operator never asked the attribution to stop."""
ptu = {
"ptu_count": 5,
"cost_per_ptu_per_hour": 2.0,
"team_id": "t",
"ptu_effective_from": "2020-01-01T00:00:00Z",
}
prisma, _ = _prisma_with_models([_model_row(model_id="dep-forever", model_info=ptu)])
alert = AsyncMock()
result = await run_scheduled_ptu_rollup(prisma, target_date=DAY, alert=alert)
assert result.lapsed == ()
alert.assert_not_awaited()
@pytest.mark.asyncio
async def test_a_broken_alert_channel_does_not_fail_the_rollup():
rows = [_model_row(model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"})]

31
tests/test_litellm/proxy/test_conftest.py Normal file
View file

@ -0,0 +1,31 @@
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@pytest.fixture
def fixture_planted_prisma_mock():
with patch("litellm.proxy.proxy_server.prisma_client", MagicMock()):
yield
def test_monkeypatch_over_fixture_patched_prisma_client(
fixture_planted_prisma_mock, monkeypatch
):
"""
Mirrors the flake in test_team_endpoints.py: an autouse fixture patches
prisma_client, the test monkeypatches the same global, and monkeypatch
records the fixture's MagicMock as the value to restore. Its undo runs
after every other finalizer, so without hook-level isolation the mock
leaks and every later no-database test on the worker fails awaiting it.
"""
import litellm.proxy.proxy_server as proxy_server
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", AsyncMock())
assert isinstance(proxy_server.prisma_client, AsyncMock)
def test_prisma_client_did_not_leak_from_previous_test():
import litellm.proxy.proxy_server as proxy_server
assert not isinstance(proxy_server.prisma_client, MagicMock)

85
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -6872,6 +6872,91 @@ async def test_update_general_settings_propagates_apply_user_budget_to_team_keys
assert ps.general_settings["apply_user_budget_to_team_keys"] is True
@pytest.mark.asyncio
async def test_update_general_settings_propagates_spend_log_cleanup_bounds():
"""The dashboard writes the cleanup bounds straight to the DB config, so
without runtime propagation the scheduled job never sees them and the knobs
do nothing until the process restarts."""
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
SPEND_LOG_CLEANUP_BOUND_SETTINGS,
)
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
db_settings = {
"maximum_spend_logs_cleanup_batch_size": 2000,
"maximum_spend_logs_cleanup_max_batches": 250,
"maximum_spend_logs_cleanup_run_budget": "90s",
"maximum_spend_logs_cleanup_batch_timeout": "10s",
}
assert set(db_settings) == set(SPEND_LOG_CLEANUP_BOUND_SETTINGS)
with patch("litellm.proxy.proxy_server.general_settings", {}):
await proxy_config._update_general_settings(db_general_settings=db_settings)
import litellm.proxy.proxy_server as ps
assert {key: ps.general_settings.get(key) for key in db_settings} == db_settings
@pytest.mark.asyncio
async def test_update_general_settings_clears_a_spend_log_cleanup_bound_dropped_from_the_db():
"""Blanking the field in the dashboard deletes the key outright, so leaving
the last value in memory would keep a bound the operator just removed."""
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
with patch(
"litellm.proxy.proxy_server.general_settings",
{"maximum_spend_logs_cleanup_run_budget": "90s", "maximum_spend_logs_cleanup_batch_timeout": "10s"},
):
await proxy_config._update_general_settings(
db_general_settings={"maximum_spend_logs_cleanup_batch_timeout": "10s"}
)
import litellm.proxy.proxy_server as ps
assert ps.general_settings["maximum_spend_logs_cleanup_run_budget"] is None
assert ps.general_settings["maximum_spend_logs_cleanup_batch_timeout"] == "10s"
@pytest.mark.asyncio
async def test_update_general_settings_keeps_a_yaml_set_spend_log_cleanup_bound():
"""A YAML-set bound never appears in the DB object, so treating its absence
as a dashboard clear would discard the deployed config on every reload."""
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
proxy_config._yaml_spend_log_cleanup_bounds = {"maximum_spend_logs_cleanup_run_budget": "90s"}
with patch("litellm.proxy.proxy_server.general_settings", {"maximum_spend_logs_cleanup_run_budget": "90s"}):
await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": True})
import litellm.proxy.proxy_server as ps
assert ps.general_settings["maximum_spend_logs_cleanup_run_budget"] == "90s"
@pytest.mark.asyncio
async def test_update_general_settings_clearing_a_db_override_falls_back_to_the_yaml_bound():
"""Clearing a dashboard override of a YAML-declared bound must restore the
YAML value. Leaving the deleted override in memory would keep enforcing the
bound the operator just removed, until the process restarted."""
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
proxy_config._yaml_spend_log_cleanup_bounds = {"maximum_spend_logs_cleanup_run_budget": "90s"}
# Memory currently holds the dashboard override, and the DB no longer carries it.
with patch("litellm.proxy.proxy_server.general_settings", {"maximum_spend_logs_cleanup_run_budget": "30s"}):
await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": True})
import litellm.proxy.proxy_server as ps
assert ps.general_settings["maximum_spend_logs_cleanup_run_budget"] == "90s"
@pytest.mark.asyncio
async def test_update_general_settings_apply_user_budget_to_team_keys_yaml_wins():
"""A DB value must not silently override an explicit YAML setting on reload."""

665
tests/test_litellm/proxy/test_spend_log_cleanup.py
View file

@ -2,12 +2,66 @@
Test cases for spend log cleanup functionality
"""
import asyncio
import math
import time
from contextlib import asynccontextmanager
from datetime import datetime, timedelta, timezone
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup
from litellm.constants import (
SPEND_LOG_CLEANUP_BATCH_SIZE,
SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP,
SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS,
)
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
SPEND_LOG_CLEANUP_BOUND_SETTINGS,
SpendLogCleanup,
TableCleanupResult,
)
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import (
SpendLogCleanupMetrics,
)
def _far_deadline() -> float:
"""A run deadline far enough out that only the other bounds can stop a batch loop."""
return time.monotonic() + 3600
def _wire_tx(db):
"""
Model the prisma seam the cleanup job actually uses.
Every statement the job issues runs inside db.tx() so it can carry a SET
LOCAL statement_timeout. Batch and probe statements are forwarded to
db.execute_raw and db.query_raw, which is what tests configure and assert
on, while the SET LOCAL statements are answered here so they neither consume
a side_effect entry nor show up in the recorded call list. Lookup is
deferred to call time so this can be wired before a test assigns its own
execute_raw.
"""
@asynccontextmanager
async def _tx():
tx = MagicMock()
async def _execute_raw(sql, *args):
if sql.lstrip().upper().startswith("SET LOCAL"):
return 0
return await db.execute_raw(sql, *args)
async def _query_raw(sql, *args):
return await db.query_raw(sql, *args)
tx.execute_raw = _execute_raw
tx.query_raw = _query_raw
yield tx
db.tx = _tx
db.query_raw = AsyncMock(return_value=[{"remaining": 0}])
def test_spend_log_cleanup_cron_scheduling():
@ -49,6 +103,7 @@ def test_spend_log_cleanup_cron_scheduler_integration():
# Mock scheduler
mock_scheduler = MagicMock()
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_cleanup_instance = MagicMock()
# Test Case 1: Cron-based scheduling
@ -155,7 +210,9 @@ async def test_cleanup_old_spend_logs_batch_deletion():
# Setup Prisma client
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_db = MagicMock()
_wire_tx(mock_db)
# Mock execute_raw to return deleted counts (3 spend-log batches, then the
# tool-index cleanup's first batch returning 0)
@ -207,7 +264,9 @@ async def test_cleanup_old_spend_logs_retention_period_cutoff():
"""
# Setup Prisma client
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_db = MagicMock()
_wire_tx(mock_db)
mock_db.execute_raw = AsyncMock(return_value=0)
mock_prisma_client.db = mock_db
@ -244,6 +303,7 @@ async def test_cleanup_drops_partitions_when_enabled_and_partitioned():
from unittest.mock import AsyncMock, MagicMock
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_prisma_client.db.execute_raw = AsyncMock(return_value=0)
partition_manager = MagicMock()
@ -285,6 +345,7 @@ async def test_cleanup_uses_delete_when_partitioning_not_enabled():
from unittest.mock import AsyncMock, MagicMock
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[10, 0, 0])
partition_manager = MagicMock()
@ -316,6 +377,7 @@ async def test_cleanup_uses_delete_when_not_partitioned():
from unittest.mock import AsyncMock, MagicMock
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[10, 0, 0])
partition_manager = MagicMock()
@ -346,6 +408,7 @@ async def test_cleanup_old_spend_logs_no_retention_period():
Test that no logs are deleted when no retention period is set
"""
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_prisma_client.db.execute_raw = AsyncMock()
cleaner = SpendLogCleanup(general_settings={}) # no retention
@ -361,6 +424,7 @@ async def test_lock_not_released_when_not_acquired():
before the lock is ever acquired.
"""
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_prisma_client.db.execute_raw = AsyncMock()
mock_redis_cache = MagicMock()
@ -418,7 +482,9 @@ async def test_delete_old_logs_aborts_on_non_int_execute_raw_return():
"""should abort deletion loop immediately when execute_raw returns a non-int
(e.g. None or dict), preventing an infinite loop."""
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_db = MagicMock()
_wire_tx(mock_db)
mock_db.execute_raw = AsyncMock(return_value=None)
mock_prisma_client.db = mock_db
@ -427,17 +493,19 @@ async def test_delete_old_logs_aborts_on_non_int_execute_raw_return():
)
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date)
result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())
assert mock_db.execute_raw.call_count == 1
assert total_deleted == 0
assert result.rows_deleted == 0
@pytest.mark.asyncio
async def test_delete_old_logs_continues_on_valid_int_return():
"""should continue deletion loop across batches when execute_raw returns valid int counts."""
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_db = MagicMock()
_wire_tx(mock_db)
mock_db.execute_raw = AsyncMock(side_effect=[500, 300, 0])
mock_prisma_client.db = mock_db
@ -446,35 +514,37 @@ async def test_delete_old_logs_continues_on_valid_int_return():
)
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date)
result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())
assert mock_db.execute_raw.call_count == 3
assert total_deleted == 800
assert result.rows_deleted == 800
@pytest.mark.asyncio
async def test_delete_old_rows_stops_at_max_batches(monkeypatch):
"""The run-loop backstop must halt a cleanup that keeps finding rows, so a
huge backlog is spread across scheduled runs instead of one unbounded loop."""
import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
monkeypatch.setattr(cleanup_module, "SPEND_LOG_RUN_LOOPS", 2)
async def test_delete_old_rows_stops_at_max_batches():
"""The batch cap must halt a cleanup that keeps finding rows, so a huge
backlog is spread across scheduled runs instead of one unbounded loop, and
the operator-facing knob must mean exactly the number of statements it names."""
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_db = MagicMock()
_wire_tx(mock_db)
mock_db.execute_raw = AsyncMock(return_value=1000)
mock_prisma_client.db = mock_db
cleaner = SpendLogCleanup(
general_settings={"maximum_spend_logs_retention_period": "7d"}
general_settings={
"maximum_spend_logs_retention_period": "7d",
"maximum_spend_logs_cleanup_max_batches": 2,
}
)
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date)
result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())
# run_count exceeds the cap only after 3 full batches (0, 1, 2)
assert mock_db.execute_raw.call_count == 3
assert total_deleted == 3000
assert mock_db.execute_raw.call_count == 2
assert result.rows_deleted == 2000
assert result.stop_reason == "batch_cap_reached"
@pytest.mark.asyncio
@ -482,7 +552,9 @@ async def test_delete_old_tool_index_rows_deletes_on_composite_key():
"""Tool index rows are derived from spend logs and expire on the same cutoff;
the delete must match on the table's composite primary key."""
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_db = MagicMock()
_wire_tx(mock_db)
mock_db.execute_raw = AsyncMock(side_effect=[5, 0])
mock_prisma_client.db = mock_db
@ -491,9 +563,9 @@ async def test_delete_old_tool_index_rows_deletes_on_composite_key():
)
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
total_deleted = await cleaner._delete_old_tool_index_rows(mock_prisma_client, cutoff_date)
result = await cleaner._delete_old_tool_index_rows(mock_prisma_client, cutoff_date, _far_deadline())
assert total_deleted == 5
assert result.rows_deleted == 5
delete_sql = mock_db.execute_raw.call_args_list[0][0][0]
assert 'DELETE FROM "LiteLLM_SpendLogToolIndex"' in delete_sql
assert 'WHERE ("request_id", "tool_name") IN' in delete_sql
@ -513,7 +585,9 @@ async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch)
)
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_db = MagicMock()
_wire_tx(mock_db)
# batch 1 succeeds, batch 2 raises (one-off DB timeout), batches 3-4 succeed,
# batch 5 returns 0 → loop exits naturally.
mock_db.execute_raw = AsyncMock(
@ -526,11 +600,11 @@ async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch)
)
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date)
result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())
# All 5 batches should have been attempted; 100 + 200 + 50 = 350 deleted.
assert mock_db.execute_raw.call_count == 5
assert total_deleted == 350
assert result.rows_deleted == 350
@pytest.mark.asyncio
@ -548,7 +622,9 @@ async def test_delete_old_logs_aborts_after_consecutive_failures(monkeypatch):
)
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_db = MagicMock()
_wire_tx(mock_db)
# Every batch raises — must abort after exactly 3 attempts, not loop forever.
mock_db.execute_raw = AsyncMock(
side_effect=ConnectionError("simulated persistent DB outage")
@ -560,10 +636,10 @@ async def test_delete_old_logs_aborts_after_consecutive_failures(monkeypatch):
)
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date)
result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())
assert mock_db.execute_raw.call_count == 3
assert total_deleted == 0
assert result.rows_deleted == 0
@pytest.mark.asyncio
@ -580,7 +656,9 @@ async def test_delete_old_logs_resets_consecutive_failures_on_success(monkeypatc
)
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_db = MagicMock()
_wire_tx(mock_db)
# Pattern: fail, fail, success (resets counter), fail, fail, success, done.
# Without reset, three of these would trip abort; with reset, they don't.
mock_db.execute_raw = AsyncMock(
@ -601,10 +679,10 @@ async def test_delete_old_logs_resets_consecutive_failures_on_success(monkeypatc
)
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date)
result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())
assert mock_db.execute_raw.call_count == 7
assert total_deleted == 150
assert result.rows_deleted == 150
@pytest.mark.asyncio
@ -617,6 +695,7 @@ async def test_cleanup_uses_logger_exception_for_full_traceback(monkeypatch):
monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger)
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
# Force the outer try/except to fire by making _should_delete_spend_logs raise.
cleaner = cleanup_module.SpendLogCleanup(
general_settings={"maximum_spend_logs_retention_period": "7d"}
@ -653,7 +732,9 @@ async def test_cleanup_releases_lock_after_persistent_batch_failures(monkeypatch
)
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_db = MagicMock()
_wire_tx(mock_db)
mock_db.execute_raw = AsyncMock(side_effect=TimeoutError("DB down"))
mock_prisma_client.db = mock_db
@ -698,6 +779,7 @@ def _mock_prisma_for_retention(side_effect: list) -> "MagicMock":
from unittest.mock import AsyncMock, MagicMock
client = MagicMock()
_wire_tx(client.db)
client.db.execute_raw = AsyncMock(side_effect=side_effect)
return client
@ -753,3 +835,536 @@ async def test_no_retention_keys_means_no_cleanup_at_all():
cleaner.pod_lock_manager = None
await cleaner.cleanup_old_spend_logs(client)
assert client.db.execute_raw.await_count == 0
@pytest.mark.asyncio
async def test_run_budget_stops_the_loop_and_leaves_the_backlog_for_the_next_run():
"""
The wall-clock budget is the bound that keeps a large backlog from turning
into one multi-hour run. With rows always available, the loop must stop on
the deadline rather than on the batch cap, and must report that reason so
operators can tell a budgeted stop from a drained table.
"""
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_db = MagicMock()
_wire_tx(mock_db)
mock_db.execute_raw = AsyncMock(return_value=1000)
mock_prisma_client.db = mock_db
cleaner = SpendLogCleanup(
general_settings={
"maximum_spend_logs_retention_period": "7d",
# Comfortably more batches than a sub-second budget can reach (each
# batch sleeps 0.1s), but small enough that a broken deadline fails
# this test in seconds instead of hanging it
"maximum_spend_logs_cleanup_max_batches": 50,
}
)
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
started_at = time.monotonic()
result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, time.monotonic() + 0.25)
elapsed = time.monotonic() - started_at
assert result.stop_reason == "budget_exhausted"
assert elapsed < 3, f"budgeted run overran its deadline: {elapsed}s"
assert mock_db.execute_raw.call_count < 50
assert result.rows_deleted > 0
@pytest.mark.asyncio
async def test_run_budget_is_shared_across_tables_not_granted_per_table():
"""
A per-table budget would let a run take N times the configured bound. The
deadline is computed once per run, so once it is spent on the first table
the later tables must stop immediately rather than each getting a fresh one.
"""
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_db = MagicMock()
_wire_tx(mock_db)
mock_db.execute_raw = AsyncMock(return_value=1000)
mock_prisma_client.db = mock_db
cleaner = SpendLogCleanup(
general_settings={
"maximum_spend_logs_retention_period": "7d",
"maximum_autorouter_session_retention_period": "365d",
# Comfortably more batches than a sub-second budget can reach (each
# batch sleeps 0.1s), but small enough that a broken deadline fails
# this test in seconds instead of hanging it
"maximum_spend_logs_cleanup_max_batches": 50,
"maximum_spend_logs_cleanup_run_budget": "1s",
}
)
cleaner.pod_lock_manager = None
started_at = time.monotonic()
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
elapsed = time.monotonic() - started_at
# three tables are eligible; a per-table budget would push this past 3s
assert elapsed < 2.5, f"budget was granted per table, not per run: {elapsed}s"
tables_touched = {call[0][0].split('"')[1] for call in mock_db.execute_raw.call_args_list}
assert "LiteLLM_SpendLogs" in tables_touched
@pytest.mark.asyncio
async def test_each_batch_carries_a_statement_and_lock_timeout():
"""
A Prisma transaction timeout cannot interrupt a statement already running,
so the Postgres statement_timeout and lock_timeout are the only things
stopping one batch from holding row locks and a pooled connection
indefinitely. Both must be set, inside the batch's own transaction, and
scoped with SET LOCAL so the pooled connection is left unchanged.
"""
recorded: list[str] = []
mock_prisma_client = MagicMock()
mock_db = MagicMock()
@asynccontextmanager
async def _tx():
tx = MagicMock()
async def _execute_raw(sql, *args):
recorded.append(sql.strip())
return 0
tx.execute_raw = _execute_raw
yield tx
mock_db.tx = _tx
mock_db.query_raw = AsyncMock(return_value=[{"remaining": 0}])
mock_prisma_client.db = mock_db
cleaner = SpendLogCleanup(
general_settings={
"maximum_spend_logs_retention_period": "7d",
"maximum_spend_logs_cleanup_batch_timeout": "12s",
}
)
await cleaner._delete_old_logs(
mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline()
)
assert "SET LOCAL statement_timeout = 12000" in recorded
assert "SET LOCAL lock_timeout = 12000" in recorded
# the timeouts must precede the delete they are meant to bound
assert recorded.index("SET LOCAL statement_timeout = 12000") < next(
i for i, sql in enumerate(recorded) if sql.startswith("DELETE")
)
@pytest.mark.parametrize(
"setting_value",
["inf", "-inf", "nan", "1e400", "0s", "-5m", "not-a-duration"],
)
def test_a_non_finite_or_non_positive_budget_falls_back_to_the_default(setting_value):
"""
The knob must not be able to remove the bound it exists to enforce.
'inf', 'nan' and '1e400' are the spellings that would turn the deadline
into no deadline at all, and '0s' and '-5m' would make every run stop before
deleting anything. All of them must land on the default rather than being
honoured, and the resulting budget must be usable arithmetic.
"""
cleaner = SpendLogCleanup(
general_settings={
"maximum_spend_logs_retention_period": "7d",
"maximum_spend_logs_cleanup_run_budget": setting_value,
}
)
assert cleaner.run_budget_seconds == SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS
assert math.isfinite(cleaner.run_budget_seconds)
assert cleaner.run_budget_seconds > 0
@pytest.mark.parametrize("setting_value", [0, -1, "abc", "", 2.9])
def test_a_bad_batch_size_falls_back_to_the_default(setting_value):
"""A zero or negative batch size would make every DELETE a no-op and the
loop spin, so unusable values must fall back rather than be honoured."""
cleaner = SpendLogCleanup(
general_settings={
"maximum_spend_logs_retention_period": "7d",
"maximum_spend_logs_cleanup_batch_size": setting_value,
}
)
assert cleaner.batch_size >= 1
def test_operator_knobs_override_the_env_defaults():
"""The knobs are meant to be reachable from general_settings (and therefore
from the admin UI), not only from environment variables."""
cleaner = SpendLogCleanup(
general_settings={
"maximum_spend_logs_retention_period": "7d",
"maximum_spend_logs_cleanup_batch_size": 250,
"maximum_spend_logs_cleanup_max_batches": 7,
"maximum_spend_logs_cleanup_run_budget": "90s",
"maximum_spend_logs_cleanup_batch_timeout": "2m",
}
)
assert cleaner.batch_size == 250
assert cleaner.max_batches == 7
assert cleaner.run_budget_seconds == 90
assert cleaner.batch_timeout_seconds == 120
_BOUND_SETTING_CASES = (
("maximum_spend_logs_cleanup_batch_size", 137, "batch_size", 137),
("maximum_spend_logs_cleanup_max_batches", 9, "max_batches", 9),
("maximum_spend_logs_cleanup_run_budget", "45s", "run_budget_seconds", 45.0),
("maximum_spend_logs_cleanup_batch_timeout", "8s", "batch_timeout_seconds", 8.0),
)
@pytest.mark.parametrize("setting_name, setting_value, attribute, expected", _BOUND_SETTING_CASES)
@pytest.mark.asyncio
async def test_a_bound_changed_after_construction_reaches_the_next_run(
setting_name, setting_value, attribute, expected
):
"""The scheduler holds one long-lived instance and the config reload mutates
general_settings in place, so a bound captured at construction would leave
every dashboard change inert until the process restarts."""
settings = {"maximum_spend_logs_retention_period": "7d"}
cleaner = SpendLogCleanup(general_settings=settings)
cleaner.pod_lock_manager = None
assert getattr(cleaner, attribute) != expected
settings[setting_name] = setting_value
await cleaner.cleanup_old_spend_logs(_mock_prisma_for_retention([0, 0]))
assert getattr(cleaner, attribute) == expected
@pytest.mark.parametrize("cleared_to_none", [True, False])
@pytest.mark.asyncio
async def test_a_bound_cleared_after_construction_falls_back_to_its_default(cleared_to_none):
"""Blanking the field in the dashboard has to restore the shipped default
rather than leave the operator's old bound in force, whether the reload
spells the clear as an explicit None or as an absent key."""
settings = {"maximum_spend_logs_retention_period": "7d", "maximum_spend_logs_cleanup_batch_size": 137}
cleaner = SpendLogCleanup(general_settings=settings)
cleaner.pod_lock_manager = None
assert cleaner.batch_size == 137
if cleared_to_none:
settings["maximum_spend_logs_cleanup_batch_size"] = None
else:
del settings["maximum_spend_logs_cleanup_batch_size"]
await cleaner.cleanup_old_spend_logs(_mock_prisma_for_retention([0, 0]))
assert cleaner.batch_size == SPEND_LOG_CLEANUP_BATCH_SIZE
def test_every_declared_bound_setting_is_covered_by_a_live_reread_case():
"""A bound added to the declared set without a live-reread case would be
propagated by the proxy and then ignored by the running job."""
assert {case[0] for case in _BOUND_SETTING_CASES} == set(SPEND_LOG_CLEANUP_BOUND_SETTINGS)
@pytest.mark.asyncio
async def test_remaining_rows_probe_is_capped_so_it_cannot_scan_the_table():
"""The remaining-eligible-rows metric must never itself become the long
scan this job exists to avoid, so its probe carries a LIMIT."""
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_db = MagicMock()
_wire_tx(mock_db)
mock_db.execute_raw = AsyncMock(return_value=0)
mock_prisma_client.db = mock_db
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
await cleaner._delete_old_logs(
mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline()
)
count_sql = mock_db.query_raw.call_args[0][0]
assert "count(*)" in count_sql
assert "LIMIT $2" in count_sql
assert mock_db.query_raw.call_args[0][2] == SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP
@pytest.mark.asyncio
async def test_a_run_skipped_because_another_pod_holds_the_lock_is_reported():
"""Operators need to tell "nothing to do" apart from "someone else is doing
it", so a lock-skipped run is recorded under its own outcome."""
recorded: list[str] = []
original_record_run = SpendLogCleanupMetrics.record_run
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
cleaner.pod_lock_manager = MagicMock()
cleaner.pod_lock_manager.redis_cache = MagicMock()
cleaner.pod_lock_manager.acquire_lock = AsyncMock(return_value=False)
cleaner.pod_lock_manager.release_lock = AsyncMock()
SpendLogCleanupMetrics.record_run = classmethod(lambda cls, outcome: recorded.append(outcome))
try:
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
finally:
SpendLogCleanupMetrics.record_run = original_record_run
assert recorded == ["skipped_locked"]
cleaner.pod_lock_manager.release_lock.assert_not_awaited()
@pytest.mark.asyncio
async def test_the_outstanding_rows_probe_carries_a_statement_timeout():
"""
The probe is a statement like any other, so if it were issued bare a slow one
would hold a connection past the budget the job advertises, which is exactly
what the bounds exist to prevent. With budget to spare it carries the same
per-statement timeout the delete batches do.
"""
recorded: list[str] = []
mock_prisma_client = MagicMock()
mock_db = MagicMock()
@asynccontextmanager
async def _tx():
tx = MagicMock()
async def _execute_raw(sql, *args):
recorded.append(sql.strip())
return 0
async def _query_raw(sql, *args):
recorded.append(sql.strip())
return [{"remaining": 7}]
tx.execute_raw = _execute_raw
tx.query_raw = _query_raw
yield tx
mock_db.tx = _tx
mock_prisma_client.db = mock_db
cleaner = SpendLogCleanup(
general_settings={
"maximum_spend_logs_retention_period": "7d",
"maximum_spend_logs_cleanup_batch_timeout": "8s",
}
)
remaining = await cleaner._count_remaining(
mock_prisma_client,
datetime.now(timezone.utc) - timedelta(days=7),
"LiteLLM_SpendLogs",
"startTime",
_far_deadline(),
)
assert remaining == 7
count_index = next(i for i, sql in enumerate(recorded) if sql.startswith("SELECT count(*)"))
assert "SET LOCAL statement_timeout = 8000" in recorded[:count_index], (
f"the probe ran without a statement timeout: {recorded}"
)
@pytest.mark.asyncio
async def test_a_statement_timeout_is_clamped_to_the_budget_that_is_left():
"""
Postgres has no 'stop at time T', only a per-statement duration, so a batch
issued just under the deadline would run a whole batch timeout past it and
the run budget would be advisory. Clamping the timeout to the remaining
budget is what makes the budget a real wall clock.
"""
recorded: list[str] = []
client = MagicMock()
@asynccontextmanager
async def _tx():
tx = MagicMock()
async def _execute_raw(sql, *args):
recorded.append(sql.strip())
return 0
tx.execute_raw = _execute_raw
tx.query_raw = AsyncMock(return_value=[{"remaining": 0}])
yield tx
client.db.tx = _tx
cleaner = SpendLogCleanup(
general_settings={
"maximum_spend_logs_retention_period": "7d",
"maximum_spend_logs_cleanup_batch_timeout": "30s",
}
)
# Only 2s of budget left against a 30s batch timeout.
await cleaner._execute_delete_batch(client, "DELETE FROM x", datetime.now(timezone.utc), time.monotonic() + 2)
timeouts = [sql for sql in recorded if "statement_timeout" in sql]
assert timeouts, f"no statement timeout was issued: {recorded}"
issued_ms = int(timeouts[0].split("=")[1].strip())
assert issued_ms <= 2000, f"the batch was given {issued_ms}ms with only 2000ms of budget left"
@pytest.mark.asyncio
async def test_no_statement_is_issued_once_the_budget_is_spent():
"""
Every table exits through _finish_table, including the ones a spent run never
started, so an unconditional probe there would put one more statement per
table past the bound.
"""
client = _mock_prisma_for_retention([0, 0])
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
result = await cleaner._finish_table(
client,
datetime.now(timezone.utc) - timedelta(days=7),
"LiteLLM_SpendLogs",
"startTime",
123,
"budget_exhausted",
time.monotonic() - 1,
)
assert result.rows_deleted == 123
assert result.stop_reason == "budget_exhausted"
client.db.query_raw.assert_not_called()
@pytest.mark.asyncio
async def test_a_batch_cancelled_by_the_deadline_is_budget_exhaustion_not_a_failure(monkeypatch):
"""
Clamping the timeout means the last batch of a budget-exhausted run is
cancelled by the deadline itself. Counting that as a batch failure would
inflate the failure metric on every such run and walk it toward the abort
threshold, so it has to be classified as the bound working.
"""
failures: list[str] = []
client = MagicMock()
_wire_tx(client.db)
# The deadline has to pass DURING the batch, not before it: a deadline
# already spent is caught by the loop's own check and no batch is ever
# issued, which would exercise none of the classification under test.
async def _cancelled_after_the_deadline(sql, *args):
await asyncio.sleep(0.05)
raise Exception("canceling statement due to statement timeout")
client.db.execute_raw = _cancelled_after_the_deadline
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
monkeypatch.setattr(SpendLogCleanupMetrics, "record_batch_failure", lambda table: failures.append(table))
result = await cleaner._delete_old_logs(
client, datetime.now(timezone.utc) - timedelta(days=7), time.monotonic() + 0.02
)
assert result.stop_reason == "budget_exhausted"
assert failures == [], f"a deadline cancellation was recorded as a batch failure: {failures}"
@pytest.mark.asyncio
async def test_partition_maintenance_is_skipped_once_the_run_budget_is_spent():
"""
Dropping a partition is DDL holding an ACCESS EXCLUSIVE lock, and unlike a
delete batch it cannot be cut short once it has started. A run whose budget is
already gone must therefore not start it at all; the next tick picks it up.
"""
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_db = MagicMock()
_wire_tx(mock_db)
mock_db.execute_raw = AsyncMock(return_value=0)
mock_prisma_client.db = mock_db
partition_manager = MagicMock()
partition_manager.is_partitioned = AsyncMock(return_value=True)
partition_manager.ensure_partitions = AsyncMock()
partition_manager.drop_partitions_older_than = AsyncMock(return_value=[])
cleaner = SpendLogCleanup(
general_settings={
"maximum_spend_logs_retention_period": "7d",
"use_spend_logs_partitioning": True,
},
partition_manager=partition_manager,
)
cleaner._should_delete_spend_logs()
# a deadline already in the past is what a run that spent its budget on an
# earlier table looks like
await cleaner._clean_spend_log_tables(mock_prisma_client, time.monotonic() - 1)
partition_manager.ensure_partitions.assert_not_awaited()
partition_manager.drop_partitions_older_than.assert_not_awaited()
@pytest.mark.asyncio
async def test_partition_maintenance_still_runs_while_the_run_has_budget():
"""The skip above must be caused by the spent budget, not by breaking the
partition path outright."""
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_db = MagicMock()
_wire_tx(mock_db)
mock_db.execute_raw = AsyncMock(return_value=0)
mock_prisma_client.db = mock_db
partition_manager = MagicMock()
partition_manager.is_partitioned = AsyncMock(return_value=True)
partition_manager.ensure_partitions = AsyncMock()
partition_manager.drop_partitions_older_than = AsyncMock(return_value=["LiteLLM_SpendLogs_p20260601"])
cleaner = SpendLogCleanup(
general_settings={
"maximum_spend_logs_retention_period": "7d",
"use_spend_logs_partitioning": True,
},
partition_manager=partition_manager,
)
cleaner._should_delete_spend_logs()
await cleaner._clean_spend_log_tables(mock_prisma_client, _far_deadline())
partition_manager.ensure_partitions.assert_awaited_once()
partition_manager.drop_partitions_older_than.assert_awaited_once()
@pytest.mark.parametrize(
"stop_reasons, expected",
[
(("exhausted",), "completed"),
(("exhausted", "exhausted"), "completed"),
(("exhausted", "batch_cap_reached"), "batch_cap_reached"),
(("batch_cap_reached", "exhausted"), "batch_cap_reached"),
(("exhausted", "budget_exhausted"), "budget_exhausted"),
(("budget_exhausted", "exhausted"), "budget_exhausted"),
(("batch_cap_reached", "budget_exhausted"), "budget_exhausted"),
(("budget_exhausted", "batch_cap_reached"), "budget_exhausted"),
(("exhausted", "aborted"), "aborted"),
(("aborted", "exhausted"), "aborted"),
(("budget_exhausted", "aborted"), "aborted"),
(("aborted", "budget_exhausted"), "aborted"),
(("aborted", "budget_exhausted", "batch_cap_reached"), "aborted"),
],
)
def test_the_reported_run_outcome_is_the_most_significant_reason_in_any_order(stop_reasons, expected):
"""
The run outcome answers "why did this run stop", so a table that merely ran
dry must never mask one that hit a bound, and an abort must outrank both.
Both orders of every pair are covered because this folds several per-table
results into one answer: a first-match-wins implementation would pass on
whichever order happened to be written and fail on its mirror.
"""
results = tuple(TableCleanupResult(rows_deleted=0, stop_reason=reason) for reason in stop_reasons)
assert SpendLogCleanup._run_outcome(results) == expected

17
tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py
View file

@ -76,6 +76,23 @@ def test_convert_mcp_to_llm_format_defaults_model(proxy_logging, make_mcp_reques
}
def test_convert_mcp_to_llm_format_exposes_headers_on_metadata(proxy_logging, make_mcp_request_obj):
"""Guardrails read the caller's HTTP headers off ``metadata.headers`` on the chat
completions path, so the MCP bridge has to put them in the same place."""
req = make_mcp_request_obj()
out = proxy_logging._convert_mcp_to_llm_format(
request_obj=req,
kwargs={"headers": {"x-nuid": "nuid-1"}},
)
assert out["metadata"]["headers"] == {"x-nuid": "nuid-1"}
def test_convert_mcp_to_llm_format_defaults_headers_to_empty(proxy_logging, make_mcp_request_obj):
req = make_mcp_request_obj()
out = proxy_logging._convert_mcp_to_llm_format(request_obj=req, kwargs={})
assert out["metadata"]["headers"] == {}
def test_convert_mcp_to_llm_format_missing_request_obj_raises(proxy_logging):
with pytest.raises(AttributeError):
proxy_logging._convert_mcp_to_llm_format(request_obj=None, kwargs={})

31
tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py
View file

@ -536,6 +536,37 @@ async def test_get_mcp_tools_from_manager_forwards_request_tags(monkeypatch):
assert mock_get_tools.await_args.kwargs["request_tags"] == ["team-a"]
@pytest.mark.asyncio
async def test_execute_tool_calls_exposes_sanitized_client_headers_to_logging(monkeypatch):
"""The Responses API MCP bridge used to log an empty header dict, hiding the caller's
headers from logging callbacks and hooks."""
_setup_proxy_logging(monkeypatch)
_setup_mcp_call_environment(monkeypatch)
captured = {}
def fake_function_setup(*_args, **kwargs):
captured.update(kwargs)
return None, None
handler_module = importlib.import_module(
"litellm.responses.mcp.litellm_proxy_mcp_handler"
)
monkeypatch.setattr(handler_module, "function_setup", fake_function_setup)
tool_name = "deepwiki-read_wiki_structure"
await LiteLLM_Proxy_MCP_Handler._execute_tool_calls(
tool_server_map={tool_name: "deepwiki"},
tool_calls=[{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}],
user_api_key_auth=None,
raw_headers={"x-nuid": "nuid-1", "x-litellm-api-key": "sk-proxy", "cookie": "s=1"},
)
expected = {"x-nuid": "nuid-1", "cookie": "***REDACTED***"}
assert captured["metadata"]["headers"] == expected
assert captured["proxy_server_request"]["headers"] == expected
@pytest.mark.asyncio
async def test_execute_tool_calls_propagates_request_tags_to_function_setup(monkeypatch):
_setup_proxy_logging(monkeypatch)

42
tests/test_litellm/test_pre_commit_lint.py
View file

@ -224,6 +224,7 @@ def test_nothing_staged_and_no_changes_is_an_explicit_no_op(tmp_path: Path) -> N
proc = _run(repo, bin_dir, {})
assert proc.returncode == 0, proc.stdout + proc.stderr
assert "nothing to check" in proc.stdout
assert "check: PASS" in proc.stdout
assert "linting Python" not in proc.stdout
@ -234,6 +235,7 @@ def test_nothing_staged_without_a_base_ref_fails_with_a_fetch_hint(tmp_path: Pat
assert proc.returncode == 1
assert "cannot resolve the merge base" in proc.stdout
assert "git fetch origin litellm_internal_staging" in proc.stdout
assert "check: FAIL" in proc.stdout
def test_partial_staging_warns_which_checks_were_skipped(tmp_path: Path) -> None:
@ -384,3 +386,43 @@ def test_a_failing_block_fails_the_whole_run(tmp_path: Path, fail: str, message:
proc = _run(repo, bin_dir, {"STUB_FAIL": fail})
assert proc.returncode == 1
assert message in proc.stdout + proc.stderr
def test_run_ends_with_a_summary_of_ran_and_skipped_blocks(tmp_path: Path) -> None:
repo, bin_dir = _sandbox(tmp_path)
proc = _run(repo, bin_dir, {})
assert proc.returncode == 0, proc.stdout + proc.stderr
assert "check: summary" in proc.stdout
assert "ran: Python lint (make lint)" in proc.stdout
assert "ran: dashboard lint (prettier + eslint + lint budgets)" in proc.stdout
assert "ran: dashboard API-type sync (npm run gen:api)" in proc.stdout
assert "skipped: tests/e2e checks (basedpyright + raw HTTP client ban) (no tests/e2e Python files in scope)" in proc.stdout
assert "check: PASS" in proc.stdout
assert "check: FAIL" not in proc.stdout
def test_staged_files_matching_no_check_print_an_explicit_noop_note_and_nonempty_log(tmp_path: Path) -> None:
repo, bin_dir = _sandbox(tmp_path)
_commit_all(repo, "base")
tests_dir = repo / "tests" / "test_litellm"
tests_dir.mkdir(parents=True)
(tests_dir / "test_x.py").write_text("def test_x() -> None: ...\n")
subprocess.run(["git", "add", "tests"], cwd=repo, check=True)
proc = _run(repo, bin_dir, {})
assert proc.returncode == 0, proc.stdout + proc.stderr
assert "no gating lint check matches the files in scope, so nothing ran" in proc.stdout
assert "tests/test_litellm/test_x.py" in proc.stdout
assert "a no-op, not a lint verdict" in proc.stdout
assert "check: PASS" in proc.stdout
assert "linting Python" not in proc.stdout
log = (repo / ".git" / "pre_commit_lint.log").read_text()
assert "check: summary" in log
assert "skipped: Python lint (make lint) (no litellm/ Python files in scope)" in log
def test_failing_run_ends_with_a_fail_verdict(tmp_path: Path) -> None:
repo, bin_dir = _sandbox(tmp_path)
proc = _run(repo, bin_dir, {"STUB_FAIL": "make-lint"})
assert proc.returncode == 1
assert "check: FAIL" in proc.stdout
assert "check: PASS" not in proc.stdout

6
type-discipline-budget.json
View file

@ -1,9 +1,9 @@
{
"LIT001": {
"limit": 22943
"limit": 22941
},
"LIT002": {
"limit": 27141
"limit": 27139
},
"LIT003": {
"limit": 269
@ -27,7 +27,7 @@
"limit": 0
},
"LIT010": {
"limit": 16722
"limit": 16716
},
"LIT011": {
"limit": 5596

173
ui/litellm-dashboard/eslint-suppressions.json
View file

@ -321,33 +321,14 @@
"src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.tsx": {
"no-nested-ternary": {
"count": 5
},
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx": {
"no-restricted-imports": {
"count": 2
}
},
"src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx": {
"no-nested-ternary": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/guardrails/_components/GuardrailTestResults.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx": {
"no-restricted-imports": {
"count": 1
},
"react-hooks/set-state-in-effect": {
"count": 1
}
@ -386,11 +367,6 @@
"count": 2
}
},
"src/app/(dashboard)/guardrails/_components/content_filter/CategoryTable.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/guardrails/_components/content_filter/CompetitorIntentConfiguration.tsx": {
"no-nested-ternary": {
"count": 1
@ -403,14 +379,8 @@
}
},
"src/app/(dashboard)/guardrails/_components/content_filter/ContentCategoryConfiguration.tsx": {
"local/no-complex-jsx-arrow": {
"count": 1
},
"no-nested-ternary": {
"count": 3
},
"no-restricted-imports": {
"count": 1
"count": 2
},
"react-hooks/set-state-in-effect": {
"count": 1
@ -419,62 +389,20 @@
"src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterConfiguration.tsx": {
"local/no-complex-jsx-arrow": {
"count": 3
},
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterDisplay.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterManager.tsx": {
"max-params": {
"count": 2
},
"no-restricted-imports": {
"count": 1
},
"react-hooks/set-state-in-effect": {
"count": 1
}
},
"src/app/(dashboard)/guardrails/_components/content_filter/CustomPatternModal.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/guardrails/_components/content_filter/KeywordModal.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/guardrails/_components/content_filter/KeywordTable.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.tsx": {
"local/no-complex-jsx-arrow": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/guardrails/_components/content_filter/PatternTable.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx": {
"no-nested-ternary": {
"count": 6
},
"no-restricted-imports": {
"count": 2
},
"react-hooks/set-state-in-effect": {
"count": 1
}
@ -487,9 +415,6 @@
"src/app/(dashboard)/guardrails/_components/guardrail_garden.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/guardrails/_components/guardrail_garden_card.tsx": {
@ -500,9 +425,6 @@
"src/app/(dashboard)/guardrails/_components/guardrail_garden_detail.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/guardrails/_components/guardrail_info.tsx": {
@ -568,23 +490,14 @@
"src/app/(dashboard)/guardrails/_components/pii_components.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/guardrails/_components/pii_configuration.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/guardrails/_components/tool_permission/ToolPermissionRulesEditor.tsx": {
"no-restricted-imports": {
"count": 2
},
"react-hooks/purity": {
"count": 1
}
@ -1055,9 +968,6 @@
"no-nested-ternary": {
"count": 2
},
"no-restricted-imports": {
"count": 1
},
"react-hooks/set-state-in-effect": {
"count": 5
}
@ -1098,9 +1008,6 @@
"no-nested-ternary": {
"count": 2
},
"no-restricted-imports": {
"count": 1
},
"react-hooks/immutability": {
"count": 2
},
@ -1115,9 +1022,6 @@
"no-nested-ternary": {
"count": 4
},
"no-restricted-imports": {
"count": 1
},
"react-hooks/set-state-in-effect": {
"count": 1
}
@ -1125,9 +1029,6 @@
"src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.tsx": {
"local/no-complex-jsx-arrow": {
"count": 2
},
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.tsx": {
@ -1135,21 +1036,11 @@
"count": 1
}
},
"src/app/(dashboard)/playground/components/compareUI/components/MessageInput.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/playground/components/compareUI/components/ModelSelector.tsx": {
"no-restricted-imports": {
"count": 2
}
},
"src/app/(dashboard)/playground/components/compareUI/components/UnifiedSelector.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx": {
"local/no-complex-jsx-arrow": {
"count": 2
@ -1603,25 +1494,12 @@
"src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageTable.tsx": {
"no-nested-ternary": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx": {
"local/no-complex-jsx-arrow": {
"count": 2
},
"no-restricted-imports": {
"count": 2
}
},
"src/app/(dashboard)/usage/_components/components/EntityUsage/SpendByProvider.tsx": {
"no-restricted-imports": {
"count": 2
}
},
"src/app/(dashboard)/usage/_components/components/EntityUsage/TopModelView.tsx": {
"no-restricted-imports": {
"count": 1
}
@ -1630,9 +1508,6 @@
"no-nested-ternary": {
"count": 1
},
"no-restricted-imports": {
"count": 1
},
"react-hooks/immutability": {
"count": 1
}
@ -1648,7 +1523,7 @@
"count": 1
},
"no-restricted-imports": {
"count": 2
"count": 1
},
"react-hooks/purity": {
"count": 1
@ -1657,14 +1532,6 @@
"count": 3
}
},
"src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx": {
"local/no-complex-jsx-arrow": {
"count": 2
},
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity.ts": {
"react-hooks/refs": {
"count": 1
@ -1927,29 +1794,14 @@
"count": 1
}
},
"src/components/EntityUsageExport/EntityUsageExportModal.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/EntityUsageExport/ExportFormatSelector.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/EntityUsageExport/ExportSummary.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/EntityUsageExport/ExportTypeSelector.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/EntityUsageExport/UsageExportHeader.tsx": {
"no-restricted-imports": {
"count": 3
"count": 1
}
},
"src/components/EntityUsageExport/types.ts": {
@ -2234,9 +2086,6 @@
},
"no-nested-ternary": {
"count": 1
},
"no-restricted-imports": {
"count": 2
}
},
"src/components/add_model/AdaptiveRoutingConfig.tsx": {
@ -2721,9 +2570,6 @@
"src/components/common_components/team_multi_select.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/components/common_components/user_search_modal.tsx": {
@ -3090,9 +2936,6 @@
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
},
"react-hooks/set-state-in-effect": {
"count": 1
}
@ -3582,6 +3425,14 @@
"count": 1
}
},
"src/components/ui/slider.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"no-nested-ternary": {
"count": 1
}
},
"src/components/ui/switch.tsx": {
"local/filename-pascal-case": {
"count": 1
@ -3625,7 +3476,7 @@
"count": 1
},
"no-restricted-imports": {
"count": 3
"count": 1
},
"react-hooks/set-state-in-effect": {
"count": 1

70
ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.test.tsx
View file

@ -1,4 +1,4 @@
import { render, screen } from "@testing-library/react";
import { render, screen, waitFor } from "@testing-library/react";
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import userEvent from "@testing-library/user-event";
import { beforeEach, describe, expect, it, vi } from "vitest";
@ -14,7 +14,7 @@ vi.mock("./ScoreChart", () => ({
}));
vi.mock("./EvaluationSettingsModal", () => ({
EvaluationSettingsModal: () => null,
EvaluationSettingsModal: ({ open }: { open: boolean }) => (open ? <div>Evaluation settings modal</div> : null),
}));
const mockGetGuardrailsUsageOverview = vi.mocked(networking.getGuardrailsUsageOverview);
@ -28,6 +28,18 @@ function wrapper({ children }: { children: React.ReactNode }) {
return <QueryClientProvider client={queryClient}>{children}</QueryClientProvider>;
}
function renderOverview(onSelectGuardrail = vi.fn()) {
return render(
<GuardrailsOverview
accessToken="test-token"
startDate="2026-08-01"
endDate="2026-08-12"
onSelectGuardrail={onSelectGuardrail}
/>,
{ wrapper },
);
}
describe("GuardrailsOverview", () => {
beforeEach(() => {
vi.clearAllMocks();
@ -92,4 +104,58 @@ describe("GuardrailsOverview", () => {
expect(onSelectGuardrail).toHaveBeenCalledWith("guardrail-low");
});
it("renders the page header and the export action", async () => {
renderOverview();
expect(await screen.findByRole("heading", { name: "Guardrails Monitor", level: 1 })).toBeInTheDocument();
expect(screen.getByText("Monitor guardrail performance across all requests")).toBeInTheDocument();
expect(screen.getByRole("button", { name: /Export Data/i })).toBeInTheDocument();
});
it("renders every summary metric card", async () => {
renderOverview();
expect(await screen.findByText("1,500")).toBeInTheDocument();
expect(screen.getByText("Total Evaluations")).toBeInTheDocument();
expect(screen.getByText("Blocked Requests")).toBeInTheDocument();
expect(screen.getByText("84")).toBeInTheDocument();
expect(screen.getByText("Pass Rate")).toBeInTheDocument();
expect(screen.getByText("94.4%")).toBeInTheDocument();
expect(screen.getByText("23ms")).toBeInTheDocument();
expect(screen.getByText("Active Guardrails")).toBeInTheDocument();
expect(screen.getByText("2")).toBeInTheDocument();
});
it("renders the table toolbar heading and its description", async () => {
renderOverview();
expect(await screen.findByRole("heading", { name: "Guardrail Performance", level: 5 })).toBeInTheDocument();
expect(screen.getByText("Click a guardrail to view details, logs, and configuration")).toBeInTheDocument();
});
it("opens the evaluation settings modal from the toolbar action", async () => {
const user = userEvent.setup();
renderOverview();
expect(screen.queryByText("Evaluation settings modal")).not.toBeInTheDocument();
await user.click(await screen.findByTitle("Evaluation settings"));
expect(await screen.findByText("Evaluation settings modal")).toBeInTheDocument();
});
it("marks the overview busy while the usage request is in flight", async () => {
mockGetGuardrailsUsageOverview.mockReturnValue(new Promise(() => {}));
renderOverview();
await waitFor(() => expect(document.querySelector('[aria-busy="true"]')).toBeInTheDocument());
});
it("shows a failure message when the usage request rejects", async () => {
mockGetGuardrailsUsageOverview.mockRejectedValue(new Error("network down"));
renderOverview();
expect(await screen.findByText("Failed to load data. Try again.")).toBeInTheDocument();
});
});

84
ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.tsx
View file

@ -1,11 +1,12 @@
import { DownloadOutlined, RiseOutlined, SafetyOutlined, SettingOutlined, WarningOutlined } from "@ant-design/icons";
import { useQuery } from "@tanstack/react-query";
import type { ColumnDef, OnChangeFn, SortingState } from "@tanstack/react-table";
import { Button, Col, Row, Spin, Typography } from "antd";
import { Download, Settings, Shield, TrendingUp, TriangleAlert } from "lucide-react";
import React, { useMemo, useState } from "react";
import { DataTable, DataTableSortHeader } from "@/components/shared/DataTable";
import { getGuardrailsUsageOverview } from "@/components/networking";
import { type PerformanceRow } from "@/components/GuardrailsMonitor/mockData";
import { Button } from "@/components/ui/button";
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
import { EvaluationSettingsModal } from "./EvaluationSettingsModal";
import { MetricCard } from "@/components/GuardrailsMonitor/MetricCard";
import { ScoreChart } from "./ScoreChart";
@ -197,51 +198,42 @@ export function GuardrailsOverview({
<div className="flex items-start justify-between mb-5">
<div>
<div className="flex items-center gap-2 mb-1">
<SafetyOutlined className="text-lg text-indigo-500" />
<Shield className="size-5 text-indigo-500" />
<h1 className="text-xl font-semibold text-gray-900">Guardrails Monitor</h1>
</div>
<p className="text-sm text-gray-500">Monitor guardrail performance across all requests</p>
</div>
<div className="flex items-center gap-3">
<Button type="default" icon={<DownloadOutlined />} title="Coming soon">
<Button variant="outline" title="Coming soon">
<Download className="size-4" />
Export Data
</Button>
</div>
</div>
<Row gutter={[16, 16]} className="mb-6">
<Col xs={12} sm={12} md={8} flex="1 0 20%">
<MetricCard label="Total Evaluations" value={metrics.totalRequests.toLocaleString()} />
</Col>
<Col xs={12} sm={12} md={8} flex="1 0 20%">
<MetricCard
label="Blocked Requests"
value={metrics.totalBlocked.toLocaleString()}
valueColor="text-red-600"
icon={<WarningOutlined className="text-red-400" />}
/>
</Col>
<Col xs={12} sm={12} md={8} flex="1 0 20%">
<MetricCard
label="Pass Rate"
value={`${metrics.passRate}%`}
valueColor="text-green-600"
icon={<RiseOutlined className="text-green-400" />}
/>
</Col>
<Col xs={12} sm={12} md={8} flex="1 0 20%">
<MetricCard
label="Avg. latency added"
value={`${metrics.avgLatency}ms`}
valueColor={
metrics.avgLatency > 150 ? "text-red-600" : metrics.avgLatency > 50 ? "text-amber-600" : "text-green-600"
}
/>
</Col>
<Col xs={12} sm={12} md={8} flex="1 0 20%">
<MetricCard label="Active Guardrails" value={metrics.count} />
</Col>
</Row>
<div className="mb-6 grid grid-cols-[repeat(auto-fit,minmax(7rem,1fr))] gap-4">
<MetricCard label="Total Evaluations" value={metrics.totalRequests.toLocaleString()} />
<MetricCard
label="Blocked Requests"
value={metrics.totalBlocked.toLocaleString()}
valueColor="text-red-600"
icon={<TriangleAlert className="size-4 text-red-400" />}
/>
<MetricCard
label="Pass Rate"
value={`${metrics.passRate}%`}
valueColor="text-green-600"
icon={<TrendingUp className="size-4 text-green-400" />}
/>
<MetricCard
label="Avg. latency added"
value={`${metrics.avgLatency}ms`}
valueColor={
metrics.avgLatency > 150 ? "text-red-600" : metrics.avgLatency > 50 ? "text-amber-600" : "text-green-600"
}
/>
<MetricCard label="Active Guardrails" value={metrics.count} />
</div>
<div className="mb-6">
<ScoreChart data={chartData} />
@ -250,7 +242,11 @@ export function GuardrailsOverview({
<div>
{(isLoading || error) && (
<div className="mb-2 flex items-center gap-2">
{isLoading && <Spin size="small" />}
{isLoading && (
<span role="status" aria-busy="true" aria-label="Loading" className="inline-flex">
<UiLoadingSpinner className="size-4 text-primary" />
</span>
)}
{error && <span className="text-sm text-red-600">Failed to load data. Try again.</span>}
</div>
)}
@ -270,20 +266,20 @@ export function GuardrailsOverview({
toolbar={() => (
<div className="flex items-start justify-between gap-4">
<div>
<Typography.Title level={5} className="mb-0! text-gray-900">
Guardrail Performance
</Typography.Title>
<h5 className="mb-0 text-base font-semibold text-gray-900">Guardrail Performance</h5>
<p className="text-xs text-gray-500 mt-0.5">
Click a guardrail to view details, logs, and configuration
</p>
</div>
<div className="flex items-center gap-2">
<Button
type="default"
icon={<SettingOutlined />}
variant="outline"
size="icon"
onClick={() => setEvaluationModalOpen(true)}
title="Evaluation settings"
/>
>
<Settings className="size-4" />
</Button>
</div>
</div>
)}

98
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx
View file

@ -1,13 +1,12 @@
import React, { useState } from "react";
import { Button } from "@tremor/react";
import { Input, Typography, Tooltip } from "antd";
import { CopyOutlined, InfoCircleOutlined } from "@ant-design/icons";
import { Copy, Info } from "lucide-react";
import NotificationsManager from "@/components/molecules/notifications_manager";
import { Button } from "@/components/ui/button";
import { Textarea } from "@/components/ui/textarea";
import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
import GuardrailTestResults from "./GuardrailTestResults";
const { TextArea } = Input;
const { Text } = Typography;
interface GuardrailTestPanelProps {
guardrailNames: string[];
onSubmit: (text: string, metadata?: Record<string, unknown> | null) => void;
@ -108,23 +107,23 @@ export function GuardrailTestPanel({
return (
<div className="space-y-4 h-full flex flex-col">
{/* Header */}
<div className="flex items-center justify-between pb-3 border-b border-gray-200">
<div className="flex items-center justify-between border-b border-border pb-3">
<div className="flex items-center space-x-3">
<div className="flex-1 min-w-0">
<div className="flex items-center space-x-2 mb-1">
<h2 className="text-lg font-semibold text-gray-900">Test Guardrails:</h2>
<div className="mb-1 flex items-center space-x-2">
<h2 className="text-lg font-semibold">Test Guardrails:</h2>
<div className="flex flex-wrap gap-2">
{guardrailNames.map((name) => (
<div
key={name}
className="inline-flex items-center space-x-1 bg-blue-50 px-3 py-1 rounded-md border border-blue-200"
className="inline-flex items-center space-x-1 rounded-md border border-blue-200 bg-blue-50 px-3 py-1"
>
<span className="font-mono text-blue-700 font-medium text-sm">{name}</span>
<span className="font-mono text-sm font-medium text-blue-700">{name}</span>
</div>
))}
</div>
</div>
<p className="text-sm text-gray-500">
<p className="text-sm text-muted-foreground">
Test {guardrailNames.length > 1 ? "guardrails" : "guardrail"} and compare results
</p>
</div>
@ -135,46 +134,63 @@ export function GuardrailTestPanel({
<div className="flex-1 overflow-auto space-y-4">
<div className="space-y-3">
<div>
<div className="flex justify-between items-center mb-2">
<div className="mb-2 flex items-center justify-between">
<div className="flex items-center gap-2">
<label className="text-sm font-medium text-gray-700">Input Text</label>
<Tooltip title="Press Enter to submit. Use Shift+Enter for new line.">
<InfoCircleOutlined className="text-gray-400 cursor-help" />
<label className="text-sm font-medium">Input Text</label>
<Tooltip>
<TooltipTrigger
render={
<span className="cursor-help text-muted-foreground">
<Info className="size-3.5" />
</span>
}
/>
<TooltipContent>Press Enter to submit. Use Shift+Enter for new line.</TooltipContent>
</Tooltip>
</div>
{inputText && (
<Button size="xs" variant="secondary" icon={CopyOutlined} onClick={handleCopyInput}>
<Button size="sm" variant="secondary" onClick={handleCopyInput}>
<Copy />
Copy Input
</Button>
)}
</div>
<TextArea
<Textarea
value={inputText}
onChange={(e) => setInputText(e.target.value)}
onKeyDown={handleKeyDown}
placeholder="Enter text to test with guardrails..."
rows={8}
className="font-mono text-sm"
className="font-mono text-sm field-sizing-fixed"
/>
<div className="flex justify-between items-center mt-1">
<Text className="text-xs text-gray-500">
Press <kbd className="px-1 py-0.5 bg-gray-100 border border-gray-300 rounded-sm text-xs">Enter</kbd> to
submit •{" "}
<kbd className="px-1 py-0.5 bg-gray-100 border border-gray-300 rounded-sm text-xs">Shift+Enter</kbd> for
new line
</Text>
<Text className="text-xs text-gray-500">Characters: {inputText.length}</Text>
<div className="mt-1 flex items-center justify-between">
<span className="text-xs text-muted-foreground">
Press <kbd className="rounded-sm border border-border bg-muted px-1 py-0.5 text-xs">Enter</kbd> to
submit • <kbd className="rounded-sm border border-border bg-muted px-1 py-0.5 text-xs">Shift+Enter</kbd>{" "}
for new line
</span>
<span className="text-xs text-muted-foreground">Characters: {inputText.length}</span>
</div>
</div>
<div>
<div className="flex items-center gap-2 mb-2">
<label className="text-sm font-medium text-gray-700">Metadata (optional)</label>
<Tooltip title="JSON object forwarded to the guardrail as request_data['metadata']. Custom guardrails can read per-request configuration from it.">
<InfoCircleOutlined className="text-gray-400 cursor-help" />
<div className="mb-2 flex items-center gap-2">
<label className="text-sm font-medium">Metadata (optional)</label>
<Tooltip>
<TooltipTrigger
render={
<span className="cursor-help text-muted-foreground">
<Info className="size-3.5" />
</span>
}
/>
<TooltipContent>
JSON object forwarded to the guardrail as request_data[&apos;metadata&apos;]. Custom guardrails can
read per-request configuration from it.
</TooltipContent>
</Tooltip>
</div>
<TextArea
<Textarea
value={metadataText}
onChange={(e) => {
setMetadataText(e.target.value);
@ -184,18 +200,20 @@ export function GuardrailTestPanel({
}}
placeholder='{"forbidden_topics": ["tax", "finance"]}'
rows={3}
className="font-mono text-sm"
status={metadataError ? "error" : undefined}
className="font-mono text-sm field-sizing-fixed"
aria-invalid={metadataError ? true : undefined}
/>
{metadataError && (
<Text type="danger" className="text-xs">
{metadataError}
</Text>
)}
{metadataError && <span className="text-xs text-destructive">{metadataError}</span>}
</div>
<div className="pt-2">
<Button onClick={handleSubmit} loading={isLoading} disabled={!inputText.trim()} className="w-full">
<Button
onClick={handleSubmit}
disabled={!inputText.trim() || isLoading}
aria-busy={isLoading}
className="w-full"
>
{isLoading && <UiLoadingSpinner className="size-4" />}
{isLoading
? `Testing ${guardrailNames.length} guardrail${guardrailNames.length > 1 ? "s" : ""}...`
: `Test ${guardrailNames.length} guardrail${guardrailNames.length > 1 ? "s" : ""}`}

207
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx
View file

@ -1,9 +1,11 @@
import React, { useState } from "react";
import { Card, List, Empty, Spin, Input, Typography } from "antd";
import { ExperimentOutlined, SearchOutlined } from "@ant-design/icons";
import { FlaskConical, Search } from "lucide-react";
import GuardrailTestPanel from "./GuardrailTestPanel";
import { applyGuardrail } from "@/components/networking";
import NotificationsManager from "@/components/molecules/notifications_manager";
import { Card, CardContent } from "@/components/ui/card";
import { InputGroup, InputGroupAddon, InputGroupInput } from "@/components/ui/input-group";
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
interface GuardrailItem {
guardrail_id?: string;
@ -112,115 +114,110 @@ const GuardrailTestPlayground: React.FC<GuardrailTestPlaygroundProps> = ({
return (
<div className="w-full h-[calc(100vh-200px)]">
<Card className="h-full" styles={{ body: { padding: 0, height: "100%" } }}>
<div className="flex h-full">
{/* Left Sidebar - Guardrails List */}
<div className="w-1/4 border-r border-gray-200 flex flex-col overflow-hidden">
<div className="p-4 border-b border-gray-200">
<div className="mb-3">
<h3 className="text-lg font-semibold mb-3">Guardrails</h3>
<Input
prefix={<SearchOutlined />}
placeholder="Search guardrails..."
value={searchQuery}
onChange={(e) => setSearchQuery(e.target.value)}
/>
<Card className="h-full overflow-hidden py-0">
<CardContent className="h-full p-0">
<div className="flex h-full">
{/* Left Sidebar - Guardrails List */}
<div className="flex w-1/4 flex-col overflow-hidden border-r border-border">
<div className="border-b border-border p-4">
<div className="mb-3">
<h3 className="mb-3 text-lg font-semibold">Guardrails</h3>
<InputGroup>
<InputGroupAddon>
<Search className="size-4 text-muted-foreground" />
</InputGroupAddon>
<InputGroupInput
placeholder="Search guardrails..."
value={searchQuery}
onChange={(e) => setSearchQuery(e.target.value)}
/>
</InputGroup>
</div>
</div>
<div className="flex-1 overflow-auto">
{isLoading ? (
<div className="flex h-32 items-center justify-center" aria-busy="true">
<UiLoadingSpinner className="size-6 text-muted-foreground" />
</div>
) : filteredGuardrails.length === 0 ? (
<div className="p-4 text-center text-muted-foreground">
{searchQuery ? "No guardrails match your search" : "No guardrails available"}
</div>
) : (
<ul className="m-0 list-none p-0">
{filteredGuardrails.map((guardrail) => (
<li
key={guardrail.guardrail_id ?? guardrail.guardrail_name}
onClick={() => {
if (guardrail.guardrail_name) {
toggleGuardrailSelection(guardrail.guardrail_name);
}
}}
className={`cursor-pointer border-b border-border py-3 pr-4 pl-6 transition-colors hover:bg-muted/40 ${
selectedGuardrails.has(guardrail.guardrail_name || "")
? "border-l-4 border-l-primary bg-accent"
: "border-l-4 border-l-transparent"
}`}
>
<div className="flex items-center space-x-2">
<FlaskConical className="size-4 text-muted-foreground" />
<span className="font-medium">{guardrail.guardrail_name}</span>
</div>
<div className="mt-1 space-y-1 text-xs">
<div>
<span className="font-medium">Type: </span>
<span className="text-muted-foreground">{guardrail.litellm_params.guardrail}</span>
</div>
<div>
<span className="font-medium">Mode: </span>
<span className="text-muted-foreground">{guardrail.litellm_params.mode}</span>
</div>
</div>
</li>
))}
</ul>
)}
</div>
<div className="border-t border-border bg-muted/40 p-3">
<span className="text-xs text-muted-foreground">
{selectedGuardrails.size} of {filteredGuardrails.length} selected
</span>
</div>
</div>
<div className="flex-1 overflow-auto">
{isLoading ? (
<div className="flex items-center justify-center h-32">
<Spin />
</div>
) : filteredGuardrails.length === 0 ? (
<div className="p-4">
<Empty description={searchQuery ? "No guardrails match your search" : "No guardrails available"} />
</div>
) : (
<List
dataSource={filteredGuardrails}
renderItem={(guardrail) => (
<List.Item
onClick={() => {
if (guardrail.guardrail_name) {
toggleGuardrailSelection(guardrail.guardrail_name);
}
}}
style={{ paddingLeft: 24, paddingRight: 16 }}
className={`cursor-pointer hover:bg-gray-50 transition-colors ${
selectedGuardrails.has(guardrail.guardrail_name || "")
? "bg-blue-50 border-l-4 border-l-blue-500"
: "border-l-4 border-l-transparent"
}`}
>
<List.Item.Meta
title={
<div className="flex items-center space-x-2">
<ExperimentOutlined className="text-gray-400" />
<span className="font-medium text-gray-900">{guardrail.guardrail_name}</span>
</div>
}
description={
<div className="text-xs space-y-1 mt-1">
<div>
<span className="font-medium">Type: </span>
<span className="text-gray-600">{guardrail.litellm_params.guardrail}</span>
</div>
<div>
<span className="font-medium">Mode: </span>
<span className="text-gray-600">{guardrail.litellm_params.mode}</span>
</div>
</div>
}
/>
</List.Item>
)}
/>
)}
</div>
{/* Right Panel - Test Area */}
<div className="flex w-3/4 flex-col">
<div className="flex items-center justify-between border-b border-border p-4">
<h2 className="mb-0 text-xl font-semibold">Guardrail Testing Playground</h2>
</div>
<div className="p-3 border-t border-gray-200 bg-gray-50">
<Typography.Text className="text-xs text-gray-600">
{selectedGuardrails.size} of {filteredGuardrails.length} selected
</Typography.Text>
<div className="flex-1 overflow-auto p-4">
{selectedGuardrails.size === 0 ? (
<div className="flex h-full flex-col items-center justify-center text-muted-foreground">
<FlaskConical className="mb-4 size-12" />
<p className="mb-2 text-lg font-medium">Select Guardrails to Test</p>
<p className="max-w-md text-center">
Choose one or more guardrails from the left sidebar to start testing and comparing results.
</p>
</div>
) : (
<div className="h-full">
<GuardrailTestPanel
guardrailNames={Array.from(selectedGuardrails)}
onSubmit={handleTestGuardrails}
results={testResults.length > 0 ? testResults : null}
errors={testErrors.length > 0 ? testErrors : null}
isLoading={isTesting}
onClose={() => setSelectedGuardrails(new Set())}
/>
</div>
)}
</div>
</div>
</div>
{/* Right Panel - Test Area */}
<div className="w-3/4 flex flex-col bg-white">
<div className="p-4 border-b border-gray-200 flex justify-between items-center">
<Typography.Title level={2} className="text-xl font-semibold mb-0">
Guardrail Testing Playground
</Typography.Title>
</div>
<div className="flex-1 overflow-auto p-4">
{selectedGuardrails.size === 0 ? (
<div className="h-full flex flex-col items-center justify-center text-gray-400">
<ExperimentOutlined style={{ fontSize: "48px", marginBottom: "16px" }} />
<Typography.Paragraph className="text-lg font-medium text-gray-600 mb-2">
Select Guardrails to Test
</Typography.Paragraph>
<Typography.Paragraph className="text-center text-gray-500 max-w-md">
Choose one or more guardrails from the left sidebar to start testing and comparing results.
</Typography.Paragraph>
</div>
) : (
<div className="h-full">
<GuardrailTestPanel
guardrailNames={Array.from(selectedGuardrails)}
onSubmit={handleTestGuardrails}
results={testResults.length > 0 ? testResults : null}
errors={testErrors.length > 0 ? testErrors : null}
isLoading={isTesting}
onClose={() => setSelectedGuardrails(new Set())}
/>
</div>
)}
</div>
</div>
</div>
</CardContent>
</Card>
</div>
);

107
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestResults.tsx
View file

@ -1,7 +1,8 @@
import React, { useState } from "react";
import { Button, Card } from "@tremor/react";
import { CopyOutlined, CheckCircleOutlined, ClockCircleOutlined, DownOutlined, RightOutlined } from "@ant-design/icons";
import { Check, ChevronDown, ChevronRight, Clock, Copy } from "lucide-react";
import NotificationsManager from "@/components/molecules/notifications_manager";
import { Button } from "@/components/ui/button";
import { Card, CardContent } from "@/components/ui/card";
interface TestResult {
guardrailName: string;
@ -66,39 +67,38 @@ export function GuardrailTestResults({ results, errors }: GuardrailTestResultsPr
}
return (
<div className="space-y-3 pt-4 border-t border-gray-200">
<h3 className="text-sm font-semibold text-gray-900">Results</h3>
<div className="space-y-3 border-t border-border pt-4">
<h3 className="text-sm font-semibold">Results</h3>
{/* Success Results */}
{results &&
results.map((result) => {
const isCollapsed = collapsedResults.has(result.guardrailName);
return (
<Card key={result.guardrailName} className="bg-green-50 border-green-200">
<div className="space-y-3">
<Card key={result.guardrailName} className="border-emerald-200 bg-emerald-50">
<CardContent className="space-y-3">
<div className="flex items-center justify-between">
<div
className="flex items-center space-x-2 cursor-pointer flex-1"
className="flex flex-1 cursor-pointer items-center space-x-2"
onClick={() => toggleResultCollapse(result.guardrailName)}
>
{isCollapsed ? (
<RightOutlined className="text-gray-500 text-xs" />
<ChevronRight className="size-3 text-muted-foreground" />
) : (
<DownOutlined className="text-gray-500 text-xs" />
<ChevronDown className="size-3 text-muted-foreground" />
)}
<CheckCircleOutlined className="text-green-600 text-lg" />
<span className="text-sm font-medium text-green-800">{result.guardrailName}</span>
<Check className="size-4 text-emerald-600" />
<span className="text-sm font-medium text-emerald-800">{result.guardrailName}</span>
</div>
<div className="flex items-center gap-3">
<div className="flex items-center space-x-1 text-xs text-gray-600">
<ClockCircleOutlined />
<div className="flex items-center space-x-1 text-xs text-muted-foreground">
<Clock className="size-3" />
<span className="font-medium">{result.latency}ms</span>
</div>
{!isCollapsed && (
<Button
size="xs"
size="sm"
variant="secondary"
icon={CopyOutlined}
onClick={async () => {
const success = await copyToClipboard(result.response_text);
if (success) {
@ -108,6 +108,7 @@ export function GuardrailTestResults({ results, errors }: GuardrailTestResultsPr
}
}}
>
<Copy />
Copy
</Button>
)}
@ -115,18 +116,18 @@ export function GuardrailTestResults({ results, errors }: GuardrailTestResultsPr
</div>
{!isCollapsed && (
<>
<div className="bg-white border border-green-200 rounded-sm p-3">
<label className="text-xs font-medium text-gray-600 mb-2 block">Output Text</label>
<div className="font-mono text-sm text-gray-900 whitespace-pre-wrap wrap-break-word">
<div className="rounded-sm border border-emerald-200 bg-background p-3">
<label className="mb-2 block text-xs font-medium text-muted-foreground">Output Text</label>
<div className="font-mono text-sm whitespace-pre-wrap wrap-break-word">
{result.response_text}
</div>
</div>
<div className="text-xs text-gray-600">
<div className="text-xs text-muted-foreground">
<span className="font-medium">Characters:</span> {result.response_text.length}
</div>
</>
)}
</div>
</CardContent>
</Card>
);
})}
@ -136,40 +137,42 @@ export function GuardrailTestResults({ results, errors }: GuardrailTestResultsPr
errors.map((errorItem) => {
const isCollapsed = collapsedResults.has(errorItem.guardrailName);
return (
<Card key={errorItem.guardrailName} className="bg-red-50 border-red-200">
<div className="flex items-start space-x-2">
<div className="cursor-pointer mt-0.5" onClick={() => toggleResultCollapse(errorItem.guardrailName)}>
{isCollapsed ? (
<RightOutlined className="text-gray-500 text-xs" />
) : (
<DownOutlined className="text-gray-500 text-xs" />
)}
</div>
<div className="text-red-600 mt-0.5">
<svg className="w-5 h-5" fill="currentColor" viewBox="0 0 20 20">
<path
fillRule="evenodd"
d="M10 18a8 8 0 100-16 8 8 0 000 16zM8.707 7.293a1 1 0 00-1.414 1.414L8.586 10l-1.293 1.293a1 1 0 101.414 1.414L10 11.414l1.293 1.293a1 1 0 001.414-1.414L11.414 10l1.293-1.293a1 1 0 00-1.414-1.414L10 8.586 8.707 7.293z"
clipRule="evenodd"
/>
</svg>
</div>
<div className="flex-1">
<div className="flex items-center justify-between mb-1">
<p
className="text-sm font-medium text-red-800 cursor-pointer"
onClick={() => toggleResultCollapse(errorItem.guardrailName)}
>
{errorItem.guardrailName} - Error
</p>
<div className="flex items-center space-x-1 text-xs text-gray-600">
<ClockCircleOutlined />
<span className="font-medium">{errorItem.latency}ms</span>
</div>
<Card key={errorItem.guardrailName} className="border-red-200 bg-red-50">
<CardContent>
<div className="flex items-start space-x-2">
<div className="mt-0.5 cursor-pointer" onClick={() => toggleResultCollapse(errorItem.guardrailName)}>
{isCollapsed ? (
<ChevronRight className="size-3 text-muted-foreground" />
) : (
<ChevronDown className="size-3 text-muted-foreground" />
)}
</div>
<div className="mt-0.5 text-destructive">
<svg className="w-5 h-5" fill="currentColor" viewBox="0 0 20 20">
<path
fillRule="evenodd"
d="M10 18a8 8 0 100-16 8 8 0 000 16zM8.707 7.293a1 1 0 00-1.414 1.414L8.586 10l-1.293 1.293a1 1 0 101.414 1.414L10 11.414l1.293 1.293a1 1 0 001.414-1.414L11.414 10l1.293-1.293a1 1 0 00-1.414-1.414L10 8.586 8.707 7.293z"
clipRule="evenodd"
/>
</svg>
</div>
<div className="flex-1">
<div className="mb-1 flex items-center justify-between">
<p
className="cursor-pointer text-sm font-medium text-red-800"
onClick={() => toggleResultCollapse(errorItem.guardrailName)}
>
{errorItem.guardrailName} - Error
</p>
<div className="flex items-center space-x-1 text-xs text-muted-foreground">
<Clock className="size-3" />
<span className="font-medium">{errorItem.latency}ms</span>
</div>
</div>
{!isCollapsed && <p className="mt-1 text-sm text-red-700">{errorItem.error.message}</p>}
</div>
{!isCollapsed && <p className="text-sm text-red-700 mt-1">{errorItem.error.message}</p>}
</div>
</div>
</CardContent>
</Card>
);
})}

39
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx
View file

@ -35,10 +35,19 @@ vi.mock("./guardrail_info", () => ({
default: () => <div>Mock Guardrail Info View</div>,
}));
vi.mock("./GuardrailTestPlayground", () => ({
__esModule: true,
default: () => <div>Mock Guardrail Test Playground</div>,
}));
vi.mock("./GuardrailTestPlayground", async () => {
const { useState } = await import("react");
const MockGuardrailTestPlayground = () => {
const [draft, setDraft] = useState("");
return (
<div>
<div>Mock Guardrail Test Playground</div>
<input aria-label="playground draft" value={draft} onChange={(e) => setDraft(e.target.value)} />
</div>
);
};
return { __esModule: true, default: MockGuardrailTestPlayground };
});
vi.mock("./TeamGuardrailsTab", () => ({
TeamGuardrailsTab: () => <div>Mock Team Guardrails Tab</div>,
@ -129,6 +138,28 @@ describe("GuardrailsPanel", () => {
expect(mockGetGuardrailsList).toHaveBeenCalledTimes(2);
});
it("should mount every tab panel up front so panel state survives tab switches", async () => {
render(<GuardrailsPanel {...defaultProps} />);
expect(await screen.findByLabelText("playground draft")).toBeInTheDocument();
expect(screen.getByText("Mock Team Guardrails Tab")).toBeInTheDocument();
});
it("should keep test playground state when switching tabs away and back", async () => {
render(<GuardrailsPanel {...defaultProps} />);
fireEvent.click(screen.getByText("Test Playground"));
const draft = await screen.findByLabelText("playground draft");
fireEvent.change(draft, { target: { value: "keep me" } });
expect(draft).toHaveValue("keep me");
fireEvent.click(screen.getByText("Guardrails"));
fireEvent.click(screen.getByText("Test Playground"));
expect(await screen.findByLabelText("playground draft")).toHaveValue("keep me");
});
it("should not delete anything when the modal is cancelled", async () => {
render(<GuardrailsPanel {...defaultProps} />);
fireEvent.click(screen.getByText("Guardrails"));

219
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx
View file

@ -1,5 +1,5 @@
import React, { useState, useEffect } from "react";
import { Tabs } from "antd";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { ChevronDown, Code, Plus } from "lucide-react";
import { getGuardrailsList, deleteGuardrailCall } from "@/components/networking";
import { buttonVariants } from "@/components/ui/button";
@ -125,118 +125,119 @@ const GuardrailsPanel: React.FC<GuardrailsPanelProps> = ({ accessToken, userRole
return (
<div className="w-full mx-auto flex-auto overflow-y-auto m-8 p-2">
<Tabs
defaultActiveKey="guardrails"
items={[
...(isAdmin
? [
{
key: "garden",
label: "Guardrail Garden",
children: <GuardrailGarden accessToken={accessToken} onGuardrailCreated={handleSuccess} />,
},
{
key: "guardrails",
label: "Guardrails",
children: (
<>
<div className="flex justify-between items-center mb-4">
<DropdownMenu>
<DropdownMenuTrigger
disabled={!accessToken}
className={cn(buttonVariants({ variant: "default" }))}
>
<Plus />
Add New Guardrail
<ChevronDown />
</DropdownMenuTrigger>
<DropdownMenuContent align="start" className="w-56">
<DropdownMenuItem onClick={handleAddGuardrail}>
<Plus />
Add Provider Guardrail
</DropdownMenuItem>
<DropdownMenuItem onClick={handleAddCustomCodeGuardrail}>
<Code />
Create Custom Code Guardrail
</DropdownMenuItem>
</DropdownMenuContent>
</DropdownMenu>
</div>
<Tabs defaultValue="guardrails">
<TabsList>
{isAdmin && (
<>
<TabsTrigger value="garden" className="flex-none">
Guardrail Garden
</TabsTrigger>
<TabsTrigger value="guardrails" className="flex-none">
Guardrails
</TabsTrigger>
<TabsTrigger value="playground" className="flex-none" disabled={!accessToken}>
Test Playground
</TabsTrigger>
</>
)}
<TabsTrigger value="submitted" className="flex-none">
Submitted Guardrails
</TabsTrigger>
</TabsList>
{selectedGuardrailId ? (
<GuardrailInfoView
guardrailId={selectedGuardrailId}
onClose={() => setSelectedGuardrailId(null)}
accessToken={accessToken}
isAdmin={isAdmin}
/>
) : (
<GuardrailTable
guardrailsList={guardrailsList}
isLoading={isLoading}
onDeleteClick={handleDeleteClick}
onGuardrailClick={(id) => setSelectedGuardrailId(id)}
/>
)}
{isAdmin && (
<>
<TabsContent value="garden" keepMounted>
<GuardrailGarden accessToken={accessToken} onGuardrailCreated={handleSuccess} />
</TabsContent>
<AddGuardrailForm
visible={isAddModalVisible}
onClose={handleCloseModal}
accessToken={accessToken}
onSuccess={handleSuccess}
/>
<TabsContent value="guardrails" keepMounted>
<div className="flex justify-between items-center mb-4">
<DropdownMenu>
<DropdownMenuTrigger disabled={!accessToken} className={cn(buttonVariants({ variant: "default" }))}>
<Plus />
Add New Guardrail
<ChevronDown />
</DropdownMenuTrigger>
<DropdownMenuContent align="start" className="w-56">
<DropdownMenuItem onClick={handleAddGuardrail}>
<Plus />
Add Provider Guardrail
</DropdownMenuItem>
<DropdownMenuItem onClick={handleAddCustomCodeGuardrail}>
<Code />
Create Custom Code Guardrail
</DropdownMenuItem>
</DropdownMenuContent>
</DropdownMenu>
</div>
<CustomCodeModal
visible={isCustomCodeModalVisible}
onClose={handleCloseCustomCodeModal}
accessToken={accessToken}
onSuccess={handleSuccess}
/>
{selectedGuardrailId ? (
<GuardrailInfoView
guardrailId={selectedGuardrailId}
onClose={() => setSelectedGuardrailId(null)}
accessToken={accessToken}
isAdmin={isAdmin}
/>
) : (
<GuardrailTable
guardrailsList={guardrailsList}
isLoading={isLoading}
onDeleteClick={handleDeleteClick}
onGuardrailClick={(id) => setSelectedGuardrailId(id)}
/>
)}
<DeleteResourceModal
isOpen={isDeleteModalOpen}
title="Delete Guardrail"
message={`Are you sure you want to delete guardrail: ${guardrailToDelete?.guardrail_name}? This action cannot be undone.`}
resourceInformationTitle="Guardrail Information"
resourceInformation={[
{ label: "Name", value: guardrailToDelete?.guardrail_name },
{ label: "ID", value: guardrailToDelete?.guardrail_id, code: true },
{ label: "Provider", value: providerDisplayName },
{ label: "Mode", value: guardrailToDelete?.litellm_params.mode },
{
label: "Default On",
value: guardrailToDelete?.litellm_params.default_on ? "Yes" : "No",
},
]}
onCancel={handleDeleteCancel}
onOk={handleDeleteConfirm}
confirmLoading={isDeleting}
/>
</>
),
},
{
key: "playground",
label: "Test Playground",
disabled: !accessToken,
children: (
<GuardrailTestPlayground
guardrailsList={guardrailsList}
isLoading={isLoading}
accessToken={accessToken}
onClose={() => {}}
/>
),
},
]
: []),
{
key: "submitted",
label: "Submitted Guardrails",
children: <TeamGuardrailsTab accessToken={accessToken} />,
},
]}
/>
<AddGuardrailForm
visible={isAddModalVisible}
onClose={handleCloseModal}
accessToken={accessToken}
onSuccess={handleSuccess}
/>
<CustomCodeModal
visible={isCustomCodeModalVisible}
onClose={handleCloseCustomCodeModal}
accessToken={accessToken}
onSuccess={handleSuccess}
/>
<DeleteResourceModal
isOpen={isDeleteModalOpen}
title="Delete Guardrail"
message={`Are you sure you want to delete guardrail: ${guardrailToDelete?.guardrail_name}? This action cannot be undone.`}
resourceInformationTitle="Guardrail Information"
resourceInformation={[
{ label: "Name", value: guardrailToDelete?.guardrail_name },
{ label: "ID", value: guardrailToDelete?.guardrail_id, code: true },
{ label: "Provider", value: providerDisplayName },
{ label: "Mode", value: guardrailToDelete?.litellm_params.mode },
{
label: "Default On",
value: guardrailToDelete?.litellm_params.default_on ? "Yes" : "No",
},
]}
onCancel={handleDeleteCancel}
onOk={handleDeleteConfirm}
confirmLoading={isDeleting}
/>
</TabsContent>
<TabsContent value="playground" keepMounted>
<GuardrailTestPlayground
guardrailsList={guardrailsList}
isLoading={isLoading}
accessToken={accessToken}
onClose={() => {}}
/>
</TabsContent>
</>
)}
<TabsContent value="submitted" keepMounted>
<TeamGuardrailsTab accessToken={accessToken} />
</TabsContent>
</Tabs>
</div>
);
};

71
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/CategoryTable.tsx
View file

@ -1,11 +1,11 @@
import React from "react";
import { Typography, Select, Tag, Button } from "antd";
import { DeleteOutlined } from "@ant-design/icons";
import { Trash2 } from "lucide-react";
import type { ColumnDef } from "@tanstack/react-table";
import { DataTable } from "@/components/shared/DataTable";
const { Text } = Typography;
const { Option } = Select;
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { ACTION_ITEMS, SEVERITY_ITEMS } from "./action_options";
interface ContentCategory {
id: string;
@ -38,14 +38,8 @@ const CategoryTable: React.FC<CategoryTableProps> = ({
const { category, display_name: displayName } = row.original;
return (
<div>
<Text strong>{displayName}</Text>
{displayName !== category && (
<div>
<Text type="secondary" style={{ fontSize: 12 }}>
{category}
</Text>
</div>
)}
<span className="font-semibold">{displayName}</span>
{displayName !== category && <div className="text-xs text-muted-foreground">{category}</div>}
</div>
);
},
@ -57,23 +51,26 @@ const CategoryTable: React.FC<CategoryTableProps> = ({
cell: ({ row }) => {
const { id, severity_threshold: severity } = row.original;
if (readOnly) {
const colorMap = {
high: "red",
medium: "orange",
low: "yellow",
} as const;
return <Tag color={colorMap[severity as keyof typeof colorMap]}>{severity.toUpperCase()}</Tag>;
return <Badge variant={severity === "high" ? "destructive" : "secondary"}>{severity.toUpperCase()}</Badge>;
}
return (
<Select
items={SEVERITY_ITEMS}
value={severity}
onChange={(value) => onSeverityChange?.(id, value as "high" | "medium" | "low")}
style={{ width: 150 }}
size="small"
onValueChange={(value: string | null) =>
value && onSeverityChange?.(id, value as "high" | "medium" | "low")
}
>
<Option value="high">High</Option>
<Option value="medium">Medium</Option>
<Option value="low">Low</Option>
<SelectTrigger size="sm" className="w-[150px]" aria-label="Severity Threshold">
<SelectValue />
</SelectTrigger>
<SelectContent alignItemWithTrigger={false}>
{SEVERITY_ITEMS.map((item) => (
<SelectItem key={item.value} value={item.value}>
{item.label}
</SelectItem>
))}
</SelectContent>
</Select>
);
},
@ -85,17 +82,24 @@ const CategoryTable: React.FC<CategoryTableProps> = ({
cell: ({ row }) => {
const { action, id } = row.original;
if (readOnly) {
return <Tag color={action === "BLOCK" ? "red" : "blue"}>{action}</Tag>;
return <Badge variant={action === "BLOCK" ? "destructive" : "secondary"}>{action}</Badge>;
}
return (
<Select
items={ACTION_ITEMS}
value={action}
onChange={(value) => onActionChange?.(id, value as "BLOCK" | "MASK")}
style={{ width: 120 }}
size="small"
onValueChange={(value: string | null) => value && onActionChange?.(id, value as "BLOCK" | "MASK")}
>
<Option value="BLOCK">Block</Option>
<Option value="MASK">Mask</Option>
<SelectTrigger size="sm" className="w-[120px]" aria-label="Action">
<SelectValue />
</SelectTrigger>
<SelectContent alignItemWithTrigger={false}>
{ACTION_ITEMS.map((item) => (
<SelectItem key={item.value} value={item.value}>
{item.label}
</SelectItem>
))}
</SelectContent>
</Select>
);
},
@ -108,7 +112,8 @@ const CategoryTable: React.FC<CategoryTableProps> = ({
id: "actions",
size: 100,
cell: ({ row }) => (
<Button type="text" danger size="small" icon={<DeleteOutlined />} onClick={() => onRemove?.(row.original.id)}>
<Button variant="ghost" size="sm" onClick={() => onRemove?.(row.original.id)}>
<Trash2 />
Delete
</Button>
),
@ -116,7 +121,7 @@ const CategoryTable: React.FC<CategoryTableProps> = ({
}
if (categories.length === 0) {
return <div style={{ textAlign: "center", padding: "40px 0", color: "#999" }}>No categories configured.</div>;
return <div className="py-10 text-center text-muted-foreground">No categories configured.</div>;
}
return <DataTable data={categories} columns={columns} getRowId={(row) => row.id} size="compact" />;

320
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentCategoryConfiguration.tsx
View file

@ -1,12 +1,22 @@
import React from "react";
import { Card, Typography, Select, Tag, Collapse, Button } from "antd";
import { DeleteOutlined, PlusOutlined, FileTextOutlined } from "@ant-design/icons";
import { ChevronRight, FileText, Plus, Trash2 } from "lucide-react";
import type { ColumnDef } from "@tanstack/react-table";
import { getCategoryYaml } from "@/components/networking";
import { DataTable } from "@/components/shared/DataTable";
const { Title, Text } = Typography;
const { Option } = Select;
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
import {
Combobox,
ComboboxContent,
ComboboxEmpty,
ComboboxInput,
ComboboxItem,
ComboboxList,
} from "@/components/ui/combobox";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { ACTION_ITEMS, SEVERITY_ITEMS } from "./action_options";
interface ContentCategory {
name: string;
@ -170,10 +180,8 @@ const ContentCategoryConfiguration: React.FC<ContentCategoryConfigurationProps>
const category = availableCategories.find((c) => c.name === row.original.category);
return (
<div>
<div style={{ fontWeight: 500 }}>{row.original.display_name}</div>
{category?.description && (
<div style={{ fontSize: "12px", color: "#888", marginTop: "4px" }}>{category.description}</div>
)}
<div className="font-medium">{row.original.display_name}</div>
{category?.description && <div className="mt-1 text-xs text-muted-foreground">{category.description}</div>}
</div>
);
},
@ -184,16 +192,20 @@ const ContentCategoryConfiguration: React.FC<ContentCategoryConfigurationProps>
size: 150,
cell: ({ row }) => (
<Select
items={ACTION_ITEMS}
value={row.original.action}
onChange={(value) => onCategoryUpdate(row.original.id, "action", value)}
style={{ width: "100%" }}
onValueChange={(value: string | null) => value && onCategoryUpdate(row.original.id, "action", value)}
>
<Option value="BLOCK">
<Tag color="red">BLOCK</Tag>
</Option>
<Option value="MASK">
<Tag color="orange">MASK</Tag>
</Option>
<SelectTrigger size="sm" className="w-full" aria-label="Action">
<SelectValue />
</SelectTrigger>
<SelectContent alignItemWithTrigger={false}>
{ACTION_ITEMS.map((item) => (
<SelectItem key={item.value} value={item.value}>
<Badge variant={item.value === "BLOCK" ? "destructive" : "secondary"}>{item.value}</Badge>
</SelectItem>
))}
</SelectContent>
</Select>
),
},
@ -203,13 +215,22 @@ const ContentCategoryConfiguration: React.FC<ContentCategoryConfigurationProps>
size: 180,
cell: ({ row }) => (
<Select
items={SEVERITY_ITEMS}
value={row.original.severity_threshold}
onChange={(value) => onCategoryUpdate(row.original.id, "severity_threshold", value)}
style={{ width: "100%" }}
onValueChange={(value: string | null) =>
value && onCategoryUpdate(row.original.id, "severity_threshold", value)
}
>
<Option value="low">Low</Option>
<Option value="medium">Medium</Option>
<Option value="high">High</Option>
<SelectTrigger size="sm" className="w-full" aria-label="Severity Threshold">
<SelectValue />
</SelectTrigger>
<SelectContent alignItemWithTrigger={false}>
{SEVERITY_ITEMS.map((item) => (
<SelectItem key={item.value} value={item.value}>
{item.label}
</SelectItem>
))}
</SelectContent>
</Select>
),
},
@ -218,7 +239,8 @@ const ContentCategoryConfiguration: React.FC<ContentCategoryConfigurationProps>
id: "actions",
size: 80,
cell: ({ row }) => (
<Button icon={<DeleteOutlined />} onClick={() => onCategoryRemove(row.original.id)} size="small">
<Button variant="outline" size="sm" onClick={() => onCategoryRemove(row.original.id)}>
<Trash2 />
Remove
</Button>
),
@ -228,172 +250,120 @@ const ContentCategoryConfiguration: React.FC<ContentCategoryConfigurationProps>
const unselectedCategories = availableCategories.filter(
(cat) => !selectedCategories.some((sel) => sel.category === cat.name),
);
const pendingCategory = availableCategories.find((c) => c.name === selectedCategoryName) ?? null;
return (
<Card
title={
<div
style={{ display: "flex", justifyContent: "space-between", alignItems: "center", flexWrap: "wrap", gap: 8 }}
>
<Title level={5} style={{ margin: 0 }}>
Blocked topics
</Title>
<Text type="secondary" style={{ fontSize: 12, fontWeight: 400 }}>
<Card>
<CardHeader>
<div className="flex flex-wrap items-center justify-between gap-2">
<CardTitle>Blocked topics</CardTitle>
<p className="text-xs font-normal text-muted-foreground">
Select topics to block using keyword and semantic analysis
</Text>
</p>
</div>
</CardHeader>
<CardContent>
<div className="mb-4 flex gap-2">
<Combobox
items={unselectedCategories}
value={pendingCategory}
onValueChange={(category: ContentCategory | null) => setSelectedCategoryName(category?.name ?? "")}
itemToStringLabel={(category: ContentCategory) => category.display_name}
>
<ComboboxInput className="w-full" placeholder="Select a content category" />
<ComboboxContent>
<ComboboxEmpty>No matching categories</ComboboxEmpty>
<ComboboxList>
{(cat: ContentCategory) => (
<ComboboxItem key={cat.name} value={cat}>
<div>
<div className="font-medium">{cat.display_name}</div>
<div className="mt-0.5 text-xs text-muted-foreground">{cat.description}</div>
</div>
</ComboboxItem>
)}
</ComboboxList>
</ComboboxContent>
</Combobox>
<Button onClick={handleAddCategory} disabled={!selectedCategoryName}>
<Plus />
Add
</Button>
</div>
}
size="small"
>
<div style={{ marginBottom: 16, display: "flex", gap: 8 }}>
<Select
placeholder="Select a content category"
value={selectedCategoryName || undefined}
onChange={setSelectedCategoryName}
style={{ flex: 1 }}
showSearch
optionLabelProp="label"
filterOption={(input, option) =>
(option?.label?.toString().toLowerCase() ?? "").includes(input.toLowerCase())
}
>
{unselectedCategories.map((cat) => (
<Option key={cat.name} value={cat.name} label={cat.display_name}>
<div>
<div style={{ fontWeight: 500 }}>{cat.display_name}</div>
<div style={{ fontSize: "12px", color: "#666", marginTop: "2px" }}>{cat.description}</div>
</div>
</Option>
))}
</Select>
<Button type="primary" onClick={handleAddCategory} disabled={!selectedCategoryName} icon={<PlusOutlined />}>
Add
</Button>
</div>
{/* Preview box - shown when category is selected but not yet added */}
{selectedCategoryName && (
<div
style={{
marginBottom: 16,
padding: "12px",
background: "#f9f9f9",
border: "1px solid #e0e0e0",
borderRadius: "4px",
}}
>
<div style={{ marginBottom: 8, fontWeight: 500, fontSize: "14px" }}>
Preview: {availableCategories.find((c) => c.name === selectedCategoryName)?.display_name}
{categoryFileTypes[selectedCategoryName] && (
<span style={{ marginLeft: 8, fontSize: "12px", color: "#888", fontWeight: 400 }}>
({categoryFileTypes[selectedCategoryName]?.toUpperCase()})
</span>
{/* Preview box - shown when category is selected but not yet added */}
{selectedCategoryName && (
<div className="mb-4 rounded-md border border-border bg-muted/40 p-3">
<div className="mb-2 text-sm font-medium">
Preview: {availableCategories.find((c) => c.name === selectedCategoryName)?.display_name}
{categoryFileTypes[selectedCategoryName] && (
<span className="ml-2 text-xs font-normal text-muted-foreground">
({categoryFileTypes[selectedCategoryName]?.toUpperCase()})
</span>
)}
</div>
{loadingPreviewYaml ? (
<div className="p-4 text-center text-muted-foreground">Loading content...</div>
) : previewYaml ? (
<pre className="m-0 max-h-[300px] max-w-full overflow-auto rounded-md border border-border bg-background p-3 text-xs leading-relaxed break-words whitespace-pre-wrap">
<code>{previewYaml}</code>
</pre>
) : (
<div className="p-2 text-center text-xs text-muted-foreground">Unable to load category content</div>
)}
</div>
{loadingPreviewYaml ? (
<div style={{ padding: "16px", textAlign: "center", color: "#888" }}>Loading content...</div>
) : previewYaml ? (
<pre
style={{
background: "#fff",
padding: "12px",
borderRadius: "4px",
overflow: "auto",
maxHeight: "300px",
maxWidth: "100%",
fontSize: "12px",
lineHeight: "1.5",
margin: 0,
border: "1px solid #e0e0e0",
whiteSpace: "pre-wrap",
wordBreak: "break-word",
}}
>
<code>{previewYaml}</code>
</pre>
) : (
<div style={{ padding: "8px", textAlign: "center", color: "#888", fontSize: "12px" }}>
Unable to load category content
</div>
)}
</div>
)}
)}
{selectedCategories.length > 0 ? (
<>
<DataTable data={selectedCategories} columns={columns} getRowId={(row) => row.id} size="compact" />
<div style={{ marginTop: 16 }}>
<Collapse
activeKey={expandedYamlCategories}
onChange={(keys) => {
const keyArray = Array.isArray(keys) ? keys : keys ? [keys] : [];
const oldExpanded = new Set(expandedYamlCategories);
// Find newly expanded categories and fetch their YAML
keyArray.forEach((key) => {
const categoryName = key as string;
if (!oldExpanded.has(categoryName) && !categoryYaml[categoryName]) {
fetchCategoryYaml(categoryName);
}
});
setExpandedYamlCategories(keyArray as string[]);
}}
ghost
items={selectedCategories.map((category) => {
{selectedCategories.length > 0 ? (
<>
<DataTable data={selectedCategories} columns={columns} getRowId={(row) => row.id} size="compact" />
<div className="mt-4 space-y-2">
{selectedCategories.map((category) => {
const fileType = categoryFileTypes[category.category] || "yaml";
const fileTypeLabel = fileType.toUpperCase();
const isExpanded = expandedYamlCategories.includes(category.category);
return {
key: category.category,
label: (
<div style={{ display: "flex", alignItems: "center", gap: 8 }}>
<FileTextOutlined />
return (
<Collapsible
key={category.category}
open={isExpanded}
onOpenChange={(open) => {
if (open && !categoryYaml[category.category]) {
fetchCategoryYaml(category.category);
}
setExpandedYamlCategories((prev) =>
open ? [...prev, category.category] : prev.filter((name) => name !== category.category),
);
}}
>
<CollapsibleTrigger className="flex items-center gap-2 text-sm">
<ChevronRight className={`size-4 transition-transform ${isExpanded ? "rotate-90" : ""}`} />
<FileText className="size-4" />
<span>
View {fileTypeLabel} for {category.display_name}
View {fileType.toUpperCase()} for {category.display_name}
</span>
</div>
),
children: loadingYaml[category.category] ? (
<div style={{ padding: "16px", textAlign: "center", color: "#888" }}>Loading content...</div>
) : categoryYaml[category.category] ? (
<pre
style={{
background: "#f5f5f5",
padding: "16px",
borderRadius: "4px",
overflow: "auto",
maxHeight: "400px",
fontSize: "12px",
lineHeight: "1.5",
margin: 0,
}}
>
<code>{categoryYaml[category.category]}</code>
</pre>
) : (
<div style={{ padding: "16px", textAlign: "center", color: "#888" }}>
Content will load when expanded
</div>
),
};
</CollapsibleTrigger>
<CollapsibleContent>
{loadingYaml[category.category] ? (
<div className="p-4 text-center text-muted-foreground">Loading content...</div>
) : categoryYaml[category.category] ? (
<pre className="m-0 max-h-[400px] overflow-auto rounded-md bg-muted p-4 text-xs leading-relaxed">
<code>{categoryYaml[category.category]}</code>
</pre>
) : (
<div className="p-4 text-center text-muted-foreground">Content will load when expanded</div>
)}
</CollapsibleContent>
</Collapsible>
);
})}
/>
</div>
</>
) : (
<div className="rounded-md border border-dashed border-border p-6 text-center text-muted-foreground">
No blocked topics selected. Add topics to detect and block harmful content.
</div>
</>
) : (
<div
style={{
textAlign: "center",
padding: "24px",
color: "#888",
border: "1px dashed #d9d9d9",
borderRadius: "4px",
}}
>
No blocked topics selected. Add topics to detect and block harmful content.
</div>
)}
)}
</CardContent>
</Card>
);
};

136
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterConfiguration.test.tsx Normal file
View file

@ -0,0 +1,136 @@
import { describe, it, expect, vi, beforeEach } from "vitest";
import { renderWithProviders, screen } from "@/../tests/test-utils";
import userEvent from "@testing-library/user-event";
import ContentFilterConfiguration from "./ContentFilterConfiguration";
vi.mock("@/components/networking", () => ({
validateBlockedWordsFile: vi.fn(),
getCategoryYaml: vi.fn(),
}));
vi.mock("@/components/molecules/notifications_manager", () => ({
default: { success: vi.fn(), error: vi.fn(), fromBackend: vi.fn() },
}));
const PREBUILT = [
{ name: "us_ssn", display_name: "US Social Security Number", category: "PII Patterns", description: "d" },
];
describe("ContentFilterConfiguration", () => {
const handlers = {
onPatternAdd: vi.fn(),
onPatternRemove: vi.fn(),
onPatternActionChange: vi.fn(),
onBlockedWordAdd: vi.fn(),
onBlockedWordRemove: vi.fn(),
onBlockedWordUpdate: vi.fn(),
};
const renderConfig = (overrides = {}) =>
renderWithProviders(
<ContentFilterConfiguration
prebuiltPatterns={PREBUILT}
categories={["PII Patterns"]}
selectedPatterns={[]}
blockedWords={[]}
accessToken="test-token"
{...handlers}
{...overrides}
/>,
);
beforeEach(() => {
vi.clearAllMocks();
});
it("should render the pattern and keyword sections", () => {
renderConfig();
expect(screen.getByText("Pattern Detection")).toBeInTheDocument();
expect(
screen.getByText("Detect sensitive information using regex patterns (SSN, credit cards, API keys, etc.)"),
).toBeInTheDocument();
expect(screen.getByText("Blocked Keywords")).toBeInTheDocument();
expect(screen.getByText("Block or mask specific sensitive terms and phrases")).toBeInTheDocument();
expect(screen.getByRole("button", { name: /add prebuilt pattern/i })).toBeInTheDocument();
expect(screen.getByRole("button", { name: /add custom regex/i })).toBeInTheDocument();
expect(screen.getByRole("button", { name: /add keyword/i })).toBeInTheDocument();
expect(screen.getByRole("button", { name: /upload yaml file/i })).toBeInTheDocument();
});
it("should show the empty states for patterns and keywords", () => {
renderConfig();
expect(screen.getByText("No patterns added.")).toBeInTheDocument();
expect(screen.getByText("No keywords added.")).toBeInTheDocument();
});
it("should open the prebuilt pattern modal", async () => {
const user = userEvent.setup();
renderConfig();
expect(screen.queryByText("Pattern type")).not.toBeInTheDocument();
await user.click(screen.getByRole("button", { name: /add prebuilt pattern/i }));
expect(await screen.findByText("Pattern type")).toBeInTheDocument();
});
it("should open the custom regex modal", async () => {
const user = userEvent.setup();
renderConfig();
expect(screen.queryByText("Add custom regex pattern")).not.toBeInTheDocument();
await user.click(screen.getByRole("button", { name: /add custom regex/i }));
expect(await screen.findByText("Add custom regex pattern")).toBeInTheDocument();
expect(screen.getByPlaceholderText("e.g., ID-[0-9]{6}")).toBeInTheDocument();
});
it("should open the keyword modal", async () => {
const user = userEvent.setup();
renderConfig();
expect(screen.queryByText("Add blocked keyword")).not.toBeInTheDocument();
await user.click(screen.getByRole("button", { name: /add keyword/i }));
expect(await screen.findByText("Add blocked keyword")).toBeInTheDocument();
expect(screen.getByPlaceholderText("Enter sensitive keyword or phrase")).toBeInTheDocument();
});
it("should list already selected patterns and keywords", () => {
renderConfig({
selectedPatterns: [
{
id: "pattern-1",
type: "prebuilt" as const,
name: "us_ssn",
display_name: "US Social Security Number",
action: "BLOCK" as const,
},
],
blockedWords: [{ id: "word-1", keyword: "secret", action: "MASK" as const, description: "Sensitive" }],
});
expect(screen.getByText("US Social Security Number")).toBeInTheDocument();
expect(screen.getByText("secret")).toBeInTheDocument();
expect(screen.queryByText("No patterns added.")).not.toBeInTheDocument();
expect(screen.queryByText("No keywords added.")).not.toBeInTheDocument();
});
it("should show only the keyword section when the keywords step is requested", () => {
renderConfig({ showStep: "keywords" });
expect(screen.getByText("Blocked Keywords")).toBeInTheDocument();
expect(screen.queryByText("Pattern Detection")).not.toBeInTheDocument();
});
it("should show only the pattern section when the patterns step is requested", () => {
renderConfig({ showStep: "patterns" });
expect(screen.getByText("Pattern Detection")).toBeInTheDocument();
expect(screen.queryByText("Blocked Keywords")).not.toBeInTheDocument();
});
});

113
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterConfiguration.tsx
View file

@ -1,8 +1,10 @@
import React, { useState } from "react";
import { Typography, Space, Upload, Card, Button } from "antd";
import { PlusOutlined, UploadOutlined } from "@ant-design/icons";
import React, { useRef, useState } from "react";
import { Plus, Upload } from "lucide-react";
import { validateBlockedWordsFile } from "@/components/networking";
import NotificationsManager from "@/components/molecules/notifications_manager";
import { Button } from "@/components/ui/button";
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
import PatternModal from "./PatternModal";
import CustomPatternModal from "./CustomPatternModal";
import KeywordModal from "./KeywordModal";
@ -11,8 +13,6 @@ import KeywordTable from "./KeywordTable";
import ContentCategoryConfiguration from "./ContentCategoryConfiguration";
import CompetitorIntentConfiguration, { CompetitorIntentConfig } from "./CompetitorIntentConfiguration";
const { Title, Text } = Typography;
interface PrebuiltPattern {
name: string;
display_name: string;
@ -115,6 +115,7 @@ const ContentFilterConfiguration: React.FC<ContentFilterConfigurationProps> = ({
const [newKeywordAction, setNewKeywordAction] = useState<"BLOCK" | "MASK">("BLOCK");
const [newKeywordDescription, setNewKeywordDescription] = useState<string>("");
const [uploadValidating, setUploadValidating] = useState(false);
const fileInputRef = useRef<HTMLInputElement>(null);
const handleAddPrebuiltPattern = () => {
if (!selectedPatternName) {
@ -201,6 +202,14 @@ const ContentFilterConfiguration: React.FC<ContentFilterConfigurationProps> = ({
return false;
};
const handleFileInputChange = (event: React.ChangeEvent<HTMLInputElement>) => {
const file = event.target.files?.[0];
event.target.value = "";
if (file) {
handleFileUpload(file);
}
};
const showPatterns = !showStep || showStep === "patterns";
const showKeywords = !showStep || showStep === "keywords";
const showCategories = !showStep || showStep === "categories";
@ -210,68 +219,78 @@ const ContentFilterConfiguration: React.FC<ContentFilterConfigurationProps> = ({
<div className="space-y-6">
{!showStep && (
<div>
<Text type="secondary">
<p className="text-muted-foreground">
Configure patterns, keywords, and content categories to detect and filter sensitive information in requests
and responses.
</Text>
</p>
</div>
)}
{showPatterns && (
<Card
title={
<div style={{ display: "flex", justifyContent: "space-between", alignItems: "center" }}>
<Title level={5} style={{ margin: 0 }}>
Pattern Detection
</Title>
<Text type="secondary" style={{ fontSize: 14, fontWeight: 400 }}>
<Card>
<CardHeader>
<div className="flex flex-wrap items-center justify-between gap-2">
<CardTitle>Pattern Detection</CardTitle>
<p className="text-sm font-normal text-muted-foreground">
Detect sensitive information using regex patterns (SSN, credit cards, API keys, etc.)
</Text>
</p>
</div>
}
size="small"
>
<div style={{ marginBottom: 16 }}>
<Space>
<Button type="primary" onClick={() => setPatternModalVisible(true)} icon={<PlusOutlined />}>
</CardHeader>
<CardContent>
<div className="mb-4 flex flex-wrap gap-2">
<Button onClick={() => setPatternModalVisible(true)}>
<Plus />
Add prebuilt pattern
</Button>
<Button onClick={() => setCustomPatternModalVisible(true)} icon={<PlusOutlined />}>
<Button variant="outline" onClick={() => setCustomPatternModalVisible(true)}>
<Plus />
Add custom regex
</Button>
</Space>
</div>
<PatternTable patterns={selectedPatterns} onActionChange={onPatternActionChange} onRemove={onPatternRemove} />
</div>
<PatternTable
patterns={selectedPatterns}
onActionChange={onPatternActionChange}
onRemove={onPatternRemove}
/>
</CardContent>
</Card>
)}
{showKeywords && (
<Card
title={
<div style={{ display: "flex", justifyContent: "space-between", alignItems: "center" }}>
<Title level={5} style={{ margin: 0 }}>
Blocked Keywords
</Title>
<Text type="secondary" style={{ fontSize: 14, fontWeight: 400 }}>
<Card>
<CardHeader>
<div className="flex flex-wrap items-center justify-between gap-2">
<CardTitle>Blocked Keywords</CardTitle>
<p className="text-sm font-normal text-muted-foreground">
Block or mask specific sensitive terms and phrases
</Text>
</p>
</div>
}
size="small"
>
<div style={{ marginBottom: 16 }}>
<Space>
<Button type="primary" onClick={() => setKeywordModalVisible(true)} icon={<PlusOutlined />}>
</CardHeader>
<CardContent>
<div className="mb-4 flex flex-wrap gap-2">
<Button onClick={() => setKeywordModalVisible(true)}>
<Plus />
Add keyword
</Button>
<Upload beforeUpload={handleFileUpload} accept=".yaml,.yml" showUploadList={false}>
<Button icon={<UploadOutlined />} loading={uploadValidating}>
Upload YAML file
</Button>
</Upload>
</Space>
</div>
<KeywordTable keywords={blockedWords} onActionChange={onBlockedWordUpdate} onRemove={onBlockedWordRemove} />
<input
ref={fileInputRef}
type="file"
accept=".yaml,.yml"
className="hidden"
onChange={handleFileInputChange}
/>
<Button
variant="outline"
disabled={uploadValidating}
aria-busy={uploadValidating}
onClick={() => fileInputRef.current?.click()}
>
{uploadValidating ? <UiLoadingSpinner className="size-4" /> : <Upload />}
Upload YAML file
</Button>
</div>
<KeywordTable keywords={blockedWords} onActionChange={onBlockedWordUpdate} onRemove={onBlockedWordRemove} />
</CardContent>
</Card>
)}

89
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterDisplay.test.tsx Normal file
View file

@ -0,0 +1,89 @@
import { describe, it, expect } from "vitest";
import { renderWithProviders, screen } from "@/../tests/test-utils";
import ContentFilterDisplay from "./ContentFilterDisplay";
const PATTERN = {
id: "pattern-1",
type: "prebuilt" as const,
name: "email",
display_name: "Email address",
action: "BLOCK" as const,
};
const KEYWORD = {
id: "word-1",
keyword: "secret",
action: "MASK" as const,
description: "Sensitive term",
};
const CATEGORY = {
id: "category-1",
category: "self_harm",
display_name: "Self Harm",
action: "BLOCK" as const,
severity_threshold: "high" as const,
};
describe("ContentFilterDisplay", () => {
it("should render nothing when there is no content filter data", () => {
const { container } = renderWithProviders(<ContentFilterDisplay patterns={[]} blockedWords={[]} categories={[]} />);
expect(container).toBeEmptyDOMElement();
});
it("should render the categories section with a configured count", () => {
renderWithProviders(<ContentFilterDisplay patterns={[]} blockedWords={[]} categories={[CATEGORY]} />);
expect(screen.getByText("Content Categories")).toBeInTheDocument();
expect(screen.getByText("1 categories configured")).toBeInTheDocument();
expect(screen.getByText("Self Harm")).toBeInTheDocument();
expect(screen.queryByText("Pattern Detection")).not.toBeInTheDocument();
expect(screen.queryByText("Blocked Keywords")).not.toBeInTheDocument();
});
it("should render the patterns section with a configured count", () => {
renderWithProviders(<ContentFilterDisplay patterns={[PATTERN]} blockedWords={[]} categories={[]} />);
expect(screen.getByText("Pattern Detection")).toBeInTheDocument();
expect(screen.getByText("1 patterns configured")).toBeInTheDocument();
expect(screen.getByText("Email address")).toBeInTheDocument();
expect(screen.queryByText("Content Categories")).not.toBeInTheDocument();
});
it("should render the keywords section with a configured count", () => {
renderWithProviders(<ContentFilterDisplay patterns={[]} blockedWords={[KEYWORD]} categories={[]} />);
expect(screen.getByText("Blocked Keywords")).toBeInTheDocument();
expect(screen.getByText("1 keywords configured")).toBeInTheDocument();
expect(screen.getByText("secret")).toBeInTheDocument();
expect(screen.getByText("Sensitive term")).toBeInTheDocument();
});
it("should render every section when all three kinds of data are present", () => {
renderWithProviders(<ContentFilterDisplay patterns={[PATTERN]} blockedWords={[KEYWORD]} categories={[CATEGORY]} />);
expect(screen.getByText("Content Categories")).toBeInTheDocument();
expect(screen.getByText("Pattern Detection")).toBeInTheDocument();
expect(screen.getByText("Blocked Keywords")).toBeInTheDocument();
});
it("should render category severity and action as static text in read-only mode", () => {
renderWithProviders(
<ContentFilterDisplay patterns={[PATTERN]} blockedWords={[KEYWORD]} categories={[CATEGORY]} readOnly={true} />,
);
expect(screen.getByText("HIGH")).toBeInTheDocument();
expect(screen.getByText("BLOCK")).toBeInTheDocument();
expect(screen.getAllByRole("button", { name: /delete/i })).toHaveLength(2);
});
it("should render category severity and action as editable controls when not read-only", () => {
renderWithProviders(
<ContentFilterDisplay patterns={[PATTERN]} blockedWords={[KEYWORD]} categories={[CATEGORY]} readOnly={false} />,
);
expect(screen.queryByText("HIGH")).not.toBeInTheDocument();
expect(screen.getAllByRole("button", { name: /delete/i })).toHaveLength(3);
});
});

67
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterDisplay.tsx
View file

@ -1,5 +1,6 @@
import React from "react";
import { Card, Text, Badge } from "@tremor/react";
import { Badge } from "@/components/ui/badge";
import { Card, CardContent } from "@/components/ui/card";
import PatternTable from "./PatternTable";
import KeywordTable from "./KeywordTable";
import CategoryTable from "./CategoryTable";
@ -66,45 +67,51 @@ const ContentFilterDisplay: React.FC<ContentFilterDisplayProps> = ({
<>
{categories.length > 0 && (
<Card className="mt-6">
<div className="flex justify-between items-center mb-4">
<Text className="text-lg font-semibold">Content Categories</Text>
<Badge color="blue">{categories.length} categories configured</Badge>
</div>
<CategoryTable
categories={categories}
onActionChange={readOnly ? undefined : onCategoryActionChange}
onSeverityChange={readOnly ? undefined : onCategorySeverityChange}
onRemove={readOnly ? undefined : onCategoryRemove}
readOnly={readOnly}
/>
<CardContent>
<div className="mb-4 flex items-center justify-between">
<p className="text-lg font-semibold">Content Categories</p>
<Badge variant="secondary">{categories.length} categories configured</Badge>
</div>
<CategoryTable
categories={categories}
onActionChange={readOnly ? undefined : onCategoryActionChange}
onSeverityChange={readOnly ? undefined : onCategorySeverityChange}
onRemove={readOnly ? undefined : onCategoryRemove}
readOnly={readOnly}
/>
</CardContent>
</Card>
)}
{patterns.length > 0 && (
<Card className="mt-6">
<div className="flex justify-between items-center mb-4">
<Text className="text-lg font-semibold">Pattern Detection</Text>
<Badge color="blue">{patterns.length} patterns configured</Badge>
</div>
<PatternTable
patterns={patterns}
onActionChange={readOnly ? noOp : onPatternActionChange || noOp}
onRemove={readOnly ? noOp : onPatternRemove || noOp}
/>
<CardContent>
<div className="mb-4 flex items-center justify-between">
<p className="text-lg font-semibold">Pattern Detection</p>
<Badge variant="secondary">{patterns.length} patterns configured</Badge>
</div>
<PatternTable
patterns={patterns}
onActionChange={readOnly ? noOp : onPatternActionChange || noOp}
onRemove={readOnly ? noOp : onPatternRemove || noOp}
/>
</CardContent>
</Card>
)}
{blockedWords.length > 0 && (
<Card className="mt-6">
<div className="flex justify-between items-center mb-4">
<Text className="text-lg font-semibold">Blocked Keywords</Text>
<Badge color="blue">{blockedWords.length} keywords configured</Badge>
</div>
<KeywordTable
keywords={blockedWords}
onActionChange={readOnly ? noOp : onBlockedWordUpdate || noOp}
onRemove={readOnly ? noOp : onBlockedWordRemove || noOp}
/>
<CardContent>
<div className="mb-4 flex items-center justify-between">
<p className="text-lg font-semibold">Blocked Keywords</p>
<Badge variant="secondary">{blockedWords.length} keywords configured</Badge>
</div>
<KeywordTable
keywords={blockedWords}
onActionChange={readOnly ? noOp : onBlockedWordUpdate || noOp}
onRemove={readOnly ? noOp : onBlockedWordRemove || noOp}
/>
</CardContent>
</Card>
)}
</>

24
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterManager.test.tsx
View file

@ -2,7 +2,6 @@ import { describe, it, expect, vi, beforeEach } from "vitest";
import { render, screen, waitFor } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import ContentFilterManager, { formatContentFilterDataForAPI } from "./ContentFilterManager";
import React from "react";
const CONTENT_FILTER_GUARDRAIL_DATA = {
litellm_params: {
@ -85,18 +84,7 @@ vi.mock("./ContentFilterDisplay", () => ({
),
}));
vi.mock("antd", async (importOriginal) => {
const actual = await importOriginal<typeof import("antd")>();
return {
...actual,
Divider: ({ children }: { children: React.ReactNode }) => <div data-testid="divider">{children}</div>,
Alert: ({ message, type }: { message: React.ReactNode; type: string }) => (
<div data-testid="unsaved-alert" data-type={type}>
{message}
</div>
),
};
});
const UNSAVED_CHANGES_TEXT = /You have unsaved changes to patterns or keywords/;
describe("ContentFilterManager", () => {
beforeEach(() => {
@ -117,7 +105,7 @@ describe("ContentFilterManager", () => {
expect(screen.getByTestId("content-filter-config")).toBeInTheDocument();
});
expect(screen.getByTestId("divider")).toHaveTextContent("Content Filter Configuration");
expect(screen.getByText("Content Filter Configuration")).toBeInTheDocument();
});
it("should return null when guardrail is not litellm_content_filter", () => {
@ -275,15 +263,15 @@ describe("ContentFilterManager", () => {
expect(screen.getByTestId("content-filter-config")).toBeInTheDocument();
});
expect(screen.queryByTestId("unsaved-alert")).not.toBeInTheDocument();
expect(screen.queryByText(UNSAVED_CHANGES_TEXT)).not.toBeInTheDocument();
await user.click(screen.getByRole("button", { name: /add pattern/i }));
await waitFor(() => {
expect(screen.getByTestId("unsaved-alert")).toBeInTheDocument();
expect(screen.getByText(UNSAVED_CHANGES_TEXT)).toBeInTheDocument();
});
expect(screen.getByTestId("unsaved-alert")).toHaveTextContent(/unsaved changes.*Save Changes/i);
expect(screen.getByText(UNSAVED_CHANGES_TEXT)).toHaveTextContent(/Save Changes/i);
});
it("should call onDataChange when patterns or keywords change", async () => {
@ -371,7 +359,7 @@ describe("ContentFilterManager", () => {
);
await waitFor(() => {
expect(screen.getByTestId("divider")).toBeInTheDocument();
expect(screen.getByText("Content Filter Configuration")).toBeInTheDocument();
});
expect(screen.queryByTestId("content-filter-config")).not.toBeInTheDocument();

28
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterManager.tsx
View file

@ -1,11 +1,11 @@
import { Alert, Divider, Typography } from "antd";
import { TriangleAlert } from "lucide-react";
import React, { useEffect, useState } from "react";
import { Alert, AlertDescription } from "@/components/shared/Alert";
import { Separator } from "@/components/ui/separator";
import ContentFilterConfiguration from "./ContentFilterConfiguration";
import ContentFilterDisplay from "./ContentFilterDisplay";
import type { CompetitorIntentConfig } from "./CompetitorIntentConfiguration";
const { Text } = Typography;
interface Pattern {
id: string;
type: "prebuilt" | "custom";
@ -233,19 +233,17 @@ const ContentFilterManager: React.FC<ContentFilterManagerProps> = ({
// Edit mode
return (
<>
<Divider orientation="left">Content Filter Configuration</Divider>
<div className="my-6 flex items-center gap-4">
<span className="shrink-0 font-medium">Content Filter Configuration</span>
<Separator className="flex-1" />
</div>
{hasUnsavedChanges && (
<Alert
type="warning"
showIcon
className="mb-4"
message={
<Text>
You have unsaved changes to patterns or keywords. Remember to click &quot;Save Changes&quot; at the
bottom.
</Text>
}
/>
<Alert variant="warning" className="mb-4">
<TriangleAlert />
<AlertDescription>
You have unsaved changes to patterns or keywords. Remember to click &quot;Save Changes&quot; at the bottom.
</AlertDescription>
</Alert>
)}
<div className="mb-6">
{guardrailSettings && guardrailSettings.content_filter_settings && (

73
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterTables.test.tsx
View file

@ -130,4 +130,77 @@ describe("content filter tables", () => {
expect(onCategoryRemove).toHaveBeenCalledWith("category-1");
});
it("should report a pattern action change", async () => {
const onActionChange = vi.fn();
const user = userEvent.setup();
renderWithProviders(
<PatternTable
patterns={[{ id: "pattern-1", type: "prebuilt", name: "email", action: "BLOCK" }]}
onActionChange={onActionChange}
onRemove={vi.fn()}
/>,
);
await user.click(screen.getByRole("combobox"));
const maskOptions = await screen.findAllByText("Mask");
await user.click(maskOptions[maskOptions.length - 1]);
expect(onActionChange).toHaveBeenCalledWith("pattern-1", "MASK");
});
it("should report a keyword action change", async () => {
const onActionChange = vi.fn();
const user = userEvent.setup();
renderWithProviders(
<KeywordTable
keywords={[{ id: "keyword-1", keyword: "secret", action: "BLOCK" }]}
onActionChange={onActionChange}
onRemove={vi.fn()}
/>,
);
await user.click(screen.getByRole("combobox"));
const maskOptions = await screen.findAllByText("Mask");
await user.click(maskOptions[maskOptions.length - 1]);
expect(onActionChange).toHaveBeenCalledWith("keyword-1", "action", "MASK");
});
it("should report category severity and action changes", async () => {
const onSeverityChange = vi.fn();
const onActionChange = vi.fn();
const user = userEvent.setup();
renderWithProviders(
<CategoryTable
categories={[
{
id: "category-1",
category: "self_harm",
display_name: "Self Harm",
action: "BLOCK",
severity_threshold: "high",
},
]}
onActionChange={onActionChange}
onSeverityChange={onSeverityChange}
onRemove={vi.fn()}
/>,
);
await user.click(screen.getAllByRole("combobox")[0]);
const lowOptions = await screen.findAllByText("Low");
await user.click(lowOptions[lowOptions.length - 1]);
expect(onSeverityChange).toHaveBeenCalledWith("category-1", "low");
await user.click(screen.getAllByRole("combobox")[1]);
const maskOptions = await screen.findAllByText("Mask");
await user.click(maskOptions[maskOptions.length - 1]);
expect(onActionChange).toHaveBeenCalledWith("category-1", "MASK");
});
});

110
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/CustomPatternModal.tsx
View file

@ -1,8 +1,10 @@
import React from "react";
import { Typography, Select, Modal, Space, Button, Input } from "antd";
const { Text } = Typography;
const { Option } = Select;
import { Button } from "@/components/ui/button";
import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
import { Input } from "@/components/ui/input";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { ACTION_ITEMS } from "./action_options";
import { ABOVE_ANTD_MODAL } from "./dialog_layering";
interface CustomPatternModalProps {
visible: boolean;
@ -28,50 +30,66 @@ const CustomPatternModal: React.FC<CustomPatternModalProps> = ({
onCancel,
}) => {
return (
<Modal title="Add custom regex pattern" open={visible} onCancel={onCancel} footer={null} width={800}>
<Space direction="vertical" style={{ width: "100%" }} size="large">
<div>
<Text strong>Pattern name</Text>
<Input
placeholder="e.g., internal_id, employee_code"
value={patternName}
onChange={(e) => onNameChange(e.target.value)}
style={{ marginTop: 8 }}
/>
<Dialog open={visible} onOpenChange={(open) => !open && onCancel()}>
<DialogContent className={`max-h-[calc(100dvh-2rem)] overflow-y-auto sm:max-w-[800px] ${ABOVE_ANTD_MODAL}`}>
<DialogHeader>
<DialogTitle>Add custom regex pattern</DialogTitle>
</DialogHeader>
<div className="space-y-6">
<div>
<p className="font-semibold">Pattern name</p>
<Input
className="mt-2"
placeholder="e.g., internal_id, employee_code"
value={patternName}
onChange={(e) => onNameChange(e.target.value)}
/>
</div>
<div>
<p className="font-semibold">Regex pattern</p>
<Input
className="mt-2"
placeholder="e.g., ID-[0-9]{6}"
value={patternRegex}
onChange={(e) => onRegexChange(e.target.value)}
/>
<p className="text-xs text-muted-foreground">Enter a valid regular expression to match sensitive data</p>
</div>
<div>
<p className="font-semibold">Action</p>
<p className="mt-1 mb-2 text-muted-foreground">
Choose what action the guardrail should take when this pattern is detected
</p>
<Select
items={ACTION_ITEMS}
value={patternAction}
onValueChange={(value: string | null) => value && onActionChange(value as "BLOCK" | "MASK")}
>
<SelectTrigger className="w-full" aria-label="Action">
<SelectValue />
</SelectTrigger>
<SelectContent alignItemWithTrigger={false}>
{ACTION_ITEMS.map((item) => (
<SelectItem key={item.value} value={item.value}>
{item.label}
</SelectItem>
))}
</SelectContent>
</Select>
</div>
</div>
<div>
<Text strong>Regex pattern</Text>
<Input
placeholder="e.g., ID-[0-9]{6}"
value={patternRegex}
onChange={(e) => onRegexChange(e.target.value)}
style={{ marginTop: 8 }}
/>
<Text type="secondary" style={{ fontSize: 12 }}>
Enter a valid regular expression to match sensitive data
</Text>
</div>
<div>
<Text strong>Action</Text>
<Text type="secondary" style={{ display: "block", marginTop: 4, marginBottom: 8 }}>
Choose what action the guardrail should take when this pattern is detected
</Text>
<Select value={patternAction} onChange={onActionChange} style={{ width: "100%" }}>
<Option value="BLOCK">Block</Option>
<Option value="MASK">Mask</Option>
</Select>
</div>
</Space>
<div style={{ display: "flex", justifyContent: "flex-end", gap: "8px", marginTop: "24px" }}>
<Button onClick={onCancel}>Cancel</Button>
<Button type="primary" onClick={onAdd}>
Add
</Button>
</div>
</Modal>
<DialogFooter>
<Button variant="outline" onClick={onCancel}>
Cancel
</Button>
<Button onClick={onAdd}>Add</Button>
</DialogFooter>
</DialogContent>
</Dialog>
);
};

87
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/KeywordModal.test.tsx Normal file
View file

@ -0,0 +1,87 @@
import { describe, it, expect, vi, beforeEach } from "vitest";
import { render, screen } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import KeywordModal from "./KeywordModal";
describe("KeywordModal", () => {
const handlers = {
onKeywordChange: vi.fn(),
onActionChange: vi.fn(),
onDescriptionChange: vi.fn(),
onAdd: vi.fn(),
onCancel: vi.fn(),
};
const renderModal = (overrides: Partial<React.ComponentProps<typeof KeywordModal>> = {}) =>
render(<KeywordModal visible keyword="" action="BLOCK" description="" {...handlers} {...overrides} />);
beforeEach(() => {
vi.clearAllMocks();
});
it("should render the keyword, action and description fields", async () => {
renderModal();
expect(await screen.findByText("Add blocked keyword")).toBeInTheDocument();
expect(screen.getByPlaceholderText("Enter sensitive keyword or phrase")).toBeInTheDocument();
expect(screen.getByPlaceholderText("Explain why this keyword is sensitive")).toBeInTheDocument();
expect(screen.getByText("Description (optional)")).toBeInTheDocument();
expect(
screen.getByText("Choose what action the guardrail should take when this keyword is detected"),
).toBeInTheDocument();
});
it("should report keyword edits", async () => {
const user = userEvent.setup();
renderModal();
await user.type(await screen.findByPlaceholderText("Enter sensitive keyword or phrase"), "s");
expect(handlers.onKeywordChange).toHaveBeenCalledWith("s");
});
it("should report description edits", async () => {
const user = userEvent.setup();
renderModal();
await user.type(await screen.findByPlaceholderText("Explain why this keyword is sensitive"), "x");
expect(handlers.onDescriptionChange).toHaveBeenCalledWith("x");
});
it("should report the chosen action", async () => {
const user = userEvent.setup();
renderModal();
await user.click(await screen.findByRole("combobox"));
const maskOptions = await screen.findAllByText("Mask");
await user.click(maskOptions[maskOptions.length - 1]);
expect(handlers.onActionChange).toHaveBeenCalled();
expect(handlers.onActionChange.mock.calls[0][0]).toBe("MASK");
});
it("should show the current keyword and description values", async () => {
renderModal({ keyword: "secret", description: "sensitive term" });
expect(await screen.findByDisplayValue("secret")).toBeInTheDocument();
expect(screen.getByDisplayValue("sensitive term")).toBeInTheDocument();
});
it("should add and cancel through the footer buttons", async () => {
const user = userEvent.setup();
renderModal();
await user.click(await screen.findByRole("button", { name: "Add" }));
expect(handlers.onAdd).toHaveBeenCalledTimes(1);
await user.click(screen.getByRole("button", { name: "Cancel" }));
expect(handlers.onCancel).toHaveBeenCalledTimes(1);
});
it("should not render its content when not visible", () => {
renderModal({ visible: false });
expect(screen.queryByText("Add blocked keyword")).not.toBeInTheDocument();
});
});

109
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/KeywordModal.tsx
View file

@ -1,8 +1,11 @@
import React from "react";
import { Typography, Select, Modal, Space, Button, Input } from "antd";
const { Text } = Typography;
const { Option } = Select;
import { Button } from "@/components/ui/button";
import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
import { Input } from "@/components/ui/input";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { Textarea } from "@/components/ui/textarea";
import { ACTION_ITEMS } from "./action_options";
import { ABOVE_ANTD_MODAL } from "./dialog_layering";
interface KeywordModalProps {
visible: boolean;
@ -28,48 +31,66 @@ const KeywordModal: React.FC<KeywordModalProps> = ({
onCancel,
}) => {
return (
<Modal title="Add blocked keyword" open={visible} onCancel={onCancel} footer={null} width={800}>
<Space direction="vertical" style={{ width: "100%" }} size="large">
<div>
<Text strong>Keyword</Text>
<Input
placeholder="Enter sensitive keyword or phrase"
value={keyword}
onChange={(e) => onKeywordChange(e.target.value)}
style={{ marginTop: 8 }}
/>
<Dialog open={visible} onOpenChange={(open) => !open && onCancel()}>
<DialogContent className={`max-h-[calc(100dvh-2rem)] overflow-y-auto sm:max-w-[800px] ${ABOVE_ANTD_MODAL}`}>
<DialogHeader>
<DialogTitle>Add blocked keyword</DialogTitle>
</DialogHeader>
<div className="space-y-6">
<div>
<p className="font-semibold">Keyword</p>
<Input
className="mt-2"
placeholder="Enter sensitive keyword or phrase"
value={keyword}
onChange={(e) => onKeywordChange(e.target.value)}
/>
</div>
<div>
<p className="font-semibold">Action</p>
<p className="mt-1 mb-2 text-muted-foreground">
Choose what action the guardrail should take when this keyword is detected
</p>
<Select
items={ACTION_ITEMS}
value={action}
onValueChange={(value: string | null) => value && onActionChange(value as "BLOCK" | "MASK")}
>
<SelectTrigger className="w-full" aria-label="Action">
<SelectValue />
</SelectTrigger>
<SelectContent alignItemWithTrigger={false}>
{ACTION_ITEMS.map((item) => (
<SelectItem key={item.value} value={item.value}>
{item.label}
</SelectItem>
))}
</SelectContent>
</Select>
</div>
<div>
<p className="font-semibold">Description (optional)</p>
<Textarea
className="mt-2 field-sizing-fixed"
placeholder="Explain why this keyword is sensitive"
value={description}
onChange={(e) => onDescriptionChange(e.target.value)}
rows={3}
/>
</div>
</div>
<div>
<Text strong>Action</Text>
<Text type="secondary" style={{ display: "block", marginTop: 4, marginBottom: 8 }}>
Choose what action the guardrail should take when this keyword is detected
</Text>
<Select value={action} onChange={onActionChange} style={{ width: "100%" }}>
<Option value="BLOCK">Block</Option>
<Option value="MASK">Mask</Option>
</Select>
</div>
<div>
<Text strong>Description (optional)</Text>
<Input.TextArea
placeholder="Explain why this keyword is sensitive"
value={description}
onChange={(e) => onDescriptionChange(e.target.value)}
rows={3}
style={{ marginTop: 8 }}
/>
</div>
</Space>
<div style={{ display: "flex", justifyContent: "flex-end", gap: "8px", marginTop: "24px" }}>
<Button onClick={onCancel}>Cancel</Button>
<Button type="primary" onClick={onAdd}>
Add
</Button>
</div>
</Modal>
<DialogFooter>
<Button variant="outline" onClick={onCancel}>
Cancel
</Button>
<Button onClick={onAdd}>Add</Button>
</DialogFooter>
</DialogContent>
</Dialog>
);
};

30
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/KeywordTable.tsx
View file

@ -1,10 +1,10 @@
import { DeleteOutlined } from "@ant-design/icons";
import { Trash2 } from "lucide-react";
import type { ColumnDef } from "@tanstack/react-table";
import { Button, Select } from "antd";
import React from "react";
import { DataTable } from "@/components/shared/DataTable";
const { Option } = Select;
import { Button } from "@/components/ui/button";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { ACTION_ITEMS } from "./action_options";
interface BlockedWord {
id: string;
@ -31,13 +31,20 @@ const KeywordTable: React.FC<KeywordTableProps> = ({ keywords, onActionChange, o
size: 150,
cell: ({ row }) => (
<Select
items={ACTION_ITEMS}
value={row.original.action}
onChange={(value) => onActionChange(row.original.id, "action", value)}
style={{ width: 120 }}
size="small"
onValueChange={(value: string | null) => value && onActionChange(row.original.id, "action", value)}
>
<Option value="BLOCK">Block</Option>
<Option value="MASK">Mask</Option>
<SelectTrigger size="sm" className="w-[120px]" aria-label="Action">
<SelectValue />
</SelectTrigger>
<SelectContent alignItemWithTrigger={false}>
{ACTION_ITEMS.map((item) => (
<SelectItem key={item.value} value={item.value}>
{item.label}
</SelectItem>
))}
</SelectContent>
</Select>
),
},
@ -51,7 +58,8 @@ const KeywordTable: React.FC<KeywordTableProps> = ({ keywords, onActionChange, o
id: "actions",
size: 100,
cell: ({ row }) => (
<Button type="text" danger size="small" icon={<DeleteOutlined />} onClick={() => onRemove(row.original.id)}>
<Button variant="ghost" size="sm" onClick={() => onRemove(row.original.id)}>
<Trash2 />
Delete
</Button>
),
@ -59,7 +67,7 @@ const KeywordTable: React.FC<KeywordTableProps> = ({ keywords, onActionChange, o
];
if (keywords.length === 0) {
return <div style={{ textAlign: "center", padding: "40px 0", color: "#999" }}>No keywords added.</div>;
return <div className="py-10 text-center text-muted-foreground">No keywords added.</div>;
}
return <DataTable data={keywords} columns={columns} getRowId={(row) => row.id} size="compact" />;

158
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.test.tsx
View file

@ -1,5 +1,5 @@
import { describe, it, expect, vi, beforeEach } from "vitest";
import { render, screen, waitFor } from "@testing-library/react";
import { render, screen } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import PatternModal from "./PatternModal";
@ -10,25 +10,25 @@ describe("PatternModal", () => {
const mockOnActionChange = vi.fn();
const mockPrebuiltPatterns = [
{ name: "us_ssn", category: "PII Patterns", description: "US Social Security Number" },
{ name: "email", category: "PII Patterns", description: "Email addresses" },
{ name: "visa", category: "Financial Patterns", description: "Visa credit card numbers" },
{ name: "aws_access_key", category: "Credential Patterns", description: "AWS Access Keys" },
{
name: "us_ssn",
display_name: "US Social Security Number",
category: "PII Patterns",
description: "US Social Security Number",
},
{ name: "email", display_name: "Email address", category: "PII Patterns", description: "Email addresses" },
{ name: "visa", display_name: "Visa card", category: "Financial Patterns", description: "Visa credit cards" },
{
name: "aws_access_key",
display_name: "AWS access key",
category: "Credential Patterns",
description: "AWS Access Keys",
},
];
const mockCategories = ["PII Patterns", "Financial Patterns", "Credential Patterns"];
beforeEach(() => {
vi.clearAllMocks();
});
it("should show dropdown with prebuilt pattern options grouped by category", async () => {
/**
* Tests that the modal displays a dropdown with prebuilt patterns
* organized by category. This verifies the pattern selection UI is working.
*/
const user = userEvent.setup();
const renderModal = () =>
render(
<PatternModal
visible={true}
@ -43,49 +43,103 @@ describe("PatternModal", () => {
/>,
);
// Wait for modal to be visible
await waitFor(() => {
expect(screen.getByText("Add prebuilt pattern")).toBeInTheDocument();
});
beforeEach(() => {
vi.clearAllMocks();
});
// Find the pattern type dropdown by looking for the first combobox input
const comboboxes = screen.getAllByRole("combobox");
const dropdown = comboboxes[0]; // First combobox is the pattern selector
expect(dropdown).toBeInTheDocument();
it("should show prebuilt pattern options grouped by category and report the picked pattern", async () => {
const user = userEvent.setup();
renderModal();
// Click to open the dropdown
await user.click(dropdown);
expect(await screen.findByText("Add prebuilt pattern")).toBeInTheDocument();
// Verify that pattern options are available in the dropdown
// Ant Design renders Select options in a portal, so we need to query the whole document
await waitFor(() => {
const options = document.querySelectorAll(".ant-select-item-option");
expect(options.length).toBeGreaterThan(0);
});
await user.click(screen.getAllByRole("combobox")[0]);
// Verify categories are shown as group labels
await waitFor(() => {
expect(document.body).toHaveTextContent("PII Patterns");
expect(document.body).toHaveTextContent("Financial Patterns");
expect(document.body).toHaveTextContent("Credential Patterns");
});
expect(await screen.findByText("PII Patterns")).toBeInTheDocument();
expect(screen.getByText("Financial Patterns")).toBeInTheDocument();
expect(screen.getByText("Credential Patterns")).toBeInTheDocument();
// Verify pattern options are available
expect(document.body).toHaveTextContent("us_ssn");
expect(document.body).toHaveTextContent("email");
expect(document.body).toHaveTextContent("visa");
expect(document.body).toHaveTextContent("aws_access_key");
expect(screen.getByText("Email address")).toBeInTheDocument();
expect(screen.getByText("Visa card")).toBeInTheDocument();
expect(screen.getByText("AWS access key")).toBeInTheDocument();
// Select a pattern by clicking on its option element
const ssnOption = Array.from(document.querySelectorAll(".ant-select-item-option")).find(
(el) => el.textContent === "us_ssn",
) as HTMLElement;
await user.click(ssnOption);
const ssnOptions = await screen.findAllByText("US Social Security Number");
await user.click(ssnOptions[ssnOptions.length - 1]);
// Verify the change handler was called with the pattern name
// Note: Ant Design Select calls onChange with (value, option), so we check if it was called
expect(mockOnPatternNameChange).toHaveBeenCalled();
const callArgs = mockOnPatternNameChange.mock.calls[0];
expect(callArgs[0]).toBe("us_ssn");
expect(mockOnPatternNameChange.mock.calls[0][0]).toBe("us_ssn");
});
it("should narrow the pattern options to the typed search text", async () => {
const user = userEvent.setup();
renderModal();
expect(await screen.findByText("Add prebuilt pattern")).toBeInTheDocument();
await user.click(screen.getAllByRole("combobox")[0]);
await user.keyboard("visa");
expect(await screen.findByText("Visa card")).toBeInTheDocument();
expect(screen.queryByText("Email address")).not.toBeInTheDocument();
expect(screen.queryByText("AWS access key")).not.toBeInTheDocument();
});
it("should match the internal pattern name when it is absent from the display name", async () => {
const user = userEvent.setup();
renderModal();
expect(await screen.findByText("Add prebuilt pattern")).toBeInTheDocument();
await user.click(screen.getAllByRole("combobox")[0]);
await user.keyboard("ssn");
expect(await screen.findByText("US Social Security Number")).toBeInTheDocument();
expect(screen.queryByText("Email address")).not.toBeInTheDocument();
expect(screen.queryByText("Visa card")).not.toBeInTheDocument();
});
it("should report the chosen action", async () => {
const user = userEvent.setup();
renderModal();
expect(await screen.findByText("Add prebuilt pattern")).toBeInTheDocument();
await user.click(screen.getAllByRole("combobox")[1]);
const maskOptions = await screen.findAllByText("Mask");
await user.click(maskOptions[maskOptions.length - 1]);
expect(mockOnActionChange).toHaveBeenCalled();
expect(mockOnActionChange.mock.calls[0][0]).toBe("MASK");
});
it("should add and cancel through the footer buttons", async () => {
const user = userEvent.setup();
renderModal();
expect(await screen.findByText("Add prebuilt pattern")).toBeInTheDocument();
await user.click(screen.getByRole("button", { name: "Add" }));
expect(mockOnAdd).toHaveBeenCalledTimes(1);
await user.click(screen.getByRole("button", { name: "Cancel" }));
expect(mockOnCancel).toHaveBeenCalledTimes(1);
});
it("should not render its content when not visible", () => {
render(
<PatternModal
visible={false}
prebuiltPatterns={mockPrebuiltPatterns}
categories={mockCategories}
selectedPatternName=""
patternAction="BLOCK"
onPatternNameChange={mockOnPatternNameChange}
onActionChange={mockOnActionChange}
onAdd={mockOnAdd}
onCancel={mockOnCancel}
/>,
);
expect(screen.queryByText("Add prebuilt pattern")).not.toBeInTheDocument();
});
});

158
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.tsx
View file

@ -1,8 +1,20 @@
import React from "react";
import { Typography, Select, Modal, Space, Button } from "antd";
const { Text } = Typography;
const { Option } = Select;
import { Button } from "@/components/ui/button";
import {
Combobox,
ComboboxCollection,
ComboboxContent,
ComboboxEmpty,
ComboboxGroup,
ComboboxInput,
ComboboxItem,
ComboboxLabel,
ComboboxList,
} from "@/components/ui/combobox";
import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { ACTION_ITEMS } from "./action_options";
import { ABOVE_ANTD_MODAL } from "./dialog_layering";
interface PrebuiltPattern {
name: string;
@ -11,6 +23,16 @@ interface PrebuiltPattern {
description: string;
}
interface PatternGroup {
category: string;
items: PrebuiltPattern[];
}
const matchesPatternQuery = (pattern: PrebuiltPattern, query: string) => {
const needle = query.toLowerCase();
return pattern.display_name.toLowerCase().includes(needle) || pattern.name.toLowerCase().includes(needle);
};
interface PatternModalProps {
visible: boolean;
prebuiltPatterns: PrebuiltPattern[];
@ -34,64 +56,84 @@ const PatternModal: React.FC<PatternModalProps> = ({
onAdd,
onCancel,
}) => {
const selectedPattern = prebuiltPatterns.find((pattern) => pattern.name === selectedPatternName) ?? null;
const patternGroups = categories
.map((category) => ({
category,
items: prebuiltPatterns.filter((pattern) => pattern.category === category),
}))
.filter((group) => group.items.length > 0);
return (
<Modal title="Add prebuilt pattern" open={visible} onCancel={onCancel} footer={null} width={800}>
<Space direction="vertical" style={{ width: "100%" }} size="large">
<div>
<Text strong>Pattern type</Text>
<Select
placeholder="Choose pattern type"
value={selectedPatternName}
onChange={onPatternNameChange}
style={{ width: "100%", marginTop: 8 }}
showSearch
filterOption={(input, option) => {
const pattern = prebuiltPatterns.find((p) => p.name === option?.value);
if (pattern) {
return (
pattern.display_name.toLowerCase().includes(input.toLowerCase()) ||
pattern.name.toLowerCase().includes(input.toLowerCase())
);
}
return false;
}}
>
{categories.map((category) => {
const categoryPatterns = prebuiltPatterns.filter((p) => p.category === category);
if (categoryPatterns.length === 0) return null;
<Dialog open={visible} onOpenChange={(open) => !open && onCancel()}>
<DialogContent className={`max-h-[calc(100dvh-2rem)] overflow-y-auto sm:max-w-[800px] ${ABOVE_ANTD_MODAL}`}>
<DialogHeader>
<DialogTitle>Add prebuilt pattern</DialogTitle>
</DialogHeader>
return (
<Select.OptGroup key={category} label={category}>
{categoryPatterns.map((pattern) => (
<Option key={pattern.name} value={pattern.name}>
{pattern.display_name}
</Option>
))}
</Select.OptGroup>
);
})}
</Select>
<div className="space-y-6">
<div>
<p className="font-semibold">Pattern type</p>
<Combobox
items={patternGroups}
value={selectedPattern}
onValueChange={(pattern: PrebuiltPattern | null) => pattern && onPatternNameChange(pattern.name)}
itemToStringLabel={(pattern: PrebuiltPattern) => pattern.display_name}
filter={matchesPatternQuery}
>
<ComboboxInput className="mt-2 w-full" placeholder="Choose pattern type" />
<ComboboxContent>
<ComboboxEmpty>No matching patterns</ComboboxEmpty>
<ComboboxList>
{(group: PatternGroup) => (
<ComboboxGroup key={group.category} items={group.items}>
<ComboboxLabel>{group.category}</ComboboxLabel>
<ComboboxCollection>
{(pattern: PrebuiltPattern) => (
<ComboboxItem key={pattern.name} value={pattern}>
{pattern.display_name}
</ComboboxItem>
)}
</ComboboxCollection>
</ComboboxGroup>
)}
</ComboboxList>
</ComboboxContent>
</Combobox>
</div>
<div>
<p className="font-semibold">Action</p>
<p className="mt-1 mb-2 text-muted-foreground">
Choose what action the guardrail should take when this pattern is detected
</p>
<Select
items={ACTION_ITEMS}
value={patternAction}
onValueChange={(value: string | null) => value && onActionChange(value as "BLOCK" | "MASK")}
>
<SelectTrigger className="w-full" aria-label="Action">
<SelectValue />
</SelectTrigger>
<SelectContent alignItemWithTrigger={false}>
{ACTION_ITEMS.map((item) => (
<SelectItem key={item.value} value={item.value}>
{item.label}
</SelectItem>
))}
</SelectContent>
</Select>
</div>
</div>
<div>
<Text strong>Action</Text>
<Text type="secondary" style={{ display: "block", marginTop: 4, marginBottom: 8 }}>
Choose what action the guardrail should take when this pattern is detected
</Text>
<Select value={patternAction} onChange={onActionChange} style={{ width: "100%" }}>
<Option value="BLOCK">Block</Option>
<Option value="MASK">Mask</Option>
</Select>
</div>
</Space>
<div style={{ display: "flex", justifyContent: "flex-end", gap: "8px", marginTop: "24px" }}>
<Button onClick={onCancel}>Cancel</Button>
<Button type="primary" onClick={onAdd}>
Add
</Button>
</div>
</Modal>
<DialogFooter>
<Button variant="outline" onClick={onCancel}>
Cancel
</Button>
<Button onClick={onAdd}>Add</Button>
</DialogFooter>
</DialogContent>
</Dialog>
);
};

42
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternTable.tsx
View file

@ -1,11 +1,11 @@
import React from "react";
import { Typography, Select, Tag, Button } from "antd";
import { DeleteOutlined } from "@ant-design/icons";
import { Trash2 } from "lucide-react";
import type { ColumnDef } from "@tanstack/react-table";
import { DataTable } from "@/components/shared/DataTable";
const { Text } = Typography;
const { Option } = Select;
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { ACTION_ITEMS } from "./action_options";
interface Pattern {
id: string;
@ -28,11 +28,7 @@ const PatternTable: React.FC<PatternTableProps> = ({ patterns, onActionChange, o
header: "Type",
accessorKey: "type",
size: 100,
cell: ({ row }) => (
<Tag color={row.original.type === "prebuilt" ? "blue" : "green"}>
{row.original.type === "prebuilt" ? "Prebuilt" : "Custom"}
</Tag>
),
cell: ({ row }) => <Badge variant="secondary">{row.original.type === "prebuilt" ? "Prebuilt" : "Custom"}</Badge>,
},
{
header: "Pattern name",
@ -44,9 +40,7 @@ const PatternTable: React.FC<PatternTableProps> = ({ patterns, onActionChange, o
accessorKey: "pattern",
cell: ({ row }) =>
row.original.pattern ? (
<Text code style={{ fontSize: 12 }}>
{row.original.pattern.substring(0, 40)}...
</Text>
<code className="rounded-sm bg-muted px-1 py-0.5 text-xs">{row.original.pattern.substring(0, 40)}...</code>
) : (
"-"
),
@ -57,13 +51,20 @@ const PatternTable: React.FC<PatternTableProps> = ({ patterns, onActionChange, o
size: 150,
cell: ({ row }) => (
<Select
items={ACTION_ITEMS}
value={row.original.action}
onChange={(value) => onActionChange(row.original.id, value as "BLOCK" | "MASK")}
style={{ width: 120 }}
size="small"
onValueChange={(value: string | null) => value && onActionChange(row.original.id, value as "BLOCK" | "MASK")}
>
<Option value="BLOCK">Block</Option>
<Option value="MASK">Mask</Option>
<SelectTrigger size="sm" className="w-[120px]" aria-label="Action">
<SelectValue />
</SelectTrigger>
<SelectContent alignItemWithTrigger={false}>
{ACTION_ITEMS.map((item) => (
<SelectItem key={item.value} value={item.value}>
{item.label}
</SelectItem>
))}
</SelectContent>
</Select>
),
},
@ -72,7 +73,8 @@ const PatternTable: React.FC<PatternTableProps> = ({ patterns, onActionChange, o
id: "actions",
size: 100,
cell: ({ row }) => (
<Button type="text" danger size="small" icon={<DeleteOutlined />} onClick={() => onRemove(row.original.id)}>
<Button variant="ghost" size="sm" onClick={() => onRemove(row.original.id)}>
<Trash2 />
Delete
</Button>
),
@ -80,7 +82,7 @@ const PatternTable: React.FC<PatternTableProps> = ({ patterns, onActionChange, o
];
if (patterns.length === 0) {
return <div style={{ textAlign: "center", padding: "40px 0", color: "#999" }}>No patterns added.</div>;
return <div className="py-10 text-center text-muted-foreground">No patterns added.</div>;
}
return <DataTable data={patterns} columns={columns} getRowId={(row) => row.id} size="compact" />;

10
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/action_options.ts Normal file
View file

@ -0,0 +1,10 @@
export const ACTION_ITEMS = [
{ value: "BLOCK", label: "Block" },
{ value: "MASK", label: "Mask" },
] as const;
export const SEVERITY_ITEMS = [
{ value: "high", label: "High" },
{ value: "medium", label: "Medium" },
{ value: "low", label: "Low" },
] as const;

1
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/dialog_layering.ts Normal file
View file

@ -0,0 +1 @@
export const ABOVE_ANTD_MODAL = "z-[1100]";

173
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.test.tsx Normal file
View file

@ -0,0 +1,173 @@
import { describe, it, expect, vi, beforeEach } from "vitest";
import { render, screen, waitFor } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import CustomCodeModal from "./CustomCodeModal";
import { createGuardrailCall, updateGuardrailCall, testCustomCodeGuardrail } from "@/components/networking";
vi.mock("@/components/networking", () => ({
createGuardrailCall: vi.fn(),
updateGuardrailCall: vi.fn(),
testCustomCodeGuardrail: vi.fn(),
}));
vi.mock("@/components/molecules/notifications_manager", () => ({
default: { success: vi.fn(), error: vi.fn(), fromBackend: vi.fn() },
}));
const mockCreate = vi.mocked(createGuardrailCall);
const mockUpdate = vi.mocked(updateGuardrailCall);
const mockTest = vi.mocked(testCustomCodeGuardrail);
describe("CustomCodeModal", () => {
const onClose = vi.fn();
const onSuccess = vi.fn();
const renderModal = (overrides = {}) =>
render(<CustomCodeModal visible onClose={onClose} onSuccess={onSuccess} accessToken="test-token" {...overrides} />);
beforeEach(() => {
vi.clearAllMocks();
mockCreate.mockResolvedValue({} as never);
mockUpdate.mockResolvedValue({} as never);
});
it("should render the create heading and the editor scaffolding", async () => {
renderModal();
expect(await screen.findByText("Create Custom Guardrail")).toBeInTheDocument();
expect(screen.getByText("Define custom logic using Python-like syntax")).toBeInTheDocument();
expect(screen.getByPlaceholderText("e.g., block-pii-custom")).toBeInTheDocument();
expect(screen.getByText("Guardrail Name")).toBeInTheDocument();
expect(screen.getByText("Mode (can select multiple)")).toBeInTheDocument();
expect(screen.getByText("Available Primitives")).toBeInTheDocument();
expect(screen.getByText("Python Logic")).toBeInTheDocument();
expect(screen.getByRole("button", { name: /save guardrail/i })).toBeInTheDocument();
});
it("should seed the editor with the empty template", async () => {
renderModal();
const editor = await screen.findByDisplayValue(/async def apply_guardrail/);
expect(editor).toBeInTheDocument();
});
it("should not render its content when not visible", () => {
renderModal({ visible: false });
expect(screen.queryByText("Create Custom Guardrail")).not.toBeInTheDocument();
});
it("should render the edit heading and existing values in edit mode", async () => {
renderModal({
editData: {
guardrail_id: "g-1",
guardrail_name: "existing-guardrail",
litellm_params: { mode: "post_call", default_on: true, custom_code: "def apply_guardrail(): pass" },
},
});
expect(await screen.findByText("Edit Custom Guardrail")).toBeInTheDocument();
expect(screen.getByDisplayValue("existing-guardrail")).toBeInTheDocument();
expect(screen.getByDisplayValue("def apply_guardrail(): pass")).toBeInTheDocument();
expect(screen.getByRole("button", { name: /update guardrail/i })).toBeInTheDocument();
});
it("should keep save disabled until a guardrail name is entered", async () => {
const user = userEvent.setup();
renderModal();
await user.click(await screen.findByRole("button", { name: /save guardrail/i }));
expect(mockCreate).not.toHaveBeenCalled();
await user.type(screen.getByPlaceholderText("e.g., block-pii-custom"), "my-guardrail");
await user.click(screen.getByRole("button", { name: /save guardrail/i }));
await waitFor(() => {
expect(mockCreate).toHaveBeenCalledTimes(1);
});
});
it("should create the guardrail with the entered name, mode and code", async () => {
const user = userEvent.setup();
renderModal();
await user.type(await screen.findByPlaceholderText("e.g., block-pii-custom"), "block-pii");
await user.click(screen.getByRole("button", { name: /save guardrail/i }));
await waitFor(() => {
expect(mockCreate).toHaveBeenCalled();
});
const [token, payload] = mockCreate.mock.calls[0] as [string, Record<string, never>];
expect(token).toBe("test-token");
expect(payload).toMatchObject({
guardrail_name: "block-pii",
litellm_params: { guardrail: "custom_code", mode: ["pre_call"], default_on: false },
});
await waitFor(() => {
expect(onSuccess).toHaveBeenCalled();
});
});
it("should switch the editor contents when a template is chosen", async () => {
const user = userEvent.setup();
renderModal();
expect(await screen.findByDisplayValue(/async def apply_guardrail/)).toBeInTheDocument();
const comboboxes = screen.getAllByRole("combobox");
await user.click(comboboxes[comboboxes.length - 1]);
const options = await screen.findAllByText("Block SSN");
await user.click(options[options.length - 1]);
expect(await screen.findByDisplayValue(/SSN detected/)).toBeInTheDocument();
});
it("should narrow the mode options to the typed search text", async () => {
const user = userEvent.setup();
renderModal();
await user.click(screen.getAllByRole("combobox")[0]);
await user.keyboard("mcp");
expect(await screen.findByText("pre_mcp_call (Before MCP Tool Call)")).toBeInTheDocument();
expect(screen.queryByText("logging_only")).not.toBeInTheDocument();
});
it("should expand the test section and run a test against the backend", async () => {
const user = userEvent.setup();
mockTest.mockResolvedValue({ success: true, result: { action: "allow" } } as never);
renderModal();
await user.click(await screen.findByText("Test Your Guardrail"));
const runButton = await screen.findByRole("button", { name: /run test/i });
await user.click(runButton);
await waitFor(() => {
expect(mockTest).toHaveBeenCalled();
});
expect(await screen.findByText("Allowed")).toBeInTheDocument();
});
it("should surface a backend test error", async () => {
const user = userEvent.setup();
mockTest.mockResolvedValue({ success: false, error: "boom", error_type: "SyntaxError" } as never);
renderModal();
await user.click(await screen.findByText("Test Your Guardrail"));
await user.click(await screen.findByRole("button", { name: /run test/i }));
expect(await screen.findByText("boom")).toBeInTheDocument();
expect(screen.getByText("[SyntaxError]")).toBeInTheDocument();
});
it("should cancel through the footer button", async () => {
const user = userEvent.setup();
renderModal();
await user.click(await screen.findByRole("button", { name: "Cancel" }));
expect(onClose).toHaveBeenCalledTimes(1);
});
});

403
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx
View file

@ -1,21 +1,34 @@
import React, { useState, useRef, useEffect } from "react";
import { Modal, Select, Switch, Collapse, Input, Divider } from "antd";
import { Button, TextInput } from "@tremor/react";
import {
CodeOutlined,
PlayCircleOutlined,
CheckCircleOutlined,
CloseCircleOutlined,
CaretRightOutlined,
SaveOutlined,
UsergroupAddOutlined,
ExportOutlined,
} from "@ant-design/icons";
import { CheckCircle2, ChevronRight, Code, ExternalLink, PlayCircle, Save, Users, XCircle } from "lucide-react";
import { createGuardrailCall, updateGuardrailCall, testCustomCodeGuardrail } from "@/components/networking";
import NotificationsManager from "@/components/molecules/notifications_manager";
const { Panel } = Collapse;
const { TextArea } = Input;
import { Button } from "@/components/ui/button";
import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
import {
Combobox,
ComboboxChip,
ComboboxChips,
ComboboxChipsInput,
ComboboxContent,
ComboboxEmpty,
ComboboxItem,
ComboboxList,
} from "@/components/ui/combobox";
import { Dialog, DialogContent, DialogDescription, DialogHeader, DialogTitle } from "@/components/ui/dialog";
import { Input } from "@/components/ui/input";
import {
Select,
SelectContent,
SelectGroup,
SelectItem,
SelectLabel,
SelectSeparator,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
import { Switch } from "@/components/ui/switch";
import { Textarea } from "@/components/ui/textarea";
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
// Code templates
const CODE_TEMPLATES = {
@ -144,6 +157,17 @@ const MODE_OPTIONS = [
{ value: "during_mcp_call", label: "during_mcp_call (During MCP Tool Call)" },
];
const TEMPLATE_ITEMS = Object.entries(CODE_TEMPLATES).map(([key, template]) => ({
value: key,
label: template.name,
}));
type ModeOption = (typeof MODE_OPTIONS)[number];
const MODE_OPTION_BY_VALUE: Record<string, ModeOption> = Object.fromEntries(
MODE_OPTIONS.map((option) => [option.value, option]),
);
// Data for editing an existing guardrail
export interface EditGuardrailData {
guardrail_id: string;
@ -470,105 +494,104 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
};
const lineCount = code.split("\n").length;
const selectedModeOptions = mode.map((value) => MODE_OPTION_BY_VALUE[value]).filter(Boolean);
return (
<Modal
open={visible}
onCancel={onClose}
footer={null}
width={1400}
className="custom-code-modal"
closable={true}
destroyOnClose
>
<div className="flex flex-col h-[80vh]">
{/* Header */}
<div className="pb-4 border-b border-gray-200">
<h2 className="text-xl font-semibold text-gray-900">
<Dialog open={visible} onOpenChange={(open) => !open && onClose()}>
<DialogContent className="max-h-[calc(100dvh-2rem)] overflow-y-auto sm:max-w-[1400px]">
<DialogHeader>
<DialogTitle className="text-xl font-semibold">
{isEditMode ? "Edit Custom Guardrail" : "Create Custom Guardrail"}
</h2>
<p className="text-sm text-gray-500 mt-1">Define custom logic using Python-like syntax</p>
</div>
</DialogTitle>
<DialogDescription>Define custom logic using Python-like syntax</DialogDescription>
</DialogHeader>
{/* Top Controls */}
<div className="flex items-center gap-4 py-4 border-b border-gray-100">
<div className="flex-1 max-w-[200px]">
<label className="block text-xs font-medium text-gray-600 mb-1">Guardrail Name</label>
<TextInput value={guardrailName} onValueChange={setGuardrailName} placeholder="e.g., block-pii-custom" />
</div>
<div className="w-[280px]">
<label className="block text-xs font-medium text-gray-600 mb-1">Mode (can select multiple)</label>
<Select
mode="multiple"
value={mode}
onChange={setMode}
options={MODE_OPTIONS}
className="w-full"
size="middle"
placeholder="Select modes"
<div className="flex items-center gap-4 border-b border-border py-4">
<div className="max-w-[200px] flex-1">
<label className="mb-1 block text-xs font-medium text-muted-foreground">Guardrail Name</label>
<Input
value={guardrailName}
onChange={(e) => setGuardrailName(e.target.value)}
placeholder="e.g., block-pii-custom"
/>
</div>
<div className="w-[180px]">
<label className="block text-xs font-medium text-gray-600 mb-1">Template</label>
<Select
value={selectedTemplate}
onChange={handleTemplateChange}
className="w-full"
size="middle"
dropdownRender={(menu) => (
<>
{menu}
<Divider style={{ margin: "8px 0" }} />
<div
style={{
padding: "8px 12px",
cursor: "pointer",
color: "#1890ff",
fontSize: "12px",
display: "flex",
alignItems: "center",
gap: "4px",
}}
onClick={(e) => {
e.preventDefault();
window.open("https://models.litellm.ai/guardrails", "_blank");
}}
onMouseEnter={(e) => {
e.currentTarget.style.backgroundColor = "#f0f0f0";
}}
onMouseLeave={(e) => {
e.currentTarget.style.backgroundColor = "transparent";
}}
>
<UsergroupAddOutlined />
<span>Browse Community templates</span>
<ExportOutlined style={{ fontSize: "10px" }} />
</div>
</>
)}
<div className="w-[280px]">
<label className="mb-1 block text-xs font-medium text-muted-foreground">Mode (can select multiple)</label>
<Combobox
items={MODE_OPTIONS}
value={selectedModeOptions}
onValueChange={(options: ModeOption[]) => setMode(options.map((option) => option.value))}
multiple
>
<Select.OptGroup label="STANDARD">
{Object.entries(CODE_TEMPLATES).map(([key, template]) => (
<Select.Option key={key} value={key}>
{template.name}
</Select.Option>
<ComboboxChips className="w-full">
{selectedModeOptions.map((option) => (
<ComboboxChip key={option.value} aria-label={option.label}>
{option.label}
</ComboboxChip>
))}
</Select.OptGroup>
<ComboboxChipsInput
className="border-0 bg-transparent"
placeholder={mode.length === 0 ? "Select modes" : undefined}
/>
</ComboboxChips>
<ComboboxContent>
<ComboboxEmpty>No matching modes</ComboboxEmpty>
<ComboboxList>
{(option: ModeOption) => (
<ComboboxItem key={option.value} value={option}>
{option.label}
</ComboboxItem>
)}
</ComboboxList>
</ComboboxContent>
</Combobox>
</div>
<div className="w-[180px]">
<label className="mb-1 block text-xs font-medium text-muted-foreground">Template</label>
<Select
items={TEMPLATE_ITEMS}
value={selectedTemplate}
onValueChange={(value: string | null) => value && handleTemplateChange(value)}
>
<SelectTrigger className="w-full" aria-label="Template">
<SelectValue />
</SelectTrigger>
<SelectContent alignItemWithTrigger={false}>
<SelectGroup>
<SelectLabel>STANDARD</SelectLabel>
{TEMPLATE_ITEMS.map((template) => (
<SelectItem key={template.value} value={template.value}>
{template.label}
</SelectItem>
))}
</SelectGroup>
<SelectSeparator />
<button
type="button"
onClick={() => window.open("https://models.litellm.ai/guardrails", "_blank")}
className="flex w-full items-center gap-1 rounded-sm px-2 py-1.5 text-xs text-primary hover:bg-accent"
>
<Users className="size-3.5" />
<span>Browse Community templates</span>
<ExternalLink className="size-2.5" />
</button>
</SelectContent>
</Select>
</div>
<div className="flex items-center gap-2 pt-5">
<span className="text-sm text-gray-600">Default On</span>
<Switch checked={defaultOn} onChange={setDefaultOn} />
<span className="text-sm text-muted-foreground">Default On</span>
<Switch checked={defaultOn} onCheckedChange={setDefaultOn} aria-label="Default On" />
</div>
</div>
{/* Main Content */}
<div className="flex flex-1 overflow-hidden mt-4 gap-6">
<div className="mt-4 flex gap-6">
{/* Code Editor */}
<div className="flex-2 flex flex-col min-w-0 overflow-y-auto">
<div className="flex items-center justify-between mb-2 shrink-0">
<span className="text-xs font-semibold text-gray-500 uppercase tracking-wide">Python Logic</span>
<span className="text-xs text-gray-400">Restricted environment (no imports)</span>
<div className="flex min-w-0 flex-1 flex-col">
<div className="mb-2 flex shrink-0 items-center justify-between">
<span className="text-xs font-semibold tracking-wide text-muted-foreground uppercase">Python Logic</span>
<span className="text-xs text-muted-foreground">Restricted environment (no imports)</span>
</div>
<div
className="relative rounded-lg overflow-hidden border border-gray-700 bg-[#1e1e1e] shrink-0"
@ -607,27 +630,23 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
</div>
{/* Test Section */}
<Collapse
activeKey={testExpanded ? ["test"] : []}
onChange={(keys) => setTestExpanded(keys.includes("test"))}
className="mt-3 bg-white border border-gray-200 rounded-lg shrink-0"
expandIcon={({ isActive }) => <CaretRightOutlined rotate={isActive ? 90 : 0} />}
<Collapsible
open={testExpanded}
onOpenChange={setTestExpanded}
className="mt-3 shrink-0 rounded-lg border border-border"
>
<Panel
header={
<span className="flex items-center gap-2 text-sm font-medium">
<PlayCircleOutlined className="text-blue-500" />
Test Your Guardrail
</span>
}
key="test"
>
<CollapsibleTrigger className="flex w-full items-center gap-2 p-3 text-sm font-medium">
<ChevronRight className={`size-4 transition-transform ${testExpanded ? "rotate-90" : ""}`} />
<PlayCircle className="size-4 text-muted-foreground" />
Test Your Guardrail
</CollapsibleTrigger>
<CollapsibleContent className="p-3 pt-0">
<div className="space-y-3">
<div>
<div className="flex items-center justify-between mb-2">
<label className="block text-xs font-medium text-gray-600">Test Input (JSON)</label>
<label className="block text-xs font-medium text-muted-foreground">Test Input (JSON)</label>
<div className="flex items-center gap-2">
<span className="text-xs text-gray-500">Load example:</span>
<span className="text-xs text-muted-foreground">Load example:</span>
<button
type="button"
onClick={() => setTestInput(JSON.stringify(TEST_INPUT_EXAMPLES.pre_call.data, null, 2))}
@ -651,7 +670,7 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
</button>
</div>
</div>
<div className="mb-2 p-2 bg-gray-50 rounded-sm text-xs text-gray-600 border border-gray-200">
<div className="mb-2 rounded-sm border border-border bg-muted/40 p-2 text-xs text-muted-foreground">
<div className="grid grid-cols-2 gap-x-4 gap-y-1">
<div>
<strong>texts</strong>: Message content (always)
@ -676,16 +695,17 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
</div>
</div>
</div>
<TextArea
<Textarea
value={testInput}
onChange={(e) => setTestInput(e.target.value)}
rows={8}
className="font-mono text-xs"
className="font-mono text-xs field-sizing-fixed"
placeholder='{"texts": ["test message"], ...}'
/>
</div>
<div className="flex items-center gap-3">
<Button size="xs" onClick={handleTest} disabled={isTesting} icon={PlayCircleOutlined}>
<Button size="sm" onClick={handleTest} disabled={isTesting} aria-busy={isTesting}>
{isTesting ? <UiLoadingSpinner className="size-4" /> : <PlayCircle />}
{isTesting ? "Running..." : "Run Test"}
</Button>
{testResult && (
@ -702,7 +722,7 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
>
{testResult.error ? (
<>
<CloseCircleOutlined />
<XCircle className="size-4" />
<span>
{testResult.error_type && <span className="font-medium">[{testResult.error_type}] </span>}
{testResult.error}
@ -710,140 +730,117 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
</>
) : testResult.action === "allow" ? (
<>
<CheckCircleOutlined /> Allowed
<CheckCircle2 className="size-4" /> Allowed
</>
) : testResult.action === "block" ? (
<>
<CloseCircleOutlined /> Blocked: {testResult.reason}
<XCircle className="size-4" /> Blocked: {testResult.reason}
</>
) : testResult.action === "modify" ? (
<>
<CheckCircleOutlined /> Modified
<CheckCircle2 className="size-4" /> Modified
{testResult.texts && testResult.texts.length > 0 && (
<span className="text-xs text-gray-500 ml-1">
→ {testResult.texts[0].substring(0, 50)}
<span className="ml-1 text-xs text-muted-foreground">
-&gt; {testResult.texts[0].substring(0, 50)}
{testResult.texts[0].length > 50 ? "..." : ""}
</span>
)}
</>
) : (
<>
<CheckCircleOutlined /> {testResult.action || "Unknown"}
<CheckCircle2 className="size-4" /> {testResult.action || "Unknown"}
</>
)}
</div>
)}
</div>
</div>
</Panel>
</Collapse>
</CollapsibleContent>
</Collapsible>
{/* Contribution CTA Banner */}
<div className="mt-3 p-4 bg-linear-to-r from-blue-50 to-indigo-50 border border-blue-200 rounded-lg flex items-center justify-between shrink-0">
<div className="mt-3 flex shrink-0 items-center justify-between rounded-lg border border-blue-200 bg-linear-to-r from-blue-50 to-indigo-50 p-4">
<div className="flex items-center gap-3">
<div className="bg-blue-100 rounded-full p-2">
<UsergroupAddOutlined className="text-blue-600 text-lg" />
<div className="rounded-full bg-blue-100 p-2">
<Users className="size-5 text-blue-600" />
</div>
<div>
<div className="text-sm font-medium text-gray-900">Built a useful guardrail?</div>
<div className="text-xs text-gray-600">Share it with the community and help others build faster</div>
<div className="text-sm font-medium">Built a useful guardrail?</div>
<div className="text-xs text-muted-foreground">
Share it with the community and help others build faster
</div>
</div>
</div>
<Button
size="xs"
onClick={() => window.open("https://github.com/BerriAI/litellm-guardrails", "_blank")}
icon={ExportOutlined}
className="bg-blue-600 hover:bg-blue-700 text-white border-0"
>
<Button size="sm" onClick={() => window.open("https://github.com/BerriAI/litellm-guardrails", "_blank")}>
<ExternalLink />
Contribute Template
</Button>
</div>
</div>
{/* Primitives Panel */}
<div className="w-[300px] shrink-0 overflow-auto border-l border-gray-200 pl-6">
<div className="flex items-center gap-2 mb-3">
<CodeOutlined className="text-blue-500" />
<span className="font-semibold text-gray-700">Available Primitives</span>
<div className="w-[300px] shrink-0 overflow-auto border-l border-border pl-6">
<div className="mb-3 flex items-center gap-2">
<Code className="size-4 text-muted-foreground" />
<span className="font-semibold">Available Primitives</span>
</div>
<p className="text-xs text-gray-500 mb-3">Click to copy functions to clipboard</p>
<p className="mb-3 text-xs text-muted-foreground">Click to copy functions to clipboard</p>
<Collapse
defaultActiveKey={["Return Values"]}
className="primitives-collapse bg-transparent border-0"
expandIconPosition="end"
>
<div className="space-y-2">
{Object.entries(PRIMITIVES).map(([category, primitives]) => (
<Panel
header={<span className="text-sm font-medium text-gray-700">{category}</span>}
<Collapsible
key={category}
className="bg-white mb-2 rounded-lg border border-gray-200"
defaultOpen={category === "Return Values"}
className="rounded-lg border border-border"
>
<div className="space-y-2">
{primitives.map((p) => (
<button
key={p.name}
onClick={() => copyPrimitive(p.name)}
className={`w-full text-left px-2 py-2 rounded transition-colors ${
copiedPrimitive === p.name ? "bg-green-100" : "bg-gray-50 hover:bg-blue-50"
}`}
>
{copiedPrimitive === p.name ? (
<span className="flex items-center gap-1 text-xs font-mono text-green-700">
<CheckCircleOutlined /> Copied!
</span>
) : (
<>
<div className="text-xs font-mono text-gray-800">{p.name}</div>
<div className="text-[10px] text-gray-500 mt-0.5">{p.desc}</div>
</>
)}
</button>
))}
</div>
</Panel>
<CollapsibleTrigger className="group flex w-full items-center justify-between px-3 py-2 text-sm font-medium">
{category}
<ChevronRight className="size-4 transition-transform group-data-panel-open:rotate-90" />
</CollapsibleTrigger>
<CollapsibleContent className="px-3 pb-3">
<div className="space-y-2">
{primitives.map((p) => (
<button
key={p.name}
onClick={() => copyPrimitive(p.name)}
className={`w-full rounded-sm px-2 py-2 text-left transition-colors ${
copiedPrimitive === p.name ? "bg-accent" : "bg-muted/40 hover:bg-accent"
}`}
>
{copiedPrimitive === p.name ? (
<span className="flex items-center gap-1 font-mono text-xs">
<CheckCircle2 className="size-3.5" /> Copied!
</span>
) : (
<>
<div className="font-mono text-xs">{p.name}</div>
<div className="mt-0.5 text-[10px] text-muted-foreground">{p.desc}</div>
</>
)}
</button>
))}
</div>
</CollapsibleContent>
</Collapsible>
))}
</Collapse>
</div>
</div>
</div>
{/* Footer */}
<div className="flex items-center justify-between pt-4 mt-4 border-t border-gray-200">
<span className="text-xs text-gray-400">Changes are auto-saved to local draft</span>
<div className="mt-4 flex items-center justify-between border-t border-border pt-4">
<span className="text-xs text-muted-foreground">Changes are auto-saved to local draft</span>
<div className="flex items-center gap-3">
<Button variant="secondary" onClick={onClose}>
Cancel
</Button>
<Button
onClick={handleSave}
loading={isSaving}
disabled={isSaving || !guardrailName.trim()}
icon={SaveOutlined}
>
<Button onClick={handleSave} disabled={isSaving || !guardrailName.trim()} aria-busy={isSaving}>
{isSaving ? <UiLoadingSpinner className="size-4" /> : <Save />}
{isEditMode ? "Update Guardrail" : "Save Guardrail"}
</Button>
</div>
</div>
</div>
<style>{`
.custom-code-modal .ant-modal-content {
padding: 24px;
}
.custom-code-modal .ant-modal-close {
top: 20px;
right: 20px;
}
.primitives-collapse .ant-collapse-item {
border: none !important;
}
.primitives-collapse .ant-collapse-header {
padding: 8px 12px !important;
}
.primitives-collapse .ant-collapse-content-box {
padding: 8px 12px !important;
}
`}</style>
</Modal>
</DialogContent>
</Dialog>
);
};

104
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.test.tsx Normal file
View file

@ -0,0 +1,104 @@
import { describe, it, expect, vi, beforeEach } from "vitest";
import { render, screen } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import GuardrailGarden from "./guardrail_garden";
import { ALL_CARDS } from "./guardrail_garden_data";
vi.mock("./guardrail_garden_detail", () => ({
__esModule: true,
default: ({ card, onBack }: { card: { name: string }; onBack: () => void }) => (
<div>
<span>Detail for {card.name}</span>
<button onClick={onBack}>Back to garden</button>
</div>
),
}));
const LITELLM_CARDS = ALL_CARDS.filter((c) => c.category === "litellm");
const PARTNER_CARDS = ALL_CARDS.filter((c) => c.category === "partner");
describe("GuardrailGarden", () => {
beforeEach(() => {
vi.clearAllMocks();
});
const renderGarden = () => render(<GuardrailGarden accessToken="test-token" onGuardrailCreated={vi.fn()} />);
it("should render both sections with their descriptions", () => {
renderGarden();
expect(screen.getByText("LiteLLM Content Filter")).toBeInTheDocument();
expect(
screen.getByText(
"Built-in guardrails powered by LiteLLM. Zero latency, no external dependencies, no additional cost.",
),
).toBeInTheDocument();
expect(screen.getByText("Partner Guardrails")).toBeInTheDocument();
expect(
screen.getByText("Third-party guardrail integrations from leading AI security providers."),
).toBeInTheDocument();
});
it("should show a capped set of litellm cards behind a show all toggle", async () => {
const user = userEvent.setup();
renderGarden();
expect(screen.getByText(`Show all (${LITELLM_CARDS.length})`)).toBeInTheDocument();
expect(screen.getByText(LITELLM_CARDS[0].name)).toBeInTheDocument();
expect(screen.queryByText(LITELLM_CARDS[LITELLM_CARDS.length - 1].name)).not.toBeInTheDocument();
await user.click(screen.getByText(`Show all (${LITELLM_CARDS.length})`));
expect(screen.getByText("Show less")).toBeInTheDocument();
expect(screen.getByText(LITELLM_CARDS[LITELLM_CARDS.length - 1].name)).toBeInTheDocument();
});
it("should always render every partner card", () => {
renderGarden();
PARTNER_CARDS.forEach((card) => {
expect(screen.getByText(card.name)).toBeInTheDocument();
});
});
it("should filter cards by the search query", async () => {
const user = userEvent.setup();
renderGarden();
const target = PARTNER_CARDS[0];
await user.type(screen.getByPlaceholderText("Search guardrails"), target.name);
expect(await screen.findByText(target.name)).toBeInTheDocument();
const otherPartner = PARTNER_CARDS.find((c) => c.name !== target.name);
if (otherPartner) {
expect(screen.queryByText(otherPartner.name)).not.toBeInTheDocument();
}
});
it("should show an empty result set for a query that matches nothing", async () => {
const user = userEvent.setup();
renderGarden();
await user.type(screen.getByPlaceholderText("Search guardrails"), "zzzzznotaguardrailzzzzz");
expect(screen.getByText("Show all (0)")).toBeInTheDocument();
PARTNER_CARDS.forEach((card) => {
expect(screen.queryByText(card.name)).not.toBeInTheDocument();
});
});
it("should open the detail view for a clicked card and return to the garden", async () => {
const user = userEvent.setup();
renderGarden();
const target = PARTNER_CARDS[0];
await user.click(screen.getByText(target.name));
expect(await screen.findByText(`Detail for ${target.name}`)).toBeInTheDocument();
expect(screen.queryByPlaceholderText("Search guardrails")).not.toBeInTheDocument();
await user.click(screen.getByRole("button", { name: "Back to garden" }));
expect(await screen.findByPlaceholderText("Search guardrails")).toBeInTheDocument();
});
});

68
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.tsx
View file

@ -1,6 +1,6 @@
import React, { useState } from "react";
import { Input } from "antd";
import { SearchOutlined, ArrowRightOutlined } from "@ant-design/icons";
import { ArrowRight, Search } from "lucide-react";
import { InputGroup, InputGroupAddon, InputGroupInput } from "@/components/ui/input-group";
import { GuardrailCardInfo, ALL_CARDS } from "./guardrail_garden_data";
import GuardrailCard from "./guardrail_garden_card";
import GuardrailDetailView from "./guardrail_garden_detail";
@ -43,72 +43,52 @@ const GuardrailGarden: React.FC<GuardrailGardenProps> = ({ accessToken, onGuardr
return (
<div>
{/* Search Bar */}
<div style={{ marginBottom: 24 }}>
<Input
size="large"
placeholder="Search guardrails"
prefix={<SearchOutlined style={{ color: "#9ca3af" }} />}
value={searchQuery}
onChange={(e) => setSearchQuery(e.target.value)}
style={{ borderRadius: 8 }}
/>
<div className="mb-6">
<InputGroup>
<InputGroupAddon>
<Search className="size-4 text-muted-foreground" />
</InputGroupAddon>
<InputGroupInput
placeholder="Search guardrails"
value={searchQuery}
onChange={(e) => setSearchQuery(e.target.value)}
/>
</InputGroup>
</div>
{/* LiteLLM Content Filter Section */}
<div style={{ marginBottom: 40 }}>
<div style={{ display: "flex", alignItems: "center", justifyContent: "space-between", marginBottom: 4 }}>
<h2 style={{ fontSize: 20, fontWeight: 600, color: "#111827", margin: 0 }}>LiteLLM Content Filter</h2>
<div className="mb-10">
<div className="mb-1 flex items-center justify-between">
<h2 className="m-0 text-xl font-semibold text-foreground">LiteLLM Content Filter</h2>
<span
style={{
display: "inline-flex",
alignItems: "center",
gap: 6,
fontSize: 14,
color: "#1a73e8",
cursor: "pointer",
}}
className="inline-flex cursor-pointer items-center gap-1.5 text-sm text-primary"
onClick={() => setShowAllLitellm(!showAllLitellm)}
>
{showAllLitellm ? (
<>Show less</>
) : (
<>
<ArrowRightOutlined style={{ fontSize: 12 }} />
<ArrowRight className="size-3" />
{`Show all (${litellmCards.length})`}
</>
)}
</span>
</div>
<p style={{ fontSize: 13, color: "#6b7280", margin: "4px 0 20px 0" }}>
<p className="mt-1 mb-5 text-[13px] text-muted-foreground">
Built-in guardrails powered by LiteLLM. Zero latency, no external dependencies, no additional cost.
</p>
<div
style={{
display: "grid",
gridTemplateColumns: "repeat(auto-fill, minmax(220px, 1fr))",
gap: 16,
}}
>
<div className="grid grid-cols-[repeat(auto-fill,minmax(220px,1fr))] gap-4">
{(showAllLitellm ? litellmCards : litellmCards.slice(0, CARDS_PER_ROW * VISIBLE_ROWS)).map((card) => (
<GuardrailCard key={card.id} card={card} onClick={() => setSelectedCard(card)} />
))}
</div>
</div>
{/* Partner Guardrails Section */}
<div style={{ marginBottom: 40 }}>
<h2 style={{ fontSize: 20, fontWeight: 600, color: "#111827", margin: "0 0 4px 0" }}>Partner Guardrails</h2>
<p style={{ fontSize: 13, color: "#6b7280", margin: "4px 0 20px 0" }}>
<div className="mb-10">
<h2 className="mt-0 mb-1 text-xl font-semibold text-foreground">Partner Guardrails</h2>
<p className="mt-1 mb-5 text-[13px] text-muted-foreground">
Third-party guardrail integrations from leading AI security providers.
</p>
<div
style={{
display: "grid",
gridTemplateColumns: "repeat(auto-fill, minmax(220px, 1fr))",
gap: 16,
}}
>
<div className="grid grid-cols-[repeat(auto-fill,minmax(220px,1fr))] gap-4">
{partnerCards.map((card) => (
<GuardrailCard key={card.id} card={card} onClick={() => setSelectedCard(card)} />
))}

38
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_card.tsx
View file

@ -1,45 +1,25 @@
import React, { useState } from "react";
import { CheckCircleFilled } from "@ant-design/icons";
import React from "react";
import { CircleCheck } from "lucide-react";
import { GuardrailCardInfo } from "./guardrail_garden_data";
import { Logo } from "@/components/molecules/logo/Logo";
const GuardrailCard: React.FC<{ card: GuardrailCardInfo; onClick: () => void }> = ({ card, onClick }) => {
const [hovered, setHovered] = useState(false);
return (
<div
onClick={onClick}
onMouseEnter={() => setHovered(true)}
onMouseLeave={() => setHovered(false)}
style={{
borderRadius: 12,
border: hovered ? "1px solid #93c5fd" : "1px solid #e5e7eb",
backgroundColor: "#ffffff",
padding: "20px 20px 16px 20px",
cursor: "pointer",
transition: "border-color 0.15s, box-shadow 0.15s",
display: "flex",
flexDirection: "column",
minHeight: 170,
boxShadow: hovered ? "0 1px 6px rgba(59,130,246,0.08)" : "none",
}}
className="flex min-h-[170px] cursor-pointer flex-col rounded-xl border border-border bg-card px-5 pt-5 pb-4 transition-[border-color,box-shadow] hover:border-primary/40 hover:shadow-sm"
>
{/* Icon + Name row */}
<div style={{ display: "flex", alignItems: "center", gap: 10, marginBottom: 10 }}>
<div className="mb-2.5 flex items-center gap-2.5">
<Logo src={card.logo} label={card.name} className="w-7 h-7 rounded-md object-contain shrink-0" />
<span style={{ fontSize: 14, fontWeight: 600, color: "#111827", lineHeight: 1.3 }}>{card.name}</span>
<span className="text-sm leading-tight font-semibold text-foreground">{card.name}</span>
</div>
{/* Description */}
<p className="line-clamp-3" style={{ fontSize: 12, color: "#6b7280", lineHeight: 1.6, margin: 0, flex: 1 }}>
{card.description}
</p>
<p className="line-clamp-3 m-0 flex-1 text-xs leading-relaxed text-muted-foreground">{card.description}</p>
{/* Eval badge */}
{card.eval && (
<div style={{ marginTop: 10, display: "flex", alignItems: "center", gap: 4 }}>
<CheckCircleFilled style={{ color: "#16a34a", fontSize: 12 }} />
<span style={{ fontSize: 11, color: "#16a34a", fontWeight: 500 }}>
<div className="mt-2.5 flex items-center gap-1 text-emerald-600">
<CircleCheck className="size-3" />
<span className="text-[11px] font-medium">
F1: {card.eval.f1}% &middot; {card.eval.testCases} test cases
</span>
</div>

31
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_detail.tsx
View file

@ -1,6 +1,6 @@
import React, { useState } from "react";
import { Button } from "antd";
import { ArrowLeftOutlined } from "@ant-design/icons";
import { ArrowLeft } from "lucide-react";
import { Button } from "@/components/ui/button";
import AddGuardrailForm from "./add_guardrail_form";
import { Logo } from "@/components/molecules/logo/Logo";
import { GUARDRAIL_PRESETS } from "./guardrail_garden_configs";
@ -44,17 +44,9 @@ const GuardrailDetailView: React.FC<GuardrailDetailViewProps> = ({ card, onBack,
{/* Back link */}
<div
onClick={onBack}
style={{
display: "inline-flex",
alignItems: "center",
gap: 6,
color: "#5f6368",
cursor: "pointer",
fontSize: 14,
marginBottom: 24,
}}
className="mb-6 inline-flex cursor-pointer items-center gap-1.5 text-sm text-muted-foreground"
>
<ArrowLeftOutlined style={{ fontSize: 11 }} />
<ArrowLeft className="size-3" />
<span>{card.name}</span>
</div>
@ -67,19 +59,8 @@ const GuardrailDetailView: React.FC<GuardrailDetailViewProps> = ({ card, onBack,
<p style={{ fontSize: 14, color: "#5f6368", margin: "0 0 20px 0", lineHeight: 1.6 }}>{card.description}</p>
{/* Action buttons — outlined style like Vertex */}
<div style={{ display: "flex", gap: 10, marginBottom: 32 }}>
<Button
onClick={() => setIsAddFormVisible(true)}
style={{
borderRadius: 20,
padding: "4px 20px",
height: 36,
borderColor: "#dadce0",
color: "#1a73e8",
fontWeight: 500,
fontSize: 14,
}}
>
<div className="mb-8 flex gap-2.5">
<Button variant="outline" className="rounded-full" onClick={() => setIsAddFormVisible(true)}>
Create Guardrail
</Button>
</div>

229
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.tsx
View file

@ -1,10 +1,21 @@
import React from "react";
import { Typography, Select, Button, Checkbox, Tooltip, Tag } from "antd";
import { CloseOutlined, EyeInvisibleOutlined, StopOutlined, FilterOutlined } from "@ant-design/icons";
import { EyeOff, Filter, Info, Ban, X } from "lucide-react";
import { PiiEntityCategory } from "@/components/guardrails/types";
const { Text } = Typography;
const { Option } = Select;
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { Checkbox } from "@/components/ui/checkbox";
import {
Combobox,
ComboboxChip,
ComboboxChips,
ComboboxChipsInput,
ComboboxContent,
ComboboxEmpty,
ComboboxItem,
ComboboxList,
} from "@/components/ui/combobox";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
// Helper functions
export const formatEntityName = (name: string) => {
@ -14,9 +25,9 @@ export const formatEntityName = (name: string) => {
export const getActionIcon = (action: string) => {
switch (action) {
case "MASK":
return <EyeInvisibleOutlined style={{ marginRight: 4 }} />;
return <EyeOff className="mr-1 size-3.5" />;
case "BLOCK":
return <StopOutlined style={{ marginRight: 4 }} />;
return <Ban className="mr-1 size-3.5" />;
default:
return null;
}
@ -30,34 +41,37 @@ export interface CategoryFilterProps {
}
export const CategoryFilter: React.FC<CategoryFilterProps> = ({ categories, selectedCategories, onChange }) => {
const categoryNames = categories.map((cat) => cat.category);
return (
<div>
<div className="flex items-center mb-2">
<FilterOutlined className="text-gray-500 mr-1" />
<Text className="text-gray-500 font-medium">Filter by category</Text>
<div className="mb-2 flex items-center">
<Filter className="mr-1 size-4 text-muted-foreground" />
<span className="font-medium text-muted-foreground">Filter by category</span>
</div>
<Select
mode="multiple"
placeholder="Select categories to filter by"
style={{ width: "100%" }}
onChange={onChange}
value={selectedCategories}
allowClear
showSearch
optionFilterProp="children"
className="mb-4"
tagRender={(props) => (
<Tag color="blue" closable={props.closable} onClose={props.onClose} className="mr-2 mb-2">
{props.label}
</Tag>
)}
>
{categories.map((cat) => (
<Option key={cat.category} value={cat.category}>
{cat.category}
</Option>
))}
</Select>
<Combobox items={categoryNames} value={selectedCategories} onValueChange={onChange} multiple>
<ComboboxChips className="mb-4 w-full">
{selectedCategories.map((category) => (
<ComboboxChip key={category} aria-label={category}>
{category}
</ComboboxChip>
))}
<ComboboxChipsInput
className="border-0 bg-transparent"
placeholder={selectedCategories.length === 0 ? "Select categories to filter by" : undefined}
/>
</ComboboxChips>
<ComboboxContent>
<ComboboxEmpty>No matching categories</ComboboxEmpty>
<ComboboxList>
{(category: string) => (
<ComboboxItem key={category} value={category}>
{category}
</ComboboxItem>
)}
</ComboboxList>
</ComboboxContent>
</Combobox>
</div>
);
};
@ -71,46 +85,34 @@ export interface QuickActionsProps {
export const QuickActions: React.FC<QuickActionsProps> = ({ onSelectAll, onUnselectAll, hasSelectedEntities }) => {
return (
<div className="bg-gray-50 p-5 rounded-lg mb-6 border border-gray-200 shadow-xs">
<div className="flex items-center justify-between mb-4">
<div className="mb-6 rounded-lg border border-border bg-muted/40 p-5 shadow-xs">
<div className="mb-4 flex items-center justify-between">
<div className="flex items-center">
<Text strong className="text-gray-700 text-base">
Quick Actions
</Text>
<Tooltip title="Apply action to all PII types at once">
<div className="ml-2 text-gray-400 cursor-help text-xs">ⓘ</div>
<span className="text-base font-semibold">Quick Actions</span>
<Tooltip>
<TooltipTrigger
render={
<span className="ml-2 cursor-help text-muted-foreground">
<Info className="size-3.5" />
</span>
}
/>
<TooltipContent>Apply action to all PII types at once</TooltipContent>
</Tooltip>
</div>
<Button
color="danger"
variant="outlined"
onClick={onUnselectAll}
disabled={!hasSelectedEntities}
icon={<CloseOutlined />}
>
<Button variant="outline" onClick={onUnselectAll} disabled={!hasSelectedEntities}>
<X />
Unselect All
</Button>
</div>
<div className="grid grid-cols-2 gap-4">
<Button
color="primary"
variant="outlined"
onClick={() => onSelectAll("MASK")}
className="h-10"
block
icon={<EyeInvisibleOutlined />}
>
Select All & Mask
<Button variant="outline" className="h-10 w-full" onClick={() => onSelectAll("MASK")}>
<EyeOff />
Select All &amp; Mask
</Button>
<Button
color="danger"
variant="outlined"
onClick={() => onSelectAll("BLOCK")}
className="h-10 hover:bg-red-100"
block
icon={<StopOutlined />}
>
Select All & Block
<Button variant="outline" className="h-10 w-full" onClick={() => onSelectAll("BLOCK")}>
<Ban />
Select All &amp; Block
</Button>
</div>
</div>
@ -138,62 +140,59 @@ export const PiiEntityList: React.FC<PiiEntityListProps> = ({
entityToCategoryMap,
}) => {
return (
<div className="border rounded-lg overflow-hidden shadow-xs">
<div className="bg-gray-50 px-5 py-3 border-b flex">
<Text strong className="flex-1 text-gray-700">
PII Type
</Text>
<Text strong className="w-32 text-right text-gray-700">
Action
</Text>
<div className="overflow-hidden rounded-lg border border-border shadow-xs">
<div className="flex border-b border-border bg-muted/40 px-5 py-3">
<span className="flex-1 font-semibold">PII Type</span>
<span className="w-32 text-right font-semibold">Action</span>
</div>
<div className="max-h-[400px] overflow-y-auto">
{entities.length === 0 ? (
<div className="py-10 text-center text-gray-500">No PII types match your filter criteria</div>
<div className="py-10 text-center text-muted-foreground">No PII types match your filter criteria</div>
) : (
entities.map((entity) => (
<div
key={entity}
className={`px-5 py-3 flex items-center justify-between hover:bg-gray-50 border-b ${
selectedEntities.includes(entity) ? "bg-blue-50" : ""
}`}
>
<div className="flex items-center flex-1">
<Checkbox
checked={selectedEntities.includes(entity)}
onChange={() => onEntitySelect(entity)}
className="mr-3"
/>
<Text className={selectedEntities.includes(entity) ? "font-medium text-gray-900" : "text-gray-700"}>
{formatEntityName(entity)}
</Text>
{entityToCategoryMap.get(entity) && (
<Tag className="ml-2 text-xs" color="blue">
{entityToCategoryMap.get(entity)}
</Tag>
)}
entities.map((entity) => {
const isSelected = selectedEntities.includes(entity);
return (
<div
key={entity}
className={`flex items-center justify-between border-b border-border px-5 py-3 hover:bg-muted/40 ${
isSelected ? "bg-accent" : ""
}`}
>
<div className="flex flex-1 items-center">
<Checkbox className="mr-3" checked={isSelected} onCheckedChange={() => onEntitySelect(entity)} />
<span className={isSelected ? "font-medium text-foreground" : "text-muted-foreground"}>
{formatEntityName(entity)}
</span>
{entityToCategoryMap.get(entity) && (
<Badge variant="secondary" className="ml-2">
{entityToCategoryMap.get(entity)}
</Badge>
)}
</div>
<div className="w-32">
<Select
value={isSelected ? selectedActions[entity] || "MASK" : "MASK"}
onValueChange={(value: string | null) => value && onActionSelect(entity, value)}
disabled={!isSelected}
>
<SelectTrigger className={`w-[120px] ${isSelected ? "" : "opacity-50"}`} aria-label="Action">
<SelectValue />
</SelectTrigger>
<SelectContent alignItemWithTrigger={false}>
{actions.map((action) => (
<SelectItem key={action} value={action}>
<span className="flex items-center">
{getActionIcon(action)}
{action}
</span>
</SelectItem>
))}
</SelectContent>
</Select>
</div>
</div>
<div className="w-32">
<Select
value={selectedEntities.includes(entity) ? selectedActions[entity] || "MASK" : "MASK"}
onChange={(value) => onActionSelect(entity, value)}
style={{ width: 120 }}
disabled={!selectedEntities.includes(entity)}
className={`${!selectedEntities.includes(entity) ? "opacity-50" : ""}`}
dropdownMatchSelectWidth={false}
>
{actions.map((action) => (
<Option key={action} value={action}>
<div className="flex items-center">
{getActionIcon(action)}
{action}
</div>
</Option>
))}
</Select>
</div>
</div>
))
);
})
)}
</div>
</div>

9
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.tsx
View file

@ -1,10 +1,7 @@
import { Typography } from "antd";
import React, { useState } from "react";
import { CategoryFilter, PiiEntityList, QuickActions } from "./pii_components";
import { PiiConfigurationProps } from "@/components/guardrails/types";
const { Title, Text } = Typography;
/**
* A reusable component for rendering PII entity selection and action configuration
* Used in both add and edit guardrail forms
@ -57,11 +54,9 @@ const PiiConfiguration: React.FC<PiiConfigurationProps> = ({
<div className="pii-configuration">
<div className="flex justify-between items-center mb-5">
<div className="flex items-center">
<Title level={4} className="m-0! font-semibold text-gray-800">
Configure PII Protection
</Title>
<h4 className="m-0 text-lg font-semibold text-foreground">Configure PII Protection</h4>
</div>
<Text className="text-gray-500">{selectedEntities.length} items selected</Text>
<span className="text-muted-foreground">{selectedEntities.length} items selected</span>
</div>
<div className="mb-6">

348
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/tool_permission/ToolPermissionRulesEditor.tsx
View file

@ -1,7 +1,12 @@
import React from "react";
import { Card, Text } from "@tremor/react";
import { Button, Divider, Empty, Input, Select, Space, Tooltip } from "antd";
import { InfoCircleOutlined, PlusOutlined, DeleteOutlined } from "@ant-design/icons";
import { Info, Plus, Trash2 } from "lucide-react";
import { Button } from "@/components/ui/button";
import { Card, CardContent } from "@/components/ui/card";
import { Input } from "@/components/ui/input";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { Separator } from "@/components/ui/separator";
import { Textarea } from "@/components/ui/textarea";
import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
export type ToolPermissionDecision = "allow" | "deny";
export type ToolPermissionDefaultAction = "allow" | "deny";
@ -28,6 +33,16 @@ interface ToolPermissionRulesEditorProps {
disabled?: boolean;
}
const DECISION_ITEMS = [
{ value: "allow", label: "Allow" },
{ value: "deny", label: "Deny" },
] as const;
const ON_DISALLOWED_ITEMS = [
{ value: "block", label: "Block" },
{ value: "rewrite", label: "Rewrite" },
] as const;
const DEFAULT_CONFIG: ToolPermissionConfig = {
rules: [],
default_action: "deny",
@ -115,8 +130,9 @@ const ToolPermissionRulesEditor: React.FC<ToolPermissionRulesEditorProps> = ({ v
if (entries.length === 0) {
return (
<Button
variant="outline"
disabled={disabled}
size="small"
size="sm"
onClick={() => updateRule(index, { allowed_param_patterns: { "": "" } })}
>
+ Restrict tool arguments (optional)
@ -126,9 +142,9 @@ const ToolPermissionRulesEditor: React.FC<ToolPermissionRulesEditorProps> = ({ v
return (
<div className="space-y-2">
<Text className="text-sm text-gray-600">Argument constraints (dot or array paths)</Text>
<p className="text-sm text-muted-foreground">Argument constraints (dot or array paths)</p>
{entries.map(([path, pattern], patternIndex) => (
<Space key={`${rule.id || index}-${patternIndex}`} align="start">
<div key={`${rule.id || index}-${patternIndex}`} className="flex items-start gap-2">
<Input
disabled={disabled}
placeholder="messages[0].content"
@ -142,20 +158,24 @@ const ToolPermissionRulesEditor: React.FC<ToolPermissionRulesEditorProps> = ({ v
onChange={(e) => updateAllowedParamPattern(index, patternIndex, e.target.value)}
/>
<Button
variant="outline"
size="icon"
aria-label="Remove constraint"
disabled={disabled}
icon={<DeleteOutlined />}
danger
onClick={() =>
updateAllowedParamEntries(index, (entries) => {
entries.splice(patternIndex, 1);
})
}
/>
</Space>
>
<Trash2 />
</Button>
</div>
))}
<Button
variant="outline"
disabled={disabled}
size="small"
size="sm"
onClick={() =>
updateRule(index, {
allowed_param_patterns: {
@ -173,148 +193,186 @@ const ToolPermissionRulesEditor: React.FC<ToolPermissionRulesEditorProps> = ({ v
return (
<Card>
<div className="flex items-center justify-between">
<div>
<Text className="text-lg font-semibold">LiteLLM Tool Permission Guardrail</Text>
<Text className="text-sm text-gray-500">
Provide regex patterns (e.g., ^mcp__github_.*$) for tool names or types and optionally constrain payload
fields.
</Text>
<CardContent>
<div className="flex items-center justify-between">
<div>
<p className="text-lg font-semibold">LiteLLM Tool Permission Guardrail</p>
<p className="text-sm text-muted-foreground">
Provide regex patterns (e.g., ^mcp__github_.*$) for tool names or types and optionally constrain payload
fields.
</p>
</div>
{!disabled && (
<Button onClick={addRule}>
<Plus />
Add Rule
</Button>
)}
</div>
{!disabled && (
<Button
icon={<PlusOutlined />}
type="primary"
onClick={addRule}
className="bg-blue-600! text-white! hover:bg-blue-500!"
>
Add Rule
</Button>
<Separator className="my-4" />
{config.rules.length === 0 ? (
<div className="py-10 text-center text-muted-foreground">No tool rules added yet</div>
) : (
<div className="space-y-4">
{config.rules.map((rule, index) => (
<Card key={rule.id || index} className="bg-muted/40">
<CardContent>
<div className="mb-3 flex items-center justify-between">
<p className="font-semibold">Rule {index + 1}</p>
<Button variant="ghost" disabled={disabled} onClick={() => removeRule(index)}>
<Trash2 />
Remove
</Button>
</div>
<div className="grid grid-cols-1 gap-4 md:grid-cols-2">
<div>
<p className="text-sm font-medium">Rule ID</p>
<Input
disabled={disabled}
placeholder="unique_rule_id"
value={rule.id}
onChange={(e) => updateRule(index, { id: e.target.value })}
/>
</div>
<div>
<p className="text-sm font-medium">Tool Name (optional)</p>
<Input
disabled={disabled}
placeholder="^mcp__github_.*$"
value={rule.tool_name ?? ""}
onChange={(e) =>
updateRule(index, {
tool_name: e.target.value.trim() === "" ? undefined : e.target.value,
})
}
/>
</div>
</div>
<div className="mt-4 grid grid-cols-1 gap-4 md:grid-cols-2">
<div>
<p className="text-sm font-medium">Tool Type (optional)</p>
<Input
disabled={disabled}
placeholder="^function$"
value={rule.tool_type ?? ""}
onChange={(e) =>
updateRule(index, {
tool_type: e.target.value.trim() === "" ? undefined : e.target.value,
})
}
/>
</div>
</div>
<div className="mt-4 flex flex-col gap-2">
<p className="text-sm font-medium">Decision</p>
<Select
items={DECISION_ITEMS}
disabled={disabled}
value={rule.decision}
onValueChange={(value: string | null) =>
value && updateRule(index, { decision: value as ToolPermissionDecision })
}
>
<SelectTrigger className="w-[200px]" aria-label="Decision">
<SelectValue />
</SelectTrigger>
<SelectContent alignItemWithTrigger={false}>
{DECISION_ITEMS.map((item) => (
<SelectItem key={item.value} value={item.value}>
{item.label}
</SelectItem>
))}
</SelectContent>
</Select>
</div>
<div className="mt-4">{renderAllowedParamPatterns(rule, index)}</div>
</CardContent>
</Card>
))}
</div>
)}
</div>
<Divider />
<Separator className="my-4" />
{config.rules.length === 0 ? (
<Empty description="No tool rules added yet" />
) : (
<div className="space-y-4">
{config.rules.map((rule, index) => (
<Card key={rule.id || index} className="bg-gray-50">
<div className="flex items-center justify-between mb-3">
<Text className="font-semibold">Rule {index + 1}</Text>
<Button
icon={<DeleteOutlined />}
danger
type="text"
disabled={disabled}
onClick={() => removeRule(index)}
>
Remove
</Button>
</div>
<div className="grid grid-cols-1 gap-4 md:grid-cols-2">
<div>
<Text className="text-sm font-medium">Rule ID</Text>
<Input
disabled={disabled}
placeholder="unique_rule_id"
value={rule.id}
onChange={(e) => updateRule(index, { id: e.target.value })}
/>
</div>
<div>
<Text className="text-sm font-medium">Tool Name (optional)</Text>
<Input
disabled={disabled}
placeholder="^mcp__github_.*$"
value={rule.tool_name ?? ""}
onChange={(e) =>
updateRule(index, {
tool_name: e.target.value.trim() === "" ? undefined : e.target.value,
})
}
/>
</div>
</div>
<div className="grid grid-cols-1 gap-4 md:grid-cols-2 mt-4">
<div>
<Text className="text-sm font-medium">Tool Type (optional)</Text>
<Input
disabled={disabled}
placeholder="^function$"
value={rule.tool_type ?? ""}
onChange={(e) =>
updateRule(index, {
tool_type: e.target.value.trim() === "" ? undefined : e.target.value,
})
}
/>
</div>
</div>
<div className="mt-4 flex flex-col gap-2">
<Text className="text-sm font-medium">Decision</Text>
<Select
disabled={disabled}
value={rule.decision}
style={{ width: 200 }}
onChange={(value) => updateRule(index, { decision: value as ToolPermissionDecision })}
>
<Select.Option value="allow">Allow</Select.Option>
<Select.Option value="deny">Deny</Select.Option>
</Select>
</div>
<div className="mt-4">{renderAllowedParamPatterns(rule, index)}</div>
</Card>
))}
<div className="grid gap-4 md:grid-cols-2">
<div>
<p className="text-sm font-medium">Default action</p>
<Select
items={DECISION_ITEMS}
disabled={disabled}
value={config.default_action}
onValueChange={(value: string | null) =>
value && updateConfig({ default_action: value as ToolPermissionDefaultAction })
}
>
<SelectTrigger className="w-full" aria-label="Default action">
<SelectValue />
</SelectTrigger>
<SelectContent alignItemWithTrigger={false}>
{DECISION_ITEMS.map((item) => (
<SelectItem key={item.value} value={item.value}>
{item.label}
</SelectItem>
))}
</SelectContent>
</Select>
</div>
<div>
<p className="flex items-center gap-1 text-sm font-medium">
On disallowed action
<Tooltip>
<TooltipTrigger
render={
<span className="cursor-help text-muted-foreground">
<Info className="size-3.5" />
</span>
}
/>
<TooltipContent>
Block returns an error when a forbidden tool is invoked. Rewrite strips the tool call but lets the
rest of the response continue.
</TooltipContent>
</Tooltip>
</p>
<Select
items={ON_DISALLOWED_ITEMS}
disabled={disabled}
value={config.on_disallowed_action}
onValueChange={(value: string | null) =>
value && updateConfig({ on_disallowed_action: value as ToolPermissionOnDisallowedAction })
}
>
<SelectTrigger className="w-full" aria-label="On disallowed action">
<SelectValue />
</SelectTrigger>
<SelectContent alignItemWithTrigger={false}>
{ON_DISALLOWED_ITEMS.map((item) => (
<SelectItem key={item.value} value={item.value}>
{item.label}
</SelectItem>
))}
</SelectContent>
</Select>
</div>
</div>
)}
<Divider />
<div className="grid gap-4 md:grid-cols-2">
<div>
<Text className="text-sm font-medium">Default action</Text>
<Select
<div className="mt-4">
<p className="text-sm font-medium">Violation message (optional)</p>
<Textarea
className="field-sizing-fixed"
disabled={disabled}
value={config.default_action}
onChange={(value) => updateConfig({ default_action: value as ToolPermissionDefaultAction })}
>
<Select.Option value="allow">Allow</Select.Option>
<Select.Option value="deny">Deny</Select.Option>
</Select>
rows={3}
placeholder="This violates our org policy..."
value={config.violation_message_template}
onChange={(e) => updateConfig({ violation_message_template: e.target.value })}
/>
</div>
<div>
<Text className="text-sm font-medium flex items-center gap-1">
On disallowed action
<Tooltip title="Block returns an error when a forbidden tool is invoked. Rewrite strips the tool call but lets the rest of the response continue.">
<InfoCircleOutlined />
</Tooltip>
</Text>
<Select
disabled={disabled}
value={config.on_disallowed_action}
onChange={(value) => updateConfig({ on_disallowed_action: value as ToolPermissionOnDisallowedAction })}
>
<Select.Option value="block">Block</Select.Option>
<Select.Option value="rewrite">Rewrite</Select.Option>
</Select>
</div>
</div>
<div className="mt-4">
<Text className="text-sm font-medium">Violation message (optional)</Text>
<Input.TextArea
disabled={disabled}
rows={3}
placeholder="This violates our org policy..."
value={config.violation_message_template}
onChange={(e) => updateConfig({ violation_message_template: e.target.value })}
/>
</div>
</CardContent>
</Card>
);
};

4
ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.ts
View file

@ -17,6 +17,10 @@ export enum ConfigType {
*/
export enum GeneralSettingsFieldName {
MAXIMUM_SPEND_LOGS_RETENTION_PERIOD = "maximum_spend_logs_retention_period",
MAXIMUM_SPEND_LOGS_CLEANUP_BATCH_SIZE = "maximum_spend_logs_cleanup_batch_size",
MAXIMUM_SPEND_LOGS_CLEANUP_MAX_BATCHES = "maximum_spend_logs_cleanup_max_batches",
MAXIMUM_SPEND_LOGS_CLEANUP_RUN_BUDGET = "maximum_spend_logs_cleanup_run_budget",
MAXIMUM_SPEND_LOGS_CLEANUP_BATCH_TIMEOUT = "maximum_spend_logs_cleanup_batch_timeout",
// Add more field names here as needed
}

12
ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts
View file

@ -6,6 +6,10 @@ import { proxyConfigKeys } from "../proxyConfig/useProxyConfig";
export interface StoreRequestInSpendLogsParams {
store_prompts_in_spend_logs: boolean;
maximum_spend_logs_retention_period?: string;
maximum_spend_logs_cleanup_batch_size?: number;
maximum_spend_logs_cleanup_max_batches?: number;
maximum_spend_logs_cleanup_run_budget?: string;
maximum_spend_logs_cleanup_batch_timeout?: string;
}
export interface StoreRequestInSpendLogsResponse {
@ -19,6 +23,8 @@ const performStoreRequestInSpendLogs = async (
const proxyBaseUrl = getProxyBaseUrl();
const url = proxyBaseUrl ? `${proxyBaseUrl}/config/update` : `/config/update`;
const { store_prompts_in_spend_logs, ...optionalSettings } = params;
const response = await fetch(url, {
method: "POST",
headers: {
@ -27,10 +33,8 @@ const performStoreRequestInSpendLogs = async (
},
body: JSON.stringify({
general_settings: {
store_prompts_in_spend_logs: params.store_prompts_in_spend_logs,
...(params.maximum_spend_logs_retention_period && {
maximum_spend_logs_retention_period: params.maximum_spend_logs_retention_period,
}),
store_prompts_in_spend_logs,
...optionalSettings,
},
}),
});

320
ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.test.tsx Normal file
View file

@ -0,0 +1,320 @@
import { render, screen, waitFor, within } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { useState } from "react";
import { beforeEach, describe, expect, it, vi } from "vitest";
import AgentBuilderView from "./AgentBuilderView";
import type { AgentModel } from "../../llm_calls/fetch_agents";
const modelCreateCall = vi.fn().mockResolvedValue({ model_id: "id-new" });
const modelPatchUpdateCall = vi.fn().mockResolvedValue({});
const modelDeleteCall = vi.fn().mockResolvedValue({});
const keyCreateCall = vi.fn().mockResolvedValue({ key: "sk-agent-key" });
const fetchMCPServers = vi.fn().mockResolvedValue([]);
const fetchAvailableAgentModels = vi.fn();
const fetchAvailableModels = vi.fn().mockResolvedValue([{ model_group: "gpt-4o" }, { model_group: "claude-sonnet-4" }]);
vi.mock("@/components/networking", () => ({
proxyBaseUrl: "https://proxy.example.com",
modelCreateCall: (...args: unknown[]) => modelCreateCall(...args),
modelPatchUpdateCall: (...args: unknown[]) => modelPatchUpdateCall(...args),
modelDeleteCall: (...args: unknown[]) => modelDeleteCall(...args),
keyCreateCall: (...args: unknown[]) => keyCreateCall(...args),
fetchMCPServers: (...args: unknown[]) => fetchMCPServers(...args),
}));
vi.mock("../../llm_calls/fetch_agents", () => ({
fetchAvailableAgentModels: (...args: unknown[]) => fetchAvailableAgentModels(...args),
}));
vi.mock("@/components/llm_calls/fetch_models", () => ({
fetchAvailableModels: (...args: unknown[]) => fetchAvailableModels(...args),
}));
vi.mock("@/components/CodeBlock", () => ({
default: ({ code }: { code: string }) => <pre data-testid="code-block">{code}</pre>,
}));
vi.mock("@/components/molecules/notifications_manager", () => ({
default: { success: vi.fn(), fromBackend: vi.fn() },
}));
const StatefulPanel = ({ label }: { label: string }) => {
const [draft, setDraft] = useState("");
return <input aria-label={label} value={draft} onChange={(event) => setDraft(event.target.value)} />;
};
vi.mock("./ChatUI", () => ({
default: () => <StatefulPanel label="chat scratch" />,
}));
vi.mock("../complianceUI/ComplianceUI", () => ({
default: () => <StatefulPanel label="batch scratch" />,
}));
const AGENTS: AgentModel[] = [
{
model_name: "support-agent",
litellm_params: { model: "litellm_agent/gpt-4o", litellm_system_prompt: "Be helpful.", temperature: 0.3 },
model_info: { id: "agent-1" },
},
{
model_name: "research-agent",
litellm_params: { model: "litellm_agent/claude-sonnet-4" },
model_info: { id: "agent-2" },
},
];
const props = {
accessToken: "sk-access",
token: "tok",
userID: "u1",
userRole: "Admin",
};
const controlUnder = (label: string): HTMLElement =>
within(screen.getByText(label).parentElement!).getByRole("combobox");
const renderView = () => render(<AgentBuilderView {...props} />);
const waitForRoster = async () => {
await screen.findByRole("button", { name: "support-agent litellm_agent" });
};
beforeEach(() => {
vi.clearAllMocks();
fetchAvailableAgentModels.mockResolvedValue(AGENTS);
fetchAvailableModels.mockResolvedValue([{ model_group: "gpt-4o" }, { model_group: "claude-sonnet-4" }]);
fetchMCPServers.mockResolvedValue([]);
modelCreateCall.mockResolvedValue({ model_id: "id-new" });
Object.defineProperty(window, "matchMedia", {
writable: true,
value: (query: string) => ({
matches: false,
media: query,
onchange: null,
addListener: () => {},
removeListener: () => {},
addEventListener: () => {},
removeEventListener: () => {},
dispatchEvent: () => false,
}),
});
});
describe("AgentBuilderView", () => {
it("asks the visitor to sign in when there is no session", () => {
render(<AgentBuilderView accessToken={null} token={null} userID={null} userRole={null} />);
expect(screen.getByText("Sign in to use Agent Builder.")).toBeInTheDocument();
expect(fetchAvailableAgentModels).not.toHaveBeenCalled();
});
it("marks itself busy while the roster loads", async () => {
let release: (agents: AgentModel[]) => void = () => {};
fetchAvailableAgentModels.mockReturnValue(
new Promise<AgentModel[]>((resolve) => {
release = resolve;
}),
);
renderView();
expect(document.querySelector('[aria-busy="true"]')).toBeInTheDocument();
release(AGENTS);
await waitForRoster();
expect(document.querySelector('[aria-busy="true"]')).not.toBeInTheDocument();
});
it("lists every agent and opens the first one's configuration", async () => {
renderView();
await waitForRoster();
expect(screen.getByRole("button", { name: "research-agent litellm_agent" })).toBeInTheDocument();
expect(screen.getByText("Agent Builder")).toBeInTheDocument();
await waitFor(() => expect(screen.getByDisplayValue("support-agent")).toBeInTheDocument());
expect(screen.getByDisplayValue("Be helpful.")).toBeInTheDocument();
expect(screen.getByDisplayValue("0.3")).toBeInTheDocument();
});
it("loads the configuration of whichever agent is picked", async () => {
const user = userEvent.setup();
renderView();
await waitForRoster();
await user.click(screen.getByRole("button", { name: "research-agent litellm_agent" }));
await waitFor(() => expect(screen.getByDisplayValue("research-agent")).toBeInTheDocument());
});
it("offers a blank draft and a save control for a new agent", async () => {
const user = userEvent.setup();
renderView();
await waitForRoster();
await user.click(screen.getByRole("button", { name: /New agent/i }));
expect(screen.getByRole("button", { name: /Save Agent/i })).toBeInTheDocument();
expect(screen.getByDisplayValue("You are a helpful assistant.")).toBeInTheDocument();
expect(screen.queryByRole("button", { name: /Update Agent/i })).not.toBeInTheDocument();
});
it("creates the agent under the litellm_agent prefix", async () => {
const user = userEvent.setup();
renderView();
await waitForRoster();
await user.click(screen.getByRole("button", { name: /New agent/i }));
await user.type(screen.getByPlaceholderText("My Agent"), "billing-agent");
await user.click(screen.getByRole("button", { name: /Save Agent/i }));
await waitFor(() => expect(modelCreateCall).toHaveBeenCalled());
const payload = modelCreateCall.mock.calls[0][1];
expect(payload.model_name).toBe("billing-agent");
expect(payload.litellm_params.model).toBe("litellm_agent/gpt-4o");
});
it("will not save a draft without a name", async () => {
const user = userEvent.setup();
renderView();
await waitForRoster();
await user.click(screen.getByRole("button", { name: /New agent/i }));
await user.click(screen.getByRole("button", { name: /Save Agent/i }));
expect(modelCreateCall).not.toHaveBeenCalled();
});
it("updates the selected agent through its model id", async () => {
const user = userEvent.setup();
renderView();
await waitForRoster();
await screen.findByDisplayValue("support-agent");
await user.click(screen.getByRole("button", { name: /Update Agent/i }));
await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
expect(modelPatchUpdateCall.mock.calls[0][2]).toBe("agent-1");
});
it("deletes only after the warning is confirmed", async () => {
const user = userEvent.setup();
renderView();
await waitForRoster();
await screen.findByDisplayValue("support-agent");
await user.click(screen.getAllByRole("button", { name: /Delete$/ })[0]);
expect(await screen.findByText(/Are you sure you want to delete "support-agent"/)).toBeInTheDocument();
expect(modelDeleteCall).not.toHaveBeenCalled();
const confirmations = screen.getAllByRole("button", { name: /Delete$/ });
await user.click(confirmations[confirmations.length - 1]);
await waitFor(() => expect(modelDeleteCall).toHaveBeenCalledWith("sk-access", "agent-1"));
});
it("abandons the delete when the warning is dismissed", async () => {
const user = userEvent.setup();
renderView();
await waitForRoster();
await screen.findByDisplayValue("support-agent");
await user.click(screen.getAllByRole("button", { name: /Delete$/ })[0]);
await screen.findByText(/Are you sure you want to delete "support-agent"/);
await user.click(screen.getAllByRole("button", { name: /Cancel$/ }).at(-1)!);
await waitFor(() =>
expect(screen.queryByText(/Are you sure you want to delete "support-agent"/)).not.toBeInTheDocument(),
);
expect(modelDeleteCall).not.toHaveBeenCalled();
});
it("shows a ready-to-run curl example on the Connect tab", async () => {
const user = userEvent.setup();
renderView();
await waitForRoster();
await user.click(screen.getByRole("tab", { name: /Connect/i }));
const snippet = await screen.findByTestId("code-block");
expect(snippet).toHaveTextContent("https://proxy.example.com/v1/chat/completions");
expect(snippet).toHaveTextContent('"model": "support-agent"');
});
it("mints a key scoped to the selected agent", async () => {
const user = userEvent.setup();
renderView();
await waitForRoster();
await user.click(screen.getByRole("tab", { name: /Connect/i }));
await user.click(await screen.findByRole("button", { name: /Create key for this agent/i }));
await waitFor(() => expect(keyCreateCall).toHaveBeenCalled());
expect(keyCreateCall.mock.calls[0][2].models).toEqual(["support-agent"]);
expect(await screen.findByTestId("code-block")).toHaveTextContent("Bearer sk-agent-key");
});
it("keeps a tab's own state alive while the user works in another tab", async () => {
const user = userEvent.setup();
renderView();
await waitForRoster();
await user.click(screen.getByRole("tab", { name: /Chat/i }));
const scratch = await screen.findByLabelText("chat scratch");
await user.type(scratch, "half a thought");
expect(scratch).toHaveValue("half a thought");
await user.click(screen.getByRole("tab", { name: /Configure/i }));
await screen.findByDisplayValue("support-agent");
await user.click(screen.getByRole("tab", { name: /Chat/i }));
expect(await screen.findByLabelText("chat scratch")).toHaveValue("half a thought");
});
it("keeps the batch tab's state alive across a round trip too", async () => {
const user = userEvent.setup();
renderView();
await waitForRoster();
await user.click(screen.getByRole("tab", { name: /Batch Test/i }));
await user.type(await screen.findByLabelText("batch scratch"), "seven cases");
await user.click(screen.getByRole("tab", { name: /Connect/i }));
await screen.findByTestId("code-block");
await user.click(screen.getByRole("tab", { name: /Batch Test/i }));
expect(await screen.findByLabelText("batch scratch")).toHaveValue("seven cases");
});
it("attaches the MCP servers the agent should reach", async () => {
const user = userEvent.setup();
fetchMCPServers.mockResolvedValue([{ server_id: "srv-1", alias: "github", server_name: "github-mcp" }]);
renderView();
await waitForRoster();
await screen.findByDisplayValue("support-agent");
await user.click(controlUnder("MCP servers"));
const options = await screen.findAllByText("github");
await user.click(options[options.length - 1]);
await user.keyboard("{Escape}");
await user.click(screen.getByRole("button", { name: /Update Agent/i }));
await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
expect(modelPatchUpdateCall.mock.calls[0][1].litellm_params.tools).toEqual([
{ type: "mcp", server_label: "litellm", server_url: "litellm_proxy/mcp/github", require_approval: "never" },
]);
});
it("warns that the builder is experimental", async () => {
renderView();
await waitForRoster();
expect(screen.getByText(/Agent Builder is experimental/)).toBeInTheDocument();
expect(within(screen.getByText(/Agent Builder is experimental/)).getByRole("link")).toHaveAttribute(
"href",
"mailto:product@berri.ai",
);
});
});

529
ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.tsx
View file

@ -1,16 +1,24 @@
"use client";
import {
CommentOutlined,
DeleteOutlined,
ExperimentOutlined,
LinkOutlined,
PlusOutlined,
RobotOutlined,
SaveOutlined,
} from "@ant-design/icons";
import { Button, Input, Modal, Select, Spin, Tabs } from "antd";
import { Bot, FlaskConical, Link as LinkIcon, MessageSquare, Plus, Save, Trash2 } from "lucide-react";
import React, { useCallback, useEffect, useState } from "react";
import {
AlertDialog,
AlertDialogAction,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { Textarea } from "@/components/ui/textarea";
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
import { MultiSelect } from "@/components/shared/MultiSelect";
import { useVisitedTabs } from "@/hooks/useVisitedTabs";
import CodeBlock from "@/components/CodeBlock";
import NotificationsManager from "@/components/molecules/notifications_manager";
import {
@ -27,8 +35,6 @@ import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_m
import ComplianceUI from "../complianceUI/ComplianceUI";
import ChatUI from "./ChatUI";
const { TextArea } = Input;
export interface AgentBuilderViewProps {
accessToken: string | null;
token: string | null;
@ -45,6 +51,8 @@ export interface AgentBuilderViewProps {
const NEW_AGENT_ID = "__new__";
type AgentTab = "configure" | "chat" | "test" | "connect";
function getConnectTabBaseUrl(
proxySettings: AgentBuilderViewProps["proxySettings"],
customProxyBaseUrl?: string,
@ -116,7 +124,7 @@ function ConnectTabContent({
Create a virtual key that can only call this agent. The key will be scoped to you (user_id) and restricted to
the model <span className="font-mono text-gray-800">{agentName}</span>.
</p>
<Button type="primary" onClick={onCreateKey} loading={creatingKey} disabled={disabledPersonalKeyCreation}>
<Button onClick={onCreateKey} disabled={creatingKey || disabledPersonalKeyCreation}>
Create key for this agent
</Button>
{disabledPersonalKeyCreation && (
@ -190,7 +198,12 @@ export default function AgentBuilderView({
const [modelGroups, setModelGroups] = useState<ModelGroup[]>([]);
const [loadingAgents, setLoadingAgents] = useState(true);
const [selectedId, setSelectedId] = useState<string | null>(null);
const [activeTab, setActiveTab] = useState<"configure" | "chat" | "test" | "connect">("configure");
const [activeTab, setActiveTab] = useState<AgentTab>("configure");
const { onTabChange, hasVisited } = useVisitedTabs("configure");
const goToTab = (tab: AgentTab) => {
setActiveTab(tab);
onTabChange(tab);
};
const [creatingKey, setCreatingKey] = useState(false);
const [createdKeyValue, setCreatedKeyValue] = useState<string | null>(null);
@ -207,6 +220,7 @@ export default function AgentBuilderView({
const [saving, setSaving] = useState(false);
const [deleting, setDeleting] = useState(false);
const [confirmingDelete, setConfirmingDelete] = useState(false);
const effectiveApiKey = apiKey || accessToken || "";
const selectedAgent =
@ -314,7 +328,7 @@ export default function AgentBuilderView({
setDraftTemperature(0.7);
setDraftMaxTokens(4096);
setDraftTools([]);
setActiveTab("configure");
goToTab("configure");
};
const handleSaveAgent = async () => {
@ -344,7 +358,7 @@ export default function AgentBuilderView({
? list.find((a) => getAgentModelId(a) === createdId) ?? list.find((a) => a.model_name === draftName.trim())
: list.find((a) => a.model_name === draftName.trim());
setSelectedId(created ? getAgentSelectionKey(created) : list[0] ? getAgentSelectionKey(list[0]) : null);
setActiveTab("chat");
goToTab("chat");
} catch (e) {
NotificationsManager.fromBackend("Failed to save agent");
} finally {
@ -411,27 +425,24 @@ export default function AgentBuilderView({
const handleDeleteAgent = () => {
if (!selectedAgent || !selectedAgentModelId || !accessToken) return;
Modal.confirm({
title: "Delete agent",
content: `Are you sure you want to delete "${selectedAgent.model_name}"? This cannot be undone.`,
okText: "Delete",
okType: "danger",
cancelText: "Cancel",
onOk: async () => {
setDeleting(true);
try {
await modelDeleteCall(accessToken, selectedAgentModelId);
NotificationsManager.success("Agent deleted");
const list = await loadAgents();
const remaining = list.filter((a) => getAgentModelId(a) !== selectedAgentModelId);
setSelectedId(remaining.length > 0 ? getAgentSelectionKey(remaining[0]) : null);
} catch (e) {
NotificationsManager.fromBackend("Failed to delete agent");
} finally {
setDeleting(false);
}
},
});
setConfirmingDelete(true);
};
const handleConfirmDelete = async () => {
if (!selectedAgent || !selectedAgentModelId || !accessToken) return;
setDeleting(true);
try {
await modelDeleteCall(accessToken, selectedAgentModelId);
NotificationsManager.success("Agent deleted");
const list = await loadAgents();
const remaining = list.filter((a) => getAgentModelId(a) !== selectedAgentModelId);
setSelectedId(remaining.length > 0 ? getAgentSelectionKey(remaining[0]) : null);
} catch (e) {
NotificationsManager.fromBackend("Failed to delete agent");
} finally {
setDeleting(false);
setConfirmingDelete(false);
}
};
if (!accessToken || !userID || !userRole) {
@ -446,13 +457,8 @@ export default function AgentBuilderView({
<div className="flex h-12 items-center justify-between px-4">
<span className="text-sm font-medium text-gray-900">Agent Builder</span>
{isNewAgent ? (
<Button
type="primary"
icon={<SaveOutlined />}
onClick={handleSaveAgent}
loading={saving}
disabled={!draftName?.trim() || !draftUnderlyingModel}
>
<Button onClick={handleSaveAgent} disabled={saving || !draftName?.trim() || !draftUnderlyingModel}>
<Save />
Save Agent
</Button>
) : (
@ -460,7 +466,7 @@ export default function AgentBuilderView({
)}
</div>
<div className="flex items-center gap-2 border-t border-amber-200 bg-amber-50 px-4 py-2 text-xs text-amber-800">
<ExperimentOutlined className="shrink-0 text-amber-600" />
<FlaskConical className="size-4 shrink-0 text-amber-600" />
<span>
Agent Builder is experimental and may change or be removed without notice. We’d love your feedback—email us
at{" "}
@ -477,12 +483,14 @@ export default function AgentBuilderView({
<div className="w-60 shrink-0 border-r border-gray-200 bg-white flex flex-col">
<div className="flex items-center justify-between border-b border-gray-200 p-3">
<span className="text-xs font-semibold uppercase tracking-wide text-gray-500">Agents</span>
<Button type="text" size="small" icon={<PlusOutlined />} onClick={handleAddAgent} aria-label="Add agent" />
<Button variant="ghost" size="icon-sm" onClick={handleAddAgent} aria-label="Add agent">
<Plus />
</Button>
</div>
<div className="flex-1 overflow-y-auto p-2">
{loadingAgents ? (
<div className="flex justify-center py-4">
<Spin size="small" />
<div className="flex justify-center py-4" aria-busy="true">
<UiLoadingSpinner className="size-4 text-gray-400" />
</div>
) : (
<>
@ -509,7 +517,7 @@ export default function AgentBuilderView({
onClick={handleAddAgent}
className="mb-1 w-full rounded-md border border-dashed border-gray-300 px-3 py-2 text-left text-sm text-gray-500 hover:border-blue-400 hover:bg-blue-50/50 hover:text-gray-700"
>
<PlusOutlined className="mr-1" /> New agent
<Plus className="mr-1 inline size-4" /> New agent
</button>
</>
)}
@ -526,227 +534,230 @@ export default function AgentBuilderView({
{(selectedId !== null || isNewAgent) && (
<>
<Tabs
activeKey={activeTab}
onChange={(k) => setActiveTab(k as "configure" | "chat" | "test" | "connect")}
className="flex-1 overflow-hidden [&_.ant-tabs-content]:h-full [&_.ant-tabs-tabpane]:h-full [&_.ant-tabs-nav]:pl-4"
items={[
{
key: "configure",
label: (
<span>
<RobotOutlined className="mr-1" /> Configure
</span>
),
children: (
<div className="h-full overflow-y-auto p-6">
{isNewAgent || selectedAgent ? (
<div className="mx-auto max-w-xl space-y-4">
{!selectedAgentModelId && selectedAgent && (
<div className="rounded-sm border border-amber-200 bg-amber-50 px-3 py-2 text-xs text-amber-800">
This agent cannot be updated or deleted here (missing model id). Manage it from Models
&amp; Endpoints.
</div>
)}
<div>
<label className="mb-1 block text-sm font-medium text-gray-700">Agent name</label>
<Input
value={draftName}
onChange={(e) => setDraftName(e.target.value)}
placeholder="My Agent"
/>
</div>
<div>
<label className="mb-1 block text-sm font-medium text-gray-700">System prompt</label>
<TextArea
value={draftSystemPrompt}
onChange={(e) => setDraftSystemPrompt(e.target.value)}
placeholder="You are a helpful assistant..."
rows={6}
/>
</div>
<div>
<label className="mb-1 block text-sm font-medium text-gray-700">Underlying LLM</label>
<Select
value={draftUnderlyingModel}
onChange={setDraftUnderlyingModel}
className="w-full"
options={modelGroups.map((m) => ({ value: m.model_group, label: m.model_group }))}
placeholder="Select model"
/>
</div>
<div className="grid grid-cols-2 gap-4">
<div>
<label className="mb-1 block text-sm font-medium text-gray-700">Temperature</label>
<Input
type="number"
min={0}
max={2}
step={0.1}
value={draftTemperature}
onChange={(e) => setDraftTemperature(Number(e.target.value))}
/>
</div>
<div>
<label className="mb-1 block text-sm font-medium text-gray-700">Max tokens</label>
<Input
type="number"
min={1}
value={draftMaxTokens}
onChange={(e) => setDraftMaxTokens(Number(e.target.value))}
/>
</div>
</div>
<div>
<label className="mb-1 block text-sm font-medium text-gray-700">MCP servers</label>
<Select
mode="multiple"
placeholder="Select MCP servers to attach (same format as chat completions API)"
value={selectedMCPServerIds}
onChange={handleMCPServerChange}
loading={loadingMCPServers}
className="w-full"
allowClear
showSearch
optionFilterProp="label"
options={mcpServers.map((s) => ({
value: s.server_id,
label: s.alias || s.server_name || s.server_id,
}))}
/>
{selectedAgent && draftTools.length > 0 && (
<p className="mt-1 text-xs text-gray-500">
{draftTools.length} MCP server{draftTools.length !== 1 ? "s" : ""} saved. Use the same{" "}
<code className="rounded-sm bg-gray-100 px-1">tools</code> array in chat completions
when calling this agent.
</p>
)}
</div>
{selectedAgent && (
<div className="flex flex-wrap items-center gap-2 pt-2">
{selectedAgentModelId && (
<>
<Button
type="primary"
icon={<SaveOutlined />}
onClick={handleUpdateAgent}
loading={saving}
disabled={!draftName?.trim() || !draftUnderlyingModel}
>
Update Agent
</Button>
<Button
type="default"
danger
icon={<DeleteOutlined />}
onClick={handleDeleteAgent}
loading={deleting}
>
Delete
</Button>
</>
)}
<Button type="primary" icon={<CommentOutlined />} onClick={() => setActiveTab("chat")}>
Test in Chat
value={activeTab}
onValueChange={(value) => goToTab(value as AgentTab)}
className="flex flex-1 flex-col overflow-hidden"
>
<TabsList variant="line" className="h-auto w-full justify-start rounded-none border-b p-0 pl-4">
<TabsTrigger value="configure" className="flex-none rounded-none px-4 py-2">
<Bot />
Configure
</TabsTrigger>
<TabsTrigger value="chat" disabled={isNewAgent} className="flex-none rounded-none px-4 py-2">
<MessageSquare />
Chat
</TabsTrigger>
<TabsTrigger value="test" disabled={isNewAgent} className="flex-none rounded-none px-4 py-2">
<FlaskConical />
Batch Test
</TabsTrigger>
<TabsTrigger value="connect" disabled={isNewAgent} className="flex-none rounded-none px-4 py-2">
<LinkIcon />
Connect
</TabsTrigger>
</TabsList>
<TabsContent
value="configure"
keepMounted={hasVisited("configure")}
className="min-h-0 overflow-hidden"
>
<div className="h-full overflow-y-auto p-6">
{isNewAgent || selectedAgent ? (
<div className="mx-auto max-w-xl space-y-4">
{!selectedAgentModelId && selectedAgent && (
<div className="rounded-sm border border-amber-200 bg-amber-50 px-3 py-2 text-xs text-amber-800">
This agent cannot be updated or deleted here (missing model id). Manage it from Models &amp;
Endpoints.
</div>
)}
<div>
<label className="mb-1 block text-sm font-medium text-gray-700">Agent name</label>
<Input
value={draftName}
onChange={(e) => setDraftName(e.target.value)}
placeholder="My Agent"
/>
</div>
<div>
<label className="mb-1 block text-sm font-medium text-gray-700">System prompt</label>
<Textarea
value={draftSystemPrompt}
onChange={(e) => setDraftSystemPrompt(e.target.value)}
placeholder="You are a helpful assistant..."
rows={6}
className="field-sizing-fixed"
/>
</div>
<div>
<label className="mb-1 block text-sm font-medium text-gray-700">Underlying LLM</label>
<Select
value={draftUnderlyingModel ?? null}
onValueChange={(model: string | null) => setDraftUnderlyingModel(model ?? undefined)}
>
<SelectTrigger className="w-full" aria-label="Underlying LLM">
<SelectValue placeholder="Select model" />
</SelectTrigger>
<SelectContent>
{modelGroups.map((m) => (
<SelectItem key={m.model_group} value={m.model_group}>
{m.model_group}
</SelectItem>
))}
</SelectContent>
</Select>
</div>
<div className="grid grid-cols-2 gap-4">
<div>
<label className="mb-1 block text-sm font-medium text-gray-700">Temperature</label>
<Input
type="number"
min={0}
max={2}
step={0.1}
value={draftTemperature}
onChange={(e) => setDraftTemperature(Number(e.target.value))}
/>
</div>
<div>
<label className="mb-1 block text-sm font-medium text-gray-700">Max tokens</label>
<Input
type="number"
min={1}
value={draftMaxTokens}
onChange={(e) => setDraftMaxTokens(Number(e.target.value))}
/>
</div>
</div>
<div>
<label className="mb-1 block text-sm font-medium text-gray-700">MCP servers</label>
<MultiSelect
placeholder="Select MCP servers to attach (same format as chat completions API)"
value={selectedMCPServerIds}
onValueChange={handleMCPServerChange}
loading={loadingMCPServers}
className="w-full"
options={mcpServers.map((s) => ({
value: s.server_id,
label: s.alias || s.server_name || s.server_id,
}))}
/>
{selectedAgent && draftTools.length > 0 && (
<p className="mt-1 text-xs text-gray-500">
{draftTools.length} MCP server{draftTools.length !== 1 ? "s" : ""} saved. Use the same{" "}
<code className="rounded-sm bg-gray-100 px-1">tools</code> array in chat completions when
calling this agent.
</p>
)}
</div>
{selectedAgent && (
<div className="flex flex-wrap items-center gap-2 pt-2">
{selectedAgentModelId && (
<>
<Button
onClick={handleUpdateAgent}
disabled={saving || !draftName?.trim() || !draftUnderlyingModel}
>
<Save />
Update Agent
</Button>
</div>
<Button variant="destructive" onClick={handleDeleteAgent} disabled={deleting}>
<Trash2 />
Delete
</Button>
</>
)}
</div>
) : null}
</div>
),
},
{
key: "chat",
label: (
<span>
<CommentOutlined className="mr-1" /> Chat
</span>
),
disabled: isNewAgent,
children: (
<div className="flex h-full flex-col min-h-0">
{selectedAgent ? (
<ChatUI
key={selectedAgent.model_name}
simplified
fixedModel={selectedAgent.model_name}
accessToken={accessToken}
token={token}
userRole={userRole}
userID={userID}
disabledPersonalKeyCreation={disabledPersonalKeyCreation}
proxySettings={proxySettings}
/>
) : (
<div className="flex flex-1 items-center justify-center text-gray-500">
Save an agent first to test in Chat.
<Button onClick={() => goToTab("chat")}>
<MessageSquare />
Test in Chat
</Button>
</div>
)}
</div>
),
},
{
key: "test",
label: (
<span>
<ExperimentOutlined className="mr-1" /> Batch Test
</span>
),
disabled: isNewAgent,
children: (
<div className="flex h-full flex-col min-h-0">
{selectedAgent ? (
<ComplianceUI
accessToken={accessToken}
disabledPersonalKeyCreation={disabledPersonalKeyCreation}
backendMode="chat_completions"
fixedModel={selectedAgent.model_name}
proxySettings={proxySettings}
/>
) : (
<div className="flex flex-1 items-center justify-center text-gray-500">
Select an agent to run batch tests.
</div>
)}
) : null}
</div>
</TabsContent>
<TabsContent value="chat" keepMounted={hasVisited("chat")} className="min-h-0 overflow-hidden">
<div className="flex h-full flex-col min-h-0">
{selectedAgent ? (
<ChatUI
key={selectedAgent.model_name}
simplified
fixedModel={selectedAgent.model_name}
accessToken={accessToken}
token={token}
userRole={userRole}
userID={userID}
disabledPersonalKeyCreation={disabledPersonalKeyCreation}
proxySettings={proxySettings}
/>
) : (
<div className="flex flex-1 items-center justify-center text-gray-500">
Save an agent first to test in Chat.
</div>
),
},
{
key: "connect",
label: (
<span>
<LinkOutlined className="mr-1" /> Connect
</span>
),
disabled: isNewAgent,
children: (
<div className="h-full overflow-y-auto p-6">
{selectedAgent ? (
<ConnectTabContent
agentName={selectedAgent.model_name}
proxySettings={proxySettings}
customProxyBaseUrl={customProxyBaseUrl}
accessToken={accessToken}
userID={userID}
disabledPersonalKeyCreation={disabledPersonalKeyCreation}
creatingKey={creatingKey}
createdKeyValue={createdKeyValue}
onCreateKey={handleCreateKeyForAgent}
/>
) : (
<div className="flex flex-1 items-center justify-center text-gray-500">
Select an agent to see how to connect.
</div>
)}
)}
</div>
</TabsContent>
<TabsContent value="test" keepMounted={hasVisited("test")} className="min-h-0 overflow-hidden">
<div className="flex h-full flex-col min-h-0">
{selectedAgent ? (
<ComplianceUI
accessToken={accessToken}
disabledPersonalKeyCreation={disabledPersonalKeyCreation}
backendMode="chat_completions"
fixedModel={selectedAgent.model_name}
proxySettings={proxySettings}
/>
) : (
<div className="flex flex-1 items-center justify-center text-gray-500">
Select an agent to run batch tests.
</div>
),
},
]}
/>
)}
</div>
</TabsContent>
<TabsContent value="connect" keepMounted={hasVisited("connect")} className="min-h-0 overflow-hidden">
<div className="h-full overflow-y-auto p-6">
{selectedAgent ? (
<ConnectTabContent
agentName={selectedAgent.model_name}
proxySettings={proxySettings}
customProxyBaseUrl={customProxyBaseUrl}
accessToken={accessToken}
userID={userID}
disabledPersonalKeyCreation={disabledPersonalKeyCreation}
creatingKey={creatingKey}
createdKeyValue={createdKeyValue}
onCreateKey={handleCreateKeyForAgent}
/>
) : (
<div className="flex flex-1 items-center justify-center text-gray-500">
Select an agent to see how to connect.
</div>
)}
</div>
</TabsContent>
</Tabs>
</>
)}
</div>
</div>
<AlertDialog open={confirmingDelete} onOpenChange={setConfirmingDelete}>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle>Delete agent</AlertDialogTitle>
<AlertDialogDescription>
Are you sure you want to delete &quot;{selectedAgent?.model_name}&quot;? This cannot be undone.
</AlertDialogDescription>
</AlertDialogHeader>
<AlertDialogFooter>
<AlertDialogAction variant="outline">Cancel</AlertDialogAction>
<Button variant="destructive" onClick={handleConfirmDelete} disabled={deleting}>
Delete
</Button>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
</div>
);
}

267
ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.test.tsx Normal file
View file

@ -0,0 +1,267 @@
import { act, render, screen, waitFor } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import RealtimePlayground from "./RealtimePlayground";
vi.mock("@/components/networking", () => ({
getProxyBaseUrl: () => "https://proxy.example.com",
}));
class FakeSocket {
static instances: FakeSocket[] = [];
static OPEN = 1;
readyState = 0;
sent: string[] = [];
onopen: (() => void) | null = null;
onmessage: ((event: { data: string }) => void) | null = null;
onerror: (() => void) | null = null;
onclose: (() => void) | null = null;
close = vi.fn(() => {
this.readyState = 3;
this.onclose?.();
});
constructor(
public url: string,
public protocols?: string[],
) {
FakeSocket.instances.push(this);
}
send(payload: string) {
this.sent.push(payload);
}
open() {
this.readyState = 1;
this.onopen?.();
}
emit(message: Record<string, unknown>) {
this.onmessage?.({ data: JSON.stringify(message) });
}
}
const latestSocket = () => FakeSocket.instances[FakeSocket.instances.length - 1];
const connect = async (user: ReturnType<typeof userEvent.setup>) => {
await user.click(screen.getByRole("button", { name: /Connect/i }));
await act(async () => {
latestSocket().open();
});
};
const props = {
accessToken: "sk-realtime",
selectedModel: "gpt-realtime",
};
beforeEach(() => {
FakeSocket.instances = [];
vi.stubGlobal("WebSocket", FakeSocket);
vi.stubGlobal(
"AudioContext",
class {
currentTime = 0;
destination = {};
close = vi.fn();
createBuffer = vi.fn(() => ({ getChannelData: () => new Float32Array(1), duration: 0 }));
createBufferSource = vi.fn(() => ({ connect: vi.fn(), start: vi.fn(), buffer: null }));
},
);
Object.defineProperty(window, "matchMedia", {
writable: true,
value: (query: string) => ({
matches: false,
media: query,
onchange: null,
addListener: () => {},
removeListener: () => {},
addEventListener: () => {},
removeEventListener: () => {},
dispatchEvent: () => false,
}),
});
Element.prototype.scrollIntoView = vi.fn();
});
afterEach(() => {
vi.unstubAllGlobals();
});
describe("RealtimePlayground", () => {
it("opens disconnected, with the invitation to connect", () => {
render(<RealtimePlayground {...props} />);
expect(screen.getByText("Realtime Voice Chat")).toBeInTheDocument();
expect(screen.getByText("Disconnected")).toBeInTheDocument();
expect(screen.getByText("Realtime Voice Playground")).toBeInTheDocument();
expect(screen.getByRole("button", { name: /Connect/i })).toBeInTheDocument();
});
it("hides the composer until a session exists", () => {
render(<RealtimePlayground {...props} />);
expect(screen.queryByPlaceholderText("Type a message or use the mic...")).not.toBeInTheDocument();
});
it("dials the realtime endpoint for the selected model, carrying the key as a protocol", async () => {
const user = userEvent.setup();
render(<RealtimePlayground {...props} />);
await user.click(screen.getByRole("button", { name: /Connect/i }));
expect(latestSocket().url).toBe("wss://proxy.example.com/v1/realtime?model=gpt-realtime");
expect(latestSocket().protocols).toEqual(["realtime", "openai-insecure-api-key.sk-realtime"]);
});
it("passes a custom proxy base url through instead of the default", async () => {
const user = userEvent.setup();
render(<RealtimePlayground {...props} customProxyBaseUrl="https://tenant.example.com" />);
await user.click(screen.getByRole("button", { name: /Connect/i }));
expect(latestSocket().url).toContain("wss://tenant.example.com/v1/realtime");
});
it("appends the selected guardrails to the session url", async () => {
const user = userEvent.setup();
render(<RealtimePlayground {...props} selectedGuardrails={["pii", "toxicity"]} />);
await user.click(screen.getByRole("button", { name: /Connect/i }));
expect(latestSocket().url).toContain("guardrails=pii%2Ctoxicity");
});
it("refuses to dial without a model and says why", async () => {
const user = userEvent.setup();
render(<RealtimePlayground {...props} selectedModel="" />);
await user.click(screen.getByRole("button", { name: /Connect/i }));
expect(FakeSocket.instances).toHaveLength(0);
expect(screen.getByText("Please select a model first")).toBeInTheDocument();
});
it("reveals the composer and the disconnect control once the session opens", async () => {
const user = userEvent.setup();
render(<RealtimePlayground {...props} />);
await connect(user);
expect(screen.getByText("Connected")).toBeInTheDocument();
expect(screen.getByText("Connected to realtime API")).toBeInTheDocument();
expect(screen.getByPlaceholderText("Type a message or use the mic...")).toBeInTheDocument();
expect(screen.getByRole("button", { name: /Disconnect/i })).toBeInTheDocument();
expect(screen.getByTitle("Start recording")).toBeInTheDocument();
});
it("configures the session against the chosen voice when it is created", async () => {
const user = userEvent.setup();
render(<RealtimePlayground {...props} />);
await connect(user);
await act(async () => {
latestSocket().emit({ type: "session.created" });
});
const update = JSON.parse(latestSocket().sent[0]);
expect(update.type).toBe("session.update");
expect(update.session.voice).toBe("alloy");
expect(update.session.type).toBe("realtime");
});
it("sends what was typed and then asks for a response", async () => {
const user = userEvent.setup();
render(<RealtimePlayground {...props} />);
await connect(user);
await user.type(screen.getByPlaceholderText("Type a message or use the mic..."), "hello there");
await user.click(screen.getByRole("button", { name: /send/i }));
const payloads = latestSocket().sent.map((raw) => JSON.parse(raw));
expect(payloads[0].item.content[0].text).toBe("hello there");
expect(payloads[1]).toEqual({ type: "response.create" });
expect(screen.getByText("hello there")).toBeInTheDocument();
expect(screen.getByPlaceholderText("Type a message or use the mic...")).toHaveValue("");
});
it("will not send an empty message", async () => {
const user = userEvent.setup();
render(<RealtimePlayground {...props} />);
await connect(user);
await user.click(screen.getByRole("button", { name: /send/i }));
expect(latestSocket().sent).toHaveLength(0);
});
it("streams assistant text deltas into a single reply", async () => {
const user = userEvent.setup();
render(<RealtimePlayground {...props} />);
await connect(user);
await act(async () => {
latestSocket().emit({ type: "response.output_text.delta", delta: "Hel" });
latestSocket().emit({ type: "response.output_text.delta", delta: "lo!" });
});
expect(screen.getByText("Hello!")).toBeInTheDocument();
});
it("falls back to the completed response when no delta arrived", async () => {
const user = userEvent.setup();
render(<RealtimePlayground {...props} />);
await connect(user);
await act(async () => {
latestSocket().emit({
type: "response.done",
response: { output: [{ content: [{ type: "output_audio", transcript: "spoken reply" }] }] },
});
});
expect(screen.getByText("spoken reply")).toBeInTheDocument();
});
it("shows what the microphone heard", async () => {
const user = userEvent.setup();
render(<RealtimePlayground {...props} />);
await connect(user);
await act(async () => {
latestSocket().emit({
type: "conversation.item.input_audio_transcription.completed",
transcript: "what is the weather",
});
});
expect(screen.getByText("what is the weather")).toBeInTheDocument();
});
it("surfaces an error frame in the transcript", async () => {
const user = userEvent.setup();
render(<RealtimePlayground {...props} />);
await connect(user);
await act(async () => {
latestSocket().emit({ type: "error", error: { message: "rate limited" } });
});
expect(screen.getByText("Error: rate limited")).toBeInTheDocument();
});
it("closes the socket and returns to the disconnected state", async () => {
const user = userEvent.setup();
render(<RealtimePlayground {...props} />);
await connect(user);
const socket = latestSocket();
await user.click(screen.getByRole("button", { name: /Disconnect/i }));
expect(socket.close).toHaveBeenCalled();
await waitFor(() => expect(screen.getByRole("button", { name: /Connect/i })).toBeInTheDocument());
expect(screen.queryByPlaceholderText("Type a message or use the mic...")).not.toBeInTheDocument();
});
});

77
ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx
View file

@ -1,13 +1,13 @@
"use client";
import { AudioMutedOutlined, AudioOutlined, CloseCircleOutlined, SendOutlined, SoundOutlined } from "@ant-design/icons";
import { Button, Input, Select, Typography } from "antd";
import { CircleX, Mic, MicOff, Send, Volume2 } from "lucide-react";
import React, { useCallback, useEffect, useRef, useState } from "react";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { getProxyBaseUrl } from "@/components/networking";
import { OPEN_AI_VOICE_SELECT_OPTIONS } from "./chatConstants";
const { Text } = Typography;
interface RealtimeMessage {
role: "user" | "assistant" | "system" | "status";
content: string;
@ -364,28 +364,37 @@ const RealtimePlayground: React.FC<RealtimePlaygroundProps> = ({
{/* Header */}
<div className="flex items-center justify-between px-4 py-3 border-b border-gray-200 bg-gray-50">
<div className="flex items-center gap-3">
<SoundOutlined className="text-lg text-blue-500" />
<Text className="font-semibold text-gray-800">Realtime Voice Chat</Text>
<Volume2 className="size-5 text-blue-500" />
<span className="font-semibold text-gray-800">Realtime Voice Chat</span>
<span className={`inline-block w-2 h-2 rounded-full ${isConnected ? "bg-green-500" : "bg-gray-300"}`} />
<Text className="text-xs text-gray-500">
<span className="text-xs text-gray-500">
{isConnected ? "Connected" : isConnecting ? "Connecting..." : "Disconnected"}
</Text>
</span>
</div>
<div className="flex items-center gap-2">
<Select
size="small"
value={selectedVoice}
onChange={setSelectedVoice}
options={OPEN_AI_VOICE_SELECT_OPTIONS}
style={{ width: 220 }}
onValueChange={(voice) => setSelectedVoice(voice ?? selectedVoice)}
disabled={isConnected}
/>
>
<SelectTrigger size="sm" className="w-[220px]" aria-label="Voice">
<SelectValue>{OPEN_AI_VOICE_SELECT_OPTIONS.find((v) => v.value === selectedVoice)?.label}</SelectValue>
</SelectTrigger>
<SelectContent>
{OPEN_AI_VOICE_SELECT_OPTIONS.map((voice) => (
<SelectItem key={voice.value} value={voice.value}>
{voice.label}
</SelectItem>
))}
</SelectContent>
</Select>
{!isConnected ? (
<Button type="primary" onClick={connect} loading={isConnecting} size="small">
<Button onClick={connect} disabled={isConnecting} size="sm">
Connect
</Button>
) : (
<Button danger onClick={disconnect} size="small" icon={<CloseCircleOutlined />}>
<Button variant="destructive" onClick={disconnect} size="sm">
<CircleX />
Disconnect
</Button>
)}
@ -396,12 +405,12 @@ const RealtimePlayground: React.FC<RealtimePlaygroundProps> = ({
<div className="flex-1 overflow-y-auto p-4 space-y-3">
{messages.length === 0 && !isConnected && (
<div className="flex flex-col items-center justify-center h-full text-gray-400 gap-3">
<SoundOutlined style={{ fontSize: 48 }} />
<Text className="text-lg text-gray-500">Realtime Voice Playground</Text>
<Text className="text-sm text-gray-400 text-center max-w-md">
<Volume2 className="size-12" />
<span className="text-lg text-gray-500">Realtime Voice Playground</span>
<p className="text-sm text-gray-400 text-center max-w-md">
Click <b>Connect</b> to start a realtime session. You can speak using your microphone or type messages.
The AI will respond with voice and text.
</Text>
</p>
</div>
)}
{messages.map((msg, i) => (
@ -433,30 +442,26 @@ const RealtimePlayground: React.FC<RealtimePlaygroundProps> = ({
<div className="border-t border-gray-200 p-3 bg-white">
<div className="flex items-center gap-2">
<Button
shape="circle"
size="large"
type={isRecording ? "primary" : "default"}
danger={isRecording}
icon={isRecording ? <AudioMutedOutlined /> : <AudioOutlined />}
size="icon-lg"
variant={isRecording ? "destructive" : "outline"}
onClick={isRecording ? stopRecording : startRecording}
title={isRecording ? "Stop recording" : "Start recording"}
className={isRecording ? "animate-pulse" : ""}
/>
className={`rounded-full ${isRecording ? "animate-pulse" : ""}`}
>
{isRecording ? <MicOff /> : <Mic />}
</Button>
<Input
placeholder="Type a message or use the mic..."
value={inputText}
onChange={(e) => setInputText(e.target.value)}
onPressEnter={sendTextMessage}
className="flex-1"
size="large"
/>
<Button
type="primary"
icon={<SendOutlined />}
onClick={sendTextMessage}
disabled={!inputText.trim()}
size="large"
onKeyDown={(e) => {
if (e.key === "Enter") sendTextMessage();
}}
className="h-10 flex-1"
/>
<Button size="icon-lg" onClick={sendTextMessage} disabled={!inputText.trim()} aria-label="Send">
<Send />
</Button>
</div>
{isRecording && (
<div className="mt-2 flex items-center gap-2 text-red-500 text-xs">

7
ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.test.tsx
View file

@ -1,4 +1,4 @@
import { render, waitFor } from "@testing-library/react";
import { render, screen, waitFor } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { beforeEach, describe, expect, it, vi } from "vitest";
import CompareUI from "./CompareUI";
@ -108,10 +108,7 @@ describe("CompareUI", () => {
let comparisonPanels = container.querySelectorAll('[data-testid^="comparison-panel-"]');
expect(comparisonPanels).toHaveLength(2);
const addButtons = Array.from(container.querySelectorAll('button[class*="ant-btn"]'));
const addComparisonButton = addButtons.find((btn) => btn.textContent?.includes("Add Comparison"));
expect(addComparisonButton).toBeInTheDocument();
await user.click(addComparisonButton!);
await user.click(screen.getByRole("button", { name: /Add Comparison/i }));
// Wait for the new comparison panel to be added (should have 3 total now)
await waitFor(() => {

Some files were not shown because too many files have changed in this diff Show more