diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index dc8f17fb665..6fe37f0aacb 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -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 diff --git a/litellm/constants.py b/litellm/constants.py index 554165f5d39..6449834d6a4 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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 diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 647ccc33a44..b8b07af59c6 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -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 diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 48bc60a07e5..4ad20772ed0 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -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: diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 8d039d95bb1..372cf110f7c 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -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) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 554b6ea952e..95a3806e8ad 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -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 diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index c8ff6e262d2..a1adda2bc95 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -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 diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py index a9a3367cd93..125dc3d773d 100644 --- a/litellm/proxy/_experimental/mcp_server/sampling_handler.py +++ b/litellm/proxy/_experimental/mcp_server/sampling_handler.py @@ -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( diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 49a1f1314f0..f237529b319 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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 diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index 5534b8b5a3c..4cf84dd0725 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -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}) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 737be8e37ce..0270bf649c2 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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).", diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 0239368ad0e..51050e62494 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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 diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 1869328c039..60a03689804 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -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) diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index 9f01c719a5f..d19023862cb 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -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 diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup_metrics.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup_metrics.py new file mode 100644 index 00000000000..340aeab938c --- /dev/null +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup_metrics.py @@ -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() diff --git a/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py b/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py index df17721d8e5..221c142d9d3 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py +++ b/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py @@ -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) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 251ed1feb10..0a5626ba0a7 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -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) diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index c00c2a5ba4c..2271501d480 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -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 ) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 116dea464ff..7e190e8b19d 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -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( diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 4bc72f554b1..912e18150b3 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -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) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index b72b5d218b2..3d7f0808fb9 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -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 diff --git a/litellm/proxy/management_helpers/access_group_key_sync.py b/litellm/proxy/management_helpers/access_group_key_sync.py new file mode 100644 index 00000000000..5d43cb29978 --- /dev/null +++ b/litellm/proxy/management_helpers/access_group_key_sync.py @@ -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, + ) diff --git a/litellm/proxy/management_helpers/access_group_team_sync.py b/litellm/proxy/management_helpers/access_group_team_sync.py new file mode 100644 index 00000000000..55c0346e375 --- /dev/null +++ b/litellm/proxy/management_helpers/access_group_team_sync.py @@ -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) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py index 132af097a55..597c2e742b3 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py @@ -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() diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index 65ebc2728c6..1c8bce28454 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -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}" diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 3dcc8257b82..34286b203c7 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -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) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 473da0afe16..bf06c6bbcda 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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", diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index fcc6aac1c14..e24e5b21583 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -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", diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index 029648f7901..efdbda47fdc 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -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("&", "&").replace("<", "<").replace(">", ">") + + +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 diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ef0376ade84..d3ca2fa64ed 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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 diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index c6e17502e5d..56818717c09 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -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, diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh index afe55603466..82498ec10cd 100755 --- a/scripts/pre_commit_lint.sh +++ b/scripts/pre_commit_lint.sh @@ -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 diff --git a/tests/proxy_admin_ui_tests/test_access_group_team_sync.py b/tests/proxy_admin_ui_tests/test_access_group_team_sync.py new file mode 100644 index 00000000000..629d77f20fc --- /dev/null +++ b/tests/proxy_admin_ui_tests/test_access_group_team_sync.py @@ -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]} diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 6d7ada17ec5..fa274324fd6 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -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() diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 02bd3535c0a..fd66667af64 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -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(): """ diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 0b95a497882..52dc91ce24d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index b56a12db5b1..4081681daef 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -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"} diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 850d01c6e34..7df83065865 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -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), diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py index 73fdee9cde3..00ed4e91efa 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py @@ -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 diff --git a/tests/test_litellm/proxy/common_utils/test_callback_utils.py b/tests/test_litellm/proxy/common_utils/test_callback_utils.py index bfd4ffe1593..0c73c3fcf22 100644 --- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -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 | | 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) diff --git a/tests/test_litellm/proxy/conftest.py b/tests/test_litellm/proxy/conftest.py index 230ccaf5fd4..61752997f0f 100644 --- a/tests/test_litellm/proxy/conftest.py +++ b/tests/test_litellm/proxy/conftest.py @@ -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) diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py index 289de707387..e949afce57b 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py @@ -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"] diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 0ebdc07d282..bdf09a95e4b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -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"] diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py index 178e9379ea7..6e670e48b6a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py @@ -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() diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index f038af4c7a1..c6960ecda5a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -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, + } diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index da805b864fc..b83b862d6b8 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -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) diff --git a/tests/test_litellm/proxy/management_helpers/test_access_group_team_sync.py b/tests/test_litellm/proxy/management_helpers/test_access_group_team_sync.py new file mode 100644 index 00000000000..eb11292cf42 --- /dev/null +++ b/tests/test_litellm/proxy/management_helpers/test_access_group_team_sync.py @@ -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"] diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_cohere_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_cohere_passthrough_logging_handler.py index 887aedaf0aa..6d7011fe10c 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_cohere_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_cohere_passthrough_logging_handler.py @@ -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__]) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py index 05051ab3745..664015003e4 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py @@ -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 diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 88dc07e741b..e99bdfb5c35 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -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 diff --git a/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py index d17f6293cc3..736fc13d137 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py @@ -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=" & ", 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 "" not in message + assert "<!channel>" 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"})] diff --git a/tests/test_litellm/proxy/test_conftest.py b/tests/test_litellm/proxy/test_conftest.py new file mode 100644 index 00000000000..6df692a67c9 --- /dev/null +++ b/tests/test_litellm/proxy/test_conftest.py @@ -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) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 57b2c874962..918d39646b0 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -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.""" diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index 03eef14dacb..87fbdd4c933 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -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 diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py b/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py index 9defb309863..56057dce7e0 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py @@ -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={}) diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index d60fff66c44..4981caa10c3 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -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) diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py index 35d98903226..33baf7474ce 100644 --- a/tests/test_litellm/test_pre_commit_lint.py +++ b/tests/test_litellm/test_pre_commit_lint.py @@ -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 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 0104733bcf3..894d99c92e0 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -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 diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index d7f71a5840d..5e322598a10 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -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 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.test.tsx index 9bfa71f77be..b616982d69b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.test.tsx @@ -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 ?
Evaluation settings modal
: null), })); const mockGetGuardrailsUsageOverview = vi.mocked(networking.getGuardrailsUsageOverview); @@ -28,6 +28,18 @@ function wrapper({ children }: { children: React.ReactNode }) { return {children}; } +function renderOverview(onSelectGuardrail = vi.fn()) { + return render( + , + { 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(); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.tsx index 8d45b4a4ee8..e1048c5322f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.tsx @@ -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({
- +

Guardrails Monitor

Monitor guardrail performance across all requests

-
- - - - - - } - /> - - - } - /> - - - 150 ? "text-red-600" : metrics.avgLatency > 50 ? "text-amber-600" : "text-green-600" - } - /> - - - - - +
+ + } + /> + } + /> + 150 ? "text-red-600" : metrics.avgLatency > 50 ? "text-amber-600" : "text-green-600" + } + /> + +
@@ -250,7 +242,11 @@ export function GuardrailsOverview({
{(isLoading || error) && (
- {isLoading && } + {isLoading && ( + + + + )} {error && Failed to load data. Try again.}
)} @@ -270,20 +266,20 @@ export function GuardrailsOverview({ toolbar={() => (
- - Guardrail Performance - +
Guardrail Performance

Click a guardrail to view details, logs, and configuration

)} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx index e310da55063..8297c1e1e3b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx @@ -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 | null) => void; @@ -108,23 +107,23 @@ export function GuardrailTestPanel({ return (
{/* Header */} -
+
-
-

Test Guardrails:

+
+

Test Guardrails:

{guardrailNames.map((name) => (
- {name} + {name}
))}
-

+

Test {guardrailNames.length > 1 ? "guardrails" : "guardrail"} and compare results

@@ -135,46 +134,63 @@ export function GuardrailTestPanel({
-
+
- - - + + + + + + } + /> + Press Enter to submit. Use Shift+Enter for new line.
{inputText && ( - )}
-