scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_fix_user_usage_filter_search

Browse source
This commit is contained in:
Devin AI 2026-08-14 22:50:23 +00:00
commit 713f16f89a
233 changed files with 15356 additions and 7785 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

84
.circleci/config.yml
View file

@ -2744,84 +2744,6 @@ jobs:
file: ./coverage.xml
flags: circleci
ui_build:
docker:
- image: cimg/node:24.19@sha256:8966565f07189a67d64d6808a2b127f31dafae566508e3547f55640e1070bfad
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
resource_class: medium+
working_directory: ~/project
steps:
- checkout
- skip_if_unrelated_changes:
category: client
- setup_google_dns
- restore_cache:
keys:
- ui-build-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
- ui-build-deps-v1-
- restore_cache:
keys:
- ui-nextjs-cache-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
- ui-nextjs-cache-v1-
- run:
name: Install dependencies
command: |
cd ui/litellm-dashboard
npm ci
- save_cache:
key: ui-build-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
paths:
- ui/litellm-dashboard/node_modules
- run:
name: Build UI
command: |
cd ui/litellm-dashboard
source ./build_ui.sh
- save_cache:
key: ui-nextjs-cache-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
paths:
- ui/litellm-dashboard/.next/cache
- persist_to_workspace:
root: .
paths:
- litellm/proxy/_experimental/out
ui_unit_tests:
docker:
- image: cimg/node:24.19@sha256:8966565f07189a67d64d6808a2b127f31dafae566508e3547f55640e1070bfad
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
resource_class: xlarge
working_directory: ~/project
steps:
- checkout
- skip_if_unrelated_changes:
category: client
- setup_google_dns
- restore_cache:
keys:
- ui-unit-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
- ui-unit-deps-v1-
- run:
name: Install dependencies
command: |
cd ui/litellm-dashboard
npm ci
- save_cache:
key: ui-unit-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
paths:
- ui/litellm-dashboard/node_modules
- run:
name: Run UI unit tests (Vitest)
command: |
cd ui/litellm-dashboard
CI=true npm run test -- --run \
--pool forks --poolOptions.forks.maxForks=6
e2e_ui_testing:
docker:
- image: cimg/python:3.12-browsers@sha256:b432899af01c9a311bf74f4f22e9ada2e5306d4b1b4383f8d29e1228a5844ef2
@ -3181,12 +3103,6 @@ workflows:
filters: *main_branches
- litellm_router_unit_testing:
filters: *main_branches
- ui_build:
filters: *main_branches
- ui_unit_tests:
requires:
- ui_build
filters: *main_branches
- auth_ui_unit_tests:
filters: *main_branches
- proxy_behavior_tests:

106
.github/workflows/test-unit-proxy-legacy.yml vendored
View file

@ -1,106 +0,0 @@
name: "Unit Tests: Proxy Legacy Tests"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
test:
runs-on: ubuntu-latest
timeout-minutes: 20
strategy:
fail-fast: false
matrix:
test-group:
- name: "auth-and-jwt"
path: "tests/proxy_unit_tests/test_[a-j]*.py"
- name: "key-generation"
path: "tests/proxy_unit_tests/test_[k-o]*.py"
- name: "proxy-config"
path: "tests/proxy_unit_tests/test_prisma*.py tests/proxy_unit_tests/test_prompt*.py tests/proxy_unit_tests/test_proxy_[c-r]*.py"
- name: "proxy-server"
path: "tests/proxy_unit_tests/test_proxy_server.py"
- name: "proxy-server-extras"
path: "tests/proxy_unit_tests/test_proxy_server_*.py tests/proxy_unit_tests/test_proxy_setting_guardrails.py"
- name: "proxy-utils"
path: "tests/proxy_unit_tests/test_proxy_utils.py"
- name: "proxy-token-counter"
path: "tests/proxy_unit_tests/test_proxy_token_counter.py"
- name: "proxy-response-and-misc"
path: "tests/proxy_unit_tests/test_[r-t]*.py"
- name: "proxy-user-auth-and-spend"
path: "tests/proxy_unit_tests/test_[u-z]*.py"
name: ${{ matrix.test-group.name }}
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Detect backend-relevant changes
id: changes
uses: ./.github/actions/detect-backend-changes
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Cache uv dependencies
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cache/uv
.venv
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-
- name: Install dependencies
if: steps.changes.outputs.decision != 'skip'
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
- name: Cache Prisma binaries
if: steps.changes.outputs.decision != 'skip'
uses: ./.github/actions/cache-prisma-binaries
- name: Generate Prisma client
if: steps.changes.outputs.decision != 'skip'
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Run tests - ${{ matrix.test-group.name }}
if: steps.changes.outputs.decision != 'skip'
env:
TEST_PATH: ${{ matrix.test-group.path }}
run: |
uv run --no-sync pytest ${TEST_PATH} \
--tb=short -vv \
--maxfail=10 \
-n 2 \
--reruns 1 \
--reruns-delay 1 \
--dist=loadscope \
--durations=20

6
litellm/constants.py
View file

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

90
litellm/integrations/langfuse/langfuse.py
View file

@ -2,8 +2,9 @@
# On success, logs events to Langfuse
import os
import traceback
from collections.abc import Callable, Iterable
from collections.abc import Callable, Iterable, Mapping
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, cast
from packaging.version import Version
@ -30,6 +31,7 @@ from litellm.types.utils import (
ImageResponse,
ModelResponse,
RerankResponse,
StandardLoggingMetadata,
StandardLoggingPayload,
StandardLoggingPromptManagementMetadata,
TextCompletionResponse,
@ -46,6 +48,11 @@ else:
Langfuse = Any
_DENIED_STEERING_KEYS: Final = frozenset({"headers", "endpoint", "caching_groups", "previous_models"})
_NO_METADATA: Final[Mapping[str, Any]] = MappingProxyType({})
_REDACTED_PROXY_HEADERS: Final[frozenset[str]] = frozenset({"authorization", "cookie", "referer"})
def _extract_cache_read_input_tokens(usage_obj) -> int:
"""
Extract cache_read_input_tokens from usage object.
@ -512,16 +519,14 @@ class LangFuseLogger:
else []
)
if standard_logging_object is None:
end_user_id = None
prompt_management_metadata: StandardLoggingPromptManagementMetadata | None = None
else:
end_user_id = standard_logging_object["metadata"].get("user_api_key_end_user_id", None)
prompt_management_metadata = cast(
StandardLoggingPromptManagementMetadata | None,
standard_logging_object["metadata"].get("prompt_management_metadata", None),
)
allowlisted_metadata: Final[StandardLoggingMetadata | dict[str, Any]] = (
standard_logging_object["metadata"] if standard_logging_object is not None else _NO_METADATA
)
end_user_id: Final = allowlisted_metadata.get("user_api_key_end_user_id", None)
prompt_management_metadata: Final[StandardLoggingPromptManagementMetadata | None] = cast(
StandardLoggingPromptManagementMetadata | None,
allowlisted_metadata.get("prompt_management_metadata", None),
)
# Clean Metadata before logging - never log raw metadata
# the raw metadata can contain circular references which leads to infinite recursion
@ -540,12 +545,7 @@ class LangFuseLogger:
tags.append(f"{key}:{value}")
# clean litellm metadata before logging
if key in [
"headers",
"endpoint",
"caching_groups",
"previous_models",
]:
if key in _DENIED_STEERING_KEYS:
continue
else:
clean_metadata[key] = value
@ -630,19 +630,18 @@ class LangFuseLogger:
trace_params["output"] = output if not mask_output else "redacted-by-litellm"
if debug is True or (isinstance(debug, str) and debug.lower() == "true"):
if "metadata" in trace_params:
# log the raw_metadata in the trace
trace_params["metadata"]["metadata_passed_to_litellm"] = metadata
else:
trace_params["metadata"] = {"metadata_passed_to_litellm": metadata}
debug_metadata: Final = {
key: value for key, value in metadata.items() if isinstance(value, (str, int, float, bool))
}
trace_params["metadata"] = {
**(trace_params.get("metadata") or _NO_METADATA),
"metadata_passed_to_litellm": debug_metadata,
}
cost: Final = kwargs.get("response_cost", None)
verbose_logger.debug("trace: %s", cost)
clean_metadata["litellm_response_cost"] = cost
if standard_logging_object is not None:
hidden_params: Final = standard_logging_object.get("hidden_params", {})
clean_metadata["hidden_params"] = filter_exceptions_from_params(hidden_params)
hidden_params: Final = standard_logging_object.get("hidden_params") if standard_logging_object else None
if (
litellm.langfuse_default_tags is not None
@ -654,22 +653,24 @@ class LangFuseLogger:
tags.append(f"proxy_base_url:{proxy_base_url}")
api_base: Final = litellm_params.get("api_base", None)
if api_base:
clean_metadata["api_base"] = api_base
vertex_location: Final = kwargs.get("vertex_location", None)
if vertex_location:
clean_metadata["vertex_location"] = vertex_location
aws_region_name: Final = kwargs.get("aws_region_name", None)
if aws_region_name:
clean_metadata["aws_region_name"] = aws_region_name
candidate_enrichments: Final = (
("litellm_response_cost", cost, True),
("hidden_params", filter_exceptions_from_params(hidden_params), hidden_params is not None),
("api_base", api_base, bool(api_base)),
("vertex_location", vertex_location, bool(vertex_location)),
("aws_region_name", aws_region_name, bool(aws_region_name)),
("cache_hit", kwargs.get("cache_hit") or False, self._supports_tags() and "cache_hit" in kwargs),
)
enrichments: Final[Mapping[str, Any]] = {
key: value for key, value, include in candidate_enrichments if include
}
if self._supports_tags():
if "cache_hit" in kwargs:
if kwargs["cache_hit"] is None:
kwargs["cache_hit"] = False
clean_metadata["cache_hit"] = kwargs["cache_hit"]
if "cache_hit" in kwargs and kwargs["cache_hit"] is None:
kwargs["cache_hit"] = False # rebind-ok: pre-existing normalization other integrations rely on
if existing_trace_id is None:
trace_params.update({"tags": tags})
@ -682,13 +683,13 @@ class LangFuseLogger:
if headers:
for key, value in headers.items():
# these headers can leak our API keys and/or JWT tokens
if key.lower() not in ["authorization", "cookie", "referer"]:
if key.lower() not in _REDACTED_PROXY_HEADERS:
clean_headers[key] = value
trace: Final[StatefulTraceClient] = self.Langfuse.trace(**trace_params)
# Log provider specific information as a span
log_provider_specific_information_as_span(trace, clean_metadata)
log_provider_specific_information_as_span(trace, enrichments)
# Log guardrail information as a span
self._log_guardrail_information_as_span(
@ -761,7 +762,10 @@ class LangFuseLogger:
"output": output if not mask_output else "redacted-by-litellm",
"usage": usage,
"usage_details": usage_details,
"metadata": log_requester_metadata(clean_metadata),
"metadata": {
**log_requester_metadata(redact_user_api_key_info(metadata=allowlisted_metadata)),
**enrichments,
},
"level": level,
"version": clean_metadata.pop("version", None),
}
@ -1058,7 +1062,7 @@ def _add_prompt_to_generation_params(
def log_provider_specific_information_as_span(
trace,
clean_metadata,
clean_metadata: Mapping[str, Any],
):
"""
Logs provider-specific information as spans.
@ -1098,7 +1102,7 @@ def log_provider_specific_information_as_span(
)
def log_requester_metadata(clean_metadata: dict):
def log_requester_metadata(clean_metadata: Mapping[str, Any]):
returned_metadata: Final = {}
requester_metadata: Final = clean_metadata.get("requester_metadata") or {}
for k, v in clean_metadata.items():

15
litellm/litellm_core_utils/exception_mapping_utils.py
View file

@ -34,12 +34,16 @@ class ExceptionCheckers:
"""
@staticmethod
def is_error_str_rate_limit(error_str: str) -> bool:
def is_error_str_rate_limit(error_str: str, status_code: int | None = None) -> bool:
"""
Check if an error string indicates a rate limit error.
Args:
error_str: The error string to check
status_code: The HTTP status the provider returned, when known. Gates only the
bare-number branch: providers echo the request back in validation errors and
429 is an ordinary token id, so an echoed prompt can put a standalone 429 in
the body of a 400. The phrase branches stay ungated (#11455).
Returns:
True if the error indicates a rate limit, False otherwise
@ -47,8 +51,9 @@ class ExceptionCheckers:
if not isinstance(error_str, str):
return False
# Only treat 429 as a rate limit signal when it appears as a standalone token
if re.search(r"\b429\b", error_str):
# A standalone 429 counts unless the provider's own status says otherwise. The
# status is read off an arbitrary exception, so a non-integer means "unknown".
if re.search(r"\b429\b", error_str) and (not isinstance(status_code, int) or status_code == 429):
return True
_error_str_lower: Final = error_str.lower()
@ -280,7 +285,9 @@ def _map_openai_exception(
else:
exception_provider = custom_llm_provider[0].upper() + custom_llm_provider[1:] + "Exception"
if ExceptionCheckers.is_error_str_rate_limit(error_str):
if ExceptionCheckers.is_error_str_rate_limit(
error_str, status_code=getattr(original_exception, "status_code", None)
):
raise RateLimitError(
message=f"RateLimitError: {exception_provider} - {message}",
model=model,

39
litellm/llms/anthropic/common_utils.py
View file

@ -5,6 +5,7 @@ This file contains common utils for anthropic calls.
import copy
import re
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from types import MappingProxyType
from typing import Any, Final, Literal
@ -12,6 +13,7 @@ import httpx
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
import litellm
from litellm.constants import DEFAULT_MODEL_CREATED_AT_TIME
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_file_ids_from_messages,
)
@ -28,6 +30,7 @@ from litellm.types.llms.anthropic import (
AnthropicMcpServerTool,
)
from litellm.types.llms.openai import AllMessageValues
from litellm.types.proxy.model_listing import ModelInfoResponse
_BEDROCK_VERSION_SUFFIX_RE: Final = re.compile(r"-v\d+(?::\d+)?$")
_INFERENCE_PROFILE_MINOR_RE: Final = re.compile(r":\d+$")
@ -1221,3 +1224,39 @@ def process_anthropic_headers(headers: httpx.Headers | dict) -> dict:
additional_headers: Final = {**llm_response_headers, **openai_headers}
return additional_headers
def _anthropic_model_entry(model: ModelInfoResponse, created_at: str) -> Mapping[str, object]:
token_limits: Final = (
("max_input_tokens", model.get("max_input_tokens")),
("max_tokens", model.get("max_output_tokens")),
)
return { # mutable-ok: JSON response body, serialized by the route and never mutated
"type": "model",
"id": model["id"],
"display_name": model["id"],
"created_at": created_at,
**{name: limit for name, limit in token_limits if limit is not None}, # mutable-ok: merged into the body above
}
def create_anthropic_model_list_response(models: Sequence[ModelInfoResponse]) -> Mapping[str, object]:
"""Build the Anthropic-native /v1/models envelope.
Clients that send an anthropic-version header parse the Anthropic Models API
shape (type/display_name/created_at plus has_more/first_id/last_id) and filter
the list themselves, so every model is returned here. The token limits carry
over from the OpenAI-shaped listing, named as the Messages API names them
"""
created_at: Final = (
datetime.fromtimestamp(DEFAULT_MODEL_CREATED_AT_TIME, tz=timezone.utc).isoformat().replace("+00:00", "Z")
)
data: Final = [ # mutable-ok: JSON response body, serialized by the route and never mutated
_anthropic_model_entry(model, created_at) for model in models
]
return { # mutable-ok: JSON response body, serialized by the route and never mutated
"data": data,
"has_more": False,
"first_id": models[0]["id"] if models else None,
"last_id": models[-1]["id"] if models else None,
}

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

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

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

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

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

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

7
litellm/llms/openai/cost_calculation.py
View file

@ -109,15 +109,16 @@ def cost_per_second(model: str, custom_llm_provider: str | None, duration: float
prompt_cost = 0.0
completion_cost = 0.0
## Speech / Audio cost calculation
if "output_cost_per_second" in model_info and model_info["output_cost_per_second"] is not None:
output_cost_per_second: Final = model_info.get("output_cost_per_second")
if output_cost_per_second is not None and output_cost_per_second > 0:
verbose_logger.debug(
"For model=%s - output_cost_per_second: %s; duration: %s",
model,
model_info.get("output_cost_per_second"),
output_cost_per_second,
duration,
)
## COST PER SECOND ##
completion_cost = model_info["output_cost_per_second"] * duration
completion_cost = output_cost_per_second * duration
elif "input_cost_per_second" in model_info and model_info["input_cost_per_second"] is not None:
verbose_logger.debug(
"For model=%s - input_cost_per_second: %s; duration: %s",

7
litellm/main.py
View file

@ -5616,7 +5616,12 @@ def completion(
elif custom_llm_provider == "hosted_vllm":
response = _complete_hosted_vllm(_dispatch_ctx)
elif (
model in litellm.open_ai_chat_completion_models
# A known OpenAI model name only decides the route when nothing else
# resolved a provider. get_llm_provider() already maps these names to
# "openai", so a different value here was asked for explicitly (or came
# from a register_model entry), and the provider config built for it
# would be handed to the OpenAI handler.
(model in litellm.open_ai_chat_completion_models and custom_llm_provider in (None, "openai"))
or custom_llm_provider == "custom_openai"
or custom_llm_provider == "deepinfra"
or custom_llm_provider == "perplexity"

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

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

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

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

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

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

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

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

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

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

16
litellm/proxy/_types.py
View file

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

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

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

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

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

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

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

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

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

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

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

14
litellm/proxy/db/routing_prisma_wrapper.py
View file

@ -103,6 +103,17 @@ class RoutingPrismaWrapper:
def reader(self) -> PrismaWrapper:
return self._reader
@property
def read_target(self) -> PrismaWrapper:
"""The wrapper `_TOP_LEVEL_READ_METHODS` dispatch to right now.
Callers that need to reason about the engine a read actually ran on
(e.g. recovering from prepared statements that went stale on it) must
consult this rather than `writer`, and `__getattr__` routes through it
so the two cannot drift apart.
"""
return self._writer if self._reader_unavailable else self._reader
@property
def reader_unavailable(self) -> bool:
return self._reader_unavailable
@ -254,8 +265,7 @@ class RoutingPrismaWrapper:
def __getattr__(self, name: str) -> Any:
if name in _TOP_LEVEL_READ_METHODS:
target: Final = self._writer if self._reader_unavailable else self._reader
return getattr(target, name)
return getattr(self.read_target, name)
writer_attr: Final = getattr(self._writer, name)
# Per-model action accessors are non-callable instances that expose
# both `find_many` and `create`. Methods like execute_raw / batch_ /

4
litellm/proxy/litellm_pre_call_utils.py
View file

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

46
litellm/proxy/proxy_server.py
View file

@ -367,7 +367,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,
@ -4079,6 +4082,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):
@ -5015,6 +5019,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)
@ -6299,6 +6309,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",
@ -9471,6 +9493,7 @@ class ProxyStartupEvent:
"/models", dependencies=[Depends(user_api_key_auth)], tags=["model management"]
) # if project requires model list
async def model_list(
request: Request = None, # pyright: ignore[reportArgumentType] # FastAPI always injects the Request; the None default only serves direct in-process callers
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
return_wildcard_routes: bool | None = False,
team_id: str | None = None,
@ -9507,6 +9530,9 @@ async def model_list(
settings: Final = cast(dict[str, object], general_settings) # any-ok: legacy settings
from litellm.llms.anthropic.common_utils import (
create_anthropic_model_list_response,
)
from litellm.proxy.management_endpoints.common_utils import (
_user_has_admin_privileges,
)
@ -9514,6 +9540,12 @@ async def model_list(
create_model_info_response,
get_available_models_for_user,
)
from litellm.types.proxy.model_listing import ModelInfoResponse
http_request: Final = cast(Request | None, request) # cast-ok: in-process callers pass no request
wants_anthropic_format: Final = (
http_request is not None and http_request.headers.get("anthropic-version") is not None
)
# Validate scope parameter if provided
if scope is not None and scope != "expand":
@ -9597,6 +9629,10 @@ async def model_list(
model_info["id"] = response_id
model_data.append(model_info)
if wants_anthropic_format:
admin_listing: Final = cast(Sequence[ModelInfoResponse], model_data) # cast-ok: rows built above
return create_anthropic_model_list_response(admin_listing)
return dict(
data=model_data,
object="list",
@ -9637,6 +9673,10 @@ async def model_list(
model_info["id"] = response_id
model_data.append(model_info)
if wants_anthropic_format:
listing: Final = cast(Sequence[ModelInfoResponse], model_data) # cast-ok: rows built above
return create_anthropic_model_list_response(listing)
return dict(
data=model_data,
object="list",
@ -15529,6 +15569,10 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro
"store_model_in_db": "Boolean",
"store_prompts_in_spend_logs": "Boolean",
"maximum_spend_logs_retention_period": "String",
"maximum_spend_logs_cleanup_batch_size": "Integer",
"maximum_spend_logs_cleanup_max_batches": "Integer",
"maximum_spend_logs_cleanup_run_budget": "String",
"maximum_spend_logs_cleanup_batch_timeout": "String",
"mcp_internal_ip_ranges": "List",
"mcp_trusted_proxy_ranges": "List",
"mcp_xff_num_trusted_hops": "Integer",

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

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

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

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

304
litellm/proxy/utils.py
View file

@ -16,6 +16,7 @@ from dataclasses import dataclass, field
from datetime import date, datetime, timedelta, timezone
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, TypeVar, Union, cast, overload
from litellm import _custom_logger_compatible_callbacks_literal
@ -679,6 +680,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
@ -3005,6 +3007,62 @@ async def prefetch_config_params(prisma_client: "PrismaClient | None", param_nam
)
class _ForcedRecreateDeclined(Exception):
"""A forced recreate was declined by the engine-generation guard.
Distinct from a reconnect *failure*: the machinery worked, it just found
that another path had already replaced the writer, so it left the engines
alone. The caller's engine may still be poisoned, so the cycle must not
report success, but it must not count as a failure either, or the record
of what could not be repaired would gate the retry that recovers.
"""
@dataclass(frozen=True, slots=True)
class _StaleReadEngine:
"""The read engine a query observed, identified rather than only counted.
`PrismaClient.read_db` resolves to the reader while it is available and to
the writer once it is not, and the two carry independent generation
counters that both start at zero and advance on the same reconnect
cadence. A bare generation compared across that switch would silently pit
one engine's counter against another's, so the wrapper is carried with the
number and a switch counts as the engine having moved.
Holding the wrapper itself rather than its `id()` is load-bearing, not
incidental: the strong reference keeps the wrapper alive, so its address
cannot be recycled under a stored observation and match an unrelated
engine later. It is only free because writer and reader both live as long
as the client does; a replaceable reader would make this a retention leak.
"""
wrapper: PrismaWrapper
generation: int
@classmethod
def observe(cls, wrapper: PrismaWrapper) -> "_StaleReadEngine":
return cls(wrapper=wrapper, generation=wrapper.engine_generation)
def is_still_live(self, current: PrismaWrapper) -> bool:
"""Whether this exact engine is still serving reads, unreplaced.
A True answer must never be the only thing standing between a poisoned
engine and its repair. The generation moves only after a replacement
connects, and a recreate whose connect raises leaves it unmoved until
some later recreate succeeds, so this can report an engine as live
after it has stopped working. What bounds that is the failed-repair
record in `_cooldown_applies`, written by a repair attempt that fails
rather than by whatever broke the engine: the two need not be the same
recreate, since the synchronous token-refresh fallback in
`PrismaWrapper.__getattr__` recreates outside the reconnect machinery
and records nothing. The record is written only for callers that named
an engine, and it collapses the rest of the burst for up to one
cooldown window rather than guaranteeing a repair, since the cooldown
conjunct underneath it still expires and lets a later caller retry.
"""
return self.wrapper is current and self.generation == current.engine_generation
class PrismaClient:
spend_log_transactions: list = []
_spend_log_transactions_lock = asyncio.Lock()
@ -3152,6 +3210,14 @@ class PrismaClient:
float(os.getenv("PRISMA_AUTH_RECONNECT_LOCK_TIMEOUT_SECONDS", "0.1")),
)
self._consecutive_reconnect_failures: int = 0
# Last generation of each read engine whose repair was attempted and
# failed. Scoped to the engine rather than counted globally so an
# unrelated reconnect failure cannot suppress a stale reader's
# recovery, and keyed per wrapper rather than held in one slot so a
# writer failure cannot evict the reader's record and hand the waiver
# back to a caller whose engine is still unrepaired. Bounded at two
# entries: a client has one writer and at most one reader.
self._failed_recreate_generations: Mapping[PrismaWrapper, int] = MappingProxyType({})
self._reconnect_escalation_threshold: int = max(1, int(os.getenv("PRISMA_RECONNECT_ESCALATION_THRESHOLD", "3")))
self._engine_pidfd: int = -1
self._engine_pid: int = 0
@ -3167,6 +3233,19 @@ class PrismaClient:
return self.db.writer
return self.db
@property
def read_db(self) -> PrismaWrapper:
"""Underlying wrapper that top-level reads are dispatched to.
Identical to `writer_db` without a read replica. With one configured
it is the reader, which is the engine `query_first` actually runs on,
so anything reasoning about the state of the connection that served a
read has to consult this rather than the writer.
"""
if isinstance(self.db, RoutingPrismaWrapper):
return self.db.read_target
return self.db
def tx(self) -> "TransactionManager":
"""Open an interactive transaction on the writer.
@ -3390,18 +3469,30 @@ class PrismaClient:
`attempt_db_reconnect`, which is singleflight: when a schema change
poisons every pooled connection at once, the first cached-plan error
recreates the client and the concurrent waiters reuse that single
recreate instead of racing to kill each other's fresh engine. We then
retry the identical query exactly once.
recreate instead of racing to kill each other's fresh engine. We pass
`force_recreate` so the reconnect skips its `SELECT 1` liveness probe:
the connection is healthy here, it is the prepared statements on it
that are stale, so a passing probe would otherwise skip the recreate
and leave the retry to hit the same error. We then retry the identical
query exactly once.
The retry reuses the original query byte-for-byte. Mutating the SQL
(e.g. injecting a unique comment) would defeat PostgreSQL's plan cache,
forcing a fresh plan on every request and pegging the database CPU.
If the reconnect is skipped because a recent reconnect is still within
its cooldown, the retry runs against the same connection and may fail
again; the get_data backoff decorator re-runs the lookup and a later
attempt reconnects once the cooldown elapses.
The reconnect cooldown must not gate the engine this query itself saw
as stale, or a migration landing within the cooldown of an earlier
reconnect leaves auth failing until it elapses. The engine observed
before the query names it, so the reconnect bypasses the cooldown only
while that same engine is still the live one.
It is observed from `read_db`, not `writer_db`: `query_first` is a
top-level read, so with a read replica configured it runs on the reader
and it is the reader's prepared statements that went stale. Naming the
writer here would let an unrelated writer reconnect re-arm the cooldown
while the reader stayed poisoned.
"""
stale_read_engine: Final = _StaleReadEngine.observe(self.read_db)
try:
return await self.db.query_first(sql_query, *args)
except Exception as e:
@ -3413,7 +3504,11 @@ class PrismaClient:
"query. This may occur during rolling deployments when schema "
"changes are applied."
)
await self.attempt_db_reconnect(reason="postgres_cached_plan_error")
await self.attempt_db_reconnect(
reason="postgres_cached_plan_error",
force_recreate=True,
stale_read_engine=stale_read_engine,
)
return await self.db.query_first(sql_query, *args)
@backoff.on_exception(
@ -4696,7 +4791,11 @@ class PrismaClient:
self._cleanup_engine_watcher()
asyncio.create_task(self._start_engine_watcher())
async def _run_reconnect_cycle(self, timeout_seconds: float | None = None) -> None:
async def _run_reconnect_cycle(
self,
timeout_seconds: float | None = None,
force_recreate: bool = False,
) -> None:
"""
Run a reconnect cycle with a single overall timeout budget.
@ -4707,6 +4806,11 @@ class PrismaClient:
the client via the non-blocking kill-then-construct flow rather than
calling disconnect(), which blocks the event loop on the synchronous
subprocess.Popen.wait() inside prisma-client-py (see issue #26191).
`force_recreate` skips the direct path's liveness probe, for callers
whose failure lives in the session state rather than the connection
(stale prepared statements after a schema change): a reachable writer
proves nothing about those, so the probe must not veto the recreate.
"""
effective_timeout: Final = (
timeout_seconds if timeout_seconds is not None else self._db_watchdog_reconnect_timeout_seconds
@ -4746,8 +4850,29 @@ class PrismaClient:
# direct path there is no SELECT 1 probe here, so the generation
# guard is the only thing standing between a crash-reconnect and
# a refresh that raced it.
await self.db.recreate_prisma_client(db_url, expected_generation=expected_generation)
recreated: Final = await self.db.recreate_prisma_client(db_url, expected_generation=expected_generation)
await self._start_engine_watcher()
# Same contract as the direct path below: a forced caller asked
# for its engine to be replaced, so a decline is not a success.
# Reachable here because the escalation threshold flips
# `_engine_confirmed_dead`, which routes the next cycle, forced
# callers included, down this branch.
if force_recreate is True and recreated is False:
# Clear the dead-engine flag first, restoring the policy the
# non-forced path already has: a decline does not raise for
# it, so it falls through to the clear below. Only the
# forced branch would strand the flag, and stranding it
# routes the next cycle back down this probe-free branch,
# where the refreshed generation now matches and the
# recreate kills the healthy engine a refresh just spawned
# (#29176). This has to stay AFTER `_start_engine_watcher`
# above: clearing the flag while the watcher is still torn
# down would be worse than either alone.
self._engine_confirmed_dead = False
raise _ForcedRecreateDeclined(
"Forced Prisma recreate declined by the generation guard; "
"the engine that failed was not replaced"
)
await asyncio.wait_for(_do_heavy_reconnect(), timeout=effective_timeout)
# Only clear the "dead engine" flag after the heavy reconnect
@ -4772,44 +4897,106 @@ class PrismaClient:
# detect a refresh that landed since cycle entry and skip the
# redundant restart.
writer: Final = self.writer_db
try:
await writer.query_raw("SELECT 1")
verbose_proxy_logger.info(
"Writer healthy on probe; skipping recreate (engine "
"likely already replaced by a token refresh)."
)
if isinstance(self.db, RoutingPrismaWrapper):
self.db.mark_writer_recovered()
await self._start_engine_watcher()
return
except Exception as probe_err:
verbose_proxy_logger.warning(
"Writer probe failed (%s); recreating Prisma client.",
probe_err,
)
if force_recreate is False:
try:
await writer.query_raw("SELECT 1")
verbose_proxy_logger.info(
"Writer healthy on probe; skipping recreate (engine "
"likely already replaced by a token refresh)."
)
if isinstance(self.db, RoutingPrismaWrapper):
self.db.mark_writer_recovered()
await self._start_engine_watcher()
return
except Exception as probe_err:
verbose_proxy_logger.warning(
"Writer probe failed (%s); recreating Prisma client.",
probe_err,
)
# Fresh Prisma client + new engine subprocess. The previous
# "lightweight" path called `disconnect()` which blocks the
# event loop on `subprocess.Popen.wait()`; since that call
# ends up killing the engine anyway, we do it non-blockingly
# via `_kill_engine_process` inside `recreate_prisma_client`.
self._cleanup_engine_watcher()
await self.db.recreate_prisma_client(db_url, expected_generation=expected_generation)
recreated: Final = await self.db.recreate_prisma_client(db_url, expected_generation=expected_generation)
await self._start_engine_watcher()
# Smoke-test the writer specifically; query_raw on the routing
# wrapper sends to the reader, which would not validate the
# newly-recreated writer engine.
# newly-recreated writer engine. The reader is left to the
# caller's own retried query, a stronger check than SELECT 1,
# and a reader that fails to come back sets `_reader_unavailable`
# so reads fall through to the writer just recreated here.
await self.writer_db.query_raw("SELECT 1")
# A recreate can decline: the optimistic-lock guard no-ops when
# the writer generation moved since cycle entry, and the routing
# wrapper then leaves the reader untouched as well. Callers that
# merely suspect a transport blip are happy either way, but a
# forced caller asked for this engine to be replaced because its
# session state is poisoned, and it was not. Do not report that
# as a success: it would reset the consecutive-failure count and
# log a repair that never happened.
if force_recreate is True and recreated is False:
raise _ForcedRecreateDeclined(
"Forced Prisma recreate declined by the generation guard; "
"the engine that failed was not replaced"
)
await asyncio.wait_for(_do_direct_reconnect(), timeout=effective_timeout)
def _cooldown_applies(self, stale_read_engine: "_StaleReadEngine | None") -> bool:
"""
Whether the reconnect cooldown should still gate this caller.
The cooldown collapses a burst of callers onto one recreate, so it
keeps gating a caller whose named engine has already been replaced:
that recreate is the one it was waiting for. While that engine is still
the live one the damage is still being served, so deferring to an
unrelated reconnect's cooldown would leave it broken until the cooldown
elapses.
A named engine always describes the one that served the failing read
(see `_query_first_with_cached_plan_fallback`), so it is compared
against `read_db`, identity included: `read_db` can resolve to a
different wrapper than it did at observation time.
The waiver is withdrawn once a repair of this same engine has been
tried and failed. A failed recreate leaves the generation where it was,
so without this every queued caller would still see its own engine live
and run its own full recreate serially instead of collapsing onto one
attempt, which is what the cooldown is for. The record is scoped to the
engine rather than to a global failure count: an unrelated reconnect
failing somewhere else says nothing about whether this engine can be
repaired, and gating on it would suppress the recovery this method
exists to allow.
The record is never cleared, and does not need to be. Generations are
monotonic per wrapper, so once the engine is repaired every later
caller names a higher one and the entry can never match again. And this
method is only ever the first half of the gate: the cooldown window
itself still expires, so an engine that can never be repaired degrades
to the plain cooldown rather than being suppressed forever.
"""
if stale_read_engine is None:
return True
if self._failed_recreate_generations.get(stale_read_engine.wrapper) == stale_read_engine.generation:
return True
return not stale_read_engine.is_still_live(self.read_db)
async def _attempt_reconnect_inside_lock(
self,
force: bool,
reason: str,
timeout_seconds: float | None,
force_recreate: bool = False,
stale_read_engine: "_StaleReadEngine | None" = None,
) -> bool:
now: Final = time.time()
if force is False and now - self._db_last_reconnect_attempt_ts < self._db_reconnect_cooldown_seconds:
if (
force is False
and self._cooldown_applies(stale_read_engine)
and now - self._db_last_reconnect_attempt_ts < self._db_reconnect_cooldown_seconds
):
verbose_proxy_logger.debug(
"Skipping DB reconnect attempt inside lock due to cooldown. reason=%s",
reason,
@ -4833,12 +5020,43 @@ class PrismaClient:
reconnect_succeeded = False
try:
await self._run_reconnect_cycle(timeout_seconds=timeout_seconds)
await self._run_reconnect_cycle(timeout_seconds=timeout_seconds, force_recreate=force_recreate)
reconnect_succeeded = True
self._consecutive_reconnect_failures = 0
verbose_proxy_logger.info("Prisma DB reconnect succeeded. reason=%s", reason)
except _ForcedRecreateDeclined as declined:
# A decline is raised only when the recreate returns False, which
# happens only at the generation guard, and the generation moves
# only after a replacement has connected. So a decline is proof
# that a replacement SUCCEEDED, and zeroing a consecutive-failure
# count on that proof is right by definition rather than by
# analogy to what a reported success used to do. Note what it
# proves is that the WRITER was replaced, not that this caller's
# engine was repaired: on a read replica the reader can still be
# poisoned, since the wrapper returns before touching it. Leaving
# the count at the threshold would let the escalation check above
# re-arm the dead-engine flag on the very next attempt and send a
# healthy replacement back down the probe-free heavy path.
self._consecutive_reconnect_failures = 0
verbose_proxy_logger.warning("Prisma DB reconnect declined. reason=%s detail=%s", reason, declined)
except Exception as reconnect_err:
self._consecutive_reconnect_failures += 1
# Remember WHICH engine could not be repaired, so the rest of this
# caller's burst collapses onto the cooldown instead of each
# retrying the recreate that just failed. Recorded only for a
# caller that named a generation: a watchdog or transport-error
# reconnect failing here is unrelated to any stale read engine and
# must not suppress its waiver.
if stale_read_engine is not None:
# Key off the wrapper the CALLER named, never a freshly resolved
# `read_db`. A failed reader recreate is itself what marks the
# reader unavailable, so re-resolving here would file the
# reader's failure under the writer: the poisoned reader would
# lose its record and the healthy writer would gain a spurious
# one, wrong in both directions at once.
self._failed_recreate_generations = MappingProxyType(
{**self._failed_recreate_generations, stale_read_engine.wrapper: stale_read_engine.generation}
)
verbose_proxy_logger.error(
"Prisma DB reconnect failed (%d consecutive). reason=%s error=%s",
self._consecutive_reconnect_failures,
@ -4856,15 +5074,35 @@ class PrismaClient:
force: bool = False,
timeout_seconds: float | None = None,
lock_timeout_seconds: float | None = None,
force_recreate: bool = False,
stale_read_engine: "_StaleReadEngine | None" = None,
) -> bool:
"""
Attempt to reconnect the Prisma client in a singleflight manner.
`force` bypasses the cooldown unconditionally; `force_recreate`
bypasses the liveness probe that would otherwise skip recreating a
reachable engine; `stale_read_engine` bypasses the cooldown only while
the engine that produced the caller's failure is still the live one
(see `_cooldown_applies`).
A `force_recreate` caller can also get False for a third reason: the
generation guard declined because another path had already replaced
the engine, which is a successful outcome reported as False. Callers
that branch on the return value (`exception_handler` raises on False,
`auth_checks` retries only on True) would misread that as a dead end,
and are safe today only because neither passes `force_recreate`. Do
not add it to one of them without revisiting how it reads the result.
Returns:
bool: True if reconnection succeeded, else False.
"""
now: Final = time.time()
if force is False and now - self._db_last_reconnect_attempt_ts < self._db_reconnect_cooldown_seconds:
if (
force is False
and self._cooldown_applies(stale_read_engine)
and now - self._db_last_reconnect_attempt_ts < self._db_reconnect_cooldown_seconds
):
verbose_proxy_logger.debug(
"Skipping DB reconnect attempt due to cooldown. reason=%s",
reason,
@ -4873,7 +5111,9 @@ class PrismaClient:
if lock_timeout_seconds is None:
async with self._db_reconnect_lock:
return await self._attempt_reconnect_inside_lock(force, reason, timeout_seconds)
return await self._attempt_reconnect_inside_lock(
force, reason, timeout_seconds, force_recreate, stale_read_engine
)
lock_acquired_by_timeout_task = False
@ -4922,7 +5162,9 @@ class PrismaClient:
return False
try:
return await self._attempt_reconnect_inside_lock(force, reason, timeout_seconds)
return await self._attempt_reconnect_inside_lock(
force, reason, timeout_seconds, force_recreate, stale_read_engine
)
finally:
self._db_reconnect_lock.release()

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

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

28
scripts/pre_commit_lint.sh
View file

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

7
tests/e2e/ui/tests/internal-user/internalUser.spec.ts
View file

@ -19,12 +19,11 @@ test.describe("Internal User", () => {
// Open the team dropdown — seeded internal user is a member of
// e2e-team-crud and e2e-team-org, so we expect at least the CRUD alias.
const teamSelect = page.locator(".ant-select", { hasText: "Search or select a team" });
const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox");
await teamSelect.click();
await page.keyboard.type(E2E_TEAM_CRUD_ALIAS);
await expect(page.locator(".ant-select-dropdown:visible").getByText(E2E_TEAM_CRUD_ALIAS).first()).toBeVisible({
timeout: 5_000,
});
const dropdown = page.locator('[data-slot="combobox-content"]:visible');
await expect(dropdown.getByText(E2E_TEAM_CRUD_ALIAS).first()).toBeVisible({ timeout: 5_000 });
});
test("Team info page omits the Settings tab for non-admin members", async ({ page }) => {

14
tests/e2e/ui/tests/internal-user/internalUserNoTeam.spec.ts
View file

@ -27,18 +27,18 @@ test.describe("Internal User with no team memberships", () => {
await page.getByRole("button", { name: /Create New Key/i }).click();
await expect(page.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 });
const teamSelect = page.locator(".ant-select", { hasText: "Search or select a team" });
const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox");
await teamSelect.click();
const dropdown = page.locator(".ant-select-dropdown:visible").first();
const dropdown = page.locator('[data-slot="combobox-content"]:visible').first();
await expect(dropdown).toBeVisible({ timeout: 5_000 });
// Wait for the settled-empty state, not a transient one. The dropdown shows
// a spinner while teams load and only swaps in "No teams found" once the
// request resolves with nothing (team_dropdown.tsx renders the spinner when
// isLoading and this copy otherwise). Asserting on it means a regression
// where teams DO load for this user fails here instead of racing a one-shot
// count() against an in-flight request.
// "Loading teams…" while teams load and only swaps in "No teams found" once
// the request resolves with nothing (team_dropdown.tsx passes both copies to
// PaginatedSearchSelect). Asserting on it means a regression where teams DO
// load for this user fails here instead of racing a one-shot count() against
// an in-flight request.
await expect(dropdown.getByText("No teams found")).toBeVisible({ timeout: 10_000 });
await expect(dropdown.getByRole("option")).toHaveCount(0);
});

4
tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts
View file

@ -18,10 +18,10 @@ test.describe("Internal User with team memberships", () => {
await page.getByRole("button", { name: /Create New Key/i }).click();
await expect(page.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 });
const teamSelect = page.locator(".ant-select", { hasText: "Search or select a team" });
const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox");
await teamSelect.click();
const dropdown = page.locator(".ant-select-dropdown:visible").first();
const dropdown = page.locator('[data-slot="combobox-content"]:visible').first();
await expect(dropdown).toBeVisible({ timeout: 5_000 });
// Both seeded memberships render, and nothing else does — proving the

6
tests/e2e/ui/tests/modelsPage/addModel.spec.ts
View file

@ -328,11 +328,11 @@ test.describe("Add Model", () => {
const teamByokRow = page.locator(".ant-form-item", { hasText: "Team-BYOK Model" });
await teamByokRow.getByRole("switch").click();
// TeamDropdown's options carry custom markup and no role="option", so match by text.
const teamDropdown = page.getByTestId("team-dropdown");
// TeamDropdown options show the alias above the team id, so match on the id line by text.
const teamDropdown = page.getByTestId("team-dropdown").getByRole("combobox");
await expect(teamDropdown).toBeVisible({ timeout: 5_000 });
await teamDropdown.click();
const teamOption = page.locator(".ant-select-dropdown:visible").getByText(E2E_TEAM_CRUD_ID).first();
const teamOption = page.locator('[data-slot="combobox-content"]:visible').getByText(E2E_TEAM_CRUD_ID).first();
await expect(teamOption).toBeVisible({ timeout: 5_000 });
await teamOption.click();

8
tests/e2e/ui/tests/proxy-admin/keys.spec.ts
View file

@ -40,11 +40,11 @@ test.describe("Proxy Admin - Keys", () => {
const keyName = `e2e-admin-key-${Date.now()}`;
await page.getByTestId("base-input").fill(keyName);
// Select team — the team dropdown has placeholder "Search or select a team"
const teamSelect = page.locator(".ant-select", { hasText: "Search or select a team" });
// Select team
const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox");
await teamSelect.click();
await page.keyboard.type(E2E_TEAM_CRUD_ALIAS);
await page.locator(".ant-select-dropdown:visible").getByText(E2E_TEAM_CRUD_ALIAS).first().click();
await page.locator('[data-slot="combobox-content"]:visible').getByText(E2E_TEAM_CRUD_ALIAS).first().click();
// Select models
await page.locator(".ant-select-selection-overflow").click();
@ -157,7 +157,7 @@ test.describe("Proxy Admin - Keys", () => {
await page.getByRole("button", { name: "More key actions" }).click();
await page.getByRole("menuitem", { name: "Delete Key" }).click();
const modal = page.locator(".ant-modal:visible");
const modal = page.getByRole("dialog", { name: "Delete Key" });
await expect(modal).toBeVisible({ timeout: 5_000 });
await modal.locator("input").fill(E2E_DELETE_KEY_ALIAS);

18
tests/e2e/ui/tests/proxy-admin/teams.spec.ts
View file

@ -47,10 +47,10 @@ test.describe("Proxy Admin - Teams", () => {
// Fill Team Name — the input has id="team_alias"
await dialog.locator("#team_alias").fill(uniqueAlias);
// Select models — the models multi-select is inside the modal
// Click to open dropdown, select "All Proxy Models"
await dialog.locator(".ant-select-selection-overflow").first().click();
await page.locator(".ant-select-dropdown:visible").getByText("All Proxy Models").click();
// Select models — the models multi-select is inside the modal. Its popup is
// portaled to the body, so scope the option lookup to the page, not the dialog.
await dialog.getByTestId("create-team-models-select").getByRole("combobox").click();
await page.getByRole("option", { name: "All Proxy Models", exact: true }).click();
await page.keyboard.press("Escape");
// Submit — click the submit button inside the dialog (not the header button)
@ -129,7 +129,7 @@ test.describe("Proxy Admin - Teams", () => {
await teamRow.locator('[data-testid^="team-actions-"]').click();
await page.getByTestId("team-action-delete").click();
const modal = page.locator(".ant-modal:visible");
const modal = page.getByRole("dialog", { name: "Delete Team?" });
await expect(modal).toBeVisible({ timeout: 5_000 });
await modal.locator("input").fill(E2E_TEAM_DELETE_ALIAS);
await modal.getByRole("button", { name: /Force Delete|Delete/i }).click();
@ -191,11 +191,11 @@ test.describe("Proxy Admin - Teams", () => {
const modelsSelect = page.locator("[data-testid='models-select']");
await expect(modelsSelect).toBeVisible({ timeout: 10_000 });
const anthropicTag = modelsSelect
.locator(".ant-select-selection-item")
const anthropicChip = modelsSelect
.locator('[data-slot="combobox-chip"]')
.filter({ hasText: "fake-anthropic-claude" });
await expect(anthropicTag).toBeVisible({ timeout: 5_000 });
await anthropicTag.locator(".ant-select-selection-item-remove").click();
await expect(anthropicChip).toBeVisible({ timeout: 5_000 });
await anthropicChip.locator('[data-slot="combobox-chip-remove"]').click();
await page.getByRole("button", { name: "Save Changes" }).click();

6
tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts
View file

@ -105,7 +105,7 @@ test.describe("Team Admin", () => {
await expect(row).toBeVisible({ timeout: 10_000 });
await row.getByTestId("delete-member").click();
const modal = page.locator(".ant-modal:visible");
const modal = page.getByRole("dialog", { name: "Delete Team Member" });
await expect(modal).toBeVisible({ timeout: 5_000 });
const remove = await captureRequestBody(page, { method: "POST", urlIncludes: "/team/member_delete" }, async () => {
@ -139,10 +139,10 @@ test.describe("Team Admin", () => {
await page.getByTestId("base-input").fill(keyName);
// Team selector — same locator pattern as the proxy-admin keys test.
const teamSelect = page.locator(".ant-select", { hasText: "Search or select a team" });
const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox");
await teamSelect.click();
await page.keyboard.type(E2E_TEAM_CRUD_ALIAS);
await page.locator(".ant-select-dropdown:visible").getByText(E2E_TEAM_CRUD_ALIAS).first().click();
await page.locator('[data-slot="combobox-content"]:visible').getByText(E2E_TEAM_CRUD_ALIAS).first().click();
// Models — pick "All Team Models"
await page.locator(".ant-select-selection-overflow").click();

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

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

282
tests/test_litellm/integrations/test_langfuse.py
View file

@ -314,7 +314,7 @@ class TestLangfuseUsageDetails(unittest.TestCase):
"litellm_params": {"metadata": {}},
"optional_params": {},
"litellm_call_id": "test-call-id-null-usage",
"standard_logging_object": None,
"standard_logging_object": self._build_standard_logging_payload(),
"response_cost": 0.0,
}
@ -382,16 +382,14 @@ class TestLangfuseUsageDetails(unittest.TestCase):
"model_id": "model-123",
"model_group": "openai",
"api_base": "https://api.openai.com",
# only real StandardLoggingMetadata fields: session_id, trace_name,
# headers and friends are request-metadata keys the allowlist drops,
# so a payload carrying them cannot occur in production
"metadata": {
"user_api_key_end_user_id": None,
"prompt_management_metadata": None,
"session_id": None,
"trace_name": None,
"trace_version": None,
"headers": None,
"endpoint": None,
"caching_groups": None,
"previous_models": None,
"user_api_key_hash": "hashed-key",
"user_api_key_alias": "canary-alias",
},
"hidden_params": {},
"request_tags": [],
@ -503,14 +501,251 @@ class TestLangfuseUsageDetails(unittest.TestCase):
# litellm_trace_id should be preferred over litellm_call_id
assert self.last_trace_kwargs.get("id") == "trace-id-from-kwargs"
def test_log_langfuse_v2_uses_litellm_trace_id_when_standard_logging_object_none(
self,
):
CANARY = "sk-lf-canary-SECRET-d4e5f6"
def _canary_request_metadata(self):
"""Raw request metadata shaped like the proxy builds it, credentials included."""
from litellm.proxy._types import UserAPIKeyAuth
team_logging = [
{
"callback_name": "langfuse",
"callback_vars": {"langfuse_secret_key": self.CANARY},
}
]
return {
"user_api_key_auth": UserAPIKeyAuth(
api_key="hashed-key",
team_metadata={"logging": team_logging},
),
"user_api_key_team_metadata": {"logging": team_logging},
"user_api_key_metadata": {"secret_manager_settings": {"vault_token": self.CANARY}},
"session_id": "canary-session",
"trace_name": "canary-trace",
"first_custom": "keep-first",
"second_custom": "keep-second",
"endpoint": "/v1/chat/completions",
"headers": {"authorization": f"Bearer {self.CANARY}"},
}
def _emitted_payload_text(self):
"""Every blob this logger handed to the langfuse SDK, as one searchable string."""
import json
blobs = [self.last_trace_kwargs]
if self.mock_langfuse_trace.generation.call_args is not None:
blobs.append(self.mock_langfuse_trace.generation.call_args.kwargs)
blobs.extend(call.kwargs for call in self.mock_langfuse_trace.span.call_args_list)
return json.dumps(blobs, default=repr)
def _drive_with_canary(self, extra_metadata=None, hidden_params=None):
metadata = {**self._canary_request_metadata(), **(extra_metadata or {})}
payload = self._build_standard_logging_payload(trace_id="canary-trace-id")
if hidden_params is not None:
payload["hidden_params"] = hidden_params
kwargs = {**self._build_langfuse_kwargs(payload), "response_cost": 0.25}
self.last_trace_kwargs = {}
self.mock_langfuse_trace.generation.reset_mock()
self.mock_langfuse_trace.span.reset_mock()
with patch(
"litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params",
side_effect=lambda generation_params, **kw: generation_params,
create=True,
):
self.logger._log_langfuse_v2(
user_id="user-1",
metadata=metadata,
litellm_params={"metadata": metadata},
output=None,
start_time=datetime.datetime(2024, 1, 1, 12, 0, 0),
end_time=datetime.datetime(2024, 1, 1, 12, 0, 1),
kwargs=kwargs,
optional_params={},
input=None,
response_obj=None,
level="INFO",
litellm_call_id="canary-call-id",
)
return self.mock_langfuse_trace.generation.call_args.kwargs["metadata"]
def test_team_callback_credentials_never_reach_langfuse(self):
"""
When standard_logging_object is None (failure case where
get_standard_logging_object_payload threw), litellm_trace_id from kwargs
should be used as the Langfuse trace_id. This matches the DB Session ID.
Regression for the credential leak: request metadata carries the whole
UserAPIKeyAuth object, whose team_metadata holds the customer's own langfuse
keys. The emitted blob is sourced from StandardLoggingPayload, so none of the
three credential carriers can ride along.
"""
generation_metadata = self._drive_with_canary()
assert self.CANARY not in self._emitted_payload_text()
for leaked_key in (
"user_api_key_auth",
"user_api_key_team_metadata",
"user_api_key_metadata",
):
assert leaked_key not in generation_metadata
def test_debug_langfuse_dump_carries_no_credentials(self):
"""
debug_langfuse dumps request metadata into the trace as a second emit site.
It must be sourced from the allowlisted payload too.
"""
self._drive_with_canary(extra_metadata={"debug_langfuse": True})
dumped = self.last_trace_kwargs["metadata"]["metadata_passed_to_litellm"]
assert "user_api_key_auth" not in dumped
assert self.CANARY not in self._emitted_payload_text()
def test_raw_request_metadata_reaches_the_emitted_blob_through_no_key(self):
"""
The emitted blob is the allowlist plus litellm enrichments, nothing else.
Nothing from raw request metadata is copied across, whatever its type, which
is what makes the credential exclusion structural rather than a filter that
has to be kept correct. Proxy callers keep their own metadata under the
allowlisted requester_metadata key.
"""
generation_metadata = self._drive_with_canary()
for caller_key in ("first_custom", "second_custom", "session_id", "trace_name"):
assert caller_key not in generation_metadata
def test_provider_specific_span_receives_the_emitted_blob(self):
"""
The provider span reads hidden_params, which is an enrichment on the emitted
blob rather than a key of request metadata. Handing it the steering dict
instead would silently stop emitting vertex grounding spans.
"""
self._drive_with_canary(hidden_params={"vertex_ai_grounding_metadata": ["ground-a", "ground-b"]})
span_inputs = [call.kwargs.get("input") for call in self.mock_langfuse_trace.span.call_args_list]
assert span_inputs == ["ground-a", "ground-b"]
assert self.CANARY not in self._emitted_payload_text()
def test_caller_cannot_spoof_an_allowlisted_identity_field(self):
"""
Request metadata never reaches the blob, so a caller naming user_api_key_alias
cannot have their value emitted in place of the proxy-resolved one.
"""
generation_metadata = self._drive_with_canary(
extra_metadata={"user_api_key_alias": "spoofed-by-caller"}
)
assert generation_metadata["user_api_key_alias"] == "canary-alias"
def test_caller_nested_metadata_cannot_erase_a_litellm_enrichment(self):
"""
log_requester_metadata drops any top-level key whose name also appears inside
requester_metadata. Sourcing the blob from the allowlist populates that nested
dict for real, so a caller naming a key litellm_response_cost would otherwise
blank out the cost litellm computed. Enrichments are layered after the dedupe.
"""
payload = self._build_standard_logging_payload(trace_id="canary-trace-id")
payload["metadata"]["requester_metadata"] = {"litellm_response_cost": "caller-value", "api_base": "caller"}
kwargs = {**self._build_langfuse_kwargs(payload), "response_cost": 0.25}
metadata = self._canary_request_metadata()
self.mock_langfuse_trace.generation.reset_mock()
with patch(
"litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params",
side_effect=lambda generation_params, **kw: generation_params,
create=True,
):
self.logger._log_langfuse_v2(
user_id="user-1",
metadata=metadata,
litellm_params={"metadata": metadata, "api_base": "https://real-api-base"},
output=None,
start_time=datetime.datetime(2024, 1, 1, 12, 0, 0),
end_time=datetime.datetime(2024, 1, 1, 12, 0, 1),
kwargs=kwargs,
optional_params={},
input=None,
response_obj=None,
level="INFO",
litellm_call_id="canary-call-id",
)
generation_metadata = self.mock_langfuse_trace.generation.call_args.kwargs["metadata"]
assert generation_metadata["litellm_response_cost"] == 0.25
assert generation_metadata["api_base"] == "https://real-api-base"
def test_denied_steering_keys_and_enrichments(self):
"""
endpoint is a plain string, so without the deny-list it would ride the
string re-injection straight into the emitted blob. The enrichments are
litellm-computed and must survive the move off clean_metadata.
"""
generation_metadata = self._drive_with_canary()
assert "endpoint" not in generation_metadata
assert "headers" not in generation_metadata
assert generation_metadata["litellm_response_cost"] == 0.25
assert "hidden_params" in generation_metadata
def test_cache_hit_is_normalized_on_the_shared_kwargs(self):
"""
kwargs here is the shared model_call_details dict. Callbacks that run after
langfuse read cache_hit off it and copy it into their own payloads, so
dropping the None to False normalization records None for datadog, logfire,
generic_api and spend tracking.
"""
metadata = self._canary_request_metadata()
payload = self._build_standard_logging_payload(trace_id="canary-trace-id")
kwargs = {**self._build_langfuse_kwargs(payload), "cache_hit": None}
with patch(
"litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params",
side_effect=lambda generation_params, **kw: generation_params,
create=True,
):
self.logger._log_langfuse_v2(
user_id="user-1",
metadata=metadata,
litellm_params={"metadata": metadata},
output=None,
start_time=datetime.datetime(2024, 1, 1, 12, 0, 0),
end_time=datetime.datetime(2024, 1, 1, 12, 0, 1),
kwargs=kwargs,
optional_params={},
input=None,
response_obj=None,
level="INFO",
litellm_call_id="canary-call-id",
)
assert kwargs["cache_hit"] is False
def test_redact_user_api_key_info_still_strips_the_emitted_blob(self):
"""
The flag used to act on the raw-derived blob. That blob is now sourced from
StandardLoggingPayload, which is where the user_api_key_* fields live, so the
redaction has to run on the assembled payload or the flag silently stops working.
"""
with patch.object(litellm, "redact_user_api_key_info", True):
generation_metadata = self._drive_with_canary()
assert not [key for key in generation_metadata if key.startswith("user_api_key")]
def test_steering_keys_still_read_from_raw_metadata(self):
"""
Only the emitted payload moves to StandardLoggingPayload. The control fields
keep reading raw metadata, which is what Braintrust's migration got wrong.
"""
self._drive_with_canary()
assert self.last_trace_kwargs.get("session_id") == "canary-session"
assert self.last_trace_kwargs.get("name") == "canary-trace"
def test_failure_trace_survives_a_missing_standard_logging_object(self):
"""
get_standard_logging_object_payload is fail-open and returns None on any
exception, which is exactly the failed-request case Langfuse most needs to
show. The trace is still emitted with the litellm_trace_id fallback, and the
blob degrades to caller strings plus enrichments rather than falling back to
raw metadata, which would ship the UserAPIKeyAuth object.
"""
metadata = self._canary_request_metadata()
kwargs = {
"standard_logging_object": None,
"model": "gpt-4",
@ -520,16 +755,17 @@ class TestLangfuseUsageDetails(unittest.TestCase):
"litellm_trace_id": "trace-id-failure",
}
self.last_trace_kwargs = {}
self.mock_langfuse_trace.generation.reset_mock()
with patch(
"litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params",
side_effect=lambda generation_params, **kwargs: generation_params,
create=True,
):
self.logger._log_langfuse_v2(
trace_id, _ = self.logger._log_langfuse_v2(
user_id="user-1",
metadata={},
litellm_params={"metadata": {}},
metadata=metadata,
litellm_params={"metadata": metadata},
output=None,
start_time=datetime.datetime.utcnow(),
end_time=datetime.datetime.utcnow(),
@ -541,8 +777,18 @@ class TestLangfuseUsageDetails(unittest.TestCase):
litellm_call_id="call-id-different",
)
# Must use litellm_trace_id, not litellm_call_id
import json
assert trace_id == "trace-id-failure"
assert self.last_trace_kwargs.get("id") == "trace-id-failure"
generation_metadata = self.mock_langfuse_trace.generation.call_args.kwargs["metadata"]
assert "user_api_key_auth" not in generation_metadata
assert self.CANARY not in self._emitted_payload_text()
assert "first_custom" not in generation_metadata
# hidden_params comes off the payload, so it is omitted rather than emitted
# as an unserializable placeholder
assert "hidden_params" not in generation_metadata
json.dumps(generation_metadata)
def test_log_langfuse_v2_session_id_passed_as_trace_session_id(self):
"""

82
tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py
View file

@ -133,6 +133,40 @@ class TestExceptionCheckers:
result = ExceptionCheckers.is_error_str_rate_limit(error_str)
assert result is True
def test_bare_429_in_body_is_ignored_when_status_code_says_otherwise(self):
"""A 429 echoed back inside a 400's body is not a rate limit.
Word boundaries don't help: 429 is an ordinary token id (" that" in several
tokenisers), so an echoed prompt_token_ids array reads as a standalone 429.
"""
error_str = (
'{"error":{"message":"`tools` must not be an empty array",'
'"type":"invalid_request_error"},'
'"prompt_token_ids":[9906,429,1234]}'
)
assert ExceptionCheckers.is_error_str_rate_limit(error_str, status_code=400) is False
def test_bare_429_still_detected_without_a_status_code(self):
"""With no status available, a standalone 429 still counts (unchanged behaviour)."""
assert ExceptionCheckers.is_error_str_rate_limit("HTTP 429 Too Many Requests") is True
assert ExceptionCheckers.is_error_str_rate_limit("HTTP 429 Too Many Requests", status_code=None) is True
assert ExceptionCheckers.is_error_str_rate_limit("HTTP 429 Too Many Requests", status_code=429) is True
def test_non_integer_status_code_does_not_suppress_bare_429(self):
"""A non-integer status counts as unknown, not as a contradiction."""
assert ExceptionCheckers.is_error_str_rate_limit("HTTP 429 Too Many Requests", status_code="not-an-int") is True
def test_rate_limit_phrase_is_honoured_under_a_non_429_status(self):
"""Phrase matching stays ungated: some providers report a real rate limit in
the text under a non-429 status (#11455)."""
assert (
ExceptionCheckers.is_error_str_rate_limit("FireworksException - rate limit exceeded", status_code=400)
is True
)
def test_is_azure_content_policy_violation_error_with_policy_violation_text(self):
"""Test detection of Azure content policy violation with explicit policy violation text"""
@ -300,6 +334,54 @@ def test_lemonade_context_window_error_mapping():
assert excinfo.value.model == model
def test_openai_compatible_400_with_bare_429_in_body_maps_to_bad_request():
"""A provider 400 whose echoed body contains a 429 must stay a 400.
``is_error_str_rate_limit`` runs before the status-code branch for
openai-compatible providers, so a validation error echoing the request back came
out as RateLimitError, which tells the caller to retry a request that cannot
succeed and books the failure against provider throttling.
"""
error_message = (
'{"error":{"message":"`tools` must not be an empty array",'
'"type":"invalid_request_error","code":400},'
'"prompt_token_ids":[9906,429,1234]}'
)
original_exception = OpenAIError(
status_code=400,
message=error_message,
headers={},
)
with pytest.raises(litellm.BadRequestError) as excinfo:
exception_type(
model="deepseek-ai/DeepSeek-V3",
original_exception=original_exception,
custom_llm_provider="deepinfra",
)
assert excinfo.value.status_code == 400
assert excinfo.value.llm_provider == "deepinfra"
def test_openai_compatible_429_still_maps_to_rate_limit():
"""A real 429 still maps to RateLimitError."""
original_exception = OpenAIError(
status_code=429,
message='{"error":{"message":"Too Many Requests","type":"rate_limit_error"}}',
headers={},
)
with pytest.raises(litellm.RateLimitError) as excinfo:
exception_type(
model="deepseek-ai/DeepSeek-V3",
original_exception=original_exception,
custom_llm_provider="deepinfra",
)
assert excinfo.value.status_code == 429
@pytest.mark.parametrize(
"error_message",
[

79
tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py
View file

@ -2028,3 +2028,82 @@ class TestCapabilityProbeUsesCallerProvider:
AnthropicModelInfo._is_adaptive_thinking_model("claude-opus-4-8", "anthropic")
is True
)
def test_create_anthropic_model_list_response_shape():
from litellm.llms.anthropic.common_utils import (
create_anthropic_model_list_response,
)
response = create_anthropic_model_list_response(
[
{"id": "claude-opus-4-6", "object": "model", "created": 0, "owned_by": "openai"},
{"id": "gpt-4o", "object": "model", "created": 0, "owned_by": "openai"},
{"id": "claude-haiku-4-5", "object": "model", "created": 0, "owned_by": "openai"},
]
)
assert "object" not in response
assert response["has_more"] is False
assert response["first_id"] == "claude-opus-4-6"
assert response["last_id"] == "claude-haiku-4-5"
assert [m["id"] for m in response["data"]] == [
"claude-opus-4-6",
"gpt-4o",
"claude-haiku-4-5",
]
for entry in response["data"]:
assert entry["type"] == "model"
assert entry["display_name"] == entry["id"]
# ISO 8601 with a Z suffix, as the Anthropic Models API returns.
assert entry["created_at"].endswith("Z")
assert "+00:00" not in entry["created_at"]
assert "max_input_tokens" not in entry
assert "max_tokens" not in entry
def test_create_anthropic_model_list_response_carries_token_limits():
from litellm.llms.anthropic.common_utils import (
create_anthropic_model_list_response,
)
response = create_anthropic_model_list_response(
[
{
"id": "claude-opus-4-6",
"object": "model",
"created": 0,
"owned_by": "openai",
"max_input_tokens": 200000,
"max_output_tokens": 64000,
},
{
"id": "input-only",
"object": "model",
"created": 0,
"owned_by": "openai",
"max_input_tokens": 8192,
},
{"id": "unknown-limits", "object": "model", "created": 0, "owned_by": "openai"},
]
)
opus, input_only, unknown = response["data"]
assert opus["max_input_tokens"] == 200000
assert opus["max_tokens"] == 64000
assert "max_output_tokens" not in opus
assert input_only["max_input_tokens"] == 8192
assert "max_tokens" not in input_only
assert "max_input_tokens" not in unknown
assert "max_tokens" not in unknown
def test_create_anthropic_model_list_response_empty():
from litellm.llms.anthropic.common_utils import (
create_anthropic_model_list_response,
)
response = create_anthropic_model_list_response([])
assert response["data"] == []
assert response["has_more"] is False
assert response["first_id"] is None
assert response["last_id"] is None

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

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

83
tests/test_litellm/llms/openai/test_cost_calculation.py Normal file
View file

@ -0,0 +1,83 @@
"""Tests for per-second transcription cost calculation."""
import pytest
import litellm
from litellm.llms.openai.cost_calculation import cost_per_second
def _register_stt(name: str, **pricing: float) -> None:
litellm.register_model(
{
name: {
"mode": "audio_transcription",
"litellm_provider": "openai",
**pricing,
}
},
persist_across_reloads=False,
)
def test_input_rate_bills_when_output_rate_is_zero():
"""A declared-but-zero output rate must not suppress the real input rate."""
_register_stt(
"test-stt-zero-output",
input_cost_per_second=5e-05,
output_cost_per_second=0.0,
)
prompt_cost, completion_cost = cost_per_second(
model="test-stt-zero-output", custom_llm_provider="openai", duration=300.0
)
assert prompt_cost == pytest.approx(0.015)
assert completion_cost == 0.0
def test_output_rate_takes_precedence_when_both_are_billable():
"""Entries duplicating one rate into both fields must not be billed twice."""
_register_stt(
"test-stt-both-rates",
input_cost_per_second=1e-04,
output_cost_per_second=1e-04,
)
prompt_cost, completion_cost = cost_per_second(
model="test-stt-both-rates", custom_llm_provider="openai", duration=10.0
)
assert prompt_cost + completion_cost == pytest.approx(1e-03)
def test_output_rate_alone_still_bills():
_register_stt("test-stt-output-only", output_cost_per_second=3e-05)
prompt_cost, completion_cost = cost_per_second(
model="test-stt-output-only", custom_llm_provider="openai", duration=60.0
)
assert prompt_cost == 0.0
assert completion_cost == pytest.approx(1.8e-03)
@pytest.mark.parametrize(
"model, provider",
[
("deepgram/nova-3", "deepgram"),
("groq/whisper-large-v3", "groq"),
("elevenlabs/scribe_v1", "elevenlabs"),
("assemblyai/best", "assemblyai"),
("whisper-1", "openai"),
],
)
def test_shipped_per_second_models_bill_a_non_zero_cost(model, provider):
prompt_cost, completion_cost = cost_per_second(model=model, custom_llm_provider=provider, duration=60.0)
assert prompt_cost + completion_cost > 0.0
def test_whisper_bills_its_documented_rate_once():
prompt_cost, completion_cost = cost_per_second(model="whisper-1", custom_llm_provider="openai", duration=30.0)
assert prompt_cost + completion_cost == pytest.approx(0.003)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

20
tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py
View file

@ -1688,9 +1688,15 @@ class TestTemporaryMCPSessionEndpoints:
expires_at=datetime.utcnow() - timedelta(seconds=30),
)
cache = {"expired": expired_entry}
with patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints._temporary_mcp_servers",
cache,
with (
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints._temporary_mcp_servers",
cache,
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none",
return_value=None,
),
):
result = await get_cached_temporary_mcp_server("expired")
@ -2274,6 +2280,10 @@ class TestTemporaryMCPSessionEndpoints:
"litellm.proxy.management_endpoints.mcp_management_endpoints._cache_temporary_mcp_server_in_redis",
AsyncMock(),
) as redis_cache_mock,
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none",
return_value=None,
),
):
response = await add_session_mcp_server(
payload=payload,
@ -3419,6 +3429,10 @@ class TestTemporaryMCPSessionEndpoints:
"litellm.proxy.management_endpoints.mcp_management_endpoints.decrypt_value_helper",
return_value=serialized,
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none",
return_value=None,
),
):
result = await get_cached_temporary_mcp_server("from-redis")
finally:

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

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

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

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

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

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

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

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

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

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

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

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

103
tests/test_litellm/proxy/proxy_server/test_routes_models.py
View file

@ -99,6 +99,62 @@ def test_get_models_happy_path(client, auth_as, patched_models, path):
}
@pytest.mark.parametrize("path", ["/v1/models", "/models"])
def test_get_models_anthropic_format_when_header_present(
client, auth_as, patched_models, path
):
"""Pins: ``GET /v1/models`` returns the Anthropic-native models shape when
the caller sends an ``anthropic-version`` header (Claude Code gateway
discovery), while the default OpenAI shape is unchanged without it."""
with auth_as():
response = client.get(path, headers={"anthropic-version": "2023-06-01"})
assert response.status_code == 200
body = response.json()
assert "object" not in body
assert body["has_more"] is False
assert body["first_id"] == "gpt-4"
assert body["last_id"] == "claude-sonnet"
assert [m["id"] for m in body["data"]] == ["gpt-4", "claude-sonnet"]
for entry in body["data"]:
assert entry["type"] == "model"
assert entry["display_name"] == entry["id"]
assert entry["created_at"].endswith("Z")
@pytest.mark.parametrize("path", ["/v1/models", "/models"])
def test_anthropic_format_exposes_token_limits(
client, auth_as, patched_models, monkeypatch, path
):
"""Claude Code sizes requests off the listing, so the Anthropic-native entries
carry the same token limits the OpenAI listing resolves, with the output budget
named max_tokens as the Messages API names it."""
from litellm.proxy import utils as proxy_utils
def _create_model_info_response(model_id, provider="openai", **kwargs):
if model_id != "claude-sonnet":
return _stub_model_info_response(model_id=model_id, provider=provider)
return {
**_stub_model_info_response(model_id=model_id, provider=provider),
"max_input_tokens": 200000,
"max_output_tokens": 64000,
}
monkeypatch.setattr(
proxy_utils, "create_model_info_response", _create_model_info_response
)
with auth_as():
response = client.get(path, headers={"anthropic-version": "2023-06-01"})
assert response.status_code == 200
gpt_4, claude = response.json()["data"]
assert claude["max_input_tokens"] == 200000
assert claude["max_tokens"] == 64000
assert "max_output_tokens" not in claude
assert "max_input_tokens" not in gpt_4
assert "max_tokens" not in gpt_4
@pytest.mark.parametrize("path", ["/v1/models", "/models"])
def test_get_models_invalid_scope_returns_400(client, auth_as, patched_models, path):
"""Pins: ``GET /v1/models``, ``GET /models`` (error path: invalid scope)."""
@ -130,3 +186,50 @@ def test_get_model_by_id_not_found(client, auth_as, patched_models, path):
response = client.get(path)
assert response.status_code == 404
assert "not found" in response.text.lower()
@pytest.mark.parametrize("params", [{}, {"scope": "expand"}])
def test_anthropic_format_returns_public_team_model_name(
client, auth_as, patched_models, monkeypatch, params
):
"""Regression: the Anthropic-native listing must go through the same team
name translation as the OpenAI listing, so a caller never sees the internal
``model_name_{team_id}_{uuid}`` routing key."""
from litellm.proxy import utils as proxy_utils
from litellm.proxy.auth import model_checks
internal_name = "model_name_team-1_c0ffee"
patched_models.get_model_list = MagicMock(
return_value=[
{
"model_name": internal_name,
"model_info": {
"team_id": "team-1",
"team_public_model_name": "gpt-4-team",
},
}
]
)
patched_models.get_model_names = MagicMock(return_value=[internal_name])
async def _fake_get_available_models_for_user(**kwargs):
return [internal_name]
monkeypatch.setattr(
proxy_utils,
"get_available_models_for_user",
_fake_get_available_models_for_user,
)
monkeypatch.setattr(
model_checks, "get_complete_model_list", lambda **kwargs: [internal_name]
)
with auth_as():
response = client.get(
"/v1/models", params=params, headers={"anthropic-version": "2023-06-01"}
)
assert response.status_code == 200
assert [m["id"] for m in response.json()["data"]] == ["gpt-4-team"]
assert internal_name not in response.text

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

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

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

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

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

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

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

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

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

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

71
tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py
View file

@ -17,13 +17,14 @@ import hashlib
import json
from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
from typing import Any
from typing import Any, Final
from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi import HTTPException
from litellm.proxy._types import LiteLLM_VerificationTokenView
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
from litellm.proxy.utils import PrismaClient
@ -270,6 +271,9 @@ async def test_query_first_with_cached_plan_fallback_reconnects_then_retries_ide
assert retry_call.args == first_call.args == (original_query, "abc")
reconnect.assert_awaited_once()
assert reconnect.await_args.kwargs.get("force", False) is False
# https://github.com/BerriAI/litellm/issues/36418: without this the healthy
# writer probe skips the recreate and the stale plans survive the retry
assert reconnect.await_args.kwargs.get("force_recreate") is True
assert [name for name, *_ in manager.mock_calls] == [
"query_first",
"attempt_db_reconnect",
@ -564,3 +568,68 @@ async def test_get_data_team_keys_forward_limit_as_take(
"where": {"team_id": "team-1"},
"include": {"litellm_budget_table": True},
}
@pytest.mark.asyncio
async def test_query_first_with_cached_plan_fallback_reports_pre_query_engine_generation(
prisma_client: PrismaClient,
) -> None:
"""The generation is snapshotted before the query, not after it fails: it
names the engine that prepared the stale statement, which is what lets the
reconnect bypass an unrelated cooldown while that engine is still live
(https://github.com/BerriAI/litellm/issues/36418). Reading it after the
failure would miss a recreate that landed in between and force a
needless second one."""
prisma_client.db.engine_generation = 3
async def _fail_then_bump(*args: Any, **kwargs: Any) -> dict[str, str]:
if prisma_client.db.engine_generation == 3:
prisma_client.db.engine_generation = 4
raise RuntimeError("cached plan must not change result type")
return {"token": "abc"}
prisma_client.db.query_first = AsyncMock(side_effect=_fail_then_bump)
prisma_client.attempt_db_reconnect = AsyncMock(return_value=True)
await prisma_client._query_first_with_cached_plan_fallback("SELECT 1")
kwargs = prisma_client.attempt_db_reconnect.await_args.kwargs
assert kwargs.get("stale_read_engine").generation == 3
@pytest.mark.asyncio
async def test_query_first_with_cached_plan_fallback_reports_the_reader_generation(
prisma_client: PrismaClient,
) -> None:
"""With a read replica configured the query runs on the READER, so the
reader's generation is the one that names the engine holding the stale
prepared statement. Snapshotting the writer's instead would let an
unrelated writer reconnect re-arm the cooldown while the reader stayed
poisoned (https://github.com/BerriAI/litellm/issues/36418). The two
generations are deliberately far apart so only the right one matches."""
writer = MagicMock(name="writer")
writer.engine_generation = 99
writer.query_first = AsyncMock(return_value={"token": "wrong-engine"})
reader = MagicMock(name="reader")
reader.engine_generation = 3
reader.query_first = AsyncMock(
side_effect=[RuntimeError("cached plan must not change result type"), {"token": "abc"}]
)
prisma_client.db = RoutingPrismaWrapper(writer=writer, reader=reader)
prisma_client.attempt_db_reconnect = AsyncMock(return_value=True)
await prisma_client._query_first_with_cached_plan_fallback("SELECT 1")
reported: Final = prisma_client.attempt_db_reconnect.await_args.kwargs.get("stale_read_engine")
pinned = {
"reported_generation": reported.generation,
"reported_the_reader_itself": reported.wrapper is reader,
"reader_served_the_query": reader.query_first.await_count,
"writer_served_the_query": writer.query_first.await_count,
}
assert pinned == {
"reported_generation": 3,
"reported_the_reader_itself": True,
"reader_served_the_query": 2,
"writer_served_the_query": 0,
}

516
tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py
View file

@ -7,17 +7,34 @@ Symbols pinned here:
- ``PrismaClient.start_db_health_watchdog_task``
- ``PrismaClient.stop_db_health_watchdog_task``
- ``PrismaClient._db_health_watchdog_loop``
Note on fixtures for the routing tests: the reader and the writer carry
independent generation counters, so a fixture that gives them far-apart values
reads clearly and proves nothing about identity, because comparing the numbers
alone already yields the right answer. Pick values so that ONLY the mechanism
under test can produce the expected result, which for identity means two
engines whose generations deliberately coincide.
Note on what to assert: pin the requirement, not the mechanism. An assertion
that restates what the implementation currently does can only ever agree with
it, including when it is wrong, so it ends up defending the defect from being
corrected. One here did exactly that, asserting that a declined heavy-path
recreate leaves the dead-engine flag set, which read as a faithful description
and was a reintroduction of #29176. "A later cycle must not kill a healthy
engine" would have failed against it whatever mechanism produced it.
"""
from __future__ import annotations
import asyncio
from typing import Any
import time
from typing import Any, Final
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy.utils import PrismaClient
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
from litellm.proxy.utils import PrismaClient, _StaleReadEngine
@pytest.mark.asyncio
@ -96,6 +113,48 @@ async def test_run_reconnect_cycle_direct_path_recreates_when_probe_fails(
}
@pytest.mark.asyncio
async def test_run_reconnect_cycle_force_recreate_skips_probe_and_recreates(
prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A healthy writer must not veto the recreate when the caller already
knows the session state is poisoned (stale prepared statements after a
schema change). Regression for
https://github.com/BerriAI/litellm/issues/36418."""
monkeypatch.setenv("DATABASE_URL", "postgres://x:y@h:5432/db")
prisma_client._engine_confirmed_dead = False
prisma_client._engine_pid = 0
prisma_client._start_engine_watcher = AsyncMock()
prisma_client._cleanup_engine_watcher = MagicMock()
writer = prisma_client.db
writer.recreate_prisma_client = AsyncMock()
writer.query_raw = AsyncMock(return_value=[{"?column?": 1}])
await prisma_client._run_reconnect_cycle(timeout_seconds=5, force_recreate=True)
pinned = {
"recreate_called": writer.recreate_prisma_client.await_count,
"writer_query_raw_calls": writer.query_raw.await_count,
}
assert pinned == {"recreate_called": 1, "writer_query_raw_calls": 1}
@pytest.mark.asyncio
async def test_attempt_db_reconnect_forwards_force_recreate_to_cycle(
prisma_client: PrismaClient,
) -> None:
"""Regression for https://github.com/BerriAI/litellm/issues/36418: the flag
has to survive both hops (attempt_db_reconnect -> inside-lock -> cycle),
otherwise the cached-plan caller silently gets a probe-gated reconnect."""
prisma_client._db_last_reconnect_attempt_ts = 0.0
prisma_client._run_reconnect_cycle = AsyncMock()
ok = await prisma_client.attempt_db_reconnect(reason="explicit", force_recreate=True)
assert ok is True
assert prisma_client._run_reconnect_cycle.await_args.kwargs.get("force_recreate") is True
@pytest.mark.asyncio
async def test_run_reconnect_cycle_passes_writer_generation_to_recreate(
prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch
@ -584,3 +643,456 @@ async def test_run_reconnect_cycle_heavy_path_forwards_entry_generation_to_recre
kwargs = prisma_client.db.recreate_prisma_client.await_args.kwargs
assert kwargs.get("expected_generation") == 4
@pytest.mark.asyncio
async def test_attempt_db_reconnect_bypasses_cooldown_for_still_live_stale_engine(
prisma_client: PrismaClient,
) -> None:
"""A schema change landing inside the cooldown of an earlier reconnect used
to leave auth failing until the cooldown elapsed. While the engine the
caller's failure came from is still the live one, the cooldown must not
gate the recreate. Regression for
https://github.com/BerriAI/litellm/issues/36418."""
prisma_client.db.engine_generation = 7
prisma_client._db_last_reconnect_attempt_ts = time.time()
prisma_client._run_reconnect_cycle = AsyncMock()
ok = await prisma_client.attempt_db_reconnect(
reason="postgres_cached_plan_error",
force_recreate=True,
stale_read_engine=_StaleReadEngine(wrapper=prisma_client.read_db, generation=7),
)
assert ok is True
assert prisma_client._run_reconnect_cycle.await_count == 1
@pytest.mark.asyncio
async def test_attempt_db_reconnect_honors_cooldown_once_stale_engine_replaced(
prisma_client: PrismaClient,
) -> None:
"""The bypass is scoped to the damaged engine: once a concurrent recreate
has replaced it, the cooldown must still collapse the rest of the burst
onto that recreate instead of killing the fresh engine."""
prisma_client.db.engine_generation = 8
prisma_client._db_last_reconnect_attempt_ts = time.time()
prisma_client._run_reconnect_cycle = AsyncMock()
ok = await prisma_client.attempt_db_reconnect(
reason="postgres_cached_plan_error",
force_recreate=True,
stale_read_engine=_StaleReadEngine(wrapper=prisma_client.read_db, generation=7),
)
assert ok is False
prisma_client._run_reconnect_cycle.assert_not_awaited()
@pytest.mark.asyncio
async def test_attempt_db_reconnect_keeps_cooldown_for_callers_without_generation(
prisma_client: PrismaClient,
) -> None:
"""Watchdog and transport-error callers name no generation, so they keep
the plain cooldown behaviour."""
prisma_client._db_last_reconnect_attempt_ts = time.time()
prisma_client._run_reconnect_cycle = AsyncMock()
ok = await prisma_client.attempt_db_reconnect(reason="watchdog_probe_failed")
assert ok is False
prisma_client._run_reconnect_cycle.assert_not_awaited()
def _routing_client(prisma_client: PrismaClient, reader_generation: int, writer_generation: int) -> tuple[Any, Any]:
"""Wire ``prisma_client.db`` to a routing wrapper with distinct engines.
Returns the (writer, reader) mocks so a test can move either generation
independently, which is the only way to tell the two counters apart.
"""
writer = MagicMock(name="writer")
writer.engine_generation = writer_generation
reader = MagicMock(name="reader")
reader.engine_generation = reader_generation
prisma_client.db = RoutingPrismaWrapper(writer=writer, reader=reader)
return writer, reader
@pytest.mark.asyncio
async def test_attempt_db_reconnect_reads_generation_from_the_reader_that_served_the_query(
prisma_client: PrismaClient,
) -> None:
"""``query_first`` is a top-level read, so with a replica configured the
stale prepared statements are on the READER. A writer reconnect that moved
the writer generation must not re-arm the cooldown while the reader the
query actually failed on is still the live, poisoned one. Regression for
https://github.com/BerriAI/litellm/issues/36418."""
_routing_client(prisma_client, reader_generation=7, writer_generation=99)
prisma_client._db_last_reconnect_attempt_ts = time.time()
prisma_client._run_reconnect_cycle = AsyncMock()
ok = await prisma_client.attempt_db_reconnect(
reason="postgres_cached_plan_error",
force_recreate=True,
stale_read_engine=_StaleReadEngine(wrapper=prisma_client.read_db, generation=7),
)
assert ok is True
assert prisma_client._run_reconnect_cycle.await_count == 1
@pytest.mark.asyncio
async def test_attempt_db_reconnect_honors_cooldown_once_the_reader_itself_was_replaced(
prisma_client: PrismaClient,
) -> None:
"""The mirror of the above: once the reader has been replaced, the recreate
the caller needed has already happened, so the cooldown collapses the rest
of the burst even though the writer generation never moved."""
_routing_client(prisma_client, reader_generation=8, writer_generation=99)
prisma_client._db_last_reconnect_attempt_ts = time.time()
prisma_client._run_reconnect_cycle = AsyncMock()
ok = await prisma_client.attempt_db_reconnect(
reason="postgres_cached_plan_error",
force_recreate=True,
stale_read_engine=_StaleReadEngine(wrapper=prisma_client.read_db, generation=7),
)
assert ok is False
prisma_client._run_reconnect_cycle.assert_not_awaited()
@pytest.mark.asyncio
async def test_attempt_db_reconnect_gates_when_reads_moved_to_an_engine_of_the_same_generation(
prisma_client: PrismaClient,
) -> None:
"""The counters are per engine, so the reader and the writer can sit on the
same number at the same time. Once the reader goes unavailable reads move to
the writer, and the caller's poisoned reader is no longer serving anything,
so the cooldown should gate it. Comparing generations alone cannot tell the
two apart and would hand out the waiver here: the generations are equal on
purpose, which is what makes this the case identity has to decide."""
writer, reader = _routing_client(prisma_client, reader_generation=5, writer_generation=5)
stale: Final = _StaleReadEngine(wrapper=reader, generation=5)
prisma_client.db._reader_unavailable = True
prisma_client._db_last_reconnect_attempt_ts = time.time()
prisma_client._run_reconnect_cycle = AsyncMock()
ok = await prisma_client.attempt_db_reconnect(
reason="postgres_cached_plan_error",
force_recreate=True,
stale_read_engine=stale,
)
pinned = {
"reads_now_served_by_the_writer": prisma_client.read_db is writer,
"generations_coincide": reader.engine_generation == writer.engine_generation,
# `_cooldown_applies` gates on the failed-repair record OR on liveness,
# and either alone produces this result. Pin that the record is empty,
# or a stray entry would make this pass while testing the other gate.
"no_failed_repair_recorded": dict(prisma_client._failed_recreate_generations) == {},
"recovered": ok,
"cycles_run": prisma_client._run_reconnect_cycle.await_count,
}
assert pinned == {
"reads_now_served_by_the_writer": True,
"generations_coincide": True,
"no_failed_repair_recorded": True,
"recovered": False,
"cycles_run": 0,
}
@pytest.mark.asyncio
async def test_failed_repair_of_one_engine_is_not_evicted_by_a_failure_on_the_other(
prisma_client: PrismaClient,
) -> None:
"""The record is kept per engine. Held in a single slot, a failed writer
repair would evict the reader's record, and the next caller naming the
reader's still-unrepaired generation would get the waiver back and run its
own redundant cycle, which is the burst the record exists to collapse."""
writer, reader = _routing_client(prisma_client, reader_generation=5, writer_generation=3)
stale_reader: Final = _StaleReadEngine(wrapper=reader, generation=5)
stale_writer: Final = _StaleReadEngine(wrapper=writer, generation=3)
prisma_client._db_last_reconnect_attempt_ts = 0.0
prisma_client._run_reconnect_cycle = AsyncMock(side_effect=RuntimeError("engine spawn failed"))
await prisma_client.attempt_db_reconnect(
reason="postgres_cached_plan_error", force_recreate=True, stale_read_engine=stale_reader
)
prisma_client.db._reader_unavailable = True
await prisma_client.attempt_db_reconnect(
reason="postgres_cached_plan_error", force_recreate=True, stale_read_engine=stale_writer
)
prisma_client.db._reader_unavailable = False
cycles_before_the_reader_returns: Final = prisma_client._run_reconnect_cycle.await_count
await prisma_client.attempt_db_reconnect(
reason="postgres_cached_plan_error", force_recreate=True, stale_read_engine=stale_reader
)
pinned = {
"cycles_before": cycles_before_the_reader_returns,
"cycles_after": prisma_client._run_reconnect_cycle.await_count,
}
assert pinned == {"cycles_before": 2, "cycles_after": 2}
@pytest.mark.asyncio
async def test_attempt_db_reconnect_withdraws_the_waiver_after_this_generation_failed_to_repair(
prisma_client: PrismaClient,
) -> None:
"""A failed recreate leaves the generation where it was, so without a record
of the failure every queued caller of the same burst would still see its own
generation live and run its own full recreate serially instead of collapsing
onto one attempt. Drives two callers rather than presetting the record, so
the record has to actually be written by the failure."""
prisma_client.db.engine_generation = 7
prisma_client._db_last_reconnect_attempt_ts = 0.0
prisma_client._run_reconnect_cycle = AsyncMock(side_effect=RuntimeError("engine spawn failed"))
first = await prisma_client.attempt_db_reconnect(
reason="postgres_cached_plan_error",
force_recreate=True,
stale_read_engine=_StaleReadEngine(wrapper=prisma_client.read_db, generation=7),
)
second = await prisma_client.attempt_db_reconnect(
reason="postgres_cached_plan_error",
force_recreate=True,
stale_read_engine=_StaleReadEngine(wrapper=prisma_client.read_db, generation=7),
)
pinned = {
"first": first,
"second": second,
"cycles_run": prisma_client._run_reconnect_cycle.await_count,
}
assert pinned == {"first": False, "second": False, "cycles_run": 1}
@pytest.mark.asyncio
async def test_attempt_db_reconnect_keeps_the_waiver_after_an_unrelated_reconnect_failure(
prisma_client: PrismaClient,
) -> None:
"""The failure record is scoped to the generation it was trying to repair.
A watchdog or transport-error reconnect names no generation, so its failure
says nothing about whether a stale read engine can be repaired and must not
gate it: gating on a global failure count would 503 authentication for the
length of the cooldown."""
prisma_client.db.engine_generation = 7
prisma_client._db_last_reconnect_attempt_ts = 0.0
prisma_client._run_reconnect_cycle = AsyncMock(side_effect=RuntimeError("watchdog reconnect failed"))
unrelated = await prisma_client.attempt_db_reconnect(reason="watchdog_probe_failed")
# Read before the second call: a global failure gate would be armed here,
# and the recovering reconnect below resets the counter either way.
failures_left_by_the_unrelated_reconnect: Final = prisma_client._consecutive_reconnect_failures
prisma_client._run_reconnect_cycle = AsyncMock()
cached_plan = await prisma_client.attempt_db_reconnect(
reason="postgres_cached_plan_error",
force_recreate=True,
stale_read_engine=_StaleReadEngine(wrapper=prisma_client.read_db, generation=7),
)
pinned = {
"unrelated_failed": unrelated,
"failures_left_by_the_unrelated_reconnect": failures_left_by_the_unrelated_reconnect,
"cached_plan_recovered": cached_plan,
"cycles_run_for_cached_plan": prisma_client._run_reconnect_cycle.await_count,
}
assert pinned == {
"unrelated_failed": False,
"failures_left_by_the_unrelated_reconnect": 1,
"cached_plan_recovered": True,
"cycles_run_for_cached_plan": 1,
}
@pytest.mark.asyncio
async def test_forced_recreate_declined_by_the_generation_guard_is_not_reported_as_success(
prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch
) -> None:
"""``recreate_prisma_client`` declines when the writer generation moved
since cycle entry, and the routing wrapper then leaves the reader untouched
too. A forced caller asked for its engine to be replaced and it was not, so
reporting success would reset the consecutive-failure count and log a repair
that never happened. The declined attempt must equally not count as a
failure, or the caller's own backoff would be gated on its next try."""
monkeypatch.setenv("DATABASE_URL", "postgres://x:y@h:5432/db")
prisma_client._engine_confirmed_dead = False
prisma_client._engine_pid = 0
prisma_client._start_engine_watcher = AsyncMock()
prisma_client._cleanup_engine_watcher = MagicMock()
prisma_client._db_last_reconnect_attempt_ts = 0.0
prisma_client._consecutive_reconnect_failures = 0
writer = prisma_client.db
writer.recreate_prisma_client = AsyncMock(return_value=False)
writer.query_raw = AsyncMock(return_value=[{"?column?": 1}])
ok = await prisma_client.attempt_db_reconnect(
reason="postgres_cached_plan_error",
force_recreate=True,
)
pinned = {
"reported_success": ok,
"recreate_attempted": writer.recreate_prisma_client.await_count,
"consecutive_failures": prisma_client._consecutive_reconnect_failures,
}
assert pinned == {
"reported_success": False,
"recreate_attempted": 1,
"consecutive_failures": 0,
}
@pytest.mark.asyncio
async def test_unforced_recreate_declined_by_the_generation_guard_still_succeeds(
prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch
) -> None:
"""The decline is only an error for a caller that forced the recreate. A
transport-blip caller is happy to learn another path already replaced the
engine, so its reconnect still reports success."""
monkeypatch.setenv("DATABASE_URL", "postgres://x:y@h:5432/db")
prisma_client._engine_confirmed_dead = False
prisma_client._engine_pid = 0
prisma_client._start_engine_watcher = AsyncMock()
prisma_client._cleanup_engine_watcher = MagicMock()
prisma_client._db_last_reconnect_attempt_ts = 0.0
writer = prisma_client.db
writer.recreate_prisma_client = AsyncMock(return_value=False)
# First call is the liveness probe, which must fail so the recreate is
# reached at all; the second is the post-recreate smoke test.
writer.query_raw = AsyncMock(side_effect=[Exception("probe fails"), [{"?column?": 1}]])
ok = await prisma_client.attempt_db_reconnect(reason="transport_blip")
pinned = {"reported_success": ok, "recreate_attempted": writer.recreate_prisma_client.await_count}
assert pinned == {"reported_success": True, "recreate_attempted": 1}
@pytest.mark.asyncio
async def test_heavy_path_forced_recreate_declined_is_not_reported_as_success(
prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A forced caller reaches the heavy branch too: the escalation threshold
flips ``_engine_confirmed_dead`` after repeated failures, and every cycle
after that takes the dead-engine path. A decline there has to be treated
exactly as it is on the direct path, or the escalation itself reintroduces
the success that never happened."""
monkeypatch.setenv("DATABASE_URL", "postgres://x:y@h:5432/db")
prisma_client._engine_confirmed_dead = True
prisma_client._engine_pid = 1234
prisma_client._start_engine_watcher = AsyncMock()
prisma_client._cleanup_engine_watcher = MagicMock()
prisma_client._db_last_reconnect_attempt_ts = 0.0
prisma_client._consecutive_reconnect_failures = 0
monkeypatch.setattr(PrismaClient, "_reap_all_zombies", staticmethod(lambda: set()))
prisma_client.db.recreate_prisma_client = AsyncMock(return_value=False)
ok = await prisma_client.attempt_db_reconnect(
reason="postgres_cached_plan_error",
force_recreate=True,
)
pinned = {
"reported_success": ok,
"recreate_attempted": prisma_client.db.recreate_prisma_client.await_count,
"consecutive_failures": prisma_client._consecutive_reconnect_failures,
# The dead-engine flag must be CLEARED. A raise normally skips the
# clear, which is right for a failure and wrong here: the guard
# declined because another path had already replaced the engine, so it
# is alive. Leaving it set routes the next cycle back down this
# probe-free branch, where the recreate would kill that healthy engine.
"engine_still_confirmed_dead": prisma_client._engine_confirmed_dead,
}
assert pinned == {
"reported_success": False,
"recreate_attempted": 1,
"consecutive_failures": 0,
"engine_still_confirmed_dead": False,
}
@pytest.mark.asyncio
async def test_declined_heavy_recreate_disarms_escalation_for_the_next_attempt(
prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Clearing the dead-engine flag on a decline is not enough on its own. The
escalation check re-arms that flag whenever the consecutive-failure count is
still at the threshold, so a decline that left the count alone would send
the very next attempt back down the probe-free heavy path and recreate over
the healthy engine another path had just installed. Drives the SECOND
attempt, because the first one alone cannot show this."""
monkeypatch.setenv("DATABASE_URL", "postgres://x:y@h:5432/db")
prisma_client._engine_pid = 1234
prisma_client._start_engine_watcher = AsyncMock()
prisma_client._cleanup_engine_watcher = MagicMock()
prisma_client._db_last_reconnect_attempt_ts = 0.0
monkeypatch.setattr(PrismaClient, "_reap_all_zombies", staticmethod(lambda: set()))
# Escalation already armed by earlier genuine failures.
prisma_client._consecutive_reconnect_failures = prisma_client._reconnect_escalation_threshold
prisma_client.db.recreate_prisma_client = AsyncMock(return_value=False)
prisma_client.db.query_raw = AsyncMock(return_value=[{"?column?": 1}])
armed: Final = prisma_client._engine_confirmed_dead is False and prisma_client._consecutive_reconnect_failures > 0
await prisma_client.attempt_db_reconnect(reason="postgres_cached_plan_error", force_recreate=True)
# Kept as its own assert, not folded into the judgement below. These are two
# claims about two moments, the first being a precondition for the second
# meaning anything, and a single combined comparison would hide which one
# failed from both the traceback and a mutation report.
assert {
"escalation_was_armed_by_the_count": armed,
"failures": prisma_client._consecutive_reconnect_failures,
"engine_confirmed_dead": prisma_client._engine_confirmed_dead,
} == {"escalation_was_armed_by_the_count": True, "failures": 0, "engine_confirmed_dead": False}
prisma_client._db_last_reconnect_attempt_ts = 0.0
await prisma_client.attempt_db_reconnect(reason="postgres_cached_plan_error", force_recreate=True)
# The requirement: a later cycle must not reclassify the healthy replacement
# as dead and restart it through the probe-free path.
assert prisma_client._engine_confirmed_dead is False
@pytest.mark.asyncio
async def test_unrelated_reconnect_failure_does_not_erase_the_burst_record(
prisma_client: PrismaClient,
) -> None:
"""The failure record names one engine, so a caller that names none must
not overwrite it. Otherwise a watchdog failure landing between two callers
of the same burst clears the record and the second caller runs its own full
recreate against the engine the first one just failed to repair."""
prisma_client.db.engine_generation = 7
prisma_client._db_last_reconnect_attempt_ts = 0.0
stale: Final = _StaleReadEngine(wrapper=prisma_client.read_db, generation=7)
prisma_client._run_reconnect_cycle = AsyncMock(side_effect=RuntimeError("engine spawn failed"))
await prisma_client.attempt_db_reconnect(
reason="postgres_cached_plan_error",
force_recreate=True,
stale_read_engine=stale,
)
# force=True the way the engine-death callers do, so this one actually
# reaches the failure branch instead of being skipped by the cooldown the
# first caller just stamped.
await prisma_client.attempt_db_reconnect(reason="engine_process_death", force=True)
cycles_before_the_second_burst_caller: Final = prisma_client._run_reconnect_cycle.await_count
await prisma_client.attempt_db_reconnect(
reason="postgres_cached_plan_error",
force_recreate=True,
stale_read_engine=stale,
)
pinned = {
"cycles_before": cycles_before_the_second_burst_caller,
"cycles_after": prisma_client._run_reconnect_cycle.await_count,
}
assert pinned == {"cycles_before": 2, "cycles_after": 2}

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

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

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

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

103
tests/test_litellm/test_main.py
View file

@ -1,3 +1,5 @@
import contextlib
import copy
import json
import os
import sys
@ -2461,3 +2463,104 @@ async def test_acompletion_forwards_aws_credentials_through_responses_bridge(
finally:
litellm.disable_aiohttp_transport = original_disable_aiohttp
litellm.in_memory_llm_clients_cache.flush_cache()
_GEMINI_RESPONSE_BODY = {
"candidates": [{"content": {"parts": [{"text": "hello"}], "role": "model"}, "finishReason": "STOP"}],
"usageMetadata": {"promptTokenCount": 2, "candidatesTokenCount": 1, "totalTokenCount": 3},
}
def _gemini_client_returning_a_reply():
"""An injected HTTP client whose post() answers like generativelanguage does."""
from litellm.llms.custom_httpx.http_handler import HTTPHandler
client = HTTPHandler()
request = httpx.Request("POST", "https://generativelanguage.googleapis.com/")
post = MagicMock(return_value=httpx.Response(200, json=_GEMINI_RESPONSE_BODY, request=request))
return client, post
@pytest.fixture
def restore_model_registry():
"""litellm.model_cost and the provider name sets are module-global.
register_model merges into the existing entry in place, hence the deep copy.
"""
model_cost = copy.deepcopy(litellm.model_cost)
openai_models = set(litellm.open_ai_chat_completion_models)
yield
litellm.model_cost.clear()
litellm.model_cost.update(model_cost)
litellm.open_ai_chat_completion_models.clear()
litellm.open_ai_chat_completion_models.update(openai_models)
def test_openai_model_name_does_not_outrank_explicit_provider():
"""`gemini/gpt-4o` goes to Google, not to litellm's OpenAI handler.
completion() checks `model in litellm.open_ai_chat_completion_models` ahead of
the gemini branch, so the call used to reach the OpenAI handler carrying
VertexGeminiConfig, whose transform_request raises NotImplementedError.
"""
assert "gpt-4o" in litellm.open_ai_chat_completion_models
client, post = _gemini_client_returning_a_reply()
with patch.object(client, "post", new=post):
response = litellm.completion(
model="gemini/gpt-4o",
messages=[{"role": "user", "content": "hello"}],
api_key="test-api-key",
client=client,
)
assert "generativelanguage.googleapis.com" in post.call_args.kwargs["url"]
assert "models/gpt-4o" in post.call_args.kwargs["url"]
assert response.choices[0].message.content == "hello"
def test_mislabelled_pricing_entry_does_not_reroute_provider(restore_model_registry):
"""register_model is the other way into the same failure.
An entry claiming litellm_provider "openai" adds its name to
open_ai_chat_completion_models, so one mislabelled price reroutes every later
call to that model in the process.
"""
litellm.register_model(
{
"gemini-2.5-pro": {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 1e-06,
"output_cost_per_token": 4e-06,
}
}
)
assert "gemini-2.5-pro" in litellm.open_ai_chat_completion_models
client, post = _gemini_client_returning_a_reply()
with patch.object(client, "post", new=post):
response = litellm.completion(
model="gemini/gemini-2.5-pro",
messages=[{"role": "user", "content": "hello"}],
api_key="test-api-key",
client=client,
)
assert "generativelanguage.googleapis.com" in post.call_args.kwargs["url"]
assert response.choices[0].message.content == "hello"
def test_openai_model_without_a_provider_still_routes_to_openai():
from openai import OpenAI
client = OpenAI(api_key="fake-key")
raw_response = client.chat.completions.with_raw_response
with patch.object(raw_response, "create") as mock_create, contextlib.suppress(Exception):
litellm.completion(
model="gpt-4o",
messages=[{"role": "user", "content": "hello"}],
client=client,
)
mock_create.assert_called()

42
tests/test_litellm/test_pre_commit_lint.py
View file

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

6
type-discipline-budget.json
View file

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

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

@ -140,9 +140,6 @@
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
},
"react-hooks/purity": {
"count": 1
},
@ -228,7 +225,7 @@
"count": 2
},
"no-restricted-imports": {
"count": 2
"count": 1
}
},
"src/app/(dashboard)/cost-tracking/_components/how_it_works.tsx": {
@ -239,9 +236,6 @@
"src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/cost-tracking/_components/pricing_calculator/multi_cost_results.test.tsx": {
@ -252,9 +246,6 @@
"src/app/(dashboard)/cost-tracking/_components/pricing_calculator/multi_cost_results.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 2
}
},
"src/app/(dashboard)/cost-tracking/_components/pricing_calculator/multi_export_dropdown.tsx": {
@ -275,17 +266,11 @@
"src/app/(dashboard)/cost-tracking/_components/provider_discount_table.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/cost-tracking/_components/provider_margin_table.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/cost-tracking/_components/use_discount_config.ts": {
@ -313,11 +298,6 @@
"count": 3
}
},
"src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.tsx": {
"no-nested-ternary": {
"count": 5
@ -1405,9 +1385,6 @@
"no-nested-ternary": {
"count": 1
},
"no-restricted-imports": {
"count": 2
},
"prefer-const": {
"count": 2
}
@ -1499,9 +1476,6 @@
"src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx": {
"local/no-complex-jsx-arrow": {
"count": 2
},
"no-restricted-imports": {
"count": 1
}
},
"src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.tsx": {
@ -1522,9 +1496,6 @@
"no-nested-ternary": {
"count": 1
},
"no-restricted-imports": {
"count": 1
},
"react-hooks/purity": {
"count": 1
},
@ -1682,21 +1653,11 @@
"count": 1
}
},
"src/app/onboarding/OnboardingErrorView.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/app/onboarding/OnboardingFormBody.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/app/onboarding/OnboardingLoadingView.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/AIHub/ModelHubTable.test.tsx": {
"max-params": {
"count": 1
@ -1709,63 +1670,30 @@
"no-nested-ternary": {
"count": 1
},
"no-restricted-imports": {
"count": 2
},
"prefer-const": {
"count": 4
}
},
"src/components/AIHub/SkillHubDashboard.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/AIHub/UsefulLinksManagement.tsx": {
"no-restricted-imports": {
"count": 1
},
"react-hooks/set-state-in-effect": {
"count": 1
}
},
"src/components/AIHub/forms/MakeAgentPublicForm.tsx": {
"no-restricted-imports": {
"count": 2
},
"react-hooks/set-state-in-effect": {
"count": 1
}
},
"src/components/AIHub/forms/MakeMCPPublicForm.test.tsx": {
"react/display-name": {
"count": 1
}
},
"src/components/AIHub/forms/MakeMCPPublicForm.tsx": {
"no-nested-ternary": {
"count": 2
},
"no-restricted-imports": {
"count": 2
},
"react-hooks/set-state-in-effect": {
"count": 1
}
},
"src/components/AIHub/forms/MakeModelPublicForm.tsx": {
"no-restricted-imports": {
"count": 2
},
"react-hooks/set-state-in-effect": {
"count": 1
}
},
"src/components/BetaBadge.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/CloudZeroCostTracking/CloudZeroCreateModal.tsx": {
"no-restricted-imports": {
"count": 1
@ -1784,42 +1712,9 @@
"count": 1
}
},
"src/components/DebugWarningBanner.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/DeprecationBanner.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/EntityUsageExport/ExportSummary.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/EntityUsageExport/UsageExportHeader.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/EntityUsageExport/types.ts": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/EntityUsageExport/utils.test.ts": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/EntityUsageExport/utils.ts": {
"max-params": {
"count": 3
},
"no-restricted-imports": {
"count": 1
}
},
"src/components/GuardrailSettingsView.tsx": {
@ -1830,9 +1725,6 @@
"src/components/GuardrailsMonitor/LogViewer.tsx": {
"no-nested-ternary": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/components/HelpLink.test.tsx": {
@ -1840,54 +1732,16 @@
"count": 1
}
},
"src/components/LicenseExpiryBanner.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/ModelSelect/ModelSelect.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/Navbar/BlogDropdown/BlogDropdown.test.tsx": {
"max-nested-callbacks": {
"count": 12
}
},
"src/components/Navbar/BlogDropdown/BlogDropdown.tsx": {
"no-restricted-imports": {
"count": 2
}
},
"src/components/Navbar/CommunityEngagementButtons/CommunityEngagementButtons.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/Navbar/NotificationsBell/NotificationsBell.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/Navbar/UserDropdown/UserDropdown.tsx": {
"no-restricted-imports": {
"count": 2
},
"react-hooks/set-state-in-effect": {
"count": 1
}
},
"src/components/Navbar/ViewSwitcher.tsx": {
"no-restricted-imports": {
"count": 2
}
},
"src/components/Navbar/WorkerDropdown/WorkerDropdown.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/SCIM.tsx": {
"no-restricted-imports": {
"count": 2
@ -1983,9 +1837,6 @@
}
},
"src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx": {
"no-restricted-imports": {
"count": 2
},
"react-hooks/set-state-in-effect": {
"count": 1
}
@ -1995,11 +1846,6 @@
"count": 1
}
},
"src/components/Settings/RouterSettings/Fallbacks/EditFallbacks.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/Settings/RouterSettings/Fallbacks/FallbackGroupConfig.tsx": {
"local/no-complex-jsx-arrow": {
"count": 1
@ -2017,9 +1863,6 @@
}
},
"src/components/Settings/RouterSettings/Fallbacks/Fallbacks.tsx": {
"no-restricted-imports": {
"count": 2
},
"prefer-const": {
"count": 2
}
@ -2295,9 +2138,6 @@
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 2
},
"react-hooks/set-state-in-effect": {
"count": 1
}
@ -2341,9 +2181,6 @@
}
},
"src/components/claude_code_plugins/MakeSkillPublicForm.tsx": {
"no-restricted-imports": {
"count": 2
},
"react-hooks/set-state-in-effect": {
"count": 1
}
@ -2372,20 +2209,7 @@
"count": 2
}
},
"src/components/common_components/AutoRotationView.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/common_components/DefaultProxyAdminTag.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/common_components/DeleteResourceModal.tsx": {
"no-restricted-imports": {
"count": 1
},
"react-hooks/set-state-in-effect": {
"count": 1
}
@ -2395,16 +2219,6 @@
"count": 1
}
},
"src/components/common_components/IconActionButton/BaseActionButton.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/common_components/KeyLifecycleSettings.tsx": {
"local/no-complex-jsx-arrow": {
"count": 1
@ -2413,16 +2227,6 @@
"count": 2
}
},
"src/components/common_components/LabeledField.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/common_components/MemberTable.tsx": {
"no-restricted-imports": {
"count": 2
}
},
"src/components/common_components/MetadataKeyValueFields.test.tsx": {
"no-restricted-imports": {
"count": 1
@ -2434,34 +2238,15 @@
}
},
"src/components/common_components/ModelAliasManager.tsx": {
"no-restricted-imports": {
"count": 1
},
"react-hooks/set-state-in-effect": {
"count": 1
}
},
"src/components/common_components/ModelSelector.tsx": {
"no-restricted-imports": {
"count": 2
},
"react-hooks/set-state-in-effect": {
"count": 1
}
},
"src/components/common_components/NewBadge.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/common_components/OrganizationDropdown.tsx": {
"local/no-complex-jsx-arrow": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/components/common_components/PassThroughGuardrailsSection.tsx": {
"no-restricted-imports": {
"count": 2
@ -2470,29 +2255,11 @@
"count": 1
}
},
"src/components/common_components/PassThroughRoutesSelector.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/common_components/PassThroughSecuritySection.tsx": {
"no-restricted-imports": {
"count": 2
}
},
"src/components/common_components/PremiumLoggingSettings.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/common_components/ProjectDropdown.tsx": {
"local/no-complex-jsx-arrow": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/components/common_components/RateLimitTypeFormItem.test.tsx": {
"no-restricted-imports": {
"count": 1
@ -2503,22 +2270,9 @@
"count": 1
}
},
"src/components/common_components/RouterSettingsAccordion.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/common_components/budget_duration_dropdown.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/components/common_components/chartUtils.test.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/common_components/chartUtils.tsx": {
@ -2527,9 +2281,6 @@
},
"no-nested-ternary": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/components/common_components/check_openapi_schema.tsx": {
@ -2554,17 +2305,11 @@
},
"no-nested-ternary": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/components/common_components/team_dropdown.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/components/common_components/team_multi_select.tsx": {
@ -2630,11 +2375,6 @@
"count": 1
}
},
"src/components/key_team_helpers/TagRateLimitEditor.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/key_team_helpers/fetch_available_models_team_key.tsx": {
"local/filename-pascal-case": {
"count": 1
@ -2696,9 +2436,6 @@
"src/components/logging_settings_view.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/components/mcp_server_management/MCPServerSelector.tsx": {
@ -2712,9 +2449,6 @@
"src/components/mcp_server_management/MCPToolPermissions.tsx": {
"local/no-complex-jsx-arrow": {
"count": 1
},
"no-restricted-imports": {
"count": 2
}
},
"src/components/mcp_tools/ByokCredentialModal.tsx": {
@ -2733,9 +2467,6 @@
"src/components/mcp_tools/McpCrudPermissionPanel.tsx": {
"no-nested-ternary": {
"count": 3
},
"no-restricted-imports": {
"count": 2
}
},
"src/components/mcp_tools/types.tsx": {
@ -2764,18 +2495,12 @@
"src/components/model_filters.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/components/model_group_alias_settings.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
},
"react-hooks/set-state-in-effect": {
"count": 1
}
@ -2837,9 +2562,6 @@
"src/components/navbar.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/components/networking.tsx": {
@ -2865,9 +2587,6 @@
"src/components/object_permissions_view.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/components/onboarding_link.tsx": {
@ -2914,9 +2633,6 @@
"src/components/organization/organization_view.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/components/page_utils.test.ts": {
@ -2940,22 +2656,9 @@
"count": 1
}
},
"src/components/permissions/AgentPermissions.tsx": {
"no-restricted-imports": {
"count": 2
}
},
"src/components/permissions/MCPServerPermissions.tsx": {
"no-nested-ternary": {
"count": 3
},
"no-restricted-imports": {
"count": 2
}
},
"src/components/permissions/VectorStorePermissions.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/policies/PolicySelector.tsx": {
@ -2982,9 +2685,6 @@
},
"max-lines": {
"count": 1
},
"no-restricted-imports": {
"count": 2
}
},
"src/components/query_param_input.tsx": {
@ -2997,17 +2697,9 @@
"count": 1
}
},
"src/components/router_settings/LatencyBasedConfiguration.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/router_settings/ReliabilityRetriesSection.tsx": {
"no-nested-ternary": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/components/router_settings/RoutingStrategySelector.tsx": {
@ -3015,18 +2707,10 @@
"count": 1
}
},
"src/components/router_settings/TagFilteringToggle.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/router_settings/index.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
},
"prefer-const": {
"count": 2
}
@ -3052,33 +2736,17 @@
"count": 1
}
},
"src/components/settings.test.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/settings.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"local/no-complex-jsx-arrow": {
"count": 4
},
"no-nested-ternary": {
"count": 2
},
"no-restricted-imports": {
"count": 3
},
"prefer-const": {
"count": 4
}
},
"src/components/shared/CreatedKeyDisplay.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/shared/advanced_date_picker.tsx": {
"local/filename-pascal-case": {
"count": 1
@ -3143,9 +2811,6 @@
"src/components/shared/numerical_input.tsx": {
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/components/shared/table_cells/cell_tooltip.tsx": {
@ -3242,9 +2907,6 @@
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 2
},
"react-hooks/set-state-in-effect": {
"count": 1
}
@ -3259,11 +2921,6 @@
"count": 1
}
},
"src/components/templates/KeyInfoHeader.tsx": {
"no-restricted-imports": {
"count": 2
}
},
"src/components/templates/key_edit_view.tsx": {
"local/filename-pascal-case": {
"count": 1
@ -3293,9 +2950,6 @@
"no-nested-ternary": {
"count": 1
},
"no-restricted-imports": {
"count": 2
},
"react-hooks/set-state-in-effect": {
"count": 1
}
@ -3475,9 +3129,6 @@
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
},
"react-hooks/set-state-in-effect": {
"count": 1
}
@ -3486,9 +3137,6 @@
"local/filename-pascal-case": {
"count": 1
},
"no-restricted-imports": {
"count": 1
},
"prefer-const": {
"count": 1
},
@ -3506,29 +3154,15 @@
"count": 1
}
},
"src/components/view_logs/CostBreakdownViewer.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/view_logs/EvalViewer/EvalViewer.tsx": {
"local/no-complex-jsx-arrow": {
"count": 1
},
"no-nested-ternary": {
"count": 1
},
"no-restricted-imports": {
"count": 1
}
},
"src/components/view_logs/GuardrailViewer/CompliancePanel.tsx": {
"no-nested-ternary": {
"count": 2
},
"no-restricted-imports": {
"count": 1
},
"react-hooks/set-state-in-effect": {
"count": 1
}
@ -3541,31 +3175,17 @@
"src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx": {
"no-nested-ternary": {
"count": 4
},
"no-restricted-imports": {
"count": 1
}
},
"src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx": {
"no-nested-ternary": {
"count": 3
},
"no-restricted-imports": {
"count": 1
}
},
"src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx": {
"no-nested-ternary": {
"count": 2
},
"no-restricted-imports": {
"count": 1
},
"react-hooks/set-state-in-effect": {
"count": 2
}
@ -3575,36 +3195,11 @@
"count": 2
}
},
"src/components/view_logs/LogDetailsDrawer/RealtimePrettyView.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/view_logs/LogDetailsDrawer/useKeyboardNavigation.ts": {
"react-hooks/immutability": {
"count": 2
}
},
"src/components/view_logs/ToolsSection/FormattedToolView.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/view_logs/ToolsSection/ToolExpandedContent.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/view_logs/ToolsSection/ToolItem.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/view_logs/VectorStoreViewer.tsx": {
"no-restricted-imports": {
"count": 1
}
},
"src/components/view_logs/columns.tsx": {
"local/filename-pascal-case": {
"count": 1

5
ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.test.tsx
View file

@ -1,4 +1,5 @@
import { renderWithProviders, screen, within } from "@/../tests/test-utils";
import { waitFor } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { beforeEach, describe, expect, it, vi } from "vitest";
import { AccessGroupsPage } from "./AccessGroupsPage";
@ -215,7 +216,9 @@ describe("AccessGroupsPage", () => {
await user.click(await openRowMenu(user, "ag-1"));
const dialog = screen.getByRole("dialog", { name: "Delete Access Group" });
await user.click(within(dialog).getByRole("button", { name: "Cancel" }));
expect(screen.queryByRole("dialog", { name: "Delete Access Group" })).not.toBeInTheDocument();
await waitFor(() => {
expect(screen.queryByRole("dialog", { name: "Delete Access Group" })).not.toBeInTheDocument();
});
expect(mockMutate).not.toHaveBeenCalled();
});

2
ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx
View file

@ -1,4 +1,4 @@
import { DateRangePickerValue } from "@tremor/react";
import type { DateRangePickerValue } from "@/components/shared/date_picker_types";
import React, { useEffect, useState } from "react";
import NotificationsManager from "@/components/molecules/notifications_manager";
import UsageDatePicker from "@/components/shared/usage_date_picker";

107
ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.test.tsx
View file

@ -1,6 +1,6 @@
import React from "react";
import { describe, it, expect, vi, beforeEach } from "vitest";
import { screen } from "@testing-library/react";
import { act, screen, waitFor } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { renderWithProviders } from "../../../../../tests/test-utils";
import CostTrackingSettings from "./cost_tracking_settings";
@ -8,25 +8,29 @@ import CostTrackingSettings from "./cost_tracking_settings";
// Mock sub-hooks so we can control their state without network calls
const mockDiscountConfig = vi.fn(() => ({}));
const mockMarginConfig = vi.fn(() => ({}));
const mockRemoveDiscount = vi.fn();
const mockRemoveMargin = vi.fn();
const stableDiscountCallbacks = {
fetchDiscountConfig: vi.fn().mockResolvedValue(undefined),
handleAddProvider: vi.fn().mockResolvedValue(true),
handleRemoveProvider: mockRemoveDiscount,
handleDiscountChange: vi.fn().mockResolvedValue(undefined),
};
const stableMarginCallbacks = {
fetchMarginConfig: vi.fn().mockResolvedValue(undefined),
handleAddMargin: vi.fn().mockResolvedValue(true),
handleRemoveMargin: mockRemoveMargin,
handleMarginChange: vi.fn().mockResolvedValue(undefined),
};
vi.mock("./use_discount_config", () => ({
useDiscountConfig: () => ({
discountConfig: mockDiscountConfig(),
fetchDiscountConfig: vi.fn().mockResolvedValue(undefined),
handleAddProvider: vi.fn().mockResolvedValue(true),
handleRemoveProvider: vi.fn().mockResolvedValue(undefined),
handleDiscountChange: vi.fn().mockResolvedValue(undefined),
}),
useDiscountConfig: () => ({ discountConfig: mockDiscountConfig(), ...stableDiscountCallbacks }),
}));
vi.mock("./use_margin_config", () => ({
useMarginConfig: () => ({
marginConfig: mockMarginConfig(),
fetchMarginConfig: vi.fn().mockResolvedValue(undefined),
handleAddMargin: vi.fn().mockResolvedValue(true),
handleRemoveMargin: vi.fn().mockResolvedValue(undefined),
handleMarginChange: vi.fn().mockResolvedValue(undefined),
}),
useMarginConfig: () => ({ marginConfig: mockMarginConfig(), ...stableMarginCallbacks }),
}));
vi.mock("./pricing_calculator/index", () => ({
@ -153,6 +157,79 @@ describe("CostTrackingSettings", () => {
});
});
describe("removing a configured provider", () => {
const expandAndRemove = async (section: string, actionName: string) => {
const user = userEvent.setup();
renderWithProviders(<CostTrackingSettings {...ADMIN_PROPS} />);
await user.click(screen.getByText(section).closest("button")!);
await user.click(await screen.findByRole("button", { name: actionName }));
return user;
};
it("should ask to confirm before removing a discount", async () => {
mockDiscountConfig.mockReturnValue({ openai: 0.05 });
await expandAndRemove("Provider Discounts", "Remove discount for openai");
expect(await screen.findByRole("button", { name: "Remove" })).toBeInTheDocument();
expect(screen.getByText(/are you sure you want to remove the discount for openai\?/i)).toBeInTheDocument();
expect(mockRemoveDiscount).not.toHaveBeenCalled();
});
it("should remove the discount once removal is confirmed", async () => {
mockDiscountConfig.mockReturnValue({ openai: 0.05 });
const user = await expandAndRemove("Provider Discounts", "Remove discount for openai");
await user.click(await screen.findByRole("button", { name: "Remove" }));
expect(mockRemoveDiscount).toHaveBeenCalledWith("openai");
});
it("should leave the discount in place when the confirmation is cancelled", async () => {
mockDiscountConfig.mockReturnValue({ openai: 0.05 });
const user = await expandAndRemove("Provider Discounts", "Remove discount for openai");
await user.click(await screen.findByRole("button", { name: "Cancel" }));
expect(mockRemoveDiscount).not.toHaveBeenCalled();
expect(screen.queryByRole("button", { name: "Remove" })).not.toBeInTheDocument();
});
it("should hold the confirmation open while the removal is still in flight", async () => {
mockDiscountConfig.mockReturnValue({ openai: 0.05 });
const { promise, resolve: settleRemoval } = Promise.withResolvers<void>();
mockRemoveDiscount.mockReturnValue(promise);
const user = await expandAndRemove("Provider Discounts", "Remove discount for openai");
await user.click(await screen.findByRole("button", { name: "Remove" }));
const removing = await screen.findByRole("button", { name: "Removing…" });
expect(removing).toBeDisabled();
expect(screen.getByRole("button", { name: "Cancel" })).toBeDisabled();
await act(async () => {
settleRemoval();
});
await waitFor(() => {
expect(screen.queryByRole("alertdialog")).not.toBeInTheDocument();
});
expect(mockRemoveDiscount).toHaveBeenCalledWith("openai");
});
it("should remove the margin once removal is confirmed", async () => {
mockMarginConfig.mockReturnValue({ openai: 0.1 });
const user = await expandAndRemove("Fee/Price Margin", "Remove margin for openai");
expect(screen.getByText(/are you sure you want to remove the margin for openai\?/i)).toBeInTheDocument();
await user.click(await screen.findByRole("button", { name: "Remove" }));
expect(mockRemoveMargin).toHaveBeenCalledWith("openai");
});
});
describe("empty state messages", () => {
it("should show the empty state message when no discount config is loaded", async () => {
mockDiscountConfig.mockReturnValue({});

290
ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx
View file

@ -1,25 +1,24 @@
import React, { useState, useEffect } from "react";
import {
Title,
Text,
Button,
Accordion,
AccordionHeader,
AccordionBody,
TabGroup,
TabList,
Tab,
TabPanels,
TabPanel,
} from "@tremor/react";
import { ChevronDown } from "lucide-react";
import { Modal, Form } from "antd";
import {
AlertDialog,
AlertDialogCancel,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { Button } from "@/components/ui/button";
import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { CostTrackingSettingsProps } from "./types";
import ProviderDiscountTable from "./provider_discount_table";
import AddProviderForm from "./add_provider_form";
import ProviderMarginTable from "./provider_margin_table";
import AddMarginForm from "./add_margin_form";
import PricingCalculator from "./pricing_calculator/index";
import { ExclamationCircleOutlined } from "@ant-design/icons";
import { DocsMenu } from "@/components/HelpLink";
import HowItWorks from "./how_it_works";
import { useDiscountConfig } from "./use_discount_config";
@ -31,6 +30,29 @@ const DOCS_LINKS = [
{ label: "Spend tracking", href: "https://docs.litellm.ai/docs/proxy/cost_tracking" },
];
const REMOVAL_COPY = {
discount: { title: "Remove Provider Discount", noun: "discount" },
margin: { title: "Remove Provider Margin", noun: "margin" },
} as const;
interface PendingRemoval {
kind: keyof typeof REMOVAL_COPY;
provider: string;
displayName: string;
}
const SECTION_HEADER_CLASS = "group/section flex w-full items-center justify-between px-6 py-4 text-left";
const SectionHeader: React.FC<{ title: string; description: string }> = ({ title, description }) => (
<CollapsibleTrigger className={SECTION_HEADER_CLASS}>
<div className="flex flex-col items-start w-full">
<span className="block text-lg font-semibold text-gray-900">{title}</span>
<span className="block text-sm text-gray-500 mt-1">{description}</span>
</div>
<ChevronDown className="size-5 shrink-0 text-gray-500 transition-transform group-data-[panel-open]/section:rotate-180" />
</CollapsibleTrigger>
);
const CostTrackingSettings: React.FC<CostTrackingSettingsProps> = ({ userID, userRole, accessToken }) => {
const [selectedProvider, setSelectedProvider] = useState<string | undefined>(undefined);
const [newDiscount, setNewDiscount] = useState<string>("");
@ -42,9 +64,10 @@ const CostTrackingSettings: React.FC<CostTrackingSettingsProps> = ({ userID, use
const [percentageValue, setPercentageValue] = useState<string>("");
const [fixedAmountValue, setFixedAmountValue] = useState<string>("");
const [models, setModels] = useState<string[]>([]);
const [pendingRemoval, setPendingRemoval] = useState<PendingRemoval | null>(null);
const [isRemoving, setIsRemoving] = useState(false);
const [form] = Form.useForm();
const [marginForm] = Form.useForm();
const [modal, contextHolder] = Modal.useModal();
const isProxyAdmin = userRole === "proxy_admin" || userRole === "Admin";
@ -104,16 +127,23 @@ const CostTrackingSettings: React.FC<CostTrackingSettingsProps> = ({ userID, use
handleAddProvider();
};
const handleRemoveProvider = async (provider: string, providerDisplayName: string) => {
modal.confirm({
title: "Remove Provider Discount",
icon: <ExclamationCircleOutlined />,
content: `Are you sure you want to remove the discount for ${providerDisplayName}?`,
okText: "Remove",
okType: "danger",
cancelText: "Cancel",
onOk: () => removeProvider(provider),
});
const handleRemoveProvider = (provider: string, providerDisplayName: string) => {
setPendingRemoval({ kind: "discount", provider, displayName: providerDisplayName });
};
const handleConfirmRemoval = async () => {
if (!pendingRemoval) return;
setIsRemoving(true);
try {
if (pendingRemoval.kind === "discount") {
await removeProvider(pendingRemoval.provider);
} else {
await removeMargin(pendingRemoval.provider);
}
} finally {
setIsRemoving(false);
setPendingRemoval(null);
}
};
const handleAddMargin = async () => {
@ -141,16 +171,8 @@ const CostTrackingSettings: React.FC<CostTrackingSettingsProps> = ({ userID, use
setMarginType("percentage");
};
const handleRemoveMargin = async (provider: string, providerDisplayName: string) => {
modal.confirm({
title: "Remove Provider Margin",
icon: <ExclamationCircleOutlined />,
content: `Are you sure you want to remove the margin for ${providerDisplayName}?`,
okText: "Remove",
okType: "danger",
cancelText: "Cancel",
onOk: () => removeMargin(provider),
});
const handleRemoveMargin = (provider: string, providerDisplayName: string) => {
setPendingRemoval({ kind: "margin", provider, displayName: providerDisplayName });
};
if (!accessToken) {
@ -159,18 +181,16 @@ const CostTrackingSettings: React.FC<CostTrackingSettingsProps> = ({ userID, use
return (
<div className="w-full p-8">
{contextHolder}
{/* Header Section - Outside the card */}
<div className="flex flex-col md:flex-row items-start md:items-center justify-between mb-6">
<div>
<div className="flex items-center gap-2">
<Title>Cost Tracking Settings</Title>
<p className="text-xl font-medium text-gray-900">Cost Tracking Settings</p>
<DocsMenu items={DOCS_LINKS} />
</div>
<Text className="text-gray-500 mt-1">
<p className="text-gray-500 mt-1">
Configure cost discounts and margins for different LLM providers. Changes are saved automatically.
</Text>
</p>
</div>
</div>
@ -178,90 +198,78 @@ const CostTrackingSettings: React.FC<CostTrackingSettingsProps> = ({ userID, use
<div className="bg-white rounded-lg shadow-sm w-full max-w-full space-y-4">
{/* Accordion 1: Provider Discounts - Only for proxy admins */}
{isProxyAdmin && (
<Accordion>
<AccordionHeader className="px-6 py-4">
<div className="flex flex-col items-start w-full">
<Text className="text-lg font-semibold text-gray-900">Provider Discounts</Text>
<Text className="text-sm text-gray-500 mt-1">
Apply percentage-based discounts to reduce costs for specific providers
</Text>
</div>
</AccordionHeader>
<AccordionBody className="px-0">
<TabGroup>
<TabList className="px-6 pt-4">
<Tab>Discounts</Tab>
<Tab>Test It</Tab>
</TabList>
<TabPanels>
<TabPanel>
<div className="p-6">
<div className="flex justify-end mb-4">
<Button onClick={() => setIsModalVisible(true)}>+ Add Provider Discount</Button>
<Collapsible className="rounded-lg border">
<SectionHeader
title="Provider Discounts"
description="Apply percentage-based discounts to reduce costs for specific providers"
/>
<CollapsibleContent className="px-0">
<Tabs defaultValue="discounts">
<TabsList className="mx-6 mt-4">
<TabsTrigger value="discounts">Discounts</TabsTrigger>
<TabsTrigger value="test-it">Test It</TabsTrigger>
</TabsList>
<TabsContent value="discounts">
<div className="p-6">
<div className="flex justify-end mb-4">
<Button onClick={() => setIsModalVisible(true)}>+ Add Provider Discount</Button>
</div>
{isFetching ? (
<div className="py-12 text-center">
<p className="text-gray-500">Loading configuration...</p>
</div>
{isFetching ? (
<div className="py-12 text-center">
<Text className="text-gray-500">Loading configuration...</Text>
</div>
) : Object.keys(discountConfig).length > 0 ? (
<ProviderDiscountTable
discountConfig={discountConfig}
onDiscountChange={handleDiscountChange}
onRemoveProvider={handleRemoveProvider}
/>
) : (
<div className="py-16 px-6 text-center">
<svg
className="mx-auto h-12 w-12 text-gray-400 mb-4"
fill="none"
stroke="currentColor"
viewBox="0 0 24 24"
>
<path
strokeLinecap="round"
strokeLinejoin="round"
strokeWidth={1.5}
d="M12 8c-1.657 0-3 .895-3 2s1.343 2 3 2 3 .895 3 2-1.343 2-3 2m0-8c1.11 0 2.08.402 2.599 1M12 8V7m0 1v8m0 0v1m0-1c-1.11 0-2.08-.402-2.599-1M21 12a9 9 0 11-18 0 9 9 0 0118 0z"
/>
</svg>
<Text className="text-gray-700 font-medium mb-2">No provider discounts configured</Text>
<Text className="text-gray-500 text-sm">
Click &quot;Add Provider Discount&quot; to get started
</Text>
</div>
)}
</div>
</TabPanel>
<TabPanel>
<div className="px-6 pb-4">
<HowItWorks />
</div>
</TabPanel>
</TabPanels>
</TabGroup>
</AccordionBody>
</Accordion>
) : Object.keys(discountConfig).length > 0 ? (
<ProviderDiscountTable
discountConfig={discountConfig}
onDiscountChange={handleDiscountChange}
onRemoveProvider={handleRemoveProvider}
/>
) : (
<div className="py-16 px-6 text-center">
<svg
className="mx-auto h-12 w-12 text-gray-400 mb-4"
fill="none"
stroke="currentColor"
viewBox="0 0 24 24"
>
<path
strokeLinecap="round"
strokeLinejoin="round"
strokeWidth={1.5}
d="M12 8c-1.657 0-3 .895-3 2s1.343 2 3 2 3 .895 3 2-1.343 2-3 2m0-8c1.11 0 2.08.402 2.599 1M12 8V7m0 1v8m0 0v1m0-1c-1.11 0-2.08-.402-2.599-1M21 12a9 9 0 11-18 0 9 9 0 0118 0z"
/>
</svg>
<p className="text-gray-700 font-medium mb-2">No provider discounts configured</p>
<p className="text-gray-500 text-sm">Click &quot;Add Provider Discount&quot; to get started</p>
</div>
)}
</div>
</TabsContent>
<TabsContent value="test-it">
<div className="px-6 pb-4">
<HowItWorks />
</div>
</TabsContent>
</Tabs>
</CollapsibleContent>
</Collapsible>
)}
{/* Accordion 2: Fee/Price Margin - Only for proxy admins */}
{isProxyAdmin && (
<Accordion>
<AccordionHeader className="px-6 py-4">
<div className="flex flex-col items-start w-full">
<Text className="text-lg font-semibold text-gray-900">Fee/Price Margin</Text>
<Text className="text-sm text-gray-500 mt-1">
Add fees or margins to LLM costs for internal billing and cost recovery
</Text>
</div>
</AccordionHeader>
<AccordionBody className="px-0">
<Collapsible className="rounded-lg border">
<SectionHeader
title="Fee/Price Margin"
description="Add fees or margins to LLM costs for internal billing and cost recovery"
/>
<CollapsibleContent className="px-0">
<div className="p-6">
<div className="flex justify-end mb-4">
<Button onClick={() => setIsMarginModalVisible(true)}>+ Add Provider Margin</Button>
</div>
{isFetching ? (
<div className="py-12 text-center">
<Text className="text-gray-500">Loading configuration...</Text>
<p className="text-gray-500">Loading configuration...</p>
</div>
) : Object.keys(marginConfig).length > 0 ? (
<ProviderMarginTable
@ -284,33 +292,49 @@ const CostTrackingSettings: React.FC<CostTrackingSettingsProps> = ({ userID, use
d="M12 8c-1.657 0-3 .895-3 2s1.343 2 3 2 3 .895 3 2-1.343 2-3 2m0-8c1.11 0 2.08.402 2.599 1M12 8V7m0 1v8m0 0v1m0-1c-1.11 0-2.08-.402-2.599-1M21 12a9 9 0 11-18 0 9 9 0 0118 0z"
/>
</svg>
<Text className="text-gray-700 font-medium mb-2">No provider margins configured</Text>
<Text className="text-gray-500 text-sm">Click &quot;Add Provider Margin&quot; to get started</Text>
<p className="text-gray-700 font-medium mb-2">No provider margins configured</p>
<p className="text-gray-500 text-sm">Click &quot;Add Provider Margin&quot; to get started</p>
</div>
)}
</div>
</AccordionBody>
</Accordion>
</CollapsibleContent>
</Collapsible>
)}
{/* Accordion 3: Pricing Calculator - Available to all roles */}
<Accordion defaultOpen={true}>
<AccordionHeader className="px-6 py-4">
<div className="flex flex-col items-start w-full">
<Text className="text-lg font-semibold text-gray-900">Pricing Calculator</Text>
<Text className="text-sm text-gray-500 mt-1">
Estimate LLM costs based on expected token usage and request volume
</Text>
</div>
</AccordionHeader>
<AccordionBody className="px-0">
<Collapsible defaultOpen={true} className="rounded-lg border">
<SectionHeader
title="Pricing Calculator"
description="Estimate LLM costs based on expected token usage and request volume"
/>
<CollapsibleContent className="px-0">
<div className="p-6">
<PricingCalculator accessToken={accessToken} models={models} />
</div>
</AccordionBody>
</Accordion>
</CollapsibleContent>
</Collapsible>
</div>
{pendingRemoval && (
<AlertDialog open onOpenChange={(open) => !open && !isRemoving && setPendingRemoval(null)}>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle>{REMOVAL_COPY[pendingRemoval.kind].title}</AlertDialogTitle>
<AlertDialogDescription>
Are you sure you want to remove the {REMOVAL_COPY[pendingRemoval.kind].noun} for{" "}
{pendingRemoval.displayName}?
</AlertDialogDescription>
</AlertDialogHeader>
<AlertDialogFooter>
<AlertDialogCancel disabled={isRemoving}>Cancel</AlertDialogCancel>
<Button variant="destructive" onClick={handleConfirmRemoval} disabled={isRemoving}>
{isRemoving ? "Removing…" : "Remove"}
</Button>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
)}
<Modal
title={
<div className="flex items-center space-x-3 pb-4 border-b border-gray-100">
@ -328,10 +352,10 @@ const CostTrackingSettings: React.FC<CostTrackingSettingsProps> = ({ userID, use
}}
>
<div className="mt-6">
<Text className="text-sm text-gray-600 mb-6">
<p className="text-sm text-gray-600 mb-6">
Select a provider and set its discount percentage. Enter a value between 0% and 100% (e.g., 5 for a 5%
discount).
</Text>
</p>
<Form form={form} onFinish={handleFormSubmit} layout="vertical" className="space-y-6">
<AddProviderForm
discountConfig={discountConfig}
@ -362,10 +386,10 @@ const CostTrackingSettings: React.FC<CostTrackingSettingsProps> = ({ userID, use
}}
>
<div className="mt-6">
<Text className="text-sm text-gray-600 mb-6">
<p className="text-sm text-gray-600 mb-6">
Select a provider (or &quot;Global&quot; for all providers) and configure the margin. You can use
percentage-based or fixed amount.
</Text>
</p>
<Form form={marginForm} layout="vertical" className="space-y-6">
<AddMarginForm
marginConfig={marginConfig}

39
ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.test.tsx
View file

@ -41,6 +41,16 @@ const DEFAULT_PROPS = {
models: ["gpt-4", "gpt-3.5-turbo", "claude-3-sonnet"],
};
const dataRows = (): HTMLElement[] =>
within(screen.getByRole("table"))
.getAllByRole("row")
.filter((row) => within(row).queryAllByRole("combobox").length > 0);
const deleteButtonIn = (row: HTMLElement): HTMLElement => {
const cells = within(row).getAllByRole("cell");
return within(cells[cells.length - 1]).getByRole("button");
};
describe("PricingCalculator", () => {
beforeEach(() => {
vi.clearAllMocks();
@ -124,8 +134,31 @@ describe("PricingCalculator", () => {
it("should render column headers for Model, Input Tokens, and Output Tokens", () => {
renderWithProviders(<PricingCalculator {...DEFAULT_PROPS} />);
expect(screen.getByText("Model")).toBeInTheDocument();
expect(screen.getByText("Input Tokens")).toBeInTheDocument();
expect(screen.getByText("Output Tokens")).toBeInTheDocument();
expect(screen.getByRole("columnheader", { name: "Model" })).toBeInTheDocument();
expect(screen.getByRole("columnheader", { name: "Input Tokens" })).toBeInTheDocument();
expect(screen.getByRole("columnheader", { name: "Output Tokens" })).toBeInTheDocument();
});
it("should render a numeric field for input tokens, output tokens and requests", () => {
renderWithProviders(<PricingCalculator {...DEFAULT_PROPS} />);
expect(screen.getAllByRole("spinbutton")).toHaveLength(3);
});
it("should offer a model picker per row", () => {
renderWithProviders(<PricingCalculator {...DEFAULT_PROPS} />);
expect(screen.getAllByRole("combobox")).toHaveLength(1);
});
it("should remove a row when its delete button is clicked", async () => {
const user = userEvent.setup();
renderWithProviders(<PricingCalculator {...DEFAULT_PROPS} />);
await user.click(screen.getByRole("button", { name: /add another model/i }));
const withTwoRows = dataRows();
expect(withTwoRows).toHaveLength(2);
await user.click(deleteButtonIn(withTwoRows[1]));
expect(dataRows()).toHaveLength(1);
});
});

227
ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.tsx
View file

@ -1,6 +1,10 @@
import React, { useState, useCallback } from "react";
import { Table, Select, InputNumber, Button, Radio } from "antd";
import { DeleteOutlined, PlusOutlined } from "@ant-design/icons";
import { Plus, Trash2 } from "lucide-react";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group";
import { Table, TableBody, TableCell, TableFooter, TableHead, TableHeader, TableRow } from "@/components/ui/table";
import { SearchSelect } from "@/components/shared/SearchSelect";
import { PricingCalculatorProps, ModelEntry } from "./types";
import MultiCostResults from "./multi_cost_results";
import { useMultiCostEstimate } from "./use_multi_cost_estimate";
@ -63,132 +67,115 @@ const PricingCalculator: React.FC<PricingCalculatorProps> = ({ accessToken, mode
const multiModelResult = getMultiModelResult(entries);
const columns = [
{
title: "Model",
dataIndex: "model",
key: "model",
width: "35%",
render: (_: string, record: ModelEntry) => (
<Select
showSearch
placeholder="Select a model"
value={record.model || undefined}
onChange={(value) => handleEntryChange(record.id, "model", value)}
optionFilterProp="label"
filterOption={(input, option) =>
String(option?.label ?? "")
.toLowerCase()
.includes(input.toLowerCase())
}
options={models.map((model) => ({
value: model,
label: model,
}))}
style={{ width: "100%" }}
size="small"
/>
),
},
{
title: "Input Tokens",
dataIndex: "input_tokens",
key: "input_tokens",
width: "18%",
render: (_: number, record: ModelEntry) => (
<InputNumber
min={0}
value={record.input_tokens}
onChange={(value) => handleEntryChange(record.id, "input_tokens", value ?? 0)}
style={{ width: "100%" }}
size="small"
formatter={(value) => `${value}`.replace(/\B(?=(\d{3})+(?!\d))/g, ",")}
/>
),
},
{
title: "Output Tokens",
dataIndex: "output_tokens",
key: "output_tokens",
width: "18%",
render: (_: number, record: ModelEntry) => (
<InputNumber
min={0}
value={record.output_tokens}
onChange={(value) => handleEntryChange(record.id, "output_tokens", value ?? 0)}
style={{ width: "100%" }}
size="small"
formatter={(value) => `${value}`.replace(/\B(?=(\d{3})+(?!\d))/g, ",")}
/>
),
},
{
title: `Requests/${timePeriod === "day" ? "Day" : "Month"}`,
dataIndex: timePeriod === "day" ? "num_requests_per_day" : "num_requests_per_month",
key: "num_requests",
width: "20%",
render: (_: number | undefined, record: ModelEntry) => (
<InputNumber
min={0}
value={timePeriod === "day" ? record.num_requests_per_day : record.num_requests_per_month}
onChange={(value) =>
handleEntryChange(
record.id,
timePeriod === "day" ? "num_requests_per_day" : "num_requests_per_month",
value ?? undefined,
)
}
style={{ width: "100%" }}
size="small"
placeholder="-"
formatter={(value) => (value ? `${value}`.replace(/\B(?=(\d{3})+(?!\d))/g, ",") : "")}
/>
),
},
{
title: "",
key: "actions",
width: 50,
render: (_: unknown, record: ModelEntry) => (
<Button
type="text"
icon={<DeleteOutlined />}
onClick={() => handleRemoveEntry(record.id)}
disabled={entries.length === 1}
danger
size="small"
/>
),
},
];
const modelOptions = models.map((model) => ({ label: model, value: model }));
const requestsField = timePeriod === "day" ? "num_requests_per_day" : "num_requests_per_month";
return (
<div className="space-y-4">
<div className="flex items-center justify-end mb-2">
<Radio.Group
<RadioGroup
value={timePeriod}
onChange={(e) => handleTimePeriodChange(e.target.value)}
size="small"
optionType="button"
buttonStyle="solid"
onValueChange={(value) => handleTimePeriodChange(value as TimePeriod)}
className="flex w-auto items-center gap-4"
>
<Radio.Button value="day">Per Day</Radio.Button>
<Radio.Button value="month">Per Month</Radio.Button>
</Radio.Group>
<label className="flex cursor-pointer items-center gap-2 text-sm">
<RadioGroupItem value="day" />
Per Day
</label>
<label className="flex cursor-pointer items-center gap-2 text-sm">
<RadioGroupItem value="month" />
Per Month
</label>
</RadioGroup>
</div>
<Table
columns={columns}
dataSource={entries}
rowKey="id"
pagination={false}
size="small"
footer={() => (
<Button type="dashed" onClick={handleAddEntry} icon={<PlusOutlined />} className="w-full">
Add Another Model
</Button>
)}
/>
<Table>
<TableHeader>
<TableRow>
<TableHead className="w-[35%]">Model</TableHead>
<TableHead className="w-[18%]">Input Tokens</TableHead>
<TableHead className="w-[18%]">Output Tokens</TableHead>
<TableHead className="w-[20%]">Requests/{timePeriod === "day" ? "Day" : "Month"}</TableHead>
<TableHead className="w-[50px]">
<span className="sr-only">Actions</span>
</TableHead>
</TableRow>
</TableHeader>
<TableBody>
{entries.map((record, index) => (
<TableRow key={record.id}>
<TableCell className="whitespace-normal">
<SearchSelect
options={modelOptions}
value={record.model || undefined}
onValueChange={(value) => handleEntryChange(record.id, "model", value)}
placeholder="Select a model"
/>
</TableCell>
<TableCell>
<Input
type="number"
min={0}
className="h-8"
value={record.input_tokens}
onChange={(e) =>
handleEntryChange(record.id, "input_tokens", e.target.value === "" ? 0 : Number(e.target.value))
}
/>
</TableCell>
<TableCell>
<Input
type="number"
min={0}
className="h-8"
value={record.output_tokens}
onChange={(e) =>
handleEntryChange(record.id, "output_tokens", e.target.value === "" ? 0 : Number(e.target.value))
}
/>
</TableCell>
<TableCell>
<Input
type="number"
min={0}
className="h-8"
placeholder="-"
value={record[requestsField] ?? ""}
onChange={(e) =>
handleEntryChange(
record.id,
requestsField,
e.target.value === "" ? undefined : Number(e.target.value),
)
}
/>
</TableCell>
<TableCell>
<Button
variant="ghost"
size="icon-sm"
aria-label={`Remove model row ${index + 1}`}
onClick={() => handleRemoveEntry(record.id)}
disabled={entries.length === 1}
className="text-destructive"
>
<Trash2 className="size-3.5" />
</Button>
</TableCell>
</TableRow>
))}
</TableBody>
<TableFooter>
<TableRow>
<TableCell colSpan={5}>
<Button variant="outline" onClick={handleAddEntry} className="w-full border-dashed">
<Plus className="size-3.5" />
Add Another Model
</Button>
</TableCell>
</TableRow>
</TableFooter>
</Table>
<MultiCostResults multiResult={multiModelResult} timePeriod={timePeriod} />
</div>

90
ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/multi_cost_results.test.tsx
View file

@ -85,6 +85,14 @@ function emptyMultiResult(): MultiModelResult {
};
}
const expandToggle = (): HTMLElement => screen.getByRole("button", { name: /cost breakdown for / });
const shownBreakdown = (): HTMLElement | null => {
const label = screen.queryByText("Total/Request");
if (label === null) return null;
return label.closest("[style*='display: none']") === null ? label : null;
};
describe("MultiCostResults", () => {
beforeEach(() => {
vi.clearAllMocks();
@ -200,40 +208,78 @@ describe("MultiCostResults", () => {
expect(screen.getByRole("button", { name: /export/i })).toBeInTheDocument();
});
it("should render a column header for each summary column", () => {
renderWithProviders(<MultiCostResults multiResult={makeMultiResult()} timePeriod="day" />);
expect(screen.getByRole("columnheader", { name: "Model" })).toBeInTheDocument();
expect(screen.getByRole("columnheader", { name: "Per Request" })).toBeInTheDocument();
expect(screen.getByRole("columnheader", { name: "Margin Fee" })).toBeInTheDocument();
expect(screen.getByRole("columnheader", { name: "Daily" })).toBeInTheDocument();
});
it("should not show the model breakdown before the row is expanded", () => {
renderWithProviders(<MultiCostResults multiResult={makeMultiResult()} timePeriod="day" />);
expect(shownBreakdown()).toBeNull();
});
it("should expand the model breakdown row when the expand button is clicked", async () => {
const user = userEvent.setup();
renderWithProviders(<MultiCostResults multiResult={makeMultiResult()} timePeriod="day" />);
// The expand column renders a button (RightOutlined icon) for rows without errors
const expandButtons = screen.getAllByRole("button");
// Find the small expand button (not the Export button)
const expandButton = expandButtons.find((btn) => !btn.textContent?.toLowerCase().includes("export"));
expect(expandButton).toBeDefined();
await user.click(expandToggle());
await user.click(expandButton!);
// After expanding, the SingleModelBreakdown should be visible
expect(screen.getByText("Total/Request")).toBeInTheDocument();
expect(shownBreakdown()).toBeVisible();
expect(screen.getByText("Daily Total (100 req)")).toBeInTheDocument();
});
it("should show the collapse icon after expanding a row", async () => {
it("should collapse the model breakdown again on a second click", async () => {
const user = userEvent.setup();
renderWithProviders(<MultiCostResults multiResult={makeMultiResult()} timePeriod="day" />);
const getExpandButton = () => {
const allButtons = screen.getAllByRole("button");
return allButtons.find((btn) => !btn.textContent?.toLowerCase().includes("export"));
};
await user.click(expandToggle());
expect(shownBreakdown()).toBeVisible();
// Before expand: button has the "down" aria-label (RightOutlined renders as down in ant icons)
// Just verify clicking works and the breakdown content appears
await user.click(getExpandButton()!);
expect(screen.getByText("Total/Request")).toBeInTheDocument();
await user.click(expandToggle());
expect(shownBreakdown()).toBeNull();
});
// After a second click, the row collapses — content may be hidden or removed
await user.click(getExpandButton()!);
// The expanded content should no longer be visible
expect(screen.queryByText("Total/Request")).not.toBeVisible();
it("should name the breakdown toggle and report its expanded state", async () => {
const user = userEvent.setup();
renderWithProviders(<MultiCostResults multiResult={makeMultiResult()} timePeriod="day" />);
const toggle = screen.getByRole("button", { name: "Show cost breakdown for gpt-4" });
expect(toggle).toHaveAttribute("aria-expanded", "false");
await user.click(toggle);
const collapseToggle = screen.getByRole("button", { name: "Hide cost breakdown for gpt-4" });
expect(collapseToggle).toHaveAttribute("aria-expanded", "true");
});
it("should not offer an expand toggle for a row that failed", () => {
renderWithProviders(
<MultiCostResults
multiResult={makeMultiResult({
entries: [
{
entry: { id: "e1", model: "gpt-4", input_tokens: 1000, output_tokens: 500 },
result: makeCostResponse(),
loading: false,
error: null,
},
{
entry: { id: "e2", model: "bad-model", input_tokens: 0, output_tokens: 0 },
result: null,
loading: false,
error: "Pricing not found",
},
],
})}
timePeriod="day"
/>,
);
expect(screen.getAllByRole("button", { name: /cost breakdown for / })).toHaveLength(1);
});
});

349
ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/multi_cost_results.tsx
View file

@ -1,7 +1,11 @@
import React, { useState } from "react";
import { Text, Button } from "@tremor/react";
import { Card, Statistic, Row, Col, Divider, Spin, Table, Tag } from "antd";
import { LoadingOutlined, DownOutlined, RightOutlined } from "@ant-design/icons";
import { ChevronDown, ChevronRight } from "lucide-react";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { Card } from "@/components/ui/card";
import { Separator } from "@/components/ui/separator";
import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table";
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
import { CostEstimateResponse } from "../types";
import { formatNumberWithCommas } from "@/utils/dataUtils";
import { MultiModelResult } from "./types";
@ -41,55 +45,57 @@ const SingleModelBreakdown: React.FC<{
<div className="space-y-3 bg-gray-50 p-4 rounded-lg">
{loading && (
<div className="flex items-center gap-2 text-gray-500 text-sm">
<Spin indicator={<LoadingOutlined spin />} size="small" />
<UiLoadingSpinner className="size-3.5" />
<span>Updating...</span>
</div>
)}
<div className="grid grid-cols-4 gap-4">
<div>
<Text className="text-xs text-gray-500 block">Total/Request</Text>
<Text className="text-base font-semibold text-blue-600">{formatCost(result.cost_per_request)}</Text>
<div className="min-w-0">
<p className="text-xs text-gray-500 block">Total/Request</p>
<p className="text-base font-semibold text-blue-600 break-words">{formatCost(result.cost_per_request)}</p>
</div>
<div>
<Text className="text-xs text-gray-500 block">Input Cost</Text>
<Text className="text-sm">{formatCost(result.input_cost_per_request)}</Text>
<div className="min-w-0">
<p className="text-xs text-gray-500 block">Input Cost</p>
<p className="text-sm break-words">{formatCost(result.input_cost_per_request)}</p>
</div>
<div>
<Text className="text-xs text-gray-500 block">Output Cost</Text>
<Text className="text-sm">{formatCost(result.output_cost_per_request)}</Text>
<div className="min-w-0">
<p className="text-xs text-gray-500 block">Output Cost</p>
<p className="text-sm break-words">{formatCost(result.output_cost_per_request)}</p>
</div>
<div>
<Text className="text-xs text-gray-500 block">Margin Fee</Text>
<Text className={`text-sm ${result.margin_cost_per_request > 0 ? "text-amber-600" : ""}`}>
<div className="min-w-0">
<p className="text-xs text-gray-500 block">Margin Fee</p>
<p className={`text-sm break-words ${result.margin_cost_per_request > 0 ? "text-amber-600" : ""}`}>
{formatCost(result.margin_cost_per_request)}
</Text>
</p>
</div>
</div>
{periodCost !== null && (
<div className="grid grid-cols-4 gap-4 pt-2 border-t border-gray-200">
<div>
<Text className="text-xs text-gray-500 block">
<div className="min-w-0">
<p className="text-xs text-gray-500 block">
{periodLabel} Total ({formatRequests(periodRequests)} req)
</Text>
<Text className={`text-base font-semibold ${timePeriod === "day" ? "text-green-600" : "text-purple-600"}`}>
</p>
<p
className={`text-base font-semibold break-words ${timePeriod === "day" ? "text-green-600" : "text-purple-600"}`}
>
{formatCost(periodCost)}
</Text>
</p>
</div>
<div>
<Text className="text-xs text-gray-500 block">{periodLabel} Input</Text>
<Text className="text-sm">{formatCost(periodInputCost)}</Text>
<div className="min-w-0">
<p className="text-xs text-gray-500 block">{periodLabel} Input</p>
<p className="text-sm break-words">{formatCost(periodInputCost)}</p>
</div>
<div>
<Text className="text-xs text-gray-500 block">{periodLabel} Output</Text>
<Text className="text-sm">{formatCost(periodOutputCost)}</Text>
<div className="min-w-0">
<p className="text-xs text-gray-500 block">{periodLabel} Output</p>
<p className="text-sm break-words">{formatCost(periodOutputCost)}</p>
</div>
<div>
<Text className="text-xs text-gray-500 block">{periodLabel} Margin Fee</Text>
<Text className={`text-sm ${(periodMarginCost ?? 0) > 0 ? "text-amber-600" : ""}`}>
<div className="min-w-0">
<p className="text-xs text-gray-500 block">{periodLabel} Margin Fee</p>
<p className={`text-sm break-words ${(periodMarginCost ?? 0) > 0 ? "text-amber-600" : ""}`}>
{formatCost(periodMarginCost)}
</Text>
</p>
</div>
</div>
)}
@ -124,7 +130,7 @@ const MultiCostResults: React.FC<MultiCostResultsProps> = ({ multiResult, timePe
if (!hasAnyResult && !isAnyLoading && !hasAnyError) {
return (
<div className="py-6 text-center border border-dashed border-gray-300 rounded-lg bg-gray-50">
<Text className="text-gray-500">Select models above to see cost estimates</Text>
<p className="text-gray-500">Select models above to see cost estimates</p>
</div>
);
}
@ -133,8 +139,8 @@ const MultiCostResults: React.FC<MultiCostResultsProps> = ({ multiResult, timePe
if (!hasAnyResult && isAnyLoading && !hasAnyError) {
return (
<div className="py-6 text-center">
<Spin indicator={<LoadingOutlined spin />} />
<Text className="text-gray-500 block mt-2">Calculating costs...</Text>
<UiLoadingSpinner className="inline-block size-5" />
<p className="text-gray-500 block mt-2">Calculating costs...</p>
</div>
);
}
@ -143,10 +149,10 @@ const MultiCostResults: React.FC<MultiCostResultsProps> = ({ multiResult, timePe
if (!hasAnyResult && hasAnyError) {
return (
<div className="space-y-4">
<Divider className="my-4" />
<Separator className="my-4" />
<div className="flex items-center justify-between">
<Text className="text-base font-semibold text-gray-900">Cost Estimates</Text>
{isAnyLoading && <Spin indicator={<LoadingOutlined spin />} size="small" />}
<p className="text-base font-semibold text-gray-900">Cost Estimates</p>
{isAnyLoading && <UiLoadingSpinner className="size-3.5" />}
</div>
{/* Error Messages */}
{errorEntries.map((e) => (
@ -174,102 +180,10 @@ const MultiCostResults: React.FC<MultiCostResultsProps> = ({ multiResult, timePe
const hasMargin = multiResult.totals.margin_per_request > 0;
const periodLabel = timePeriod === "day" ? "Daily" : "Monthly";
const periodCostKey = timePeriod === "day" ? "daily_cost" : "monthly_cost";
const summaryColumns = [
{
title: "Model",
dataIndex: "model",
key: "model",
render: (
text: string,
record: {
id: string;
provider?: string | null;
error?: string | null;
loading?: boolean;
hasZeroCost?: boolean | null;
},
) => (
<div className="flex flex-col gap-1">
<div className="flex items-center gap-2">
<span className="font-medium text-sm">{text}</span>
{record.provider && (
<Tag color="blue" className="text-xs">
{record.provider}
</Tag>
)}
{record.loading && <Spin indicator={<LoadingOutlined spin />} size="small" />}
</div>
{record.error && <div className="text-xs text-red-600 bg-red-50 px-2 py-1 rounded-sm">⚠️ {record.error}</div>}
{record.hasZeroCost && !record.error && (
<div className="text-xs text-amber-600 bg-amber-50 px-2 py-1 rounded-sm">
⚠️ No pricing data found for this model. Set base_model in config.
</div>
)}
</div>
),
},
{
title: "Per Request",
dataIndex: "cost_per_request",
key: "cost_per_request",
align: "right" as const,
render: (value: number | null, record: { error?: string | null }) =>
record.error ? (
<span className="text-gray-400">-</span>
) : (
<span className="font-mono text-sm">{formatCost(value)}</span>
),
},
{
title: "Margin Fee",
dataIndex: "margin_cost_per_request",
key: "margin_cost_per_request",
align: "right" as const,
render: (value: number | null, record: { error?: string | null }) =>
record.error ? (
<span className="text-gray-400">-</span>
) : (
<span className={`font-mono text-sm ${(value ?? 0) > 0 ? "text-amber-600" : "text-gray-400"}`}>
{formatCost(value)}
</span>
),
},
{
title: periodLabel,
dataIndex: periodCostKey,
key: "period_cost",
align: "right" as const,
render: (value: number | null, record: { error?: string | null }) =>
record.error ? (
<span className="text-gray-400">-</span>
) : (
<span className="font-mono text-sm">{formatCost(value)}</span>
),
},
{
title: "",
key: "expand",
width: 40,
render: (_: unknown, record: { id: string; error?: string | null }) =>
record.error ? null : (
<Button
size="xs"
variant="light"
onClick={() => toggleExpanded(record.id)}
className="text-gray-400 hover:text-gray-600"
>
{expandedModels.has(record.id) ? <DownOutlined /> : <RightOutlined />}
</Button>
),
},
];
// Include both valid results and errors in the table data
const allEntriesWithModels = multiResult.entries.filter((e) => e.entry.model);
const summaryData = allEntriesWithModels.map((e) => ({
key: e.entry.id,
id: e.entry.id,
model: e.result?.model || e.entry.model,
provider: e.result?.provider,
@ -284,78 +198,153 @@ const MultiCostResults: React.FC<MultiCostResultsProps> = ({ multiResult, timePe
return (
<div className="space-y-4">
<Divider className="my-4" />
<Separator className="my-4" />
<div className="flex items-center justify-between">
<Text className="text-base font-semibold text-gray-900">Cost Estimates</Text>
<p className="text-base font-semibold text-gray-900">Cost Estimates</p>
<div className="flex items-center gap-2">
{isAnyLoading && <Spin indicator={<LoadingOutlined spin />} size="small" />}
{isAnyLoading && <UiLoadingSpinner className="size-3.5" />}
<MultiExportDropdown multiResult={multiResult} />
</div>
</div>
{/* Combined Totals - Always show when there are results */}
<Card size="small" className="bg-linear-to-r from-slate-50 to-blue-50 border-slate-200">
<Row gutter={[16, 8]}>
<Col xs={24} sm={12}>
<Statistic
title={<span className="text-xs">Total Per Request</span>}
value={formatCost(multiResult.totals.cost_per_request)}
valueStyle={{ color: "#1890ff", fontSize: "18px", fontFamily: "monospace" }}
/>
</Col>
<Col xs={24} sm={12}>
<Statistic
title={<span className="text-xs">Total {periodLabel}</span>}
value={formatCost(timePeriod === "day" ? multiResult.totals.daily_cost : multiResult.totals.monthly_cost)}
valueStyle={{
color: timePeriod === "day" ? "#52c41a" : "#722ed1",
fontSize: "18px",
fontFamily: "monospace",
}}
/>
</Col>
</Row>
<Card size="sm" className="px-4 bg-linear-to-r from-slate-50 to-blue-50">
<div className="grid grid-cols-1 sm:grid-cols-2 gap-x-4 gap-y-2">
<div className="min-w-0">
<span className="text-xs text-gray-500">Total Per Request</span>
<div className="text-lg font-mono text-blue-600 break-words">
{formatCost(multiResult.totals.cost_per_request)}
</div>
</div>
<div className="min-w-0">
<span className="text-xs text-gray-500">Total {periodLabel}</span>
<div
className={`text-lg font-mono break-words ${timePeriod === "day" ? "text-green-600" : "text-purple-600"}`}
>
{formatCost(timePeriod === "day" ? multiResult.totals.daily_cost : multiResult.totals.monthly_cost)}
</div>
</div>
</div>
{hasMargin && (
<Row gutter={[16, 8]} className="mt-3 pt-3 border-t border-slate-200">
<Col xs={24} sm={12}>
<div className="grid grid-cols-1 sm:grid-cols-2 gap-x-4 gap-y-2 mt-3 pt-3 border-t border-slate-200">
<div className="min-w-0">
<div className="text-xs text-gray-500">Margin Fee/Request</div>
<div className="text-sm font-mono text-amber-600">
<div className="text-sm font-mono text-amber-600 break-words">
{formatCost(multiResult.totals.margin_per_request)}
</div>
</Col>
<Col xs={24} sm={12}>
</div>
<div className="min-w-0">
<div className="text-xs text-gray-500">{periodLabel} Margin Fee</div>
<div className="text-sm font-mono text-amber-600">
<div className="text-sm font-mono text-amber-600 break-words">
{formatCost(timePeriod === "day" ? multiResult.totals.daily_margin : multiResult.totals.monthly_margin)}
</div>
</Col>
</Row>
</div>
</div>
)}
</Card>
{/* Per-Model Table */}
{summaryData.length > 0 && (
<Table
columns={summaryColumns}
dataSource={summaryData}
pagination={false}
size="small"
className="border border-gray-200 rounded-lg"
expandable={{
expandedRowKeys: Array.from(expandedModels),
expandedRowRender: (record) => {
const entry = validEntries.find((e) => e.entry.id === record.id);
if (!entry?.result) return null;
<Table className="border border-gray-200 rounded-lg">
<TableHeader>
<TableRow>
<TableHead>Model</TableHead>
<TableHead className="text-right">Per Request</TableHead>
<TableHead className="text-right">Margin Fee</TableHead>
<TableHead className="text-right">{periodLabel}</TableHead>
<TableHead className="w-10">
<span className="sr-only">Cost breakdown</span>
</TableHead>
</TableRow>
</TableHeader>
<TableBody>
{summaryData.map((record) => {
const isExpanded = expandedModels.has(record.id);
const periodCost = timePeriod === "day" ? record.daily_cost : record.monthly_cost;
const breakdownEntry = validEntries.find((e) => e.entry.id === record.id);
return (
<div className="py-2">
<SingleModelBreakdown result={entry.result} loading={entry.loading} timePeriod={timePeriod} />
</div>
<React.Fragment key={record.id}>
<TableRow>
<TableCell className="whitespace-normal">
<div className="flex min-w-0 flex-col gap-1">
<div className="flex items-center gap-2">
<span className="font-medium text-sm break-words">{record.model}</span>
{record.provider && (
<Badge variant="secondary" className="text-xs">
{record.provider}
</Badge>
)}
{record.loading && <UiLoadingSpinner className="size-3.5" />}
</div>
{record.error && (
<div className="text-xs text-red-600 bg-red-50 px-2 py-1 rounded-sm">⚠️ {record.error}</div>
)}
{record.hasZeroCost && !record.error && (
<div className="text-xs text-amber-600 bg-amber-50 px-2 py-1 rounded-sm">
⚠️ No pricing data found for this model. Set base_model in config.
</div>
)}
</div>
</TableCell>
<TableCell className="text-right">
{record.error ? (
<span className="text-gray-400">-</span>
) : (
<span className="font-mono text-sm">{formatCost(record.cost_per_request)}</span>
)}
</TableCell>
<TableCell className="text-right">
{record.error ? (
<span className="text-gray-400">-</span>
) : (
<span
className={`font-mono text-sm ${(record.margin_cost_per_request ?? 0) > 0 ? "text-amber-600" : "text-gray-400"}`}
>
{formatCost(record.margin_cost_per_request)}
</span>
)}
</TableCell>
<TableCell className="text-right">
{record.error ? (
<span className="text-gray-400">-</span>
) : (
<span className="font-mono text-sm">{formatCost(periodCost)}</span>
)}
</TableCell>
<TableCell className="text-right">
{!record.error && (
<Button
variant="ghost"
size="icon-xs"
aria-expanded={isExpanded}
aria-label={`${isExpanded ? "Hide" : "Show"} cost breakdown for ${record.model}`}
onClick={() => toggleExpanded(record.id)}
className="text-gray-400 hover:text-gray-600"
>
{isExpanded ? <ChevronDown className="size-3" /> : <ChevronRight className="size-3" />}
</Button>
)}
</TableCell>
</TableRow>
{isExpanded && breakdownEntry?.result && (
<TableRow>
<TableCell colSpan={5} className="whitespace-normal">
<div className="py-2">
<SingleModelBreakdown
result={breakdownEntry.result}
loading={breakdownEntry.loading}
timePeriod={timePeriod}
/>
</div>
</TableCell>
</TableRow>
)}
</React.Fragment>
);
},
showExpandColumn: false,
}}
/>
})}
</TableBody>
</Table>
)}
</div>
);

204
ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_discount_table.test.tsx
View file

@ -5,49 +5,21 @@ import userEvent from "@testing-library/user-event";
import { renderWithProviders } from "../../../../../tests/test-utils";
import ProviderDiscountTable from "./provider_discount_table";
vi.mock("@heroicons/react/outline", () => ({
TrashIcon: function TrashIcon() {
return null;
},
PencilAltIcon: function PencilAltIcon() {
return null;
},
CheckIcon: function CheckIcon() {
return null;
},
XIcon: function XIcon() {
return null;
},
}));
vi.mock("@tremor/react", () => ({
Table: ({ children }: any) => <table>{children}</table>,
TableHead: ({ children }: any) => <thead>{children}</thead>,
TableRow: ({ children }: any) => <tr>{children}</tr>,
TableHeaderCell: ({ children }: any) => <th>{children}</th>,
TableBody: ({ children }: any) => <tbody>{children}</tbody>,
TableCell: ({ children }: any) => <td>{children}</td>,
Text: ({ children }: any) => <span>{children}</span>,
TextInput: ({ value, onValueChange, onKeyDown, placeholder, ...rest }: any) => (
<input
value={value}
onChange={(e) => onValueChange?.(e.target.value)}
onKeyDown={onKeyDown}
placeholder={placeholder}
{...rest}
/>
),
Icon: ({ icon: IconComponent, onClick }: any) => {
const name = IconComponent?.displayName ?? IconComponent?.name ?? "icon";
return <button onClick={onClick} aria-label={name} />;
},
}));
const DEFAULT_DISCOUNT_CONFIG = {
openai: 0.05,
anthropic: 0.1,
};
const ROW_ACTION_NAME = {
edit: /^Edit discount for /,
save: /^Save discount for /,
cancel: /^Cancel editing discount for /,
remove: /^Remove discount for /,
} as const;
const rowAction = (action: keyof typeof ROW_ACTION_NAME): HTMLElement =>
screen.getByRole("button", { name: ROW_ACTION_NAME[action] });
describe("ProviderDiscountTable", () => {
const onDiscountChange = vi.fn();
const onRemoveProvider = vi.fn();
@ -75,9 +47,9 @@ describe("ProviderDiscountTable", () => {
onRemoveProvider={onRemoveProvider}
/>,
);
expect(screen.getByText("Provider")).toBeInTheDocument();
expect(screen.getByText("Discount Percentage")).toBeInTheDocument();
expect(screen.getByText("Actions")).toBeInTheDocument();
expect(screen.getByRole("columnheader", { name: "Provider" })).toBeInTheDocument();
expect(screen.getByRole("columnheader", { name: "Discount Percentage" })).toBeInTheDocument();
expect(screen.getByRole("columnheader", { name: "Actions" })).toBeInTheDocument();
});
it("should display provider display names in the table", () => {
@ -91,6 +63,21 @@ describe("ProviderDiscountTable", () => {
expect(screen.getByText("OpenAI")).toBeInTheDocument();
});
it("should sort rows by provider display name", () => {
renderWithProviders(
<ProviderDiscountTable
discountConfig={DEFAULT_DISCOUNT_CONFIG}
onDiscountChange={onDiscountChange}
onRemoveProvider={onRemoveProvider}
/>,
);
const rows = screen.getAllByRole("row").slice(1);
expect(rows.map((row) => row.textContent)).toEqual([
expect.stringContaining("Anthropic"),
expect.stringContaining("OpenAI"),
]);
});
it("should display the formatted discount percentage", () => {
renderWithProviders(
<ProviderDiscountTable
@ -102,6 +89,17 @@ describe("ProviderDiscountTable", () => {
expect(screen.getByText("5.0%")).toBeInTheDocument();
});
it("should render the provider logo alongside the display name", () => {
renderWithProviders(
<ProviderDiscountTable
discountConfig={{ openai: 0.05 }}
onDiscountChange={onDiscountChange}
onRemoveProvider={onRemoveProvider}
/>,
);
expect(screen.getByRole("img", { name: "OpenAI logo" })).toBeInTheDocument();
});
it("should show a text input when the edit icon is clicked", async () => {
const user = userEvent.setup();
renderWithProviders(
@ -112,8 +110,7 @@ describe("ProviderDiscountTable", () => {
/>,
);
const pencilButton = screen.getByRole("button", { name: /PencilAltIcon/i });
await user.click(pencilButton);
await user.click(rowAction("edit"));
expect(screen.getByPlaceholderText("5")).toBeInTheDocument();
});
@ -128,11 +125,26 @@ describe("ProviderDiscountTable", () => {
/>,
);
await user.click(screen.getByRole("button", { name: /PencilAltIcon/i }));
await user.click(rowAction("edit"));
expect(screen.queryByText("5.0%")).not.toBeInTheDocument();
});
it("should seed the edit input with the current discount as a percentage", async () => {
const user = userEvent.setup();
renderWithProviders(
<ProviderDiscountTable
discountConfig={{ openai: 0.05 }}
onDiscountChange={onDiscountChange}
onRemoveProvider={onRemoveProvider}
/>,
);
await user.click(rowAction("edit"));
expect(screen.getByPlaceholderText("5")).toHaveValue("5");
});
it("should call onDiscountChange with the new value when the save icon is clicked", async () => {
const user = userEvent.setup();
renderWithProviders(
@ -143,17 +155,57 @@ describe("ProviderDiscountTable", () => {
/>,
);
await user.click(screen.getByRole("button", { name: /PencilAltIcon/i }));
await user.click(rowAction("edit"));
const input = screen.getByPlaceholderText("5");
await user.clear(input);
await user.type(input, "10");
await user.click(screen.getByRole("button", { name: /CheckIcon/i }));
await user.click(rowAction("save"));
expect(onDiscountChange).toHaveBeenCalledWith("openai", "0.1");
});
it("should save the edited discount when Enter is pressed", async () => {
const user = userEvent.setup();
renderWithProviders(
<ProviderDiscountTable
discountConfig={{ openai: 0.05 }}
onDiscountChange={onDiscountChange}
onRemoveProvider={onRemoveProvider}
/>,
);
await user.click(rowAction("edit"));
const input = screen.getByPlaceholderText("5");
await user.clear(input);
await user.type(input, "10{Enter}");
expect(onDiscountChange).toHaveBeenCalledWith("openai", "0.1");
expect(screen.queryByPlaceholderText("5")).not.toBeInTheDocument();
});
it("should abandon the edit when Escape is pressed", async () => {
const user = userEvent.setup();
renderWithProviders(
<ProviderDiscountTable
discountConfig={{ openai: 0.05 }}
onDiscountChange={onDiscountChange}
onRemoveProvider={onRemoveProvider}
/>,
);
await user.click(rowAction("edit"));
const input = screen.getByPlaceholderText("5");
await user.clear(input);
await user.type(input, "10{Escape}");
expect(onDiscountChange).not.toHaveBeenCalled();
expect(screen.getByText("5.0%")).toBeInTheDocument();
});
it("should restore the display view after saving", async () => {
const user = userEvent.setup();
renderWithProviders(
@ -164,8 +216,8 @@ describe("ProviderDiscountTable", () => {
/>,
);
await user.click(screen.getByRole("button", { name: /PencilAltIcon/i }));
await user.click(screen.getByRole("button", { name: /CheckIcon/i }));
await user.click(rowAction("edit"));
await user.click(rowAction("save"));
expect(screen.queryByPlaceholderText("5")).not.toBeInTheDocument();
});
@ -180,30 +232,14 @@ describe("ProviderDiscountTable", () => {
/>,
);
await user.click(screen.getByRole("button", { name: /PencilAltIcon/i }));
await user.click(screen.getByRole("button", { name: /XIcon/i }));
await user.click(rowAction("edit"));
await user.click(rowAction("cancel"));
expect(screen.queryByPlaceholderText("5")).not.toBeInTheDocument();
expect(onDiscountChange).not.toHaveBeenCalled();
expect(screen.getByText("5.0%")).toBeInTheDocument();
});
it("should not call onDiscountChange when canceling edit", async () => {
const user = userEvent.setup();
renderWithProviders(
<ProviderDiscountTable
discountConfig={{ openai: 0.05 }}
onDiscountChange={onDiscountChange}
onRemoveProvider={onRemoveProvider}
/>,
);
await user.click(screen.getByRole("button", { name: /PencilAltIcon/i }));
await user.click(screen.getByRole("button", { name: /XIcon/i }));
expect(onDiscountChange).not.toHaveBeenCalled();
});
it("should call onRemoveProvider with the provider key and display name when the trash icon is clicked", async () => {
const user = userEvent.setup();
renderWithProviders(
@ -214,7 +250,7 @@ describe("ProviderDiscountTable", () => {
/>,
);
await user.click(screen.getByRole("button", { name: /TrashIcon/i }));
await user.click(rowAction("remove"));
expect(onRemoveProvider).toHaveBeenCalledWith("openai", "OpenAI");
});
@ -229,12 +265,42 @@ describe("ProviderDiscountTable", () => {
/>,
);
await user.click(screen.getByRole("button", { name: /PencilAltIcon/i }));
await user.click(rowAction("edit"));
const input = screen.getByPlaceholderText("5");
await user.clear(input);
await user.type(input, "150");
await user.click(screen.getByRole("button", { name: /CheckIcon/i }));
await user.click(rowAction("save"));
expect(onDiscountChange).not.toHaveBeenCalled();
});
it("should expose each row action as a button named for its provider", async () => {
const user = userEvent.setup();
renderWithProviders(
<ProviderDiscountTable
discountConfig={{ openai: 0.05 }}
onDiscountChange={onDiscountChange}
onRemoveProvider={onRemoveProvider}
/>,
);
expect(screen.getByRole("button", { name: "Edit discount for OpenAI" })).toBeInTheDocument();
expect(screen.getByRole("button", { name: "Remove discount for OpenAI" })).toBeInTheDocument();
await user.click(screen.getByRole("button", { name: "Edit discount for OpenAI" }));
expect(screen.getByRole("button", { name: "Save discount for OpenAI" })).toBeInTheDocument();
expect(screen.getByRole("button", { name: "Cancel editing discount for OpenAI" })).toBeInTheDocument();
});
it("should render the empty message when no discounts are configured", () => {
renderWithProviders(
<ProviderDiscountTable
discountConfig={{}}
onDiscountChange={onDiscountChange}
onRemoveProvider={onRemoveProvider}
/>,
);
expect(screen.getByText("No provider discounts configured")).toBeInTheDocument();
});
});

106
ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_discount_table.tsx
View file

@ -1,6 +1,7 @@
import React, { useState } from "react";
import { TextInput, Icon, Text } from "@tremor/react";
import { TrashIcon, PencilAltIcon, CheckIcon, XIcon } from "@heroicons/react/outline";
import { Check, SquarePen, Trash2, X } from "lucide-react";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { SimpleTable } from "@/components/common_components/simple_table";
import { DiscountConfig } from "./types";
import { getProviderLogoAndName } from "@/components/provider_info_helpers";
@ -79,45 +80,57 @@ const ProviderDiscountTable: React.FC<ProviderDiscountTableProps> = ({
},
{
header: "Discount Percentage",
cell: (row) => (
<div className="flex items-center gap-2">
{editingProvider === row.provider ? (
<>
<TextInput
value={editValue}
onValueChange={setEditValue}
onKeyDown={(e) => handleKeyDown(e, row.provider)}
placeholder="5"
className="w-20"
autoFocus
/>
<span className="text-gray-600">%</span>
<Icon
icon={CheckIcon}
size="sm"
onClick={() => handleSaveEdit(row.provider)}
className="cursor-pointer text-green-600 hover:text-green-700"
/>
<Icon
icon={XIcon}
size="sm"
onClick={handleCancelEdit}
className="cursor-pointer text-gray-600 hover:text-gray-700"
/>
</>
) : (
<>
<Text className="font-medium">{(row.discount * 100).toFixed(1)}%</Text>
<Icon
icon={PencilAltIcon}
size="sm"
onClick={() => handleStartEdit(row.provider, row.discount)}
className="cursor-pointer text-blue-600 hover:text-blue-700"
/>
</>
)}
</div>
),
cell: (row) => {
const { displayName } = getProviderLogoAndName(row.provider);
return (
<div className="flex items-center gap-2">
{editingProvider === row.provider ? (
<>
<Input
value={editValue}
onChange={(e) => setEditValue(e.target.value)}
onKeyDown={(e) => handleKeyDown(e, row.provider)}
placeholder="5"
className="w-20"
autoFocus
/>
<span className="text-gray-600">%</span>
<Button
variant="ghost"
size="icon-sm"
aria-label={`Save discount for ${displayName}`}
onClick={() => handleSaveEdit(row.provider)}
className="cursor-pointer text-green-600 hover:text-green-700"
>
<Check className="size-5" />
</Button>
<Button
variant="ghost"
size="icon-sm"
aria-label={`Cancel editing discount for ${displayName}`}
onClick={handleCancelEdit}
className="cursor-pointer text-gray-600 hover:text-gray-700"
>
<X className="size-5" />
</Button>
</>
) : (
<>
<p className="font-medium">{(row.discount * 100).toFixed(1)}%</p>
<Button
variant="ghost"
size="icon-sm"
aria-label={`Edit discount for ${displayName}`}
onClick={() => handleStartEdit(row.provider, row.discount)}
className="cursor-pointer text-blue-600 hover:text-blue-700"
>
<SquarePen className="size-5" />
</Button>
</>
)}
</div>
);
},
width: "250px",
},
{
@ -125,12 +138,15 @@ const ProviderDiscountTable: React.FC<ProviderDiscountTableProps> = ({
cell: (row) => {
const { displayName } = getProviderLogoAndName(row.provider);
return (
<Icon
icon={TrashIcon}
size="sm"
<Button
variant="ghost"
size="icon-sm"
aria-label={`Remove discount for ${displayName}`}
onClick={() => onRemoveProvider(row.provider, displayName)}
className="cursor-pointer hover:text-red-600"
/>
>
<Trash2 className="size-5" />
</Button>
);
},
width: "80px",

158
ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_margin_table.test.tsx
View file

@ -6,43 +6,15 @@ import { renderWithProviders } from "../../../../../tests/test-utils";
import ProviderMarginTable from "./provider_margin_table";
import { Providers, providerLogoMap } from "@/components/provider_info_helpers";
vi.mock("@heroicons/react/outline", () => ({
TrashIcon: function TrashIcon() {
return null;
},
PencilAltIcon: function PencilAltIcon() {
return null;
},
CheckIcon: function CheckIcon() {
return null;
},
XIcon: function XIcon() {
return null;
},
}));
const ROW_ACTION_NAME = {
edit: /^Edit margin for /,
save: /^Save margin for /,
cancel: /^Cancel editing margin for /,
remove: /^Remove margin for /,
} as const;
vi.mock("@tremor/react", () => ({
Table: ({ children }: any) => <table>{children}</table>,
TableHead: ({ children }: any) => <thead>{children}</thead>,
TableRow: ({ children }: any) => <tr>{children}</tr>,
TableHeaderCell: ({ children }: any) => <th>{children}</th>,
TableBody: ({ children }: any) => <tbody>{children}</tbody>,
TableCell: ({ children }: any) => <td>{children}</td>,
Text: ({ children }: any) => <span>{children}</span>,
TextInput: ({ value, onValueChange, placeholder, autoFocus, className }: any) => (
<input
value={value}
onChange={(e) => onValueChange?.(e.target.value)}
placeholder={placeholder}
autoFocus={autoFocus}
className={className}
/>
),
Icon: ({ icon: IconComponent, onClick }: any) => {
const name = IconComponent?.displayName ?? IconComponent?.name ?? "icon";
return <button onClick={onClick} aria-label={name} />;
},
}));
const rowAction = (action: keyof typeof ROW_ACTION_NAME): HTMLElement =>
screen.getByRole("button", { name: ROW_ACTION_NAME[action] });
describe("ProviderMarginTable", () => {
const onMarginChange = vi.fn();
@ -71,9 +43,9 @@ describe("ProviderMarginTable", () => {
onRemoveProvider={onRemoveProvider}
/>,
);
expect(screen.getByText("Provider")).toBeInTheDocument();
expect(screen.getByText("Margin")).toBeInTheDocument();
expect(screen.getByText("Actions")).toBeInTheDocument();
expect(screen.getByRole("columnheader", { name: "Provider" })).toBeInTheDocument();
expect(screen.getByRole("columnheader", { name: "Margin" })).toBeInTheDocument();
expect(screen.getByRole("columnheader", { name: "Actions" })).toBeInTheDocument();
});
it("should display the provider display name", () => {
@ -122,6 +94,21 @@ describe("ProviderMarginTable", () => {
expect(screen.getByText("Global (All Providers)")).toBeInTheDocument();
});
it("should sort the global row above provider rows", () => {
renderWithProviders(
<ProviderMarginTable
marginConfig={{ openai: 0.1, global: 0.05 }}
onMarginChange={onMarginChange}
onRemoveProvider={onRemoveProvider}
/>,
);
const rows = screen.getAllByRole("row").slice(1);
expect(rows.map((row) => row.textContent)).toEqual([
expect.stringContaining("Global (All Providers)"),
expect.stringContaining("OpenAI"),
]);
});
it("should display a numeric margin as a percentage", () => {
renderWithProviders(
<ProviderMarginTable
@ -165,12 +152,28 @@ describe("ProviderMarginTable", () => {
/>,
);
await user.click(screen.getByRole("button", { name: /PencilAltIcon/i }));
await user.click(rowAction("edit"));
expect(screen.getByPlaceholderText("10")).toBeInTheDocument();
expect(screen.getByPlaceholderText("0.001")).toBeInTheDocument();
});
it("should seed the percentage input from a numeric margin and leave the fixed amount blank", async () => {
const user = userEvent.setup();
renderWithProviders(
<ProviderMarginTable
marginConfig={{ openai: 0.1 }}
onMarginChange={onMarginChange}
onRemoveProvider={onRemoveProvider}
/>,
);
await user.click(rowAction("edit"));
expect(screen.getByPlaceholderText("10")).toHaveValue("10");
expect(screen.getByPlaceholderText("0.001")).toHaveValue("");
});
it("should call onMarginChange with a percentage value when save is clicked", async () => {
const user = userEvent.setup();
renderWithProviders(
@ -181,17 +184,37 @@ describe("ProviderMarginTable", () => {
/>,
);
await user.click(screen.getByRole("button", { name: /PencilAltIcon/i }));
await user.click(rowAction("edit"));
const percentInput = screen.getByPlaceholderText("10");
await user.clear(percentInput);
await user.type(percentInput, "20");
await user.click(screen.getByRole("button", { name: /CheckIcon/i }));
await user.click(rowAction("save"));
expect(onMarginChange).toHaveBeenCalledWith("openai", 0.2);
});
it("should call onMarginChange with a fixed-amount-only object when the percentage is cleared", async () => {
const user = userEvent.setup();
renderWithProviders(
<ProviderMarginTable
marginConfig={{ openai: 0.1 }}
onMarginChange={onMarginChange}
onRemoveProvider={onRemoveProvider}
/>,
);
await user.click(rowAction("edit"));
await user.clear(screen.getByPlaceholderText("10"));
await user.type(screen.getByPlaceholderText("0.001"), "0.002");
await user.click(rowAction("save"));
expect(onMarginChange).toHaveBeenCalledWith("openai", { fixed_amount: 0.002 });
});
it("should cancel edit mode without calling onMarginChange when X is clicked", async () => {
const user = userEvent.setup();
renderWithProviders(
@ -202,8 +225,8 @@ describe("ProviderMarginTable", () => {
/>,
);
await user.click(screen.getByRole("button", { name: /PencilAltIcon/i }));
await user.click(screen.getByRole("button", { name: /XIcon/i }));
await user.click(rowAction("edit"));
await user.click(rowAction("cancel"));
expect(onMarginChange).not.toHaveBeenCalled();
expect(screen.queryByPlaceholderText("10")).not.toBeInTheDocument();
@ -219,7 +242,7 @@ describe("ProviderMarginTable", () => {
/>,
);
await user.click(screen.getByRole("button", { name: /TrashIcon/i }));
await user.click(rowAction("remove"));
expect(onRemoveProvider).toHaveBeenCalledWith("openai", "OpenAI");
});
@ -234,11 +257,50 @@ describe("ProviderMarginTable", () => {
/>,
);
await user.click(screen.getByRole("button", { name: /TrashIcon/i }));
await user.click(rowAction("remove"));
expect(onRemoveProvider).toHaveBeenCalledWith("global", "Global");
});
it("should expose each row action as a button named for its provider", async () => {
const user = userEvent.setup();
renderWithProviders(
<ProviderMarginTable
marginConfig={{ openai: 0.1 }}
onMarginChange={onMarginChange}
onRemoveProvider={onRemoveProvider}
/>,
);
expect(screen.getByRole("button", { name: "Edit margin for OpenAI" })).toBeInTheDocument();
expect(screen.getByRole("button", { name: "Remove margin for OpenAI" })).toBeInTheDocument();
await user.click(screen.getByRole("button", { name: "Edit margin for OpenAI" }));
expect(screen.getByRole("button", { name: "Save margin for OpenAI" })).toBeInTheDocument();
expect(screen.getByRole("button", { name: "Cancel editing margin for OpenAI" })).toBeInTheDocument();
});
it("should name the global row's actions after the global provider", () => {
renderWithProviders(
<ProviderMarginTable
marginConfig={{ global: 0.05 }}
onMarginChange={onMarginChange}
onRemoveProvider={onRemoveProvider}
/>,
);
expect(screen.getByRole("button", { name: "Edit margin for Global" })).toBeInTheDocument();
expect(screen.getByRole("button", { name: "Remove margin for Global" })).toBeInTheDocument();
});
it("should render the empty message when no margins are configured", () => {
renderWithProviders(
<ProviderMarginTable marginConfig={{}} onMarginChange={onMarginChange} onRemoveProvider={onRemoveProvider} />,
);
expect(screen.getByText("No provider margins configured")).toBeInTheDocument();
});
describe("when both percentage and fixed amount are entered", () => {
it("should call onMarginChange with an object containing both values", async () => {
const user = userEvent.setup();
@ -250,7 +312,7 @@ describe("ProviderMarginTable", () => {
/>,
);
await user.click(screen.getByRole("button", { name: /PencilAltIcon/i }));
await user.click(rowAction("edit"));
const percentInput = screen.getByPlaceholderText("10");
await user.clear(percentInput);
@ -259,7 +321,7 @@ describe("ProviderMarginTable", () => {
const fixedInput = screen.getByPlaceholderText("0.001");
await user.type(fixedInput, "0.002");
await user.click(screen.getByRole("button", { name: /CheckIcon/i }));
await user.click(rowAction("save"));
expect(onMarginChange).toHaveBeenCalledWith("openai", {
percentage: 0.05,

129
ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_margin_table.tsx
View file

@ -1,6 +1,7 @@
import React, { useState } from "react";
import { TextInput, Icon, Text } from "@tremor/react";
import { TrashIcon, PencilAltIcon, CheckIcon, XIcon } from "@heroicons/react/outline";
import { Check, SquarePen, Trash2, X } from "lucide-react";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { SimpleTable } from "@/components/common_components/simple_table";
import { MarginConfig } from "./types";
import { getProviderLogoAndName } from "@/components/provider_info_helpers";
@ -17,6 +18,9 @@ interface ProviderMarginRow {
margin: number | { percentage?: number; fixed_amount?: number };
}
const marginRowDisplayName = (provider: string): string =>
provider === "global" ? "Global" : getProviderLogoAndName(provider).displayName;
const ProviderMarginTable: React.FC<ProviderMarginTableProps> = ({
marginConfig,
onMarginChange,
@ -119,67 +123,82 @@ const ProviderMarginTable: React.FC<ProviderMarginTableProps> = ({
},
{
header: "Margin",
cell: (row) => (
<div className="flex items-center gap-2">
{editingProvider === row.provider ? (
<>
<div className="flex items-center gap-2">
<TextInput
value={editPercentage}
onValueChange={setEditPercentage}
placeholder="10"
className="w-20"
autoFocus
/>
<span className="text-gray-600">%</span>
<span className="text-gray-400">+</span>
<span className="text-gray-600">$</span>
<TextInput
value={editFixedAmount}
onValueChange={setEditFixedAmount}
placeholder="0.001"
className="w-24"
/>
</div>
<Icon
icon={CheckIcon}
size="sm"
onClick={() => handleSaveEdit(row.provider)}
className="cursor-pointer text-green-600 hover:text-green-700"
/>
<Icon
icon={XIcon}
size="sm"
onClick={handleCancelEdit}
className="cursor-pointer text-gray-600 hover:text-gray-700"
/>
</>
) : (
<>
<Text className="font-medium">{formatMargin(row.margin)}</Text>
<Icon
icon={PencilAltIcon}
size="sm"
onClick={() => handleStartEdit(row.provider, row.margin)}
className="cursor-pointer text-blue-600 hover:text-blue-700"
/>
</>
)}
</div>
),
cell: (row) => {
const displayName = marginRowDisplayName(row.provider);
return (
<div className="flex items-center gap-2">
{editingProvider === row.provider ? (
<>
<div className="flex items-center gap-2">
<Input
value={editPercentage}
onChange={(e) => setEditPercentage(e.target.value)}
placeholder="10"
className="w-20"
autoFocus
/>
<span className="text-gray-600">%</span>
<span className="text-gray-400">+</span>
<span className="text-gray-600">$</span>
<Input
value={editFixedAmount}
onChange={(e) => setEditFixedAmount(e.target.value)}
placeholder="0.001"
className="w-24"
/>
</div>
<Button
variant="ghost"
size="icon-sm"
aria-label={`Save margin for ${displayName}`}
onClick={() => handleSaveEdit(row.provider)}
className="cursor-pointer text-green-600 hover:text-green-700"
>
<Check className="size-5" />
</Button>
<Button
variant="ghost"
size="icon-sm"
aria-label={`Cancel editing margin for ${displayName}`}
onClick={handleCancelEdit}
className="cursor-pointer text-gray-600 hover:text-gray-700"
>
<X className="size-5" />
</Button>
</>
) : (
<>
<p className="font-medium">{formatMargin(row.margin)}</p>
<Button
variant="ghost"
size="icon-sm"
aria-label={`Edit margin for ${displayName}`}
onClick={() => handleStartEdit(row.provider, row.margin)}
className="cursor-pointer text-blue-600 hover:text-blue-700"
>
<SquarePen className="size-5" />
</Button>
</>
)}
</div>
);
},
width: "350px",
},
{
header: "Actions",
cell: (row) => {
const displayName = row.provider === "global" ? "Global" : getProviderLogoAndName(row.provider).displayName;
const displayName = marginRowDisplayName(row.provider);
return (
<Icon
icon={TrashIcon}
size="sm"
<Button
variant="ghost"
size="icon-sm"
aria-label={`Remove margin for ${displayName}`}
onClick={() => onRemoveProvider(row.provider, displayName)}
className="cursor-pointer hover:text-red-600"
/>
>
<Trash2 className="size-5" />
</Button>
);
},
width: "80px",

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

@ -1,4 +1,4 @@
import type { DateRangePickerValue } from "@tremor/react";
import type { DateRangePickerValue } from "@/components/shared/date_picker_types";
import React, { useCallback, useMemo, useState } from "react";
import { formatDate } from "@/components/networking";
import AdvancedDatePicker from "@/components/shared/advanced_date_picker";

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

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

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

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

6
ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.test.tsx
View file

@ -62,6 +62,8 @@ const settingsRow = async (fieldName: string) => {
return row as HTMLElement;
};
const numericValueIn = (row: HTMLElement) => Number((within(row).getByRole("spinbutton") as HTMLInputElement).value);
describe("GeneralSettings General tab", () => {
beforeEach(() => {
vi.mocked(getGeneralSettingsCall).mockResolvedValue([...SETTINGS_FIXTURE.map((s) => ({ ...s }))]);
@ -87,7 +89,7 @@ describe("GeneralSettings General tab", () => {
await user.click(screen.getByText("General"));
const row = await settingsRow("max_ui_session_budget");
expect(within(row).getByRole("spinbutton")).toHaveValue("7.50");
expect(numericValueIn(row)).toBe(7.5);
const actionCell = row.querySelectorAll("td")[3];
const resetIcon = actionCell.querySelector("svg");
@ -95,7 +97,7 @@ describe("GeneralSettings General tab", () => {
await user.click(resetIcon as unknown as Element);
expect(deleteConfigFieldSetting).toHaveBeenCalledWith("token", "max_ui_session_budget");
expect(within(row).getByRole("spinbutton")).toHaveValue("1.00");
expect(numericValueIn(row)).toBe(1);
});
});

246
ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx
View file

@ -1,22 +1,14 @@
import React, { useState, useEffect } from "react";
import {
Card,
Table,
TableHead,
TableRow,
TableHeaderCell,
TableCell,
TableBody,
Title,
Text,
Button,
Icon,
Switch,
} from "@tremor/react";
import { Button } from "@/components/ui/button";
import { Card, CardContent, CardTitle } from "@/components/ui/card";
import { Input } from "@/components/ui/input";
import { InputGroup, InputGroupAddon, InputGroupInput } from "@/components/ui/input-group";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { Switch } from "@/components/ui/switch";
import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { getGeneralSettingsCall, updateConfigFieldSetting, deleteConfigFieldSetting } from "@/components/networking";
import { InputNumber, Select as AntdSelect } from "antd";
import { TrashIcon } from "@heroicons/react/outline";
import { Trash2 } from "lucide-react";
import { StatusBadge } from "@/components/shared/table_cells";
import RouterSettings from "@/components/router_settings";
@ -44,16 +36,22 @@ export interface generalSettingsItem {
field_default_value?: any;
}
const NUMERIC_INPUT_WIDTH = "w-36";
const toNumericValue = (raw: string): number | null => (raw === "" ? null : Number(raw));
const SettingValueEditor: React.FC<{
setting: generalSettingsItem;
onChange: (fieldName: string, newValue: any) => void;
}> = ({ setting, onChange }) => {
if (setting.field_type === "Integer") {
return (
<InputNumber
<Input
type="number"
step={1}
value={setting.field_value}
onChange={(newValue) => onChange(setting.field_name, newValue)}
className={NUMERIC_INPUT_WIDTH}
value={setting.field_value ?? ""}
onChange={(event) => onChange(setting.field_name, toNumericValue(event.target.value))}
/>
);
}
@ -61,42 +59,55 @@ const SettingValueEditor: React.FC<{
return (
<Switch
checked={setting.field_value === true || setting.field_value === "true"}
onChange={(checked) => onChange(setting.field_name, checked)}
onCheckedChange={(checked) => onChange(setting.field_name, checked)}
/>
);
}
if (setting.field_type === "Float") {
return (
<InputNumber
<Input
type="number"
min={0}
max={1}
step={0.05}
value={setting.field_value}
onChange={(newValue) => onChange(setting.field_name, newValue)}
className={NUMERIC_INPUT_WIDTH}
value={setting.field_value ?? ""}
onChange={(event) => onChange(setting.field_name, toNumericValue(event.target.value))}
/>
);
}
if (setting.field_type === "Dollar") {
return (
<InputNumber
min={0.01}
step={0.25}
prefix="$"
value={setting.field_value}
onChange={(newValue) => onChange(setting.field_name, newValue)}
/>
<InputGroup className={NUMERIC_INPUT_WIDTH}>
<InputGroupAddon>$</InputGroupAddon>
<InputGroupInput
type="number"
min={0.01}
step={0.25}
value={setting.field_value ?? ""}
onChange={(event) => onChange(setting.field_name, toNumericValue(event.target.value))}
/>
</InputGroup>
);
}
if (setting.field_type === "Select") {
return (
<AntdSelect
allowClear
style={{ minWidth: "8rem" }}
placeholder="Default"
value={setting.field_value || undefined}
options={(setting.field_options ?? []).map((option) => ({ label: option, value: option }))}
onChange={(newValue) => onChange(setting.field_name, newValue ?? "")}
/>
<Select
value={setting.field_value || null}
onValueChange={(newValue) => onChange(setting.field_name, newValue ?? "")}
>
<SelectTrigger className="min-w-32">
<SelectValue placeholder="Default" />
</SelectTrigger>
<SelectContent>
<SelectItem value={null}>Default</SelectItem>
{(setting.field_options ?? []).map((option) => (
<SelectItem key={option} value={option}>
{option}
</SelectItem>
))}
</SelectContent>
</Select>
);
}
return null;
@ -131,33 +142,43 @@ export const PromptCachingPanel: React.FC<{
return (
<Card>
<Title>Prompt Caching</Title>
<CardContent>
<CardTitle>Prompt Caching</CardTitle>
<div className="mt-6 flex items-start justify-between gap-8">
<div className="max-w-2xl">
<Text className="font-medium">Automatic Anthropic prompt caching</Text>
<p className="mt-1 text-xs text-gray-500">{enableSetting.field_description}</p>
</div>
<Switch checked={enabled} onChange={(checked) => persist(ENABLE_ANTHROPIC_PROMPT_CACHING, checked)} />
</div>
{ttlSetting && (
<div className="mt-6 flex items-start justify-between gap-8">
<div className="max-w-2xl">
<Text className={`font-medium ${enabled ? "" : "text-gray-400"}`}>Cache lifetime (TTL)</Text>
<p className="mt-1 text-xs text-gray-500">{ttlSetting.field_description}</p>
<div className="min-w-0 max-w-2xl">
<p className="font-medium">Automatic Anthropic prompt caching</p>
<p className="mt-1 break-words text-xs text-gray-500">{enableSetting.field_description}</p>
</div>
<AntdSelect
allowClear
disabled={!enabled}
style={{ minWidth: "10rem" }}
placeholder="5m (default)"
value={ttlSetting.field_value || undefined}
options={(ttlSetting.field_options ?? []).map((option) => ({ label: option, value: option }))}
onChange={(newValue) => persist(ANTHROPIC_PROMPT_CACHING_TTL, newValue ?? "")}
/>
<Switch checked={enabled} onCheckedChange={(checked) => persist(ENABLE_ANTHROPIC_PROMPT_CACHING, checked)} />
</div>
)}
{ttlSetting && (
<div className="mt-6 flex items-start justify-between gap-8">
<div className="min-w-0 max-w-2xl">
<p className={`font-medium ${enabled ? "" : "text-gray-400"}`}>Cache lifetime (TTL)</p>
<p className="mt-1 break-words text-xs text-gray-500">{ttlSetting.field_description}</p>
</div>
<Select
disabled={!enabled}
value={ttlSetting.field_value || null}
onValueChange={(newValue) => persist(ANTHROPIC_PROMPT_CACHING_TTL, newValue ?? "")}
>
<SelectTrigger className="min-w-40">
<SelectValue placeholder="5m (default)" />
</SelectTrigger>
<SelectContent>
<SelectItem value={null}>5m (default)</SelectItem>
{(ttlSetting.field_options ?? []).map((option) => (
<SelectItem key={option} value={option}>
{option}
</SelectItem>
))}
</SelectContent>
</Select>
</div>
)}
</CardContent>
</Card>
);
};
@ -254,55 +275,60 @@ const GeneralSettings: React.FC<GeneralSettingsPageProps> = ({ accessToken, user
</TabsContent>
<TabsContent value="general" className="px-8 py-6">
<Card>
<Table>
<TableHead>
<TableRow>
<TableHeaderCell>Setting</TableHeaderCell>
<TableHeaderCell>Value</TableHeaderCell>
<TableHeaderCell>Status</TableHeaderCell>
<TableHeaderCell>Action</TableHeaderCell>
</TableRow>
</TableHead>
<TableBody>
{generalSettings
.filter((value) => value.field_type !== "TypedDictionary" && value.field_tab !== PROMPT_CACHING_TAB)
.map((value, index) => (
<TableRow key={index}>
<TableCell>
<Text>{value.field_name}</Text>
<p
style={{
fontSize: "0.65rem",
color: "#808080",
fontStyle: "italic",
}}
className="mt-1"
>
{value.field_description}
</p>
</TableCell>
<TableCell>
<SettingValueEditor setting={value} onChange={handleInputChange} />
</TableCell>
<TableCell>
{value.stored_in_db == true ? (
<StatusBadge tone="success" label="In DB" />
) : value.stored_in_db == false ? (
<StatusBadge tone="neutral" label="In Config" />
) : (
<StatusBadge tone="neutral" label="Not Set" />
)}
</TableCell>
<TableCell>
<Button onClick={() => handleUpdateField(value.field_name)}>Update</Button>
<Icon icon={TrashIcon} color="red" onClick={() => handleResetField(value.field_name)}>
Reset
</Icon>
</TableCell>
</TableRow>
))}
</TableBody>
</Table>
<CardContent>
<Table>
<TableHeader>
<TableRow>
<TableHead>Setting</TableHead>
<TableHead>Value</TableHead>
<TableHead>Status</TableHead>
<TableHead>Action</TableHead>
</TableRow>
</TableHeader>
<TableBody>
{generalSettings
.filter((value) => value.field_type !== "TypedDictionary" && value.field_tab !== PROMPT_CACHING_TAB)
.map((value, index) => (
<TableRow key={index}>
<TableCell className="whitespace-normal">
<p className="break-words">{value.field_name}</p>
<p
style={{
fontSize: "0.65rem",
color: "#808080",
fontStyle: "italic",
}}
className="mt-1 break-words"
>
{value.field_description}
</p>
</TableCell>
<TableCell>
<SettingValueEditor setting={value} onChange={handleInputChange} />
</TableCell>
<TableCell>
{value.stored_in_db == true ? (
<StatusBadge tone="success" label="In DB" />
) : value.stored_in_db == false ? (
<StatusBadge tone="neutral" label="In Config" />
) : (
<StatusBadge tone="neutral" label="Not Set" />
)}
</TableCell>
<TableCell>
<Button onClick={() => handleUpdateField(value.field_name)}>Update</Button>
<span
onClick={() => handleResetField(value.field_name)}
className="inline-flex shrink-0 cursor-pointer items-center justify-center px-1.5 py-1.5 text-red-500"
>
<Trash2 className="h-5 w-5 shrink-0" />
</span>
</TableCell>
</TableRow>
))}
</TableBody>
</Table>
</CardContent>
</Card>
</TabsContent>
</Tabs>

2
ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx
View file

@ -14,7 +14,7 @@ import { MoneyCell } from "@/components/shared/table_cells";
import { Card as ShadcnCard, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
import { hasCapability, type Capability } from "@/utils/capabilities";
import { formatNumberWithCommas } from "@/utils/dataUtils";
import type { DateRangePickerValue } from "@tremor/react";
import type { DateRangePickerValue } from "@/components/shared/date_picker_types";
import { ChevronDown, ChevronRight, ExternalLink, Info, Loader2 } from "lucide-react";
import type { ColumnDef } from "@tanstack/react-table";
import { Alert, AlertDescription } from "@/components/shared/Alert";

2
ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx
View file

@ -7,7 +7,7 @@
*/
import { ChevronDown, ChevronRight, Download, ExternalLink, Info, Loader2, Sparkles, X } from "lucide-react";
import type { DateRangePickerValue } from "@tremor/react";
import type { DateRangePickerValue } from "@/components/shared/date_picker_types";
import React, { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { BarChart } from "@/components/shared/charts";

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