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

chore: merge litellm_internal_staging into litellm_lit_7022_azure_ai_passthrough_config

Browse source
This commit is contained in:
mateo-berri 2026-09-09 15:36:06 -07:00
commit 13745e89ac
176 changed files with 13599 additions and 2289 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
basedpyright-code-budget.json
View file

@ -105,7 +105,7 @@
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38271
"limit": 38269
},
"reportUnknownParameterType": {
"limit": 19584

13
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -1801,7 +1801,16 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Remove conflicting keys from data to avoid duplicate keyword arguments
filtered_data = {k: v for k, v in data.items() if k not in ("model", "file_id")}
for model_id, model_file_id in specific_model_file_id_mapping.items():
delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **filtered_data) # type: ignore
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
delete_data = {
**{k: v for k, v in filtered_data.items() if k != "_litellm_internal_model_credentials"},
**(
{"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
if credentials is not None
else {}
),
}
delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **delete_data)
stored_file_object = await self.delete_unified_file_id(file_id, litellm_parent_otel_span)
@ -1812,7 +1821,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
prom_logger.record_managed_file_deleted(result="success")
if stored_file_object:
return stored_file_object
return OpenAIFileObject.model_validate(stored_file_object).model_copy(update={"id": file_id})
elif delete_response:
delete_response.id = file_id
return delete_response

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.65"
version = "0.1.66"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.65"
version = "0.1.66"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

4
litellm/__init__.py
View file

@ -546,7 +546,7 @@ _key_management_system: Optional["KeyManagementSystem"] = None
#### PII MASKING ####
output_parse_pii: bool = False
#############################################
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map, mark_litellm_import_complete
model_cost = get_model_cost_map(url=model_cost_map_url)
cost_discount_config: Dict[str, float] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount
@ -2405,3 +2405,5 @@ def __getattr__(name: str) -> Any:
# ALL_LITELLM_RESPONSE_TYPES is lazy-loaded via __getattr__ to avoid loading utils at import time
mark_litellm_import_complete()

24
litellm/_internal_context.py
View file

@ -6,9 +6,33 @@ be settable from user input. Context variables are scoped to the current
asyncio task and cannot be injected via HTTP request bodies.
"""
from collections.abc import Generator
from contextlib import contextmanager
from contextvars import ContextVar
from datetime import datetime, timezone
from typing import Final
# When True, suppresses async logging and billing for internal sub-calls
# (e.g., emulated file-search steps that make nested LLM calls).
is_internal_call: Final[ContextVar[bool]] = ContextVar("is_internal_call", default=False)
# One request prices its totals, its per-token-type lines and the rates it reports on
# separate code paths. Each reads the clock for off-peak pricing, so without a pinned
# moment they can land on either side of a window boundary and disagree with each other.
_billing_time: Final[ContextVar[datetime | None]] = ContextVar("billing_time", default=None)
@contextmanager
def pinned_billing_time(moment: datetime) -> Generator[None]:
"""Price every rate lookup inside this block at ``moment`` rather than at each one's own clock read."""
token: Final = _billing_time.set(moment)
try:
yield
finally:
_billing_time.reset(token)
def current_billing_time() -> datetime:
"""The pinned billing moment, or now in UTC outside a pinned block."""
pinned: Final = _billing_time.get()
return pinned if pinned is not None else datetime.now(timezone.utc)

7
litellm/constants.py
View file

@ -143,6 +143,7 @@ DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD: Final = float(
os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD", 0.3)
)
MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH: Final = int(os.getenv("MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH", 150))
MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH: Final = 2048
DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS: Final = 2000
@ -197,6 +198,7 @@ LITELLM_UI_ALLOW_HEADERS: Final = [
"x-litellm-adaptive-router-model",
"x-litellm-applied-guardrails",
"x-litellm-guardrail-scan-id",
"x-litellm-guardrail-scan-metadata",
"x-litellm-cache-key",
]
@ -333,6 +335,7 @@ DEFAULT_SSL_CIPHERS: Final = os.getenv(
########### v2 Architecture constants for managing writing updates to the database ###########
REDIS_UPDATE_BUFFER_KEY: Final = "litellm_spend_update_buffer"
REDIS_GATEWAY_REQUESTS_BUFFER_KEY: Final = "litellm_gateway_requests_buffer"
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_spend_update_buffer"
REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_team_spend_update_buffer"
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_org_spend_update_buffer"
@ -1767,6 +1770,10 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
SPECIAL_LITELLM_AUTH_TOKEN: Final = ["ui-token"]
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60))
DEFAULT_ACCESS_GROUP_CACHE_TTL: Final = int(os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600))
SPEND_LOG_KEY_METADATA_CACHE_TTL: Final = 600
SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL: Final = 30
SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS: Final = 10000
SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS: Final = 5000
# Short TTL for negative MCP access-group existence lookups. Keeps unauthenticated
# callers from forcing a DB query per request for unknown names, while bounding
# staleness so a transient DB error (which surfaces as an empty list) cannot

13
litellm/cost_calculator.py
View file

@ -25,6 +25,7 @@ from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import
TranscriptionUsageObjectTransformation,
)
from litellm.litellm_core_utils.llm_cost_calc.utils import (
BilledTokenRates,
CostCalculatorUtils,
_generic_cost_per_character,
_get_regional_uplift_multiplier,
@ -45,6 +46,9 @@ from litellm.llms.azure.cost_calculation import (
from litellm.llms.azure_ai.cost_calculator import (
cost_per_token as azure_ai_cost_per_token,
)
from litellm.llms.azure_ai.cost_calculator import (
is_azure_model_router as azure_ai_is_model_router_name,
)
from litellm.llms.base_llm.search.transformation import SearchResponse
from litellm.llms.bedrock.cost_calculation import (
cost_per_token as bedrock_cost_per_token,
@ -1122,6 +1126,7 @@ def _store_cost_breakdown_in_logging_obj(
service_tier: str | None = None,
data_residency: str | None = None,
vertex_location: str | None = None,
billed_token_rates: BilledTokenRates | None = None,
) -> None:
"""
Helper function to store cost breakdown in the logging object.
@ -1166,6 +1171,7 @@ def _store_cost_breakdown_in_logging_obj(
service_tier=service_tier,
data_residency=data_residency,
vertex_location=vertex_location,
billed_token_rates=billed_token_rates,
)
except Exception as breakdown_error:
@ -1659,11 +1665,10 @@ def completion_cost(
data_residency=data_residency,
vertex_location=vertex_location,
response=completion_response,
request_model=request_model_for_cost,
)
# Get additional costs from provider (e.g., routing fees, infrastructure costs)
if custom_llm_provider == "azure_ai":
if custom_llm_provider == "azure_ai" and not azure_ai_is_model_router_name(model):
model_for_additional_costs = request_model_for_cost
if completion_response is not None:
hidden_params = getattr(completion_response, "_hidden_params", None) or {}
@ -1735,6 +1740,7 @@ def completion_cost(
_reasoning_cost: float | None = None
_cache_read_cost: float | None = None
_cache_creation_cost: float | None = None
_billed_token_rates: BilledTokenRates | None = None
if cost_per_token_usage_object is not None and model:
_breakdown_provider: str | None = (
custom_llm_provider if isinstance(custom_llm_provider, str) else None
@ -1746,10 +1752,12 @@ def completion_cost(
service_tier=service_tier,
data_residency=data_residency,
vertex_location=vertex_location,
custom_cost_per_token=custom_cost_per_token,
)
_reasoning_cost = _token_type_breakdown.reasoning_cost
_cache_read_cost = _token_type_breakdown.cache_read_cost
_cache_creation_cost = _token_type_breakdown.cache_creation_cost
_billed_token_rates = _token_type_breakdown.rates
_store_cost_breakdown_in_logging_obj(
litellm_logging_obj=litellm_logging_obj,
prompt_tokens_cost_usd_dollar=prompt_tokens_cost_usd_dollar,
@ -1769,6 +1777,7 @@ def completion_cost(
service_tier=service_tier,
data_residency=data_residency,
vertex_location=vertex_location,
billed_token_rates=_billed_token_rates,
)
return _final_cost

30
litellm/experimental_mcp_client/client.py
View file

@ -18,6 +18,7 @@ from mcp import ClientSession, McpError, ReadResourceResult, Resource, StdioServ
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
from mcp.shared.message import SessionMessage
from mcp.shared.session import RequestResponder
from typing_extensions import Unpack
_TransportStreams: TypeAlias = tuple[
@ -56,10 +57,13 @@ def missing_streamable_http_client_error() -> ImportError:
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
from mcp.types import CallToolResult as MCPCallToolResult
from mcp.types import (
ClientResult,
GetPromptRequestParams,
GetPromptResult,
Prompt,
ResourceTemplate,
ServerNotification,
ServerRequest,
TextContent,
)
from mcp.types import Tool as MCPTool
@ -146,8 +150,8 @@ _SDK_READ_TIMEOUT_CODE: Final = int(httpx.codes.REQUEST_TIMEOUT)
otherwise carries JSON-RPC error codes."""
def _as_read_timeout(exc: BaseException) -> TimeoutError | None:
"""The session read timeout elapsing, re-expressed as a ``TimeoutError``, or ``None``.
def as_mcp_read_timeout(exc: BaseException) -> TimeoutError | None:
"""Normalize an MCP SDK read timeout for client and gateway diagnostics, or return ``None``.
The SDK reports its own elapsed read timeout as ``McpError`` carrying an HTTP status code in a
field that otherwise holds JSON-RPC error codes, and it relays an upstream's JSON-RPC error
@ -442,6 +446,18 @@ class MCPClient:
in_flight_error: BaseException | None = None
try:
read_stream, write_stream = transport[0], transport[1]
stream_error: Final[asyncio.Future[Exception]] = asyncio.get_running_loop().create_future()
async def receive_message(
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
) -> None:
if not isinstance(message, (ValueError, httpx.RequestError, OSError)):
return
if not stream_error.done():
stream_error.set_result(message)
# The SDK closes pending requests when its message handler raises.
raise RuntimeError("MCP response stream failed")
# Build session kwargs with optional callbacks
session_kwargs: Final[dict[str, Any]] = {}
if self._sampling_callback is not None:
@ -456,6 +472,7 @@ class MCPClient:
read_stream,
write_stream,
read_timeout_seconds=timedelta(seconds=self.timeout),
message_handler=receive_message,
**session_kwargs,
)
session: Final = await session_ctx.__aenter__()
@ -467,6 +484,10 @@ class MCPClient:
if isinstance(ins, str) and ins.strip():
self._last_initialize_instructions = ins.strip()
return await operation(session)
except McpError:
if stream_error.done():
raise stream_error.result()
raise
finally:
try:
await session_ctx.__aexit__(None, None, None)
@ -501,11 +522,10 @@ class MCPClient:
transport_ctx, http_client = self._create_transport_context()
return await self._execute_session_operation(transport_ctx, operation)
except Exception as e:
read_timeout: Final = _as_read_timeout(e)
read_timeout: Final = as_mcp_read_timeout(e)
if read_timeout is not None:
verbose_logger.warning(
"MCP client timed out after %ss waiting for %s to answer; the server accepted the "
"request and ended its response stream without a JSON-RPC reply",
"MCP client timed out after %ss waiting for a valid MCP response from %s",
self.timeout,
self.server_url or "stdio",
)

2
litellm/files/main.py
View file

@ -31,7 +31,7 @@ FileCreateProvider = Literal[
FileRetrieveProvider = Literal[
"openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic"
]
FileDeleteProvider = Literal["openai", "azure", "gemini", "litellm_proxy", "manus", "anthropic"]
FileDeleteProvider = Literal["openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic"]
FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic"]
import litellm
from litellm import get_secret_str

4
litellm/integrations/azure_sentinel/azure_sentinel.py
View file

@ -21,6 +21,8 @@ from types import MappingProxyType
from typing import Final, TypeVar
from urllib.parse import urlparse
import httpx
from litellm._logging import verbose_logger
from litellm.integrations.batch_utils import (
BatchSendCancelled,
@ -418,7 +420,7 @@ class AzureSentinelLogger(CustomBatchLogger):
"Content-Type": "application/json",
}
async def _send_batch(batch: Sequence[_QueuedPayload]):
async def _send_batch(batch: Sequence[_QueuedPayload]) -> httpx.Response:
body: Final = safe_dumps(batch)
return await self.async_httpx_client.post(
url=api_endpoint,

30
litellm/integrations/custom_guardrail.py
View file

@ -850,20 +850,24 @@ class CustomGuardrail(CustomLogger):
if self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True:
return None
# CHECK IF GUARDRAIL REJECTS THE REQUEST
target: Final = self._deployment_hook_target()
hook_request_data: Final = {**request_data, "guardrail_to_apply": self} if target is not self else request_data
result: Final = await target.async_post_call_success_hook(
user_api_key_dict=UserAPIKeyAuth(
user_id=request_data.get("user_api_key_user_id"),
team_id=request_data.get("user_api_key_team_id"),
end_user_id=request_data.get("user_api_key_end_user_id"),
api_key=request_data.get("user_api_key_hash"),
request_route=request_data.get("user_api_key_request_route"),
),
data=hook_request_data,
response=response,
)
try:
if target is not self:
request_data["guardrail_to_apply"] = self # rebind-ok: dispatch consumes this key
result: Final = await target.async_post_call_success_hook(
user_api_key_dict=UserAPIKeyAuth(
user_id=request_data.get("user_api_key_user_id"),
team_id=request_data.get("user_api_key_team_id"),
end_user_id=request_data.get("user_api_key_end_user_id"),
api_key=request_data.get("user_api_key_hash"),
request_route=request_data.get("user_api_key_request_route"),
),
data=request_data,
response=response,
)
finally:
if target is not self:
request_data.pop("guardrail_to_apply", None)
if not self._is_valid_response_type(result):
return None

208
litellm/litellm_core_utils/get_model_cost_map.py
View file

@ -2,6 +2,7 @@
Pulls the cost + context window + provider route for known models from https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json
This can be disabled by setting the LITELLM_LOCAL_MODEL_COST_MAP environment variable to True.
The ``lite`` and ``litellm-proxy`` CLI entry points also use the bundled map without fetching.
```
export LITELLM_LOCAL_MODEL_COST_MAP=True
@ -9,17 +10,22 @@ export LITELLM_LOCAL_MODEL_COST_MAP=True
"""
import asyncio
import hashlib
import json
import os
import random
import sys
import threading
import time
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from dataclasses import dataclass, replace
from datetime import datetime, timezone
from importlib.resources import files
from pathlib import Path
from typing import Final, Protocol
import httpx
from typing_extensions import ReadOnly, TypedDict
from litellm import verbose_logger
from litellm.constants import (
@ -31,6 +37,12 @@ from litellm.litellm_core_utils.fallback_generalizations import (
)
FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations"
_CLI_ENTRYPOINT_NAMES: Final = frozenset({"lite", "litellm-proxy"})
def _is_cli_process() -> bool:
return Path(sys.argv[0]).stem in _CLI_ENTRYPOINT_NAMES
# Reserved top-level keys that are not model entries. They must be excluded
# from the model-count integrity check so a real upstream shrink can't be masked.
@ -42,6 +54,10 @@ def _count_model_entries(model_cost: dict) -> int:
return sum(1 for key in model_cost if key not in RESERVED_TOP_LEVEL_KEYS)
def git_blob_id(body: bytes) -> str:
return hashlib.sha1(b"blob %d\0" % len(body) + body, usedforsecurity=False).hexdigest()
class GetModelCostMap:
"""
Handles fetching, validating, and loading the model cost map.
@ -53,15 +69,24 @@ class GetModelCostMap:
_backup_model_count: int = -1 # -1 = not yet loaded
@staticmethod
def read_local_model_cost_map_bytes() -> bytes:
return files("litellm").joinpath("model_prices_and_context_window_backup.json").read_bytes()
@staticmethod
def read_local_model_cost_map_text() -> str:
return files("litellm").joinpath("model_prices_and_context_window_backup.json").read_text(encoding="utf-8")
return GetModelCostMap.read_local_model_cost_map_bytes().decode("utf-8")
@staticmethod
def load_local_model_cost_map_with_revision() -> "ModelCostMapReloaded":
body: Final = GetModelCostMap.read_local_model_cost_map_bytes()
content: Final = json.loads(body)
return ModelCostMapReloaded(model_cost_map=content, revision=git_blob_id(body))
@staticmethod
def load_local_model_cost_map() -> dict:
"""Load the local backup model cost map bundled with the package."""
content: Final = json.loads(GetModelCostMap.read_local_model_cost_map_text())
return content
return GetModelCostMap.load_local_model_cost_map_with_revision().model_cost_map
@classmethod
def _get_backup_model_count(cls) -> int:
@ -161,11 +186,18 @@ class GetModelCostMap:
RETRYABLE_FETCH_STATUS_CODES: Final = frozenset({429, 500, 502, 503, 504})
MODEL_COST_MAP_FETCH_MAX_ATTEMPTS: Final = 3
MODEL_COST_MAP_FETCH_MAX_WAIT_SECONDS: Final = 30.0
_litellm_import_complete = threading.Event()
def mark_litellm_import_complete() -> None:
_litellm_import_complete.set()
@dataclass(frozen=True, slots=True)
class ModelCostMapReloaded:
model_cost_map: dict # mutable-ok: adopted as litellm.model_cost, whose consumer contract is a plain mutable dict
revision: str | None = None
etag: str | None = None
@dataclass(frozen=True, slots=True)
@ -254,7 +286,9 @@ def _classify_fetch_response(response: httpx.Response, url: str) -> _FetchAttemp
return ModelCostMapReloadUnavailable(reason=f"invalid JSON from {url}: {e}")
if not isinstance(parsed, dict):
return ModelCostMapReloadUnavailable(reason=f"expected a JSON object from {url}, got {type(parsed).__name__}")
return ModelCostMapReloaded(model_cost_map=parsed)
return ModelCostMapReloaded(
model_cost_map=parsed, revision=git_blob_id(response.content), etag=response.headers.get("etag")
)
def _next_retry_wait(
@ -295,12 +329,13 @@ async def _fetch_remote_model_cost_map_with_retry(
def _fetch_remote_model_cost_map_with_retry_sync(
url: str,
timeout: int,
max_attempts: int,
attempts: range,
sleep: Callable[[float], None],
rng: random.Random,
client: _SyncGetClient,
) -> ModelCostMapReloadResult:
for attempt in range(1, max_attempts + 1):
max_attempts: Final = attempts.stop - 1
for attempt in attempts:
outcome = _attempt_fetch_sync(client=client, url=url, timeout=timeout)
if not isinstance(outcome, _FetchAttemptRetryable):
return outcome
@ -328,13 +363,12 @@ async def refetch_model_cost_map(
map they already have.
"""
if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true":
_cost_map_source_info.loaded_at = datetime.now(timezone.utc)
_cost_map_source_info.source = "local"
_cost_map_source_info.url = None
_cost_map_source_info.is_env_forced = True
_cost_map_source_info.fallback_reason = None
return ModelCostMapReloaded(
model_cost_map=_finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map())
)
return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision())
result: Final = await _fetch_remote_model_cost_map_with_retry(
url=url,
@ -355,11 +389,12 @@ async def refetch_model_cost_map(
backup_model_count=GetModelCostMap._get_backup_model_count(),
):
return ModelCostMapReloadUnavailable(reason=f"model cost map from {url} failed integrity validation")
_cost_map_source_info.loaded_at = datetime.now(timezone.utc)
_cost_map_source_info.source = "remote"
_cost_map_source_info.url = url
_cost_map_source_info.is_env_forced = False
_cost_map_source_info.fallback_reason = None
return ModelCostMapReloaded(model_cost_map=_finalize_model_cost_map(result.model_cost_map))
return _finalize_loaded_model_cost_map(result)
class ModelCostMapSourceInfo:
@ -370,13 +405,35 @@ class ModelCostMapSourceInfo:
is_env_forced: bool = False
fallback_reason: str | None = None
loaded_at: "datetime | None" = None
source_revision: str | None = None
etag: str | None = None
# Module-level singleton tracking the source of the current cost map
_cost_map_source_info: Final = ModelCostMapSourceInfo()
def get_model_cost_map_source_info() -> dict:
class CostMapProvenance(TypedDict):
source_revision: ReadOnly[str | None]
etag: ReadOnly[str | None]
class CostMapSourceInfo(CostMapProvenance):
source: ReadOnly[str]
url: ReadOnly[str | None]
is_env_forced: ReadOnly[bool]
fallback_reason: ReadOnly[str | None]
loaded_at: ReadOnly[str | None]
def get_model_cost_map_provenance() -> CostMapProvenance:
return {
"source_revision": _cost_map_source_info.source_revision,
"etag": _cost_map_source_info.etag,
}
def get_model_cost_map_source_info() -> CostMapSourceInfo:
"""
Return metadata about where the current model cost map was loaded from.
@ -385,12 +442,19 @@ def get_model_cost_map_source_info() -> dict:
- url: the remote URL attempted (or None for local-only)
- is_env_forced: True if LITELLM_LOCAL_MODEL_COST_MAP=True forced local usage
- fallback_reason: human-readable reason if remote failed and local was used
- loaded_at: ISO 8601 time this process last loaded the map
- source_revision: git blob id of the loaded file's bytes
- etag: the ETag of the remote fetch (None for the bundled backup)
"""
loaded_at: Final = _cost_map_source_info.loaded_at
return {
"source": _cost_map_source_info.source,
"url": _cost_map_source_info.url,
"is_env_forced": _cost_map_source_info.is_env_forced,
"fallback_reason": _cost_map_source_info.fallback_reason,
"loaded_at": loaded_at.isoformat() if loaded_at is not None else None,
"source_revision": _cost_map_source_info.source_revision,
"etag": _cost_map_source_info.etag,
}
@ -466,6 +530,74 @@ def _finalize_model_cost_map(model_cost: dict) -> dict:
return _expand_model_aliases(model_cost)
def _finalize_loaded_model_cost_map(loaded: ModelCostMapReloaded) -> ModelCostMapReloaded:
_cost_map_source_info.source_revision = loaded.revision
_cost_map_source_info.etag = loaded.etag
return replace(loaded, model_cost_map=_finalize_model_cost_map(loaded.model_cost_map))
def adopt_model_cost_map(
new_model_cost_map: dict, # mutable-ok: public API preserves the mutable cost-map contract
) -> int:
import litellm
from litellm import utils
litellm.model_cost = new_model_cost_map
utils._invalidate_model_cost_lowercase_map() # pyright: ignore[reportPrivateUsage] # required cache invalidation
litellm.add_known_models(model_cost_map=new_model_cost_map)
fetched_model_count: Final = len(new_model_cost_map) if new_model_cost_map else 0
utils.reapply_runtime_model_cost_registrations()
return fetched_model_count
def _retry_remote_fetch_in_background(
url: str,
timeout: int,
max_attempts: int,
sleep: Callable[[float], None],
rng: random.Random,
client: _SyncGetClient,
first_outcome: _FetchAttemptRetryable,
) -> None:
try:
first_wait: Final = _next_retry_wait(outcome=first_outcome, attempt=1, max_attempts=max_attempts, rng=rng)
if isinstance(first_wait, ModelCostMapReloadUnavailable):
return
sleep(first_wait)
result: Final = _fetch_remote_model_cost_map_with_retry_sync(
url=url,
timeout=timeout,
attempts=range(2, max_attempts + 1),
sleep=sleep,
rng=rng,
client=client,
)
if isinstance(result, ModelCostMapReloadUnavailable):
verbose_logger.warning(
"LiteLLM: Failed to fetch remote model cost map from %s after %d attempts; keeping local backup",
url,
max_attempts,
)
return
_litellm_import_complete.wait()
if not GetModelCostMap.validate_model_cost_map(
fetched_map=result.model_cost_map,
backup_model_count=GetModelCostMap._get_backup_model_count(), # pyright: ignore[reportPrivateUsage] # integrity cache
):
verbose_logger.warning(
"LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s",
url,
)
return
finalized: Final = _finalize_loaded_model_cost_map(result).model_cost_map
_cost_map_source_info.source = "remote"
_cost_map_source_info.fallback_reason = None
_cost_map_source_info.loaded_at = datetime.now(timezone.utc)
adopt_model_cost_map(finalized)
except Exception as e: # noqa: BLE001 # a failed background retry must not kill the task; the backup stays
verbose_logger.warning("LiteLLM: Background model cost map retry failed: %s", e)
def get_model_cost_map(
url: str,
timeout: int = 5,
@ -477,10 +609,12 @@ def get_model_cost_map(
"""
Public entry point — returns the model cost map dict.
1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set, uses the local backup only.
1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set or this is a ``lite`` /
``litellm-proxy`` CLI process, uses the local backup only.
2. Otherwise fetches from ``url``, retrying transient HTTP errors
(429/5xx/transport) with Retry-After-aware backoff, validates
integrity, and falls back to the local backup on any failure.
(429/5xx/transport) with Retry-After-aware backoff in a background
thread, validates integrity, and falls back to the local backup on any
failure.
Only the backup model count is cached (a single int) for validation.
The full backup dict is only parsed when it must be *returned* as a
@ -489,34 +623,44 @@ def get_model_cost_map(
_cost_map_source_info.loaded_at = datetime.now(timezone.utc)
# Note: can't use get_secret_bool here — this runs during litellm.__init__
# before litellm._key_management_settings is set.
if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true":
if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true" or _is_cli_process():
_cost_map_source_info.source = "local"
_cost_map_source_info.url = None
_cost_map_source_info.is_env_forced = True
_cost_map_source_info.fallback_reason = None
return _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map())
return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map
_cost_map_source_info.url = url
_cost_map_source_info.is_env_forced = False
result: Final = _fetch_remote_model_cost_map_with_retry_sync(
url=url,
timeout=timeout,
max_attempts=max_attempts,
sleep=sleep,
rng=rng if rng is not None else random.Random(),
client=client if client is not None else httpx,
)
if isinstance(result, ModelCostMapReloadUnavailable):
fetch_client: Final = client if client is not None else httpx
fetch_rng: Final = rng if rng is not None else random.Random()
outcome: Final = _attempt_fetch_sync(client=fetch_client, url=url, timeout=timeout)
if isinstance(outcome, _FetchAttemptRetryable) and max_attempts > 1:
threading.Thread(
target=_retry_remote_fetch_in_background,
kwargs={ # mutable-ok: threading requires a mutable keyword-arguments mapping
"url": url,
"timeout": timeout,
"max_attempts": max_attempts,
"sleep": sleep,
"rng": fetch_rng,
"client": fetch_client,
"first_outcome": outcome,
},
name="litellm-model-cost-map-retry",
daemon=True,
).start()
if not isinstance(outcome, ModelCostMapReloaded):
verbose_logger.warning(
"LiteLLM: Failed to fetch remote model cost map from %s: %s. Falling back to local backup.",
url,
result.reason,
outcome.reason,
)
_cost_map_source_info.source = "local"
_cost_map_source_info.fallback_reason = f"Remote fetch failed: {result.reason}"
return _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map())
content: Final = result.model_cost_map
_cost_map_source_info.fallback_reason = f"Remote fetch failed: {outcome.reason}"
return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map
content: Final = outcome.model_cost_map
# Validate using cached count (cheap int comparison, no file I/O)
if not GetModelCostMap.validate_model_cost_map(
@ -529,8 +673,8 @@ def get_model_cost_map(
)
_cost_map_source_info.source = "local"
_cost_map_source_info.fallback_reason = "Remote data failed integrity validation"
return _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map())
return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map
_cost_map_source_info.source = "remote"
_cost_map_source_info.fallback_reason = None
return _finalize_model_cost_map(content)
return _finalize_loaded_model_cost_map(outcome).model_cost_map

5
litellm/litellm_core_utils/litellm_logging.py
View file

@ -202,6 +202,7 @@ if TYPE_CHECKING:
from litellm.integrations.otel.logger import OpenTelemetryV2
from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
from litellm.litellm_core_utils.llm_cost_calc.utils import BilledTokenRates
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, LoggedRelayResponse
try:
from litellm_enterprise.enterprise_callbacks.callback_controls import (
@ -589,6 +590,7 @@ class Logging(LiteLLMLoggingBaseClass):
# Initialize cost breakdown field
self.cost_breakdown: CostBreakdown | None = None
self.billed_token_rates: BilledTokenRates | None = None
# Init Caching related details
self.caching_details: CachingDetails | None = None
@ -1586,6 +1588,7 @@ class Logging(LiteLLMLoggingBaseClass):
service_tier: str | None = None,
data_residency: str | None = None,
vertex_location: str | None = None,
billed_token_rates: "BilledTokenRates | None" = None,
) -> None:
"""
Helper method to store cost breakdown in the logging object.
@ -1605,8 +1608,10 @@ class Logging(LiteLLMLoggingBaseClass):
service_tier: Tier the costs above were priced on, already resolved
data_residency: Region uplift the costs above were priced on, already resolved
vertex_location: Vertex AI location the costs above were priced on, already resolved
billed_token_rates: Per-token rates the costs above were billed at, already resolved
"""
self.billed_token_rates = billed_token_rates
self.cost_breakdown = CostBreakdown(
input_cost=input_cost,
output_cost=output_cost,

247
litellm/litellm_core_utils/llm_cost_calc/utils.py
View file

@ -10,6 +10,7 @@ from typing import Any, Final, Literal, TypedDict, cast
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
import litellm
from litellm._internal_context import current_billing_time
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import (
select_tier_for_input,
@ -19,6 +20,7 @@ from litellm.types.utils import (
CacheCreationTokenDetails,
CallTypes,
CompletionTokensDetailsWrapper,
CostPerToken,
DataResidency,
ImageResponse,
ModelInfo,
@ -305,7 +307,7 @@ def _is_within_off_peak_window(off_peak_hours_utc: str | Sequence[str], current_
than being localised, so callers must pass datetime.now(timezone.utc), never datetime.now(),
or every window shifts by the host's offset.
"""
reference: Final = current_time if current_time is not None else datetime.now(timezone.utc)
reference: Final = current_time if current_time is not None else current_billing_time()
now: Final = (reference.astimezone(timezone.utc) if reference.tzinfo is not None else reference).time()
windows: Final = (off_peak_hours_utc,) if isinstance(off_peak_hours_utc, str) else off_peak_hours_utc
for window in windows:
@ -392,7 +394,7 @@ def _is_off_peak(off_peak: Mapping[str, object], current_time: datetime | None =
rules: the flat hours_utc windows, which apply every day, or any entry in windows, whose
hours apply only on its weekdays.
"""
reference: Final = current_time if current_time is not None else datetime.now(timezone.utc)
reference: Final = current_time if current_time is not None else current_billing_time()
reference_utc: Final = (
reference.astimezone(timezone.utc) if reference.tzinfo is not None else reference.replace(tzinfo=timezone.utc)
)
@ -1195,7 +1197,7 @@ def generic_cost_per_token(
usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens, 0
)
billing_time: Final = current_time if current_time is not None else datetime.now(timezone.utc)
billing_time: Final = current_time if current_time is not None else current_billing_time()
(
prompt_base_cost,
completion_base_cost,
@ -1309,42 +1311,90 @@ def _coerce_token_count(value: object) -> int:
return value if isinstance(value, int) and value > 0 else 0
@dataclass(frozen=True, slots=True)
class BilledTokenRates:
"""Per-token rates one request's usage bills at, after token tiers, off-peak windows and the
regional multipliers the totals apply, so each cost line equals its token count times its rate."""
input_cost_per_token: float
output_cost_per_token: float
cache_read_input_token_cost: float
cache_creation_input_token_cost: float
cache_creation_input_token_cost_above_1hr: float
output_cost_per_reasoning_token: float
def scaled(self, multiplier: float) -> "BilledTokenRates":
if multiplier == 1.0:
return self
return BilledTokenRates(
input_cost_per_token=self.input_cost_per_token * multiplier,
output_cost_per_token=self.output_cost_per_token * multiplier,
cache_read_input_token_cost=self.cache_read_input_token_cost * multiplier,
cache_creation_input_token_cost=self.cache_creation_input_token_cost * multiplier,
cache_creation_input_token_cost_above_1hr=self.cache_creation_input_token_cost_above_1hr * multiplier,
output_cost_per_reasoning_token=self.output_cost_per_reasoning_token * multiplier,
)
@dataclass(frozen=True, slots=True)
class TokenTypeCostBreakdown:
reasoning_cost: float
cache_read_cost: float
cache_creation_cost: float
rates: BilledTokenRates | None = None
"""Rates these lines were billed at, so a caller reporting both cannot resolve them a second,
differently-argued way. None when the model's pricing could not be resolved."""
def get_token_type_cost_breakdown(
model: str,
custom_llm_provider: str | None,
def _reasoning_token_count(usage: Usage) -> int:
parsed: Final = (
parse_completion_tokens_details(usage)["reasoning_tokens"] if usage.completion_tokens_details is not None else 0
)
return parsed or _coerce_token_count(getattr(usage, "reasoning_tokens", 0))
def _cache_token_counts(usage: Usage) -> tuple[int, int, CacheCreationTokenDetails | None]:
"""(cache read tokens, cache creation tokens, cache creation details): read from prompt_tokens_details
first, then the private top-level counters the Usage constructor mirrors cache tokens onto for
providers/callers that bypass the details."""
parsed: Final = parse_prompt_tokens_details(usage) if usage.prompt_tokens_details is not None else None
parsed_read: Final = parsed["cache_hit_tokens"] if parsed is not None else 0
parsed_creation: Final = parsed["cache_creation_tokens"] if parsed is not None else 0
return (
parsed_read or _coerce_token_count(getattr(usage, "_cache_read_input_tokens", 0)),
parsed_creation or _coerce_token_count(getattr(usage, "_cache_creation_input_tokens", 0)),
parsed["cache_creation_token_details"] if parsed is not None else None,
)
def _custom_pricing_rates(custom_cost_per_token: CostPerToken) -> BilledTokenRates:
"""Flat custom pricing has no tiers, uplifts or reasoning rate: cache tokens bill at the configured
cache rates (else the input rate) and reasoning at the output rate, as _cost_per_token_custom_pricing_helper does."""
input_rate: Final = custom_cost_per_token["input_cost_per_token"]
output_rate: Final = custom_cost_per_token["output_cost_per_token"]
cache_creation_rate: Final = custom_cost_per_token.get("cache_creation_input_token_cost", input_rate)
return BilledTokenRates(
input_cost_per_token=input_rate,
output_cost_per_token=output_rate,
cache_read_input_token_cost=custom_cost_per_token.get("cache_read_input_token_cost", input_rate),
cache_creation_input_token_cost=cache_creation_rate,
cache_creation_input_token_cost_above_1hr=cache_creation_rate,
output_cost_per_reasoning_token=output_rate,
)
def _cost_map_billed_rates(
model_info: ModelInfo,
usage: Usage,
service_tier: str | None = None,
data_residency: str | None = None,
vertex_location: str | None = None,
current_time: datetime | None = None,
) -> TokenTypeCostBreakdown:
"""
Provider-agnostic cost of reasoning and cache tokens, derived from the usage
object and model pricing alone.
This works for every provider, including Perplexity/Cerebras/Dashscope whose
cost calculators bypass ``generic_cost_per_token``, because cache tokens always
land on ``prompt_tokens_details`` (via the Usage constructor and provider
transformations) and reasoning tokens on ``completion_tokens_details``. It reuses
the same rate-resolution primitives as the total-cost path so the breakdown can
never drift from the totals. Returns zeros (never raises) when the model or its
pricing cannot be resolved.
"""
try:
model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider)
except Exception:
return TokenTypeCostBreakdown(0.0, 0.0, 0.0)
billing_time: Final = current_time if current_time is not None else datetime.now(timezone.utc)
custom_llm_provider: str | None,
service_tier: str | None,
data_residency: str | None,
vertex_location: str | None,
current_time: datetime | None,
) -> BilledTokenRates:
billing_time: Final = current_time if current_time is not None else current_billing_time()
(
_prompt_base_cost,
prompt_base_cost,
completion_base_cost,
cache_creation_cost_rate,
cache_creation_cost_above_1hr_rate,
@ -1356,13 +1406,6 @@ def get_token_type_cost_breakdown(
current_time=billing_time,
threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider),
)
reasoning_tokens = (
parse_completion_tokens_details(usage)["reasoning_tokens"] if usage.completion_tokens_details is not None else 0
)
if not reasoning_tokens:
reasoning_tokens = _coerce_token_count(getattr(usage, "reasoning_tokens", 0))
reasoning_rate: Final = _resolve_billed_reasoning_rate(
model_info=model_info,
usage=usage,
@ -1370,57 +1413,103 @@ def get_token_type_cost_breakdown(
completion_base_cost=completion_base_cost,
current_time=billing_time,
)
reasoning_cost = float(reasoning_tokens) * reasoning_rate
multiplier: Final = (
_get_regional_uplift_multiplier(model_info, data_residency)
* get_vertex_regional_endpoint_uplift(model_info, vertex_location)
* get_provider_specific_geo_multiplier(model_info=model_info, usage=usage)
)
return BilledTokenRates(
input_cost_per_token=prompt_base_cost,
output_cost_per_token=completion_base_cost,
cache_read_input_token_cost=cache_read_cost_rate,
cache_creation_input_token_cost=cache_creation_cost_rate,
cache_creation_input_token_cost_above_1hr=cache_creation_cost_above_1hr_rate,
output_cost_per_reasoning_token=reasoning_rate,
).scaled(multiplier)
cache_read_tokens = 0
cache_creation_tokens = 0
cache_creation_token_details: CacheCreationTokenDetails | None = None
if usage.prompt_tokens_details is not None:
prompt_tokens_details: Final = parse_prompt_tokens_details(usage)
cache_read_tokens = prompt_tokens_details["cache_hit_tokens"]
cache_creation_tokens = prompt_tokens_details["cache_creation_tokens"]
cache_creation_token_details = prompt_tokens_details["cache_creation_token_details"]
# Fall back to the private top-level counters the Usage constructor mirrors cache
# tokens onto, so providers/callers that bypass prompt_tokens_details are covered.
if not cache_read_tokens:
cache_read_tokens = _coerce_token_count(getattr(usage, "_cache_read_input_tokens", 0))
if not cache_creation_tokens:
cache_creation_tokens = _coerce_token_count(getattr(usage, "_cache_creation_input_tokens", 0))
cache_read_cost = float(cache_read_tokens) * cache_read_cost_rate
cache_creation_cost = calculate_cache_writing_cost(
cache_creation_tokens=cache_creation_tokens,
cache_creation_token_details=cache_creation_token_details,
cache_creation_cost_above_1hr=cache_creation_cost_above_1hr_rate,
cache_creation_cost=cache_creation_cost_rate,
def get_billed_token_rates(
model: str,
custom_llm_provider: str | None,
usage: Usage,
service_tier: str | None = None,
data_residency: str | None = None,
vertex_location: str | None = None,
current_time: datetime | None = None,
custom_cost_per_token: CostPerToken | None = None,
) -> BilledTokenRates | None:
"""Rates the cost calculator bills ``usage`` at, resolved exactly as the totals and the token-type
breakdown resolve them. None when the model's pricing cannot be resolved."""
if custom_cost_per_token is not None:
return _custom_pricing_rates(custom_cost_per_token)
try:
model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider)
except Exception: # noqa: BLE001 # get_model_info raises a bare Exception for an unmapped model: no rates
return None
return _cost_map_billed_rates(
model_info=model_info,
usage=usage,
custom_llm_provider=custom_llm_provider,
service_tier=service_tier,
data_residency=data_residency,
vertex_location=vertex_location,
current_time=current_time,
)
# Apply the same flat regional-processing uplift the totals get, so per-type
# costs stay reconciled with input_cost/output_cost for regionalized OpenAI hosts.
uplift: Final = _get_regional_uplift_multiplier(model_info, data_residency)
if uplift != 1.0:
reasoning_cost *= uplift
cache_read_cost *= uplift
cache_creation_cost *= uplift
vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location)
if vertex_uplift != 1.0:
reasoning_cost *= vertex_uplift
cache_read_cost *= vertex_uplift
cache_creation_cost *= vertex_uplift
def get_token_type_cost_breakdown(
model: str,
custom_llm_provider: str | None,
usage: Usage,
service_tier: str | None = None,
data_residency: str | None = None,
vertex_location: str | None = None,
current_time: datetime | None = None,
custom_cost_per_token: CostPerToken | None = None,
) -> TokenTypeCostBreakdown:
"""
Provider-agnostic cost of reasoning and cache tokens, derived from the usage
object and model pricing alone.
# Mirror the provider-specific geo uplift (e.g. Anthropic us: 1.1) the totals
# apply, so cache and reasoning line items stay reconciled with them.
geo_multiplier: Final = get_provider_specific_geo_multiplier(model_info=model_info, usage=usage)
if geo_multiplier != 1.0:
reasoning_cost *= geo_multiplier
cache_read_cost *= geo_multiplier
cache_creation_cost *= geo_multiplier
This works for every provider, including Perplexity/Cerebras/Dashscope whose
cost calculators bypass ``generic_cost_per_token``, because cache tokens always
land on ``prompt_tokens_details`` (via the Usage constructor and provider
transformations) and reasoning tokens on ``completion_tokens_details``. It reuses
the same rate resolution as the total-cost path (``get_billed_token_rates``) so the
breakdown can never drift from the totals. A deployment billed by
``custom_cost_per_token`` is priced from those flat rates instead of the cost map and,
like its totals, bills cache writes flat rather than by their 5m/1h split.
Returns zeros (never raises) when the model or its pricing cannot be resolved.
"""
rates: Final = get_billed_token_rates(
model=model,
custom_llm_provider=custom_llm_provider,
usage=usage,
service_tier=service_tier,
data_residency=data_residency,
vertex_location=vertex_location,
current_time=current_time,
custom_cost_per_token=custom_cost_per_token,
)
if rates is None:
return TokenTypeCostBreakdown(0.0, 0.0, 0.0)
cache_read_tokens, cache_creation_tokens, cache_creation_token_details = _cache_token_counts(usage)
cache_creation_cost: Final = (
float(cache_creation_tokens) * rates.cache_creation_input_token_cost
if custom_cost_per_token is not None
else calculate_cache_writing_cost(
cache_creation_tokens=cache_creation_tokens,
cache_creation_token_details=cache_creation_token_details,
cache_creation_cost_above_1hr=rates.cache_creation_input_token_cost_above_1hr,
cache_creation_cost=rates.cache_creation_input_token_cost,
)
)
return TokenTypeCostBreakdown(
reasoning_cost=reasoning_cost,
cache_read_cost=cache_read_cost,
reasoning_cost=float(_reasoning_token_count(usage)) * rates.output_cost_per_reasoning_token,
cache_read_cost=float(cache_read_tokens) * rates.cache_read_input_token_cost,
cache_creation_cost=cache_creation_cost,
rates=rates,
)

16
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -24,6 +24,7 @@ from litellm.types.utils import (
Choices,
CompletionTokensDetails,
CompletionTokensDetailsWrapper,
Delta,
Function,
FunctionCall,
ModelResponse,
@ -326,6 +327,18 @@ class ChunkProcessor:
return chunk_id
return ""
@staticmethod
def _get_role_from_chunks(chunks: Sequence["_BaseChunk"]) -> str:
return ChunkProcessor._role_of_choice(next((c["choices"][0] for c in chunks if c.get("choices")), None))
@staticmethod
def _role_of_choice(choice: object) -> str:
match choice:
case StreamingChoices(delta=Delta(role=str() as role)) | {"delta": {"role": str() as role}} if role:
return role
case _:
return "assistant"
@staticmethod
def _get_model_from_chunks(chunks: Sequence["_BaseChunk"], first_chunk_model: str) -> str:
"""
@ -353,8 +366,7 @@ class ChunkProcessor:
model: Final = ChunkProcessor._get_model_from_chunks(chunks, first_chunk_model)
system_fingerprint: Final = chunk.get("system_fingerprint", None)
first_chunk_with_choices: Final = next((c for c in chunks if c.get("choices")), chunk)
role: Final = first_chunk_with_choices["choices"][0]["delta"]["role"]
role: Final = ChunkProcessor._get_role_from_chunks(chunks)
finish_reason = "stop"
for chunk in chunks:
if "choices" in chunk and len(chunk["choices"]) > 0:

23
litellm/llms/anthropic/experimental_pass_through/utils.py
View file

@ -3,6 +3,8 @@ from collections.abc import Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from pydantic import BaseModel, ConfigDict, ValidationError
import litellm
from litellm.types.utils import ModelInfo
@ -21,10 +23,27 @@ _EFFORT_DEGRADATION_CHAIN: Final[Mapping[str, tuple[str, ...]]] = MappingProxyTy
_THINKING_OFF: Final = "none"
class _ClaudeCodeUserId(BaseModel):
"""The JSON Claude Code packs into ``metadata.user_id``; only ``session_id`` is per conversation."""
model_config = ConfigDict(frozen=True)
session_id: str
def prompt_cache_key_from_user_id(user_id: object) -> str | None:
if user_id is None:
"""The per-session key Claude Code carries inside ``metadata.user_id``, or nothing.
Anthropic defines ``user_id`` as an opaque end-user id, so a plain string names a person, not
a conversation. Keying the provider cache on it pins every parallel session and subagent of that
person to one slot, which caches worse than the provider's own prompt-prefix hashing does.
"""
if not isinstance(user_id, str):
return None
try:
return _ClaudeCodeUserId.model_validate_json(user_id).session_id[:OPENAI_MAX_PROMPT_CACHE_KEY_LENGTH] or None
except ValidationError:
return None
return str(user_id)[:OPENAI_MAX_PROMPT_CACHE_KEY_LENGTH] or None
def litellm_logging_obj_from_kwargs(kwargs: Mapping[str, object]) -> "LiteLLMLoggingObject | None":

7
litellm/llms/azure/audio_transcriptions.py
View file

@ -37,6 +37,7 @@ class AzureAudioTranscription(AzureChatCompletion):
azure_ad_token: str | None = None,
atranscription: bool = False,
litellm_params: dict | None = None,
custom_llm_provider: str = "azure",
) -> TranscriptionResponse | Coroutine[Any, Any, TranscriptionResponse]:
data: Final = {"model": model, "file": audio_file, **optional_params}
@ -53,6 +54,7 @@ class AzureAudioTranscription(AzureChatCompletion):
logging_obj=logging_obj,
model=model,
litellm_params=litellm_params,
custom_llm_provider=custom_llm_provider,
)
azure_client: Final = self.get_azure_openai_client(
@ -99,7 +101,7 @@ class AzureAudioTranscription(AzureChatCompletion):
additional_args={"complete_input_dict": data},
original_response=stringified_response,
)
hidden_params: Final = {"model": model, "custom_llm_provider": "azure"}
hidden_params: Final = {"model": model, "custom_llm_provider": custom_llm_provider}
final_response: Final[TranscriptionResponse] = convert_to_model_response_object(
response_object=stringified_response,
model_response_object=model_response,
@ -122,6 +124,7 @@ class AzureAudioTranscription(AzureChatCompletion):
client=None,
max_retries=None,
litellm_params: dict | None = None,
custom_llm_provider: str = "azure",
) -> TranscriptionResponse:
response = None
try:
@ -178,7 +181,7 @@ class AzureAudioTranscription(AzureChatCompletion):
},
original_response=stringified_response,
)
hidden_params: Final = {"model": model, "custom_llm_provider": "azure"}
hidden_params: Final = {"model": model, "custom_llm_provider": custom_llm_provider}
response = convert_to_model_response_object(
_response_headers=headers,
response_object=stringified_response,

116
litellm/llms/azure_ai/cost_calculator.py
View file

@ -11,7 +11,7 @@ from litellm.types.utils import Usage
from litellm.utils import get_model_info
def _is_azure_model_router(model: str) -> bool:
def is_azure_model_router(model: str) -> bool:
"""
Check if the model is Azure AI Foundry Model Router.
@ -31,6 +31,18 @@ def _is_azure_model_router(model: str) -> bool:
return "model-router" in model_lower or "model_router" in model_lower or model_lower == "azure-model-router"
ROUTER_FEE_ENTRY_NAMES: Final = frozenset({"model-router", "model_router"})
def is_router_fee_entry(model: str) -> bool:
return model.lower().removeprefix("azure_ai/") in ROUTER_FEE_ENTRY_NAMES
def _router_fee_entry_name(model: str) -> str:
entry_name: Final = model.lower().removeprefix("azure_ai/")
return entry_name if entry_name in ROUTER_FEE_ENTRY_NAMES else "model_router"
def calculate_azure_model_router_flat_cost(model: str, prompt_tokens: int) -> float:
"""
Calculate the flat cost for Azure AI Foundry Model Router.
@ -42,20 +54,39 @@ def calculate_azure_model_router_flat_cost(model: str, prompt_tokens: int) -> fl
Returns:
float: The flat cost in USD, or 0.0 if not applicable
"""
if not _is_azure_model_router(model):
if not is_azure_model_router(model):
return 0.0
# Get the model router pricing from model_prices_and_context_window.json
# Use "model_router" as the key (without actual model name suffix)
model_info: Final = get_model_info(model="model_router", custom_llm_provider="azure_ai")
model_info: Final = get_model_info(model=_router_fee_entry_name(model), custom_llm_provider="azure_ai")
router_flat_cost_per_token: Final = model_info.get("input_cost_per_token", 0)
if router_flat_cost_per_token and router_flat_cost_per_token > 0:
return prompt_tokens * router_flat_cost_per_token
return 0.0
def _response_model_cost(model: str, usage: Usage, service_tier: str | None) -> tuple[float, float]:
try:
return generic_cost_per_token(
model=model, usage=usage, custom_llm_provider="azure_ai", service_tier=service_tier
)
except Exception as e:
if not is_azure_model_router(model):
raise
verbose_logger.debug(
"Azure AI Model Router: model '%s' not in cost map, only the routing fee applies. Error: %s", model, e
)
return 0.0, 0.0
def _router_fee_name(model: str, request_model: str | None) -> str | None:
if is_router_fee_entry(model):
return None
if is_azure_model_router(model):
return model
if request_model is not None and is_azure_model_router(request_model):
return request_model
return None
def cost_per_token(
model: str,
usage: Usage,
@ -64,68 +95,31 @@ def cost_per_token(
service_tier: str | None = None,
) -> tuple[float, float]:
"""
Calculate the cost per token for Azure AI models.
Price the response model's own tokens for Azure AI, plus the Model Router fee exactly once when either the
priced name or request_model is a Model Router name.
For Azure AI Foundry Model Router:
- Adds a flat cost of $0.14 per million input tokens (from model_prices_and_context_window.json)
- Plus the cost of the actual model used (handled by generic_cost_per_token)
A response priced as the router entry itself already carries the fee, so nothing is added on top of it. A
router deployment name that is missing from the cost map prices at the fee alone.
completion_cost passes only the priced name: when that name is a routed model it adds the fee itself through
AzureModelRouterConfig.calculate_additional_costs as the "Azure Model Router Flat Cost" line of the cost
breakdown, and when the name is router-shaped the fee is already in the prompt cost returned here.
Args:
model: str, the model name without provider prefix (from response)
usage: LiteLLM Usage block
response_time_ms: Optional response time in milliseconds
request_model: Optional[str], the original request model name (to detect router usage)
request_model: Optional[str], the original request model name; a Model Router name adds the routing fee
service_tier: Optional service tier the request was priced on
Returns:
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
Raises:
ValueError: If the model is not found in the cost map and cost cannot be calculated
(except for Model Router models where we return just the routing flat cost)
ValueError: If a model that is not a Model Router name is missing from the cost map
"""
prompt_cost = 0.0
completion_cost = 0.0
# Determine if this was a model router request
# Check both the response model and the request model
is_router_request: Final = _is_azure_model_router(model) or (
request_model is not None and _is_azure_model_router(request_model)
)
# Calculate base cost using generic cost calculator
# This may raise an exception if the model is not in the cost map
try:
prompt_cost, completion_cost = generic_cost_per_token(
model=model,
usage=usage,
custom_llm_provider="azure_ai",
service_tier=service_tier,
)
except Exception as e:
# For Model Router, the model name (e.g., "azure-model-router") may not be in the cost map
# because it's a routing service, not an actual model. In this case, we continue
# to calculate just the routing flat cost.
if not _is_azure_model_router(model):
# Re-raise for non-router models - they should have pricing defined
raise
verbose_logger.debug(
"Azure AI Model Router: model '%s' not in cost map, calculating routing flat cost only. Error: %s", model, e
)
# Add flat cost for Azure Model Router
# The flat cost is defined in model_prices_and_context_window.json for azure_ai/model_router
if is_router_request:
# Use the request model for flat cost calculation if available, otherwise use response model
router_model_for_calc: Final = request_model if request_model else model
router_flat_cost: Final = calculate_azure_model_router_flat_cost(router_model_for_calc, usage.prompt_tokens)
if router_flat_cost > 0:
verbose_logger.debug(
f"Azure AI Model Router flat cost: ${router_flat_cost:.6f} "
f"({usage.prompt_tokens} tokens × ${router_flat_cost / usage.prompt_tokens:.9f}/token)"
)
# Add flat cost to prompt cost
prompt_cost += router_flat_cost
return prompt_cost, completion_cost
prompt_cost, completion_cost = _response_model_cost(model=model, usage=usage, service_tier=service_tier)
fee_name: Final = _router_fee_name(model=model, request_model=request_model)
if fee_name is None:
return prompt_cost, completion_cost
return prompt_cost + calculate_azure_model_router_flat_cost(fee_name, usage.prompt_tokens), completion_cost

72
litellm/llms/bedrock/files/transformation.py
View file

@ -7,13 +7,13 @@ from contextlib import suppress
from functools import cache
from itertools import chain
from types import MappingProxyType
from typing import Any, Final, TypeAlias, TypedDict
from typing import Any, Final, Literal, TypeAlias, TypedDict
from urllib.parse import unquote
import httpx
from httpx import Headers, Response
from openai.types.file_deleted import FileDeleted
from pydantic import BaseModel, ConfigDict, TypeAdapter
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
from typing_extensions import ReadOnly
from litellm._logging import verbose_logger
@ -60,11 +60,12 @@ from litellm.utils import get_llm_provider
from ..base_aws_llm import BaseAWSLLM
from ..common_utils import BedrockError, merge_bedrock_aws_request_params, resolve_s3_encryption_key_id
# litellm_params key used to hand the SigV4-signed GET headers from
# `transform_file_content_request` to `validate_environment` (the only hook
# the shared file-content HTTP handler exposes for setting request headers).
# Same pattern as the `upload_url` handoff in `transform_create_file_request`.
S3_SIGNED_GET_HEADERS_PARAM: Final = "_s3_signed_get_headers"
S3_SIGNED_REQUEST_HEADERS_PARAM: Final = "_s3_signed_request_headers"
class _S3DeleteContext(BaseModel):
file_id: str = Field(min_length=1)
# litellm_params key carrying the size of the body uploaded to S3, handed from
# `transform_create_file_request` to `transform_create_file_response`.
@ -291,7 +292,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
) -> dict:
result: Final[dict[str, object]] = {}
result.update(headers)
signed_headers: Final = litellm_params.pop(S3_SIGNED_GET_HEADERS_PARAM, None)
signed_headers: Final = litellm_params.pop(S3_SIGNED_REQUEST_HEADERS_PARAM, None)
if isinstance(signed_headers, Mapping):
result.update(signed_headers) # any-ok: untyped handoff headers
# otherwise no extra headers - AWS credentials are handled by BaseAWSLLM
@ -1187,18 +1188,27 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
def transform_delete_file_request(
self,
file_id: str,
optional_params: dict,
litellm_params: dict,
) -> tuple[str, dict]:
raise NotImplementedError("BedrockFilesConfig does not support file deletion")
optional_params: Mapping[str, object],
litellm_params: MutableMapping[str, object],
) -> tuple[str, dict[str, str]]:
return self._transform_s3_file_request(
file_id=file_id, method="DELETE", optional_params=optional_params, litellm_params=litellm_params
)
def transform_delete_file_response(
self,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
litellm_params: Mapping[str, object],
) -> FileDeleted:
raise NotImplementedError("BedrockFilesConfig does not support file deletion")
if raw_response.status_code != 204:
raise BedrockError(
status_code=raw_response.status_code if raw_response.status_code >= 400 else 502,
message=raw_response.text or f"S3 file deletion returned HTTP {raw_response.status_code}",
headers=raw_response.headers,
)
context: Final = _S3DeleteContext.model_validate(logging_obj.model_call_details.get("additional_args"))
return FileDeleted(id=context.file_id, deleted=True, object="file")
def transform_list_files_request(
self,
@ -1233,6 +1243,18 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
if not file_id:
raise ValueError("file_id is required for Bedrock file content retrieval")
return self._transform_s3_file_request(
file_id=file_id, method="GET", optional_params=optional_params, litellm_params=litellm_params
)
def _transform_s3_file_request(
self,
*,
file_id: str,
method: Literal["GET", "DELETE"],
optional_params: Mapping[str, object],
litellm_params: MutableMapping[str, object],
) -> tuple[str, dict[str, str]]:
s3_uri: Final = extract_s3_uri_from_file_id(file_id)
bucket_name, object_key = _validate_file_id_against_configured_buckets(
s3_uri=s3_uri,
@ -1240,40 +1262,32 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params),
)
# The shared file-content handler passes optional_params={}, so AWS
# credentials/region arrive via litellm_params here (unlike the upload
# path). s3_region_name wins over aws_region_name, same priority as
# get_complete_file_url above.
merged_params: Final[dict[str, object]] = {}
merged_params.update(litellm_params)
merged_params.update(optional_params)
request_params: Final = _BedrockS3RequestParams.model_validate(merged_params)
request_params: Final = _BedrockS3RequestParams.model_validate({**litellm_params, **optional_params})
region_preference: Final = request_params.s3_region_name or request_params.aws_region_name
region_params: Final[dict[str, str | None]] = {"aws_region_name": region_preference}
aws_region_name: Final = self._get_aws_region_name(optional_params=region_params, model="")
s3_endpoint_url = (
s3_endpoint_url: Final = (
request_params.s3_endpoint_url or f"https://s3.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}"
).rstrip("/")
url: Final = f"{s3_endpoint_url}/{bucket_name}/{encode_s3_object_key_for_url(object_key)}"
litellm_params[S3_SIGNED_GET_HEADERS_PARAM] = self._sign_s3_get_request(
litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] = self._sign_s3_request_without_body(
api_base=url,
aws_region_name=aws_region_name,
request_params=request_params,
method=method,
)
return url, {}
def _sign_s3_get_request(
def _sign_s3_request_without_body(
self,
api_base: str,
aws_region_name: str,
request_params: _BedrockS3RequestParams,
method: Literal["GET", "DELETE"] = "GET",
) -> dict[str, str]:
"""
SigV4-sign an S3 GetObject request, mirroring `_sign_s3_request` (PUT).
"""
try:
import hashlib
@ -1297,7 +1311,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
empty_body_hash: Final = hashlib.sha256(b"").hexdigest()
aws_request: Final = AWSRequest( # any-ok: botocore AWSRequest is untyped
method="GET",
method=method,
url=api_base,
headers={"x-amz-content-sha256": empty_body_hash},
)

4
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -579,7 +579,7 @@ class BaseLLMHTTPHandler:
data: dict[str, object], # mutable-ok: async_completion takes dict
signed_headers: dict[str, object], # mutable-ok: async_completion takes dict
signed_json_body: bytes | None,
):
) -> Coroutine[object, object, ModelResponse | CustomStreamWrapper]:
async_client: Final = client if isinstance(client, AsyncHTTPHandler) else None
if stream is True:
return self.acompletion_stream_function(
@ -626,7 +626,7 @@ class BaseLLMHTTPHandler:
if acompletion is True and provider_config.uses_async_transform_request:
async def transform_then_dispatch():
async def transform_then_dispatch() -> ModelResponse | CustomStreamWrapper:
transformed: Final = cast( # cast-ok: async_transform_request is declared as a bare dict
"dict[str, object]",
await provider_config.async_transform_request(

100
litellm/llms/fireworks_ai/responses/transformation.py
View file

@ -1,10 +1,10 @@
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from urllib.parse import unquote
import httpx
from openai.types.responses import EasyInputMessageParam, ResponseInputItemParam
from openai.types.responses import EasyInputMessageParam, ResponseInputContentParam, ResponseInputItemParam
from litellm.llms.fireworks_ai.common_utils import (
resolve_fireworks_api_key,
@ -31,6 +31,17 @@ def _session_params(litellm_params: GenericLiteLLMParams) -> Mapping[str, object
)
_INSTRUCTION_ROLES: Final = frozenset({"system", "developer"})
def _role(item: ResponseInputItemParam) -> str | None:
match item:
case {"role": str(role)}:
return role
case _:
return None
def _developer_item_as_system(item: ResponseInputItemParam) -> ResponseInputItemParam:
if "role" not in item or item["role"] != "developer":
return item
@ -43,6 +54,69 @@ def _developer_items_as_system(input: str | ResponseInputParam) -> str | Respons
return [_developer_item_as_system(item) for item in input]
def _text_part(part: ResponseInputContentParam) -> str | None:
match part:
case {"type": "input_text", "text": str(text)}:
return text
case _:
return None
def _text_only_content(item: ResponseInputItemParam) -> str | None:
match item:
case {"role": "system" | "developer", "content": str(text)}:
return text
case {"role": "system" | "developer", "content": [*parts]}:
texts: Final = tuple(map(_text_part, parts))
return None if any(text is None for text in texts) else "\n\n".join(text for text in texts if text)
case _:
return None
def _leading_instruction_block_length(roles: Sequence[str | None]) -> int:
return next((index for index, role in enumerate(roles) if role not in _INSTRUCTION_ROLES), len(roles))
def _closing_instruction_block_start(roles: Sequence[str | None], leading_length: int) -> int:
last_conversation_index: Final = next(
(index for index in range(len(roles) - 1, leading_length - 1, -1) if roles[index] not in _INSTRUCTION_ROLES),
None,
)
if last_conversation_index is None or roles[last_conversation_index] != "assistant":
return len(roles)
return last_conversation_index + 1
def _hoisted_indices(roles: Sequence[str | None]) -> tuple[int, ...]:
leading_length: Final = _leading_instruction_block_length(roles)
closing_start: Final = _closing_instruction_block_start(roles, leading_length)
return tuple(
index for index, role in enumerate(roles[:closing_start]) if index < leading_length or role == "developer"
)
def _with_instruction_items_folded(
input: str | ResponseInputParam, instructions: str | None
) -> tuple[str | None, str | ResponseInputParam]:
if isinstance(input, str):
return instructions, input
items: Final = tuple(input)
folded: Final = MappingProxyType(
{
index: text
for index in _hoisted_indices(tuple(map(_role, items)))
if (text := _text_only_content(items[index])) is not None
}
)
joined: Final = "\n\n".join(chunk for chunk in (instructions, *folded.values()) if chunk)
return (
instructions if not folded else joined or None,
[ # mutable-ok: the base class takes the input items as a list
_developer_item_as_system(item) for index, item in enumerate(items) if index not in folded
],
)
class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
@property
def custom_llm_provider(self) -> LlmProviders:
@ -68,9 +142,6 @@ class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
base: Final = (api_base or get_secret_str("FIREWORKS_API_BASE") or FIREWORKS_AI_DEFAULT_API_BASE).rstrip("/")
return f"{base}/responses"
def _validate_input_param(self, input: str | ResponseInputParam) -> str | ResponseInputParam:
return _developer_items_as_system(super()._validate_input_param(input))
def transform_responses_api_request(
self,
model: str,
@ -79,10 +150,25 @@ class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
litellm_params: GenericLiteLLMParams,
headers: dict, # mutable-ok: overrides the base class signature
) -> dict: # mutable-ok: overrides the base class signature
instructions_param: Final[object] = response_api_optional_request_params.get("instructions")
validated_input: Final = self._validate_input_param(input)
instructions, folded_input = (
_with_instruction_items_folded(validated_input, instructions_param)
if isinstance(instructions_param, str | None)
else (instructions_param, _developer_items_as_system(validated_input))
)
instruction_entries: Final = () if instructions is None else (("instructions", instructions),)
folded_params: Final = { # mutable-ok: the base class takes the optional params as a dict
key: value
for key, value in (
*((key, value) for key, value in response_api_optional_request_params.items() if key != "instructions"),
*instruction_entries,
)
}
return super().transform_responses_api_request(
model=resolve_fireworks_resource_name(model),
input=input,
response_api_optional_request_params=response_api_optional_request_params,
input=folded_input,
response_api_optional_request_params=folded_params,
litellm_params=litellm_params,
headers=headers,
)

1
litellm/main.py
View file

@ -7805,6 +7805,7 @@ def transcription(
azure_ad_token=azure_ad_token,
max_retries=max_retries,
litellm_params=litellm_params_dict,
custom_llm_provider=custom_llm_provider,
)
elif custom_llm_provider == "openai" or (custom_llm_provider in litellm.openai_compatible_providers):
api_base = (

136
litellm/model_prices_and_context_window_backup.json
View file

@ -3626,6 +3626,79 @@
"supports_xhigh_reasoning_effort": true,
"supports_minimal_reasoning_effort": false
},
"azure_ai/gpt-chat-latest": {
"cache_read_input_token_cost": 5e-07,
"deprecation_date": "2026-12-02",
"input_cost_per_token": 5e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 3e-05,
"source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"azure_ai/codex-mini": {
"cache_read_input_token_cost": 3.75e-07,
"deprecation_date": "2026-11-15",
"input_cost_per_token": 1.5e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "responses",
"output_cost_per_token": 6e-06,
"source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/",
"supported_endpoints": [
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true
},
"azure_ai/whisper": {
"deprecation_date": "2026-12-15",
"input_cost_per_second": 0.0001,
"litellm_provider": "azure_ai",
"mode": "audio_transcription",
"output_cost_per_second": 0.0001,
"source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/"
},
"azure_ai/gpt-5.5-2026-04-23": {
"cache_read_input_token_cost": 5e-07,
"cache_read_input_token_cost_above_272k_tokens": 1e-06,
@ -4029,13 +4102,29 @@
"supports_minimal_reasoning_effort": false
},
"azure_ai/model_router": {
"deprecation_date": "2027-05-20",
"input_cost_per_token": 1.4e-07,
"output_cost_per_token": 0,
"litellm_provider": "azure_ai",
"max_input_tokens": 200000,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/aoai/",
"comment": "Flat cost of $0.14 per M input tokens for Azure AI Foundry Model Router infrastructure. Use pattern: azure_ai/model_router/<deployment-name> where deployment-name is your Azure deployment (e.g., azure-model-router)"
},
"azure_ai/model-router": {
"deprecation_date": "2027-05-20",
"input_cost_per_token": 1.4e-07,
"output_cost_per_token": 0,
"litellm_provider": "azure_ai",
"max_input_tokens": 200000,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/aoai/",
"comment": "Catalog-name twin of azure_ai/model_router: the flat $0.14 per M input tokens is the router's own fee, the routed model is priced on top of it"
},
"azure/eu/gpt-4o-2024-08-06": {
"deprecation_date": "2027-04-14",
"cache_read_input_token_cost": 1.375e-06,
@ -10347,6 +10436,18 @@
"/v1/ocr"
]
},
"azure_ai/cohere-command-a": {
"input_cost_per_token": 2.5e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 131072,
"max_output_tokens": 8182,
"max_tokens": 8182,
"mode": "chat",
"output_cost_per_token": 1e-05,
"source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/cohere/",
"supports_function_calling": true,
"supports_tool_choice": true
},
"azure_ai/doc-intelligence/prebuilt-read": {
"litellm_provider": "azure_ai",
"ocr_cost_per_page": 0.0015,
@ -10698,6 +10799,41 @@
"supports_vision": true,
"supports_web_search": true
},
"azure_ai/grok-4-20-reasoning": {
"cache_read_input_token_cost": 1.25e-06,
"deprecation_date": "2027-04-06",
"input_cost_per_token": 1.25e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 262000,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/grok/",
"supports_function_calling": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true,
"supports_reasoning": true
},
"azure_ai/grok-4-20-non-reasoning": {
"cache_read_input_token_cost": 1.25e-06,
"deprecation_date": "2027-04-06",
"input_cost_per_token": 1.25e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 262000,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/grok/",
"supports_function_calling": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"azure_ai/grok-4-fast-non-reasoning": {
"deprecation_date": "2026-05-01",
"input_cost_per_token": 2e-07,

3
litellm/proxy/_experimental/mcp_server/mcp_context.py
View file

@ -21,3 +21,6 @@ _mcp_gateway_initialize_instructions: Final[ContextVar[str | None]] = ContextVar
# Per-request scoped server name; set in MCP HTTP/SSE handlers when the path
# identifies exactly one upstream server. Never populated from client-supplied headers.
_mcp_gateway_server_name: Final[ContextVar[str | None]] = ContextVar("_mcp_gateway_server_name", default=None)
# Set server-side by the /mcp/proxy route. Never populated from client-supplied headers.
_mcp_proxy_mode: Final[ContextVar[bool]] = ContextVar("_mcp_proxy_mode", default=False)

82
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -3,12 +3,15 @@ import importlib
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from traceback import walk_tb
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal
from uuid import uuid4
import anyio
import httpx
from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
from pydantic import ValidationError
from starlette.datastructures import Headers
from litellm._logging import verbose_logger
@ -30,6 +33,8 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
list_fault_http_status,
outcome_wire_value,
)
from litellm.proxy._experimental.mcp_server.faults.traversal import iter_exception_tree
from litellm.proxy._experimental.mcp_server.oauth_utils import _redact_mcp_resource_url
from litellm.proxy._experimental.mcp_server.ui_session_utils import (
acting_user_auth,
build_effective_auth_contexts,
@ -78,11 +83,39 @@ _MCP_GUARDRAIL_REJECTIONS: Final = (
def _connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str:
reference: Final = uuid4().hex
verbose_logger.error(
"MCP connection test failed (reference=%s): %s",
reference,
tuple(
(
type(cause).__name__,
tuple(
(frame.f_code.co_filename, lineno, frame.f_code.co_name)
for frame, lineno in walk_tb(cause.__traceback__)
),
)
for cause in iter_exception_tree(exc)
),
)
return next(
(
message
for cause in iter_exception_tree(exc)
if (message := _known_connection_error_message(cause, url, timeout_seconds)) is not None
),
"An unexpected error occurred while testing the MCP connection. "
f"Retry; if it persists, share reference {reference} with your gateway administrator.",
)
def _known_connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str | None:
if isinstance(exc, MCPServerURLCredentialsError):
return str(exc.detail)
if isinstance(exc, TimeoutError):
return (
f"Failed to connect to MCP server: no response from {url or 'the server'} "
"Failed to connect to MCP server: no valid MCP response received from "
f"{_redact_mcp_resource_url(url) or 'the server'} "
f"within {timeout_seconds:.0f}s. Check that the LiteLLM proxy can reach this URL "
"from its network (DNS, egress rules, firewalls) and that the server answers MCP requests."
)
@ -99,13 +132,45 @@ def _connection_error_message(exc: BaseException, url: str | None, timeout_secon
return "Failed to connect to MCP server: the connection timed out."
if isinstance(exc, httpx.HTTPStatusError):
return f"Failed to connect to MCP server: it returned HTTP {exc.response.status_code}."
return "Failed to connect to MCP server. Check proxy logs for details."
if isinstance(exc, (httpx.NetworkError, httpx.RemoteProtocolError, ConnectionError)):
return (
"Failed to connect to MCP server: the connection was interrupted. "
"Check the server and network connection, then retry."
)
if isinstance(exc, ValueError) and str(exc).startswith("Unexpected content type:"):
return (
"Failed to connect to MCP server: the endpoint returned an unsupported content type. "
"Check that the URL is an MCP endpoint, not a web page, and matches the selected transport."
)
if isinstance(exc, ValidationError) and exc.title in ("JSONRPCMessage", "InitializeResult", "ListToolsResult"):
return (
"Failed to connect to MCP server: the endpoint returned invalid JSON or an invalid MCP response. "
"Check the MCP endpoint URL and the server's protocol implementation."
)
if MCP_AVAILABLE and isinstance(exc, McpError):
if exc.error.code == -32000 and exc.error.message == "Connection closed":
return (
"Failed to connect to MCP server: the connection was closed before the request completed. "
"Check that the server stays running and returns a complete MCP response, then retry."
)
if exc.error.code == 32600 and exc.error.message == "Session terminated":
return (
"Failed to connect to MCP server: the MCP session was terminated. "
"Check that the URL points to an MCP endpoint and matches the selected transport, "
"then retry to start a new session."
)
return (
f"Failed to connect to MCP server: the MCP request failed (JSON-RPC code {exc.error.code}). "
"Check that the endpoint supports MCP initialization and tool listing, and check the upstream server logs."
)
return None
if MCP_AVAILABLE:
from mcp.shared.exceptions import McpError
from mcp.types import Tool as MCPTool
from litellm.experimental_mcp_client.client import MCPClient
from litellm.experimental_mcp_client.client import MCPClient, as_mcp_read_timeout
from litellm.llms.litellm_proxy.skills.skill_search import (
DEFAULT_SKILL_SEARCH_TOP_K,
)
@ -1342,11 +1407,18 @@ if MCP_AVAILABLE:
except (KeyboardInterrupt, SystemExit, asyncio.CancelledError):
raise
except BaseException as e:
verbose_logger.error("Error in MCP operation: %s", e, exc_info=True)
effective_timeout: Final = (
min(request.timeout if request.timeout is not None else MCP_CLIENT_TIMEOUT, timeout_seconds)
if any(
isinstance(cause, McpError) and as_mcp_read_timeout(cause) is not None
for cause in iter_exception_tree(e)
)
else timeout_seconds
)
return {
"status": "error",
"error": True,
"message": _connection_error_message(e, request.url, timeout_seconds),
"message": _connection_error_message(e, request.url, effective_timeout),
}
async def _preview_openapi_tools(spec_path: str) -> dict:

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

@ -15,7 +15,7 @@ import types
import uuid
from collections.abc import AsyncIterator, Callable, Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, Protocol
from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol
import httpx
from fastapi import FastAPI, HTTPException
@ -47,6 +47,7 @@ from litellm.proxy._experimental.mcp_server.mcp_context import (
_mcp_active_toolset_id,
_mcp_gateway_initialize_instructions,
_mcp_gateway_server_name,
_mcp_proxy_mode, # pyright: ignore[reportPrivateUsage] # server-owned request mode
)
from litellm.proxy._experimental.mcp_server.mcp_debug import MCPDebug
from litellm.proxy._experimental.mcp_server.oauth_utils import (
@ -537,11 +538,22 @@ if MCP_AVAILABLE:
notification_options: NotificationOptions | None = None,
experimental_capabilities: dict[str, dict[str, object]] | None = None,
) -> InitializationOptions:
opts: Final = Server.create_initialization_options(
base_options: Final = Server.create_initialization_options(
self,
notification_options=notification_options,
experimental_capabilities=experimental_capabilities or {},
)
opts: Final = (
base_options.model_copy(
update={ # mutable-ok: Pydantic update payload
"capabilities": base_options.capabilities.model_copy(
update={"prompts": None, "resources": None} # mutable-ok: Pydantic update payload
)
}
)
if _mcp_proxy_mode.get()
else base_options
)
updates: Final[dict[str, str]] = {}
merged: Final = _mcp_gateway_initialize_instructions.get()
if merged is not None:
@ -755,12 +767,12 @@ if MCP_AVAILABLE:
_stateful_auth_context_cleanup_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await _stateful_auth_context_cleanup_task
if _session_manager_cm:
await _session_manager_cm.__aexit__(None, None, None)
if _session_manager_stateful_cm:
await _session_manager_stateful_cm.__aexit__(None, None, None)
if _sse_session_manager_cm:
await _sse_session_manager_cm.__aexit__(None, None, None)
if _session_manager_stateful_cm:
await _session_manager_stateful_cm.__aexit__(None, None, None)
if _session_manager_cm:
await _session_manager_cm.__aexit__(None, None, None)
except Exception as e:
verbose_logger.exception("Error during session manager shutdown: %s", e)
@ -822,17 +834,20 @@ if MCP_AVAILABLE:
"MCP list_tools - MCP server auth headers: %s",
list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None,
)
from mcp.types import Tool
from litellm.proxy._experimental.mcp_server.tool_search import (
get_mcp_proxy_tool_definitions,
get_virtual_tool_definitions,
)
if _mcp_proxy_mode.get():
return [Tool.model_validate(d) for d in get_mcp_proxy_tool_definitions()] # mutable-ok: MCP SDK list
if getattr(
getattr(user_api_key_auth, "object_permission", None),
"mcp_tool_search_enabled",
False,
):
from mcp.types import Tool
from litellm.proxy._experimental.mcp_server.tool_search import (
get_virtual_tool_definitions,
)
return [Tool.model_validate(d) for d in get_virtual_tool_definitions()]
# Get mcp_servers from context variable
@ -906,6 +921,12 @@ if MCP_AVAILABLE:
verbose_logger.debug("Host progressToken captured: %s...", str(host_token)[:8])
return forward_progress
def _reject_mcp_proxy_operation() -> NoReturn:
from mcp.shared.exceptions import McpError
from mcp.types import METHOD_NOT_FOUND, ErrorData
raise McpError(ErrorData(code=METHOD_NOT_FOUND, message="Operation unavailable on /mcp/proxy"))
async def _build_virtual_call_logging_obj(
name: str,
arguments: dict[str, object],
@ -961,16 +982,91 @@ if MCP_AVAILABLE:
from litellm.proxy._experimental.mcp_server.tool_search import (
AGENT_SEARCH_TOOL_NAME,
DEFAULT_AGENT_SEARCH_TOP_K,
MCP_PROXY_CALL_TOOL_NAME,
MCP_PROXY_TOOL_NAMES,
MCP_TOOL_SEARCH_TOOL_NAME,
SKILL_SEARCH_TOOL_NAME,
VIRTUAL_TOOL_NAMES,
coerce_top_k,
handle_agent_search,
handle_mcp_proxy_tool,
handle_mcp_tool_call,
handle_mcp_tool_search,
handle_skill_search,
)
if _mcp_proxy_mode.get() and name not in MCP_PROXY_TOOL_NAMES:
return CallToolResult(
content=[ # mutable-ok: MCP result content
TextContent(type="text", text=f"Tool {name} is unavailable on /mcp/proxy")
],
isError=True,
)
if _mcp_proxy_mode.get() and name in MCP_PROXY_TOOL_NAMES:
assert user_api_key_auth is not None
proxy_call_start: Final = datetime.now() # noqa: DTZ005 # logging pipeline uses naive datetimes
proxy_logging_obj: Final = (
await _build_virtual_call_logging_obj(
name=name,
arguments=arguments or {}, # mutable-ok: logging pipeline payload
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
client_ip=client_ip,
)
if name == MCP_PROXY_CALL_TOOL_NAME
else None
)
try:
proxy_result: Final = await handle_mcp_proxy_tool(
name=name,
arguments=arguments or {}, # mutable-ok: proxy handler payload
user_api_key_dict=user_api_key_auth,
client_ip=client_ip,
mcp_servers=mcp_servers,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
litellm_logging_obj=proxy_logging_obj,
)
except Exception as exc:
if proxy_logging_obj is not None:
from litellm.proxy.proxy_server import proxy_logging_obj as request_logging_obj
failure_end: Final = datetime.now() # noqa: DTZ005 # matches the logging pipeline start time
failure_traceback: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG)
try:
proxy_logging_obj.failure_handler(exc, failure_traceback, proxy_call_start, failure_end)
await proxy_logging_obj.async_failure_handler(
exc, failure_traceback, proxy_call_start, failure_end
)
if not isinstance(exc, MCPUpstreamAuthError):
await request_logging_obj.post_call_failure_hook(
request_data={ # mutable-ok: failure hook mutates its request payload
"name": name,
"arguments": arguments,
"litellm_logging_obj": proxy_logging_obj,
},
original_exception=exc,
user_api_key_dict=user_api_key_auth,
route="/mcp/call_tool",
traceback_str=failure_traceback,
)
except Exception: # noqa: BLE001 # a failing failure hook must not mask the tool call's own error
verbose_logger.exception("Error logging failed MCP proxy tool call")
raise
if proxy_logging_obj is not None:
return await _fire_mcp_tool_call_logging(
logging_obj=proxy_logging_obj,
result=proxy_result,
start_time=proxy_call_start,
end_time=datetime.now(), # noqa: DTZ005 # matches the logging pipeline start time
user_api_key_auth=user_api_key_auth,
request_data=types.MappingProxyType({"name": name, "arguments": arguments}),
)
return proxy_result
if name not in VIRTUAL_TOOL_NAMES:
return None
@ -1216,6 +1312,8 @@ if MCP_AVAILABLE:
"""
List all available prompts
"""
if _mcp_proxy_mode.get():
_reject_mcp_proxy_operation()
from mcp.server.lowlevel.server import request_ctx
req_ctx: Final = request_ctx.get(None)
@ -1273,8 +1371,8 @@ if MCP_AVAILABLE:
Returns:
GetPromptResult: Getting prompt execution results
"""
# Validate arguments
if _mcp_proxy_mode.get():
_reject_mcp_proxy_operation()
from mcp.server.lowlevel.server import request_ctx
req_ctx: Final = request_ctx.get(None)
@ -1311,6 +1409,8 @@ if MCP_AVAILABLE:
@server.list_resources()
async def list_resources() -> list[Resource]:
"""List all available resources."""
if _mcp_proxy_mode.get():
_reject_mcp_proxy_operation()
from mcp.server.lowlevel.server import request_ctx
req_ctx: Final = request_ctx.get(None)
@ -1355,6 +1455,8 @@ if MCP_AVAILABLE:
@server.list_resource_templates()
async def list_resource_templates() -> list[ResourceTemplate]:
"""List all available resource templates."""
if _mcp_proxy_mode.get():
_reject_mcp_proxy_operation()
from mcp.server.lowlevel.server import request_ctx
req_ctx: Final = request_ctx.get(None)
@ -1400,6 +1502,8 @@ if MCP_AVAILABLE:
@server.read_resource()
async def read_resource(url: AnyUrl) -> list[ReadResourceContents]:
if _mcp_proxy_mode.get():
_reject_mcp_proxy_operation()
from mcp.server.lowlevel.server import request_ctx
req_ctx: Final = request_ctx.get(None)
@ -1998,6 +2102,7 @@ if MCP_AVAILABLE:
litellm_trace_id: str | None = None,
request_tags: list[str] | None = None,
client_ip: str | None = None,
mcp_proxy_mode: bool = False,
) -> AggregateToolListing:
"""
Helper method to fetch tools from MCP servers based on server filtering criteria.
@ -2177,9 +2282,14 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_auth,
)
# Apply display-name/description overrides last so that
# permission filtering always works against original names.
filtered_tools = apply_tool_overrides(filtered_tools, server)
if mcp_proxy_mode:
from litellm.proxy._experimental.mcp_server.tool_search import with_mcp_proxy_identity
filtered_tools = [ # mutable-ok: MCP tool pipeline
with_mcp_proxy_identity(tool, server.server_id) for tool in filtered_tools
]
else:
filtered_tools = apply_tool_overrides(filtered_tools, server)
verbose_logger.debug(
"Successfully fetched %s tools from server %s, %s after filtering",
@ -2491,6 +2601,7 @@ if MCP_AVAILABLE:
log_list_tools_to_spendlogs: bool = False,
list_tools_log_source: str | None = None,
client_ip: str | None = None,
mcp_proxy_mode: bool = False,
) -> AggregateToolListing:
"""
List all available MCP tools.
@ -2520,6 +2631,7 @@ if MCP_AVAILABLE:
log_list_tools_to_spendlogs=log_list_tools_to_spendlogs,
list_tools_log_source=list_tools_log_source,
client_ip=client_ip,
mcp_proxy_mode=mcp_proxy_mode,
)
verbose_logger.debug("Successfully fetched %s tools from managed MCP servers", len(listing.tools))
return listing
@ -3419,7 +3531,9 @@ if MCP_AVAILABLE:
server_name: str | None,
session_id: str | None = None,
) -> StandardLoggingMCPToolCall:
mcp_server: Final = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
mcp_server: Final = global_mcp_server_manager._get_mcp_server_from_tool_name(
add_server_prefix_to_name(name, server_name) if server_name else name
)
namespaced_tool_name: Final = f"{server_name}/{name}" if server_name else name
if mcp_server:
mcp_info: Final = mcp_server.mcp_info or {}

265
litellm/proxy/_experimental/mcp_server/tool_search.py
View file

@ -1,5 +1,6 @@
from __future__ import annotations
import hashlib
import json
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
@ -30,6 +31,12 @@ if TYPE_CHECKING:
MCP_TOOL_SEARCH_SETTINGS_KEY: Final[str] = "mcp_tool_search"
MCP_TOOL_SEARCH_TOOL_NAME: Final[str] = "mcp_tool_search"
MCP_TOOL_CALL_TOOL_NAME: Final[str] = "mcp_tool_call"
MCP_PROXY_SEARCH_TOOL_NAME: Final[str] = "search_tools"
MCP_PROXY_SCHEMA_TOOL_NAME: Final[str] = "get_tool_schema"
MCP_PROXY_CALL_TOOL_NAME: Final[str] = "call_tool"
MCP_PROXY_TOOL_NAMES: Final = frozenset(
(MCP_PROXY_SEARCH_TOOL_NAME, MCP_PROXY_SCHEMA_TOOL_NAME, MCP_PROXY_CALL_TOOL_NAME)
)
AGENT_SEARCH_TOOL_NAME: Final[str] = "agent_search"
SKILL_SEARCH_TOOL_NAME: Final[str] = "skill_search"
VIRTUAL_TOOL_NAMES: Final = frozenset(
@ -51,6 +58,29 @@ class ToolSearchResult(TypedDict, total=False):
score: ReadOnly[float]
class MCPProxySearchResult(TypedDict, total=False):
tool_id: Required[ReadOnly[str]]
name: Required[ReadOnly[str]]
description: Required[ReadOnly[str]]
score: ReadOnly[float]
class MCPProxySchemaResult(MCPProxySearchResult, total=False):
inputSchema: Required[ReadOnly[Mapping[str, object]]]
outputSchema: ReadOnly[Mapping[str, object]]
class MCPProxyToolIdentity(TypedDict):
server_id: ReadOnly[str]
tool_name: ReadOnly[str]
@dataclass(frozen=True, slots=True)
class MCPToolSearchHit:
tool: Tool
score: float | None = None
@dataclass(frozen=True, slots=True)
class SemanticToolRanker:
embed: Embedder
@ -76,6 +106,55 @@ def _scored_result(tool: Tool, score: float) -> ToolSearchResult:
return {"name": tool.name, "description": tool.description or "", "inputSchema": tool.inputSchema, "score": score}
_MCP_PROXY_IDENTITY_META_KEY: Final[str] = "litellm.ai/proxy_tool_identity"
def with_mcp_proxy_identity(tool: Tool, server_id: str) -> Tool:
identity: Final[MCPProxyToolIdentity] = {"server_id": server_id, "tool_name": tool.name}
return tool.model_copy( # mutable-ok: Pydantic requires mutable update and metadata mappings
update={ # mutable-ok: Pydantic update payload
"meta": {**(tool.meta or {}), _MCP_PROXY_IDENTITY_META_KEY: identity} # mutable-ok: metadata mapping
}
)
def _mcp_proxy_identity(tool: Tool) -> MCPProxyToolIdentity:
identity: Final = (tool.meta or {}).get(_MCP_PROXY_IDENTITY_META_KEY) # mutable-ok: absent metadata default
if not isinstance(identity, Mapping):
raise TypeError("MCP proxy tool identity is missing")
server_id: Final = identity.get("server_id")
tool_name: Final = identity.get("tool_name")
if not isinstance(server_id, str) or not isinstance(tool_name, str):
raise TypeError("MCP proxy tool identity is invalid")
return {"server_id": server_id, "tool_name": tool_name} # mutable-ok: TypedDict identity payload
def mcp_proxy_tool_id(tool: Tool) -> str:
identity: Final = _mcp_proxy_identity(tool)
return hashlib.sha256(f"{identity['server_id']}\0{identity['tool_name']}".encode()).hexdigest()[:32]
def _proxy_search_result(hit: MCPToolSearchHit) -> MCPProxySearchResult:
base: Final[MCPProxySearchResult] = {
"tool_id": mcp_proxy_tool_id(hit.tool),
"name": hit.tool.name,
"description": hit.tool.description or "",
}
return {**base, "score": hit.score} if hit.score is not None else base # mutable-ok: wire result payload
def _proxy_schema_result(tool: Tool) -> MCPProxySchemaResult:
base: Final[MCPProxySchemaResult] = {
"tool_id": mcp_proxy_tool_id(tool),
"name": tool.name,
"description": tool.description or "",
"inputSchema": tool.inputSchema,
}
if tool.outputSchema is None:
return base
return {**base, "outputSchema": tool.outputSchema} # mutable-ok: wire schema payload
def _tool_text(tool: Tool) -> str:
return "\n".join(part for part in (tool.name, tool.description or "") if part)
@ -107,6 +186,38 @@ def search_tools(query: str, tools: Sequence[Tool], top_k: int = 5) -> tuple[Too
return tuple(_tool_result(tool) for _, tool in _top_hits(tools, scores, minimum=1.0, limit=top_k))
async def rank_mcp_tools(
query: str,
tools: Sequence[Tool],
top_k: int,
settings: MCPToolSearchSettings,
ranker: SemanticToolRanker | None,
) -> tuple[MCPToolSearchHit, ...] | EmbeddingFailed:
core, rest = _split_core_tools(tools, settings.core_tools)
core_hits: Final = tuple(MCPToolSearchHit(tool) for tool in core)
if not query:
return core_hits
limit: Final = min(top_k, settings.top_k)
if ranker is None:
scores: Final = tuple(_keyword_score(query, tool) for tool in rest)
return (
*core_hits,
*(MCPToolSearchHit(tool) for _, tool in _top_hits(rest, scores, minimum=1.0, limit=limit)),
)
semantic_scores: Final = await ranker.index.scores(
query, tuple(_tool_text(tool) for tool in rest), ranker.embed, ranker.embedding_model
)
if isinstance(semantic_scores, EmbeddingFailed):
return semantic_scores
return (
*core_hits,
*(
MCPToolSearchHit(tool, score)
for score, tool in _top_hits(rest, semantic_scores, settings.similarity_threshold, limit)
),
)
async def search_mcp_tools(
query: str,
tools: Sequence[Tool],
@ -114,21 +225,12 @@ async def search_mcp_tools(
settings: MCPToolSearchSettings,
ranker: SemanticToolRanker | None,
) -> tuple[ToolSearchResult, ...] | EmbeddingFailed:
"""Core tools the caller can access come first, then up to `top_k` ranked matches from the remaining tools."""
core, rest = _split_core_tools(tools, settings.core_tools)
limit: Final = min(top_k, settings.top_k)
core_results: Final = tuple(_tool_result(tool) for tool in core)
if ranker is None:
return (*core_results, *search_tools(query, rest, limit))
if not query:
return core_results
scores: Final = await ranker.index.scores(
query, tuple(_tool_text(tool) for tool in rest), ranker.embed, ranker.embedding_model
hits: Final = await rank_mcp_tools(query, tools, top_k, settings, ranker)
if isinstance(hits, EmbeddingFailed):
return hits
return tuple(
_scored_result(hit.tool, hit.score) if hit.score is not None else _tool_result(hit.tool) for hit in hits
)
if isinstance(scores, EmbeddingFailed):
return scores
hits: Final = _top_hits(rest, scores, minimum=settings.similarity_threshold, limit=limit)
return (*core_results, *(_scored_result(tool, score) for score, tool in hits))
class _ToolParamSchema(TypedDict, total=False):
@ -223,10 +325,48 @@ _SKILL_SEARCH_DEFINITION: Final[VirtualToolDefinition] = {
}
_MCP_PROXY_SEARCH_DEFINITION: Final[VirtualToolDefinition] = {
"name": MCP_PROXY_SEARCH_TOOL_NAME,
"description": "Search accessible MCP tools by describing what you need. Returns opaque tool IDs.",
"inputSchema": {
"type": "object",
"properties": {"query": {"type": "string", "description": "What the tool should do."}},
"required": _json_array("query"),
},
}
_MCP_PROXY_SCHEMA_DEFINITION: Final[VirtualToolDefinition] = {
"name": MCP_PROXY_SCHEMA_TOOL_NAME,
"description": "Return the complete schema for an accessible MCP tool ID.",
"inputSchema": {
"type": "object",
"properties": {"tool_id": {"type": "string", "description": "Opaque ID from search_tools."}},
"required": _json_array("tool_id"),
},
}
_MCP_PROXY_CALL_DEFINITION: Final[VirtualToolDefinition] = {
"name": MCP_PROXY_CALL_TOOL_NAME,
"description": "Call an accessible MCP tool by opaque ID with schema-valid arguments.",
"inputSchema": {
"type": "object",
"properties": {
"tool_id": {"type": "string", "description": "Opaque ID from search_tools."},
"arguments": {"type": "object", "description": "Arguments validated against the selected tool schema."},
},
"required": _json_array("tool_id"),
},
}
def get_virtual_tool_definitions() -> tuple[VirtualToolDefinition, ...]:
return (_MCP_TOOL_SEARCH_DEFINITION, _MCP_TOOL_CALL_DEFINITION, _AGENT_SEARCH_DEFINITION, _SKILL_SEARCH_DEFINITION)
def get_mcp_proxy_tool_definitions() -> tuple[VirtualToolDefinition, ...]:
return (_MCP_PROXY_SEARCH_DEFINITION, _MCP_PROXY_SCHEMA_DEFINITION, _MCP_PROXY_CALL_DEFINITION)
def _text_tool_result(text: str, is_error: bool) -> CallToolResult:
from mcp.types import CallToolResult, TextContent
@ -314,7 +454,9 @@ async def handle_mcp_tool_search(
oauth2_headers: dict[str, str] | None = None,
raw_headers: dict[str, str] | None = None,
) -> CallToolResult:
from litellm.proxy._experimental.mcp_server.server import _list_mcp_tools
from litellm.proxy._experimental.mcp_server.server import (
_list_mcp_tools, # pyright: ignore[reportPrivateUsage] # shared catalog owner
)
from litellm.proxy.proxy_server import llm_router, proxy_logging_obj
settings: Final = mcp_tool_search_settings()
@ -351,6 +493,97 @@ async def handle_mcp_tool_search(
return _text_tool_result(json.dumps(results), is_error=False)
async def handle_mcp_proxy_tool(
name: str,
arguments: dict[str, object], # mutable-ok: MCP dispatcher passes mutable call arguments
user_api_key_dict: UserAPIKeyAuth,
client_ip: str | None = None,
mcp_servers: list[str] | None = None, # mutable-ok: preserve MCP scope container for existing resolver
mcp_auth_header: str | None = None,
mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, # mutable-ok: preserve forwarded headers
oauth2_headers: dict[str, str] | None = None, # mutable-ok: preserve forwarded headers
raw_headers: dict[str, str] | None = None, # mutable-ok: preserve request headers
litellm_logging_obj: LiteLLMLoggingObj | None = None,
) -> CallToolResult:
from fastapi import HTTPException
from jsonschema import ValidationError as JsonSchemaValidationError
from jsonschema import validate
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server.server import ( # pyright: ignore[reportPrivateUsage] # shared catalog owner
_list_mcp_tools, # pyright: ignore[reportPrivateUsage] # shared catalog owner
)
listing: Final = await _list_mcp_tools(
user_api_key_auth=user_api_key_dict,
mcp_servers=mcp_servers,
client_ip=client_ip,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
mcp_proxy_mode=True,
)
tools_by_id: Final = {mcp_proxy_tool_id(tool): tool for tool in listing.tools} # mutable-ok: lookup index
if name == MCP_PROXY_SEARCH_TOOL_NAME:
llm_router: Final = proxy_server.llm_router
proxy_logging_obj: Final = proxy_server.proxy_logging_obj
settings: Final = mcp_tool_search_settings()
if isinstance(settings, ValidationError):
return _text_tool_result(str(settings), is_error=True)
if settings.embedding_model is not None and llm_router is None:
return _text_tool_result(
f"litellm_settings.{MCP_TOOL_SEARCH_SETTINGS_KEY}.embedding_model needs a model_list so it can be called",
is_error=True,
)
ranker: Final = (
SemanticToolRanker(
embed=router_embedder(llm_router, settings.embedding_model, user_api_key_dict, proxy_logging_obj),
embedding_model=settings.embedding_model,
index=global_mcp_tool_search_index,
)
if settings.embedding_model is not None and llm_router is not None
else None
)
results: Final = await rank_mcp_tools(str(arguments.get("query", "")), listing.tools, 5, settings, ranker)
if isinstance(results, EmbeddingFailed):
return _text_tool_result(results.reason, is_error=True)
return _text_tool_result(json.dumps(tuple(_proxy_search_result(hit) for hit in results)), is_error=False)
tool_id: Final = arguments.get("tool_id")
tool: Final = tools_by_id.get(tool_id) if isinstance(tool_id, str) else None
if tool is None:
return _text_tool_result("Unknown or unauthorized tool_id", is_error=True)
if name == MCP_PROXY_SCHEMA_TOOL_NAME:
return _text_tool_result(json.dumps(_proxy_schema_result(tool)), is_error=False)
if name != MCP_PROXY_CALL_TOOL_NAME:
raise HTTPException(status_code=400, detail=f"Unknown MCP proxy tool: {name}")
tool_arguments: Final = arguments.get("arguments", {}) # mutable-ok: JSON Schema validator consumes mapping
if not isinstance(tool_arguments, dict):
return _text_tool_result("arguments must be an object", is_error=True)
try:
validate(instance=tool_arguments, schema=tool.inputSchema)
except JsonSchemaValidationError as exc:
return _text_tool_result(f"Invalid arguments: {exc.message}", is_error=True)
return await handle_mcp_tool_call(
tool_name=_mcp_proxy_identity(tool)["tool_name"],
arguments=tool_arguments,
user_api_key_dict=user_api_key_dict,
requested_server_id=_mcp_proxy_identity(tool)["server_id"],
client_ip=client_ip,
mcp_servers=mcp_servers,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
litellm_logging_obj=litellm_logging_obj,
)
async def handle_mcp_tool_call(
tool_name: str,
arguments: dict[str, Any],
@ -362,6 +595,7 @@ async def handle_mcp_tool_call(
oauth2_headers: dict[str, str] | None = None,
raw_headers: dict[str, str] | None = None,
litellm_logging_obj: LiteLLMLoggingObj | None = None,
requested_server_id: str | None = None,
) -> CallToolResult:
from litellm.proxy._experimental.mcp_server.server import (
_get_allowed_mcp_servers,
@ -400,4 +634,5 @@ async def handle_mcp_tool_call(
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
litellm_logging_obj=litellm_logging_obj,
requested_server_id=requested_server_id,
)

16
litellm/proxy/_experimental/mcp_server/toolset_db.py
View file

@ -132,10 +132,18 @@ async def update_mcp_toolset(
data: UpdateMCPToolsetRequest,
touched_by: str,
) -> MCPToolset | None:
data_dict: Final = data.model_dump(exclude_none=True, exclude={"toolset_id"})
if "tools" in data_dict:
data_dict["tools"] = json.dumps(data_dict["tools"])
data_dict["updated_by"] = touched_by
"""A partial update: absent keeps, null clears. A toolset always has a name and a
tool list, so a null ``toolset_name`` or ``tools`` is a no-op rather than a clear;
emptying the tool selection is an explicit ``[]``, which cannot be mistaken for a
caller that left the field out."""
data_dict: Final = dict( # mutable-ok: Prisma requires a plain dict for JSON query serialization
(
(field, json.dumps(value) if field == "tools" else value)
for field, value in data.model_dump(exclude_unset=True).items()
if field != "toolset_id" and (field not in ("toolset_name", "tools") or value is not None)
),
updated_by=touched_by,
)
try:
row: Final = await _toolset_table(prisma_client).update(
where={"toolset_id": data.toolset_id},

128
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -17027,6 +17027,134 @@
"mcp_app"
]
}
},
"/mcp/proxy": {
"delete": {
"description": "Serve the fixed three-tool MCP proxy surface.",
"operationId": "proxy_mcp_route_mcp_proxy_delete",
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
}
},
"summary": "Proxy Mcp Route",
"tags": [
"mcp_app"
]
},
"get": {
"description": "Serve the fixed three-tool MCP proxy surface.",
"operationId": "proxy_mcp_route_mcp_proxy_get",
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
}
},
"summary": "Proxy Mcp Route",
"tags": [
"mcp_app"
]
},
"head": {
"description": "Serve the fixed three-tool MCP proxy surface.",
"operationId": "proxy_mcp_route_mcp_proxy_head",
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
}
},
"summary": "Proxy Mcp Route",
"tags": [
"mcp_app"
]
},
"options": {
"description": "Serve the fixed three-tool MCP proxy surface.",
"operationId": "proxy_mcp_route_mcp_proxy_options",
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
}
},
"summary": "Proxy Mcp Route",
"tags": [
"mcp_app"
]
},
"patch": {
"description": "Serve the fixed three-tool MCP proxy surface.",
"operationId": "proxy_mcp_route_mcp_proxy_patch",
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
}
},
"summary": "Proxy Mcp Route",
"tags": [
"mcp_app"
]
},
"post": {
"description": "Serve the fixed three-tool MCP proxy surface.",
"operationId": "proxy_mcp_route_mcp_proxy_post",
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
}
},
"summary": "Proxy Mcp Route",
"tags": [
"mcp_app"
]
},
"put": {
"description": "Serve the fixed three-tool MCP proxy surface.",
"operationId": "proxy_mcp_route_mcp_proxy_put",
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
}
},
"summary": "Proxy Mcp Route",
"tags": [
"mcp_app"
]
}
}
}
},

51
litellm/proxy/_types.py
View file

@ -3,6 +3,7 @@ import json
import os
from collections.abc import Callable, Mapping
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple
import httpx
@ -502,6 +503,7 @@ class LiteLLMRoutes(enum.Enum):
mcp_inference_routes = [
"/mcp",
"/mcp/",
"/mcp/proxy",
"/mcp/{subpath}",
"/mcp/tools",
"/mcp/tools/list",
@ -1293,6 +1295,13 @@ class UpdateKeyRequest(KeyRequestBase):
rotation_interval: str | None = None
organization_id: str | None = None
@model_validator(mode="before")
@classmethod
def drop_blank_team_id(cls, values: object) -> object:
if isinstance(values, Mapping) and values.get("team_id") == "":
return MappingProxyType({k: v for k, v in values.items() if k != "team_id"})
return values
@field_validator("organization_id", mode="before")
@classmethod
def treat_cleared_organization_id_as_unset(cls, v: object) -> object:
@ -5147,9 +5156,26 @@ class CostEstimateRequest(LiteLLMPydanticObjectBase):
model: str = Field(description="Model name (from /model_group/info)")
input_tokens: int = Field(description="Expected input tokens per request", ge=0)
output_tokens: int = Field(description="Expected output tokens per request", ge=0)
cache_read_input_tokens: int = Field(
default=0, description="Input tokens read from the prompt cache; counted within input_tokens", ge=0
)
cache_creation_input_tokens: int = Field(
default=0, description="Input tokens written to the prompt cache; counted within input_tokens", ge=0
)
reasoning_tokens: int = Field(
default=0, description="Reasoning tokens the model emits; counted within output_tokens", ge=0
)
num_requests_per_day: int | None = Field(default=None, description="Number of requests per day", ge=0)
num_requests_per_month: int | None = Field(default=None, description="Number of requests per month", ge=0)
@model_validator(mode="after")
def validate_token_subsets(self) -> "CostEstimateRequest":
if self.cache_read_input_tokens + self.cache_creation_input_tokens > self.input_tokens:
raise ValueError("cache_read_input_tokens plus cache_creation_input_tokens cannot exceed input_tokens")
if self.reasoning_tokens > self.output_tokens:
raise ValueError("reasoning_tokens cannot exceed output_tokens")
return self
class CostEstimateResponse(LiteLLMPydanticObjectBase):
"""Response body for /cost/estimate endpoint."""
@ -5157,6 +5183,9 @@ class CostEstimateResponse(LiteLLMPydanticObjectBase):
model: str
input_tokens: int
output_tokens: int
cache_read_input_tokens: int = 0
cache_creation_input_tokens: int = 0
reasoning_tokens: int = 0
num_requests_per_day: int | None = None
num_requests_per_month: int | None = None
# Per-request costs
@ -5164,17 +5193,33 @@ class CostEstimateResponse(LiteLLMPydanticObjectBase):
input_cost_per_request: float = Field(description="Input token cost per request (before margin)")
output_cost_per_request: float = Field(description="Output token cost per request (before margin)")
margin_cost_per_request: float = Field(default=0.0, description="Margin/fee added per request")
cache_read_cost_per_request: float = Field(default=0.0, description="Cache-read share of input_cost_per_request")
cache_creation_cost_per_request: float = Field(
default=0.0, description="Cache-write share of input_cost_per_request"
)
reasoning_cost_per_request: float = Field(default=0.0, description="Reasoning share of output_cost_per_request")
# Daily costs (if num_requests_per_day provided)
daily_cost: float | None = Field(default=None, description="Total daily cost (includes margin)")
daily_input_cost: float | None = Field(default=None, description="Daily input token cost")
daily_output_cost: float | None = Field(default=None, description="Daily output token cost")
daily_margin_cost: float | None = Field(default=None, description="Daily margin/fee")
daily_cache_read_cost: float | None = Field(default=None, description="Cache-read share of daily_input_cost")
daily_cache_creation_cost: float | None = Field(default=None, description="Cache-write share of daily_input_cost")
daily_reasoning_cost: float | None = Field(default=None, description="Reasoning share of daily_output_cost")
# Monthly costs (if num_requests_per_month provided)
monthly_cost: float | None = Field(default=None, description="Total monthly cost (includes margin)")
monthly_input_cost: float | None = Field(default=None, description="Monthly input token cost")
monthly_output_cost: float | None = Field(default=None, description="Monthly output token cost")
monthly_margin_cost: float | None = Field(default=None, description="Monthly margin/fee")
# Pricing info
input_cost_per_token: float | None = None
output_cost_per_token: float | None = None
monthly_cache_read_cost: float | None = Field(default=None, description="Cache-read share of monthly_input_cost")
monthly_cache_creation_cost: float | None = Field(
default=None, description="Cache-write share of monthly_input_cost"
)
monthly_reasoning_cost: float | None = Field(default=None, description="Reasoning share of monthly_output_cost")
# Pricing info: the rates this request's usage bills at, after token tiers and regional multipliers
input_cost_per_token: float | None = Field(default=None, description="Rate billed per input token")
output_cost_per_token: float | None = Field(default=None, description="Rate billed per output token")
cache_read_input_token_cost: float | None = Field(default=None, description="Rate billed per cache-read token")
cache_creation_input_token_cost: float | None = Field(default=None, description="Rate billed per cache-write token")
output_cost_per_reasoning_token: float | None = Field(default=None, description="Rate billed per reasoning token")
provider: str | None = None

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

@ -475,6 +475,7 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None
_NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({})
_TEAM_GRANT_RELATIONS: Final[Mapping[str, object]] = MappingProxyType({"litellm_model_table": True})
def _has_ptu_flat_cost(model: str, llm_router: "Router") -> bool:
@ -2858,7 +2859,9 @@ class TeamNotFoundError(HTTPException):
async def _get_team_db_check(
team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None
) -> "_PrismaTeamRow | None":
response = await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id})
response = await _team_table(TeamRepository(prisma_client)).find_unique(
where={"team_id": team_id}, include=_TEAM_GRANT_RELATIONS
)
if response is None and team_id_upsert:
from litellm.proxy.management_endpoints.team_endpoints import new_team
@ -3158,7 +3161,9 @@ async def get_team_object_by_alias(
# Query database by team_alias
try:
teams: Final = await _team_table(TeamRepository(prisma_client)).find_many(where={"team_alias": team_alias})
teams: Final = await _team_table(TeamRepository(prisma_client)).find_many(
where={"team_alias": team_alias}, include=_TEAM_GRANT_RELATIONS
)
if not teams:
raise HTTPException(

5
litellm/proxy/auth/handle_jwt.py
View file

@ -53,6 +53,7 @@ from litellm.proxy._types import (
)
from litellm.proxy.auth.auth_checks import can_team_access_model
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.auth.team_grants import team_model_aliases
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
get_management_object_ttl,
@ -1595,7 +1596,7 @@ class JWTAuthManager:
model=requested_model,
team_object=team_object,
llm_router=llm_router,
team_model_aliases=None,
team_model_aliases=team_model_aliases(team_object),
)
):
is_allowed = allowed_routes_check(
@ -2132,7 +2133,7 @@ class JWTAuthManager:
model=requested_model,
team_object=team_object,
llm_router=llm_router,
team_model_aliases=None,
team_model_aliases=team_model_aliases(team_object),
)
except ProxyException:
continue

122
litellm/proxy/auth/team_grants.py Normal file
View file

@ -0,0 +1,122 @@
"""Project a team row (plus the caller's membership in it) onto the ``team_*`` fields of ``UserAPIKeyAuth``.
The virtual-key path gets these fields for free from the combined-view SQL join. Every other auth path
starts from a ``LiteLLM_TeamTable`` object instead and has to copy them over by hand, which is how JWT
callers kept losing grants (aliases, permissions, limits) one field at a time. Build the badge through
``team_grants`` and the two paths cannot drift.
"""
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Annotated, Final
from pydantic import BaseModel, BeforeValidator, ConfigDict, TypeAdapter, ValidationError
from pydantic.main import IncEx
from typing_extensions import ReadOnly, TypedDict
from litellm.proxy._types import (
LiteLLM_ObjectPermissionTable,
LiteLLM_TeamMembership,
LiteLLM_TeamTable,
Member,
)
_MODEL_ALIASES_ADAPTER: Final = TypeAdapter(dict[str, str])
_JSON_COLUMNS: Final[Mapping[str, IncEx | bool]] = MappingProxyType(
{"metadata": True, "litellm_model_table": MappingProxyType({"model_aliases": True})}
)
def _decode_model_aliases(value: object) -> object:
"""``LiteLLM_ModelTable.model_aliases`` is typed ``str | dict``; writers hand Prisma ``json.dumps(...)``, so take both."""
if not isinstance(value, str):
return value
try:
return _MODEL_ALIASES_ADAPTER.validate_json(value)
except ValidationError:
return None
class TeamModelAliasTable(BaseModel):
model_config = ConfigDict(protected_namespaces=())
model_aliases: Annotated[Mapping[str, str] | None, BeforeValidator(_decode_model_aliases)] = None
class _TeamJsonColumns(BaseModel):
"""The two loosely typed columns on ``LiteLLM_TeamTable``, re-read with the shape the badge needs."""
metadata: Mapping[str, object] | None = None
litellm_model_table: TeamModelAliasTable | None = None
class TeamGrants(TypedDict, total=False):
"""Keyword arguments for ``UserAPIKeyAuth``. Empty when the caller has no team, so the model's own defaults apply."""
team_alias: ReadOnly[str | None]
team_tpm_limit: ReadOnly[int | None]
team_rpm_limit: ReadOnly[int | None]
team_max_budget: ReadOnly[float | None]
team_soft_budget: ReadOnly[float | None]
team_spend: ReadOnly[float | None]
team_models: ReadOnly[Sequence[str]]
team_blocked: ReadOnly[bool]
team_metadata: ReadOnly[Mapping[str, object] | None]
team_model_aliases: ReadOnly[Mapping[str, str] | None]
team_object_permission_id: ReadOnly[str | None]
team_object_permission: ReadOnly[LiteLLM_ObjectPermissionTable | None]
team_member: ReadOnly[Member | None]
team_member_spend: ReadOnly[float | None]
team_member_tpm_limit: ReadOnly[int | None]
team_member_rpm_limit: ReadOnly[int | None]
def _json_columns(team_object: LiteLLM_TeamTable) -> _TeamJsonColumns:
try:
return _TeamJsonColumns.model_validate(team_object.model_dump(include=_JSON_COLUMNS))
except ValidationError:
return _TeamJsonColumns()
def team_model_aliases(team_object: LiteLLM_TeamTable | None) -> Mapping[str, str] | None:
if team_object is None:
return None
alias_table: Final = _json_columns(team_object).litellm_model_table
return alias_table.model_aliases if alias_table is not None else None
def team_grants(
team_object: LiteLLM_TeamTable | None,
team_membership: LiteLLM_TeamMembership | None,
user_id: str | None,
) -> TeamGrants:
if team_object is None:
return TeamGrants()
json_columns: Final = _json_columns(team_object)
return TeamGrants(
team_alias=team_object.team_alias,
team_tpm_limit=team_object.tpm_limit,
team_rpm_limit=team_object.rpm_limit,
team_max_budget=team_object.max_budget,
team_soft_budget=team_object.soft_budget,
team_spend=team_object.spend,
team_models=tuple(team_object.models),
team_blocked=team_object.blocked,
team_metadata=json_columns.metadata,
team_model_aliases=(
json_columns.litellm_model_table.model_aliases if json_columns.litellm_model_table is not None else None
),
team_object_permission_id=team_object.object_permission_id,
team_object_permission=team_object.object_permission,
team_member=next(
(m for m in team_object.members_with_roles if user_id is not None and m.user_id == user_id),
None,
),
team_member_spend=team_membership.spend if team_membership is not None else None,
team_member_tpm_limit=(
team_membership.safe_get_team_member_tpm_limit() if team_membership is not None else None
),
team_member_rpm_limit=(
team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None
),
)

24
litellm/proxy/auth/user_api_key_auth.py
View file

@ -82,6 +82,7 @@ from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request
from litellm.proxy.auth.resolvers import CredentialRef, Principal
from litellm.proxy.auth.resolvers.store import IdentityStore
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.auth.team_grants import team_grants
from litellm.proxy.auth.trusted_proxy_utils import get_trusted_proxy_cidrs
from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator
from litellm.proxy.common_utils.http_parsing_utils import (
@ -1476,24 +1477,16 @@ async def _user_api_key_auth_builder(
user_id=user_id,
user_email=user_email,
team_id=team_id,
team_alias=(team_object.team_alias if team_object is not None else None),
team_tpm_limit=(team_object.tpm_limit if team_object is not None else None),
team_rpm_limit=(team_object.rpm_limit if team_object is not None else None),
team_models=(team_object.models if team_object is not None else []),
team_metadata=(team_object.metadata if team_object is not None else None),
org_id=org_id,
end_user_id=end_user_id,
parent_otel_span=parent_otel_span,
jwt_claims=jwt_claims,
**team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id),
)
valid_token = UserAPIKeyAuth(
api_key=None,
team_id=team_id,
team_alias=(team_object.team_alias if team_object is not None else None),
team_tpm_limit=(team_object.tpm_limit if team_object is not None else None),
team_rpm_limit=(team_object.rpm_limit if team_object is not None else None),
team_models=(team_object.models if team_object is not None else []),
user_role=(
LitellmUserRoles(user_object.user_role)
if user_object is not None and user_object.user_role is not None
@ -1507,17 +1500,8 @@ async def _user_api_key_auth_builder(
user_tpm_limit=(user_object.tpm_limit if user_object is not None else None),
user_rpm_limit=(user_object.rpm_limit if user_object is not None else None),
user_model_max_budget=(user_object.model_max_budget if user_object is not None else None),
team_member_rpm_limit=(
team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None
),
team_member_tpm_limit=(
team_membership.safe_get_team_member_tpm_limit() if team_membership is not None else None
),
team_metadata=(team_object.metadata if team_object is not None else None),
jwt_claims=jwt_claims,
)
valid_token.team_object_permission = (
team_object.object_permission if team_object is not None else None
**team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id),
)
# AUTO_REGISTER deferred from _resolve_jwt_to_virtual_key.
@ -2836,7 +2820,6 @@ async def _authorize_authenticated_request(
return None
@tracer.wrap()
def _seed_request_destinations(user_api_key_dict: UserAPIKeyAuth, request: Request | None = None) -> None:
"""Anchor the OTLP destinations this key or team overrides its traces to.
@ -2874,6 +2857,7 @@ def _seed_request_destinations(user_api_key_dict: UserAPIKeyAuth, request: Reque
verbose_proxy_logger.debug("OTel V2: tenant destination resolution failed: %s", exc)
@tracer.wrap()
async def user_api_key_auth(
request: Request,
api_key: str = fastapi.Security(api_key_header),

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

@ -1,10 +1,12 @@
import copy
import json
import os
from collections.abc import Callable, Iterable, Mapping
from dataclasses import dataclass
from itertools import accumulate
from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias
from typing_extensions import assert_never
from typing_extensions import ReadOnly, TypedDict, assert_never
import litellm
from litellm import get_secret
@ -12,6 +14,7 @@ from litellm._logging import verbose_proxy_logger
from litellm.constants import (
CLIENT_OUTPUT_CEILING_METADATA_KEY,
CONSUMED_REQUEST_TAGS_METADATA_KEY,
MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH,
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
ROUTING_REQUEST_TAGS_METADATA_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
@ -28,6 +31,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
encrypt_value_helper,
)
from litellm.proxy.types_utils.utils import get_instance_fn
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import (
StandardLoggingGuardrailInformation,
StandardLoggingPayload,
@ -52,6 +56,15 @@ reset_color_code: Final = "\033[0m"
TRUSTED_PILLAR_RESPONSE_HEADERS_METADATA_KEY: Final = "_pillar_response_headers_trusted"
GUARDRAIL_SCAN_IDS_METADATA_KEY: Final = "guardrail_scan_ids"
GUARDRAIL_SCAN_METADATA_METADATA_KEY: Final = "guardrail_scan_metadata"
class GuardrailScanMetadata(TypedDict):
guardrail: ReadOnly[str | None]
stage: ReadOnly[str]
provider: ReadOnly[str]
scan_id: ReadOnly[str]
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
@ -450,6 +463,16 @@ def get_remaining_tokens_and_requests_from_request_data(data: dict) -> dict[str,
return headers
def _serialize_scan_metadata_header(entries: Iterable[object], *, max_length: int) -> str | None:
"""Compact JSON list of scan metadata entries, dropping trailing entries so the header fits in max_length."""
encoded: Final = tuple(json.dumps(entry, separators=(",", ":")) for entry in entries)
lengths: Final = tuple(accumulate(len(item) + 1 for item in encoded))
kept: Final = sum(1 for length in lengths if length + 1 <= max_length)
if kept == 0:
return None
return f"[{','.join(encoded[:kept])}]"
def get_logging_caching_headers(request_data: dict) -> dict | None:
_metadata: Final[dict] = {}
metadata_bucket: Final = request_data.get("metadata")
@ -468,6 +491,15 @@ def get_logging_caching_headers(request_data: dict) -> dict | None:
if scan_ids:
headers["x-litellm-guardrail-scan-id"] = ",".join(scan_ids)
scan_metadata: Final = _metadata.get(GUARDRAIL_SCAN_METADATA_METADATA_KEY)
scan_metadata_header: Final = (
_serialize_scan_metadata_header(scan_metadata, max_length=MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH)
if isinstance(scan_metadata, (list, tuple))
else None
)
if scan_metadata_header:
headers["x-litellm-guardrail-scan-metadata"] = scan_metadata_header
if "applied_policies" in _metadata:
headers["x-litellm-applied-policies"] = ",".join(_metadata["applied_policies"])
@ -501,6 +533,7 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset(
"applied_policies",
"applied_guardrails",
GUARDRAIL_SCAN_IDS_METADATA_KEY,
GUARDRAIL_SCAN_METADATA_METADATA_KEY,
"policy_sources",
"guardrails",
"guardrail_config",
@ -565,21 +598,40 @@ def add_guardrail_to_applied_guardrails_header(request_data: dict, guardrail_nam
_metadata["applied_guardrails"] = [guardrail_name]
def add_guardrail_scan_id(request_data: dict, scan_id: str | None) -> None:
def add_guardrail_scan_id(
request_data: dict[str, object],
scan_id: str | None,
*,
guardrail_name: str | None,
provider: str,
stage: GuardrailEventHooks,
) -> None:
"""
Record a provider scan id so it can be surfaced to the caller.
Record a provider scan id, keyed to the guardrail execution that produced it, so it can be surfaced to the caller.
Guardrails only return scan details to the client when they block, so allowed requests carry no
audit trail. Ids recorded here become the x-litellm-guardrail-scan-id response header.
audit trail. Ids recorded here become the x-litellm-guardrail-scan-id response header, and the
(guardrail, stage, provider, scan_id) entries become the x-litellm-guardrail-scan-metadata header.
"""
if not scan_id:
return
_, _metadata = get_or_create_metadata_bucket(request_data)
existing: Final = _metadata.get(GUARDRAIL_SCAN_IDS_METADATA_KEY)
scan_ids: Final = tuple(existing) if isinstance(existing, (list, tuple)) else ()
scan_ids: Final[tuple[object, ...]] = tuple(existing) if isinstance(existing, (list, tuple)) else ()
if scan_id not in scan_ids:
_metadata[GUARDRAIL_SCAN_IDS_METADATA_KEY] = (*scan_ids, scan_id)
entry: Final[GuardrailScanMetadata] = {
"guardrail": guardrail_name,
"stage": stage.value,
"provider": provider,
"scan_id": scan_id,
}
existing_entries: Final = _metadata.get(GUARDRAIL_SCAN_METADATA_METADATA_KEY)
entries: Final[tuple[object, ...]] = tuple(existing_entries) if isinstance(existing_entries, (list, tuple)) else ()
if entry not in entries:
_metadata[GUARDRAIL_SCAN_METADATA_METADATA_KEY] = (*entries, entry)
def add_policy_to_applied_policies_header(request_data: dict, policy_name: str | None):
"""

107
litellm/proxy/db/db_url_settings.py
View file

@ -34,12 +34,20 @@ writer's connection params (pool size, timeouts, pgbouncer mode) for the
ones the reader URL does not pin itself.
"""
import _ssl
import hashlib
import os
import socket
import ssl
import struct
import sys
import tempfile
import urllib.parse
from collections.abc import Mapping
from collections.abc import Callable, Mapping, Sequence
from functools import partial
from pathlib import Path
from types import MappingProxyType
from typing import Annotated, Final, cast
from typing import Annotated, Final, Protocol, TypeAlias, cast
from pydantic import AliasChoices, BeforeValidator, Field
from pydantic_settings import BaseSettings, SettingsConfigDict
@ -126,21 +134,100 @@ def add_missing_query_params(url: str, params: Mapping[str, str | int | float])
LIBPQ_VERIFY_SSLMODES: Final[frozenset[str]] = frozenset({"verify-ca", "verify-full"})
PEM_CERT_HEADER: Final = b"-----BEGIN CERTIFICATE-----"
PG_SSL_REQUEST: Final = struct.pack("!ii", 8, 80877103)
TLS_PROBE_TIMEOUT_SECONDS: Final = 10.0
RootCertResolver: TypeAlias = Callable[[str, str, int], str] # mutable-ok: Callable parameter syntax
def translate_libpq_ssl_params(url: str) -> str:
class _VerifiedChainSource(Protocol):
def get_verified_chain(self) -> Sequence[_ssl.Certificate] | None: ...
def _verified_chain_der(tls: ssl.SSLSocket) -> tuple[bytes, ...]:
if sys.version_info >= (3, 13):
return tuple(tls.get_verified_chain())
legacy: Final = cast( # cast-ok: the stub omits _sslobj, the C object has get_verified_chain since 3.10
"_VerifiedChainSource | None",
tls._sslobj, # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # public API only from 3.13
)
chain: Final = () if legacy is None else legacy.get_verified_chain() or ()
return tuple(cert.public_bytes(_ssl.ENCODING_DER) for cert in chain)
def _server_trust_anchor(cafile: str, host: str, port: int) -> bytes | None:
try:
context: Final = ssl.create_default_context(cafile=cafile)
with socket.create_connection((host, port), timeout=TLS_PROBE_TIMEOUT_SECONDS) as raw:
raw.sendall(PG_SSL_REQUEST)
if raw.recv(1) != b"S":
return None
with context.wrap_socket(raw, server_hostname=host) as tls:
chain: Final = _verified_chain_der(tls)
except (OSError, ValueError):
return None
return chain[-1] if chain else None
def pin_bundle_root(cert_path: str, host: str, port: int) -> str:
"""Reduce a multi-root CA bundle to the one root that verifies ``host``.
Prisma's ``sslcert`` loads a single PEM certificate (native-tls
``Certificate::from_pem``), so pointing it at a bundle such as the AWS RDS
global bundle trusts only the first of its 108 regional roots and the
handshake fails with "unable to get local issuer certificate" for every
other region. A single-certificate file is returned as is. For a bundle,
one verifying handshake (chain and hostname, whole bundle as trust store)
identifies the trust anchor the server actually chains to, which is
written to a single-certificate file for Prisma. If the probe fails the
bundle path is returned unchanged, so Prisma fails closed exactly as
before rather than trusting anything the bundle would not.
"""
try:
if Path(cert_path).read_bytes().count(PEM_CERT_HEADER) < 2:
return cert_path
except OSError:
return cert_path
root: Final = _server_trust_anchor(cert_path, host, port)
if root is None:
return cert_path
pinned: Final = Path(tempfile.gettempdir()) / f"litellm-sslcert-{hashlib.sha256(root).hexdigest()[:16]}.pem"
return str(pinned) if _replace_file(pinned, ssl.DER_cert_to_PEM_cert(root)) else cert_path
def _replace_file(target: Path, content: str) -> bool:
"""Write ``content`` to a private temp file and rename it over ``target``, so
readers never see a partial file and a symlink planted at ``target`` is
replaced rather than followed."""
try:
fd, staged = tempfile.mkstemp(dir=target.parent, prefix=f"{target.name}.")
except OSError:
return False
try:
with os.fdopen(fd, "w") as handle:
handle.write(content)
os.replace(staged, target)
except OSError:
Path(staged).unlink(missing_ok=True)
return False
return True
def translate_libpq_ssl_params(url: str, resolve_root_cert: RootCertResolver = pin_bundle_root) -> str:
"""Rewrite libpq's certificate-verification params into Prisma's dialect.
Prisma's engine only knows ``sslmode=disable|prefer|require``, ``sslcert``
(the CA bundle) and ``sslaccept=strict``. It silently discards
(a single CA certificate) and ``sslaccept=strict``. It silently discards
``sslrootcert`` and downgrades ``sslmode=verify-ca`` / ``verify-full`` to
``prefer``, so a URL copied from libpq / RDS docs connects over TLS with no
certificate check at all. ``verify-ca`` and ``verify-full`` both become
``require`` (Prisma has no CA-only mode), ``sslrootcert`` becomes
``sslcert``, and either one turns on ``sslaccept=strict`` (chain and
hostname), matching libpq where a root cert makes ``require`` verify.
Prisma params the operator pinned themselves win; anything else is left
untouched.
``sslcert`` (run through ``resolve_root_cert``, which pins a multi-root
bundle down to the server's root), and either one turns on
``sslaccept=strict`` (chain and hostname), matching libpq where a root
cert makes ``require`` verify. Prisma params the operator pinned
themselves win; anything else is left untouched.
"""
parsed: Final = urllib.parse.urlsplit(url)
pairs: Final = tuple(urllib.parse.parse_qsl(parsed.query, keep_blank_values=True))
@ -154,7 +241,9 @@ def translate_libpq_ssl_params(url: str) -> str:
if key != "sslrootcert"
)
root_cert: Final = tuple(
("sslcert", value) for key, value in pairs if key == "sslrootcert" and "sslcert" not in keys
("sslcert", resolve_root_cert(value, parsed.hostname or "", parsed.port or int(DEFAULT_POSTGRES_PORT)))
for key, value in pairs
if key == "sslrootcert" and "sslcert" not in keys
)
strict: Final = () if "sslaccept" in keys else (("sslaccept", "strict"),)
query: Final = urllib.parse.urlencode(translated + root_cert + strict)

220
litellm/proxy/db/gateway_request_tracking.py
View file

@ -10,13 +10,28 @@ strings rather than passing the raw path through. Nothing a caller sends can
add a key, so the fold and the table it commits to are bounded by (days x
routes) however much traffic arrives, and the response path carries no
unbounded queue that would block once full.
A flush commits its whole snapshot as one multi-row ``INSERT ... ON CONFLICT DO
UPDATE`` rather than one upsert per key, so a worker costs the primary one
statement per interval however many routes it served. With
``use_redis_transaction_buffer`` on, workers instead push their snapshot to a
Redis list and one lock-holding pod folds every entry and writes the table, so
the deployment as a whole costs the primary one statement per interval.
"""
from dataclasses import asdict
import json
from collections.abc import AsyncIterator, Iterable
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Final
from itertools import chain
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, TypeAlias
from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.caching import RedisCache
from litellm.constants import MAX_REDIS_BUFFER_DEQUEUE_COUNT, REDIS_GATEWAY_REQUESTS_BUFFER_KEY
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
from litellm.proxy.middleware.billable_request_metrics_middleware import BillableCategory
from litellm.types.proxy.gateway_requests import (
GatewayRequestCounts,
@ -28,6 +43,15 @@ if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
_EMPTY: Final = GatewayRequestCounts(successful_requests=0, failed_requests=0)
_TABLE: Final = '"LiteLLM_DailyGatewayRequests"'
_COLUMNS_PER_ROW: Final = 5
_UTC_NOW: Final = "(NOW() AT TIME ZONE 'UTC')"
GATEWAY_REQUESTS_JOB_NAME: Final = "update_gateway_requests_job"
_BufferedRows: TypeAlias = tuple[tuple[str, str, str, int, int], ...]
_BUFFERED_ROWS: Final = TypeAdapter(_BufferedRows)
_BUFFERED_ENTRIES: Final = TypeAdapter(tuple[str | bytes, ...])
_NO_COUNTS: Final[GatewayRequestSnapshot] = MappingProxyType({})
def _utc_date() -> str:
@ -59,20 +83,54 @@ class GatewayRequestAccumulator:
route) however long the database is unreachable.
This buys at-least-once, not exactly-once, and the cost is worth stating.
The batch commits inside its context manager's ``__aexit__``, so a failure
raised after the transaction committed (a connection dropped while reading
the acknowledgement) restores counts that are already persisted, and the
next flush increments them a second time. Exactly-once would need a dedup
key the upserts could ignore on replay. For a traffic-volume metric a rare
The statement commits on the server before its acknowledgement is read, so
a failure raised after the commit (a connection dropped while reading the
acknowledgement) restores counts that are already persisted, and the next
flush increments them a second time. Exactly-once would need a dedup key
the upsert could ignore on replay. For a traffic-volume metric a rare
overcount on a dropped acknowledgement beats losing a whole interval to
every database blip, so the trade is deliberate.
"""
for key, counts in snapshot.items():
existing = self._counts.get(key, _EMPTY)
self._counts[key] = GatewayRequestCounts(
successful_requests=existing.successful_requests + counts.successful_requests,
failed_requests=existing.failed_requests + counts.failed_requests,
)
self._counts = dict(fold_counts(chain(self._counts.items(), snapshot.items()))) # mutable-ok: fold replaced
def fold_counts(items: Iterable[tuple[GatewayRequestKey, GatewayRequestCounts]]) -> GatewayRequestSnapshot:
"""Sum counts key-wise; the result stays bounded by (date x category x route)."""
folded: Final[dict[GatewayRequestKey, GatewayRequestCounts]] = {} # mutable-ok: local fold returned once
for key, counts in items:
existing = folded.get(key, _EMPTY)
folded[key] = GatewayRequestCounts(
successful_requests=existing.successful_requests + counts.successful_requests,
failed_requests=existing.failed_requests + counts.failed_requests,
)
return folded
def build_gateway_requests_upsert(snapshot: GatewayRequestSnapshot) -> tuple[str, tuple[str | int, ...]]:
"""
One ``INSERT ... ON CONFLICT DO UPDATE`` that increments every (date, category,
route) in the snapshot. Rows are ordered by the conflict key so concurrent
writers lock rows in the same order and cannot deadlock.
"""
ordered: Final = sorted(snapshot.items(), key=lambda item: (item[0].date, item[0].category, item[0].route))
rows: Final = ", ".join(
f"(${base + 1}::text, ${base + 2}::text, ${base + 3}::text, ${base + 4}::bigint, ${base + 5}::bigint, {_UTC_NOW})"
for base in range(0, len(ordered) * _COLUMNS_PER_ROW, _COLUMNS_PER_ROW)
)
sql: Final = (
f'INSERT INTO {_TABLE} ("date", "category", "route", "successful_requests", "failed_requests", "updated_at")\n'
f"VALUES {rows}\n"
'ON CONFLICT ("date", "category", "route") DO UPDATE SET\n'
f' "successful_requests" = {_TABLE}."successful_requests" + EXCLUDED."successful_requests",\n'
f' "failed_requests" = {_TABLE}."failed_requests" + EXCLUDED."failed_requests",\n'
f' "updated_at" = {_UTC_NOW}'
)
params: Final[tuple[str | int, ...]] = tuple(
value
for key, counts in ordered
for value in (key.date, key.category, key.route, counts.successful_requests, counts.failed_requests)
)
return sql, params
async def commit_gateway_requests_to_db(
@ -80,50 +138,130 @@ async def commit_gateway_requests_to_db(
prisma_client: "PrismaClient",
snapshot: GatewayRequestSnapshot,
) -> None:
"""Upsert one incrementing row per (date, category, route)."""
"""Increment every (date, category, route) in the snapshot with a single statement."""
if not snapshot:
return
ordered: Final = sorted(snapshot.items(), key=lambda item: (item[0].date, item[0].category, item[0].route))
sql, params = build_gateway_requests_upsert(snapshot)
await prisma_client.db.execute_raw(sql, *params) # pyright: ignore[reportAny] # untyped prisma client
# pyright: ignore[reportAny] on both lines -- prisma's generated client is untyped,
# so .db and every table action off it resolve to Any at this boundary. The dict
# literals below are the shape prisma's generated inputs require.
async with prisma_client.db.batch_() as batcher: # pyright: ignore[reportAny] # untyped prisma client
for key, counts in ordered:
columns = asdict(key)
batcher.litellm_dailygatewayrequests.upsert( # pyright: ignore[reportAny] # untyped prisma client
where={"date_category_route": columns}, # mutable-ok: prisma input is dict-shaped
data={ # mutable-ok: prisma input is dict-shaped
"create": { # mutable-ok: prisma input is dict-shaped
**columns,
"successful_requests": counts.successful_requests,
"failed_requests": counts.failed_requests,
},
"update": { # mutable-ok: prisma input is dict-shaped
"successful_requests": {"increment": counts.successful_requests}, # mutable-ok: as above
"failed_requests": {"increment": counts.failed_requests}, # mutable-ok: as above
},
},
verbose_proxy_logger.debug(
"Gateway request tracking - committed %d aggregated rows in one statement", len(snapshot)
)
class GatewayRequestRedisBuffer:
"""
Folds every worker's snapshot through one Redis list so a single pod per
interval writes the table, mirroring the spend writer's transaction buffer.
Each entry is one worker's snapshot as JSON rows; the lock holder pops them,
sums them, and commits one statement. A commit failure pushes the summed
rows back so the next holder retries, keeping the at-least-once guarantee.
If that push fails too, the rows go back to the holder's own accumulator so
they ride along with its next flush instead of vanishing with the pop.
"""
def __init__(self, *, redis_cache: RedisCache, pod_lock_manager: PodLockManager) -> None:
self._redis_cache: Final = redis_cache
self._pod_lock_manager: Final = pod_lock_manager
async def push(self, snapshot: GatewayRequestSnapshot) -> None:
if not snapshot:
return
rows: Final[_BufferedRows] = tuple(
(key.date, key.category, key.route, counts.successful_requests, counts.failed_requests)
for key, counts in snapshot.items()
)
await self._redis_cache.async_rpush(key=REDIS_GATEWAY_REQUESTS_BUFFER_KEY, values=(json.dumps(rows),))
async def _pop_batch(self) -> tuple[str | bytes, ...]:
popped: Final[object] = await self._redis_cache.async_lpop( # pyright: ignore[reportAny] # redis returns Any
key=REDIS_GATEWAY_REQUESTS_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT
)
if not popped:
return ()
return _BUFFERED_ENTRIES.validate_python(popped if isinstance(popped, list) else (popped,))
async def _pop_all(self) -> AsyncIterator[str | bytes]:
while True:
batch = await self._pop_batch()
for entry in batch:
yield entry
if len(batch) < MAX_REDIS_BUFFER_DEQUEUE_COUNT:
return
async def pop(self) -> GatewayRequestSnapshot:
entries: Final = tuple([entry async for entry in self._pop_all()])
return fold_counts(
(
GatewayRequestKey(date=date, category=category, route=route),
GatewayRequestCounts(successful_requests=succeeded, failed_requests=failed),
)
for entry in entries
for date, category, route, succeeded, failed in _BUFFERED_ROWS.validate_json(entry)
)
verbose_proxy_logger.debug("Gateway request tracking - committed %d aggregated rows", len(ordered))
async def commit_if_leader(self, prisma_client: "PrismaClient") -> GatewayRequestSnapshot:
"""
Drain the list and write it as one statement, but only on the pod holding the job lock.
The lock is a lease, never released: the holder re-enters it on every flush and
keeps committing alone until the TTL lapses, so the primary sees one statement
per flush interval deployment-wide instead of one per worker.
Returns the popped rows that could be neither committed nor re-queued, for the
caller to keep in memory. Empty on success.
"""
if not await self._pod_lock_manager.acquire_lock(cronjob_id=GATEWAY_REQUESTS_JOB_NAME):
return _NO_COUNTS
buffered: Final = await self.pop()
try:
await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=buffered)
except Exception: # noqa: BLE001 -- a failed commit must not stop the scheduler
verbose_proxy_logger.warning(
"Gateway request tracking - failed to commit %d buffered rows, re-queuing to Redis for the next flush",
len(buffered),
exc_info=True,
)
return await self._requeue(buffered)
return _NO_COUNTS
async def _requeue(self, snapshot: GatewayRequestSnapshot) -> GatewayRequestSnapshot:
try:
await self.push(snapshot)
except Exception: # noqa: BLE001 -- the rows go back to the caller's accumulator instead
verbose_proxy_logger.warning(
"Gateway request tracking - Redis re-queue failed, keeping %d rows in memory for the next flush",
len(snapshot),
exc_info=True,
)
return snapshot
return _NO_COUNTS
async def flush_gateway_requests(
prisma_client: "PrismaClient",
accumulator: GatewayRequestAccumulator,
redis_buffer: GatewayRequestRedisBuffer | None = None,
) -> None:
"""
Scheduler entrypoint. Never raises: a metering failure must not kill the job.
With ``redis_buffer`` the snapshot goes to Redis and only the lease holder
writes to Postgres. Shutdown passes no buffer so a departing worker writes its
own counts directly instead of parking them behind a lease it may not hold.
``CancelledError`` is deliberately not caught, so a flush cancelled during
shutdown drops its snapshot rather than restoring counts onto an accumulator
the process is about to discard.
"""
snapshot: Final = accumulator.drain()
try:
await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=snapshot)
if redis_buffer is None:
await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=snapshot)
else:
await redis_buffer.push(snapshot)
except Exception: # noqa: BLE001 -- a failed flush must not stop the scheduler
accumulator.restore(snapshot)
verbose_proxy_logger.warning(
@ -131,3 +269,13 @@ async def flush_gateway_requests(
len(snapshot),
exc_info=True,
)
return
if redis_buffer is None:
return
try:
accumulator.restore(await redis_buffer.commit_if_leader(prisma_client))
except Exception: # noqa: BLE001 -- entries still in Redis are drained by the next flush
verbose_proxy_logger.warning(
"Gateway request tracking - leader drain failed, buffered rows stay in Redis for the next flush",
exc_info=True,
)

10
litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py
View file

@ -17,7 +17,8 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.proxy.common_utils.callback_utils import add_guardrail_scan_id
from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations
from litellm.types.utils import (
GenericGuardrailAPIInputs,
GuardrailStatus,
@ -218,6 +219,13 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail):
metadata: Final = request_data.get("metadata") or {}
request_data["metadata"] = metadata
metadata["_openai_moderation_response"] = moderation_response.model_dump()
add_guardrail_scan_id(
request_data=request_data,
scan_id=moderation_response.id,
guardrail_name=self.guardrail_name,
provider=SupportedGuardrailIntegrations.OPENAI_MODERATION.value,
stage=GuardrailEventHooks.post_call if input_type == "response" else GuardrailEventHooks.pre_call,
)
# Check if content is flagged and raise exception if needed
self._check_moderation_result(moderation_response)

34
litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py
View file

@ -721,10 +721,18 @@ class PanwPrismaAirsHandler(CustomGuardrail):
}
}
def _record_scan_id(self, request_data: dict[str, object], scan_result: Mapping[str, object]) -> None:
def _record_scan_id(
self, request_data: dict[str, object], scan_result: Mapping[str, object], stage: GuardrailEventHooks
) -> None:
"""Surface the AIRS scan id on the response, so allowed calls are auditable too."""
scan_id: Final = scan_result.get("scan_id")
add_guardrail_scan_id(request_data=request_data, scan_id=str(scan_id) if scan_id else None)
add_guardrail_scan_id(
request_data=request_data,
scan_id=str(scan_id) if scan_id else None,
guardrail_name=self.guardrail_name,
provider=self._PROVIDER_NAME,
stage=stage,
)
def _handle_api_error_with_logging(
self,
@ -948,7 +956,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
event_type=GuardrailEventHooks.post_call,
)
add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name)
self._record_scan_id(request_data, scan_result)
self._record_scan_id(request_data, scan_result, GuardrailEventHooks.post_call)
def _check_and_mark_scanned(self, data: dict, scan_type: str) -> bool:
"""
@ -1078,7 +1086,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
duration=(end_time - start_time).total_seconds(),
event_type=GuardrailEventHooks.pre_call,
)
self._record_scan_id(data, scan_result)
self._record_scan_id(data, scan_result, GuardrailEventHooks.pre_call)
action: Final = scan_result.get("action", "block")
category: Final = scan_result.get("category", "unknown")
@ -1199,7 +1207,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
duration=(end_time - start_time).total_seconds(),
event_type=GuardrailEventHooks.post_call,
)
self._record_scan_id(data, scan_result)
self._record_scan_id(data, scan_result, GuardrailEventHooks.post_call)
action: Final = scan_result.get("action", "block")
category: Final = scan_result.get("category", "unknown")
@ -1401,7 +1409,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
duration=(end_time - start_time).total_seconds(),
event_type=GuardrailEventHooks.post_call,
)
self._record_scan_id(request_data, scan_result)
self._record_scan_id(request_data, scan_result, GuardrailEventHooks.post_call)
# Add guardrail to applied guardrails header for observability
add_guardrail_to_applied_guardrails_header(
@ -1475,7 +1483,11 @@ class PanwPrismaAirsHandler(CustomGuardrail):
)
continue
self._record_scan_id(request_data, scan_result)
self._record_scan_id(
request_data,
scan_result,
GuardrailEventHooks.post_call if is_response else GuardrailEventHooks.pre_call,
)
action = scan_result.get("action", "block")
masked_args = self._masked_tool_call_arguments(
@ -1829,7 +1841,11 @@ class PanwPrismaAirsHandler(CustomGuardrail):
new_texts.append(text)
continue
self._record_scan_id(request_data, scan_result)
self._record_scan_id(
request_data,
scan_result,
GuardrailEventHooks.post_call if is_response else GuardrailEventHooks.pre_call,
)
action = scan_result.get("action", "block")
masked_text = self._get_masked_text(scan_result, is_response=is_response)
@ -1901,7 +1917,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
)
# If we reach here, fallback_on_error="allow"
else:
self._record_scan_id(request_data, mcp_scan_result)
self._record_scan_id(request_data, mcp_scan_result, GuardrailEventHooks.pre_call)
action = mcp_scan_result.get("action", "block")
masked_text = self._get_masked_text(mcp_scan_result, is_response=False)
if action == "allow":

11
litellm/proxy/litellm_pre_call_utils.py
View file

@ -235,6 +235,7 @@ _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = (
"applied_policies",
"policy_sources",
"guardrail_scan_ids",
"guardrail_scan_metadata",
"routing_decision",
GATEWAY_INJECTED_CACHE_METADATA_KEY,
"pillar_response_headers",
@ -291,6 +292,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = (
"applied_policies",
"policy_sources",
"guardrail_scan_ids",
"guardrail_scan_metadata",
"routing_decision",
GATEWAY_INJECTED_CACHE_METADATA_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
@ -1041,7 +1043,9 @@ def resolve_tenant_otel_destinations(
the request has an outcome, so honouring the filter would mean holding every span
back until the call finishes. Those entries keep today's behaviour instead, where
the tenant's credentials reach the backend through per-request tracer routing and
the operator's exporter is left alone.
the operator's exporter is left alone. Its ``callback_vars`` still take part in the
merge for a backend another entry made eligible, so the destination carries the
same credentials the runtime parser resolves for that request.
A backend the request disabled dynamically, through the key's
``litellm_disabled_callbacks`` or the ``x-litellm-disable-callbacks`` header in
@ -1069,12 +1073,13 @@ def resolve_tenant_otel_destinations(
callback
for item in entries
if (callback := _get_validated_callback_metadata(item=item, source="otel-destination")) is not None
if callback.callback_type != "failure"
if callback.callback_name.lower() not in disabled
)
return tuple(
destination
for name in dict.fromkeys(callback.callback_name for callback in callbacks)
for name in dict.fromkeys(
callback.callback_name for callback in callbacks if callback.callback_type != "failure"
)
if (
destination := destination_for(
name,

47
litellm/proxy/management_endpoints/common_daily_activity.py
View file

@ -14,6 +14,7 @@ from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.spend_tracking.key_metadata_recovery import (
attach_user_emails,
recover_double_hashed_key_metadata,
recover_key_metadata_from_spend_logs,
)
from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled
from litellm.proxy.utils import PrismaClient
@ -433,9 +434,29 @@ def update_breakdown_metrics(
return breakdown
def _spend_logs_window(dates: AbstractSet[str | None]) -> tuple[datetime, datetime] | None:
parsed: Final = sorted(day for day in (_parse_spend_date(raw) for raw in dates) if day is not None)
if not parsed:
return None
return (parsed[0] - timedelta(days=1), parsed[-1] + timedelta(days=2))
def _parse_spend_date(raw: str | None) -> datetime | None:
if not isinstance(raw, str):
return None
try:
return datetime.fromisoformat(raw)
except ValueError:
return None
_EMPTY_KEY_METADATA: Final[Mapping[str, _KeyMetadataDict]] = MappingProxyType({})
async def get_api_key_metadata(
prisma_client: PrismaClient,
api_keys: AbstractSet[str],
spend_logs_window: tuple[datetime, datetime] | None = None,
) -> Mapping[str, _KeyMetadataDict]:
"""Get api key metadata, falling back to deleted keys table for keys not found in active table.
@ -481,11 +502,17 @@ async def get_api_key_metadata(
)
still_missing: Final = api_keys - frozenset(result)
combined: Final = (
result
if not still_missing
else MappingProxyType({**result, **(await recover_double_hashed_key_metadata(prisma_client, still_missing))})
from_reverse_hash: Final = (
await recover_double_hashed_key_metadata(prisma_client, still_missing) if still_missing else _EMPTY_KEY_METADATA
)
after_token_recovery: Final = MappingProxyType({**result, **from_reverse_hash})
unresolved: Final = api_keys - frozenset(after_token_recovery)
from_spend_logs: Final = (
await recover_key_metadata_from_spend_logs(prisma_client, unresolved, spend_logs_window)
if unresolved and spend_logs_window is not None
else _EMPTY_KEY_METADATA
)
combined: Final = MappingProxyType({**after_token_recovery, **from_spend_logs})
return await attach_user_emails(prisma_client, combined)
@ -898,7 +925,9 @@ async def _aggregate_spend_records(
api_key_metadata: dict[str, _KeyMetadataDict] = {}
if api_keys:
api_key_metadata = await get_api_key_metadata(prisma_client, api_keys)
api_key_metadata = await get_api_key_metadata(
prisma_client, api_keys, _spend_logs_window(frozenset(record.date for record in records))
)
return await asyncio.to_thread(
_aggregate_spend_records_sync,
@ -1094,7 +1123,9 @@ async def _aggregate_grouping_sets_records(
api_key_metadata: dict[str, _KeyMetadataDict] = {}
if api_keys:
api_key_metadata = await get_api_key_metadata(prisma_client, api_keys)
api_key_metadata = await get_api_key_metadata(
prisma_client, api_keys, _spend_logs_window(frozenset(r.date for r in records))
)
return await asyncio.to_thread(
_aggregate_grouping_sets_records_sync,
@ -1357,7 +1388,9 @@ async def get_daily_activity_aggregated(
r.api_key for r in entity_records if r.api_key and r.api_key != PTU_SENTINEL_API_KEY
)
entity_key_metadata: Final = (
await get_api_key_metadata(prisma_client, entity_api_keys)
await get_api_key_metadata(
prisma_client, entity_api_keys, _spend_logs_window(frozenset(r.date for r in entity_records))
)
if entity_api_keys
else {} # mutable-ok: matches the helper's dict return
)

258
litellm/proxy/management_endpoints/cost_tracking_settings.py
View file

@ -18,6 +18,7 @@ from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel
import litellm
from litellm._internal_context import current_billing_time, pinned_billing_time
from litellm._logging import verbose_proxy_logger
from litellm.cost_calculator import completion_cost
from litellm.proxy._types import (
@ -27,7 +28,15 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.types.utils import CostPerToken, LlmProvidersSet, ModelInfo
from litellm.types.utils import (
CostBreakdown,
CostPerToken,
LlmProvidersSet,
ModelInfo,
ModelResponse,
PromptTokensDetailsWrapper,
Usage,
)
router: Final = APIRouter()
@ -46,13 +55,15 @@ def _configured_price(key: str, sources: tuple[Mapping[str, object], ...]) -> fl
def _extract_custom_pricing(
litellm_params: Mapping[str, object], model_info: Mapping[str, object]
litellm_params: Mapping[str, object], model_info: Mapping[str, object], builtin: ModelInfo | None
) -> CostPerToken | None:
"""
Pull per-token pricing configured on a deployment so on-prem / self-hosted
models (absent from the public cost map) still estimate a real cost.
Pricing may live on ``litellm_params`` or ``model_info``; ``litellm_params``
wins, matching the router's cost-map registration precedence.
wins, matching the router's cost-map registration precedence. Cache rates the
deployment leaves unset come from the backend model's built-in entry, then its
own input rate, again matching what the router registers for live billing.
"""
sources: Final = (litellm_params, model_info)
input_price: Final = _configured_price("input_cost_per_token", sources)
@ -61,15 +72,21 @@ def _extract_custom_pricing(
if input_price is None and output_price is None:
return None
input_rate: Final = input_price or 0.0
cache_sources: Final = sources if builtin is None else (*sources, builtin)
cache_read_price: Final = _configured_price("cache_read_input_token_cost", cache_sources)
cache_creation_price: Final = _configured_price("cache_creation_input_token_cost", cache_sources)
return CostPerToken(
input_cost_per_token=input_price or 0.0,
input_cost_per_token=input_rate,
output_cost_per_token=output_price or 0.0,
cache_read_input_token_cost=input_rate if cache_read_price is None else cache_read_price,
cache_creation_input_token_cost=input_rate if cache_creation_price is None else cache_creation_price,
)
def _lookup_model_info(model: str) -> ModelInfo | None:
def _lookup_model_info(model: str, custom_llm_provider: str | None = None) -> ModelInfo | None:
try:
return litellm.get_model_info(model=model)
return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
except Exception:
return None
@ -98,17 +115,14 @@ def _resolve_model_for_cost_lookup(model: str) -> ResolvedCostModel:
model_info: Final = first_deployment.get("model_info", {})
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
provider: Final = str(custom_llm_provider) if custom_llm_provider is not None else None
custom_cost_per_token: Final = _extract_custom_pricing(litellm_params, model_info)
# Check base_model first (needed for Azure custom deployment names)
# base_model wins (needed for Azure custom deployment names)
base_model: Final = model_info.get("base_model") or litellm_params.get("base_model")
if base_model:
verbose_proxy_logger.debug("Resolved model '%s' to base_model '%s' from router", model, base_model)
return ResolvedCostModel(str(base_model), provider, custom_cost_per_token)
resolved_model: Final = litellm_params.get("model")
resolved_model: Final = base_model or litellm_params.get("model")
if resolved_model:
verbose_proxy_logger.debug("Resolved model '%s' to '%s' from router", model, resolved_model)
custom_cost_per_token: Final = _extract_custom_pricing(
litellm_params, model_info, _lookup_model_info(str(resolved_model), provider)
)
return ResolvedCostModel(str(resolved_model), provider, custom_cost_per_token)
except Exception as e:
verbose_proxy_logger.debug("Could not resolve model '%s' from router: %s", model, e)
@ -117,19 +131,59 @@ def _resolve_model_for_cost_lookup(model: str) -> ResolvedCostModel:
return ResolvedCostModel(model, None, None)
def _calculate_period_costs(num_requests, cost_per_request, input_cost, output_cost, margin_cost):
"""
Calculate costs for a given number of requests.
@dataclass(frozen=True, slots=True)
class CostLines:
"""Cost of one request split the way the spend logs split it: the cache lines are
shares of input_cost and the reasoning line is a share of output_cost."""
Returns tuple of (total_cost, input_cost, output_cost, margin_cost) or all None if num_requests is None/0.
"""
if not num_requests:
return None, None, None, None
return (
cost_per_request * num_requests,
input_cost * num_requests,
output_cost * num_requests,
margin_cost * num_requests,
total_cost: float
input_cost: float
output_cost: float
margin_cost: float
cache_read_cost: float
cache_creation_cost: float
reasoning_cost: float
def times(self, num_requests: int | None) -> "CostLines | None":
if not num_requests:
return None
return CostLines(
total_cost=self.total_cost * num_requests,
input_cost=self.input_cost * num_requests,
output_cost=self.output_cost * num_requests,
margin_cost=self.margin_cost * num_requests,
cache_read_cost=self.cache_read_cost * num_requests,
cache_creation_cost=self.cache_creation_cost * num_requests,
reasoning_cost=self.reasoning_cost * num_requests,
)
def _cost_lines(cost_per_request: float, cost_breakdown: CostBreakdown | None) -> CostLines:
breakdown: Final = cost_breakdown if cost_breakdown is not None else CostBreakdown()
return CostLines(
total_cost=cost_per_request,
input_cost=breakdown.get("input_cost", 0.0),
output_cost=breakdown.get("output_cost", 0.0),
margin_cost=breakdown.get("margin_total_amount", 0.0),
cache_read_cost=breakdown.get("cache_read_cost", 0.0),
cache_creation_cost=breakdown.get("cache_creation_cost", 0.0),
reasoning_cost=breakdown.get("reasoning_cost", 0.0),
)
def _usage_for_estimate(request: CostEstimateRequest) -> Usage:
cache_tokens: Final = request.cache_read_input_tokens + request.cache_creation_input_tokens
return Usage(
prompt_tokens=request.input_tokens,
completion_tokens=request.output_tokens,
total_tokens=request.input_tokens + request.output_tokens,
reasoning_tokens=request.reasoning_tokens,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=request.cache_read_input_tokens,
cache_creation_tokens=request.cache_creation_input_tokens,
)
if cache_tokens
else None,
)
@ -530,11 +584,14 @@ async def estimate_cost(
- model: Model name (e.g., "gpt-4", "claude-3-opus")
- input_tokens: Expected input tokens per request
- output_tokens: Expected output tokens per request
- cache_read_input_tokens: Cache-read tokens per request, counted within input_tokens (optional)
- cache_creation_input_tokens: Cache-write tokens per request, counted within input_tokens (optional)
- reasoning_tokens: Reasoning tokens per request, counted within output_tokens (optional)
- num_requests_per_day: Number of requests per day (optional)
- num_requests_per_month: Number of requests per month (optional)
Returns cost breakdown including:
- Per-request costs (input, output, margin)
- Per-request costs (input, output, margin, plus the cache-read, cache-write and reasoning shares)
- Daily costs (if num_requests_per_day provided)
- Monthly costs (if num_requests_per_month provided)
@ -543,14 +600,15 @@ async def estimate_cost(
{
"model": "gpt-4",
"input_tokens": 1000,
"cache_read_input_tokens": 800,
"output_tokens": 500,
"reasoning_tokens": 200,
"num_requests_per_day": 100,
"num_requests_per_month": 3000
}
```
"""
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.utils import ModelResponse, Usage
# Resolve model name (handles router aliases like 'e-model-router' -> 'azure_ai/gpt-4')
resolved: Final = _resolve_model_for_cost_lookup(request.model)
@ -559,15 +617,8 @@ async def estimate_cost(
verbose_proxy_logger.debug("Cost estimate: request.model='%s' resolved to '%s'", request.model, resolved_model)
# Create a mock response with usage for completion_cost
mock_response: Final = ModelResponse(
model=resolved_model,
usage=Usage(
prompt_tokens=request.input_tokens,
completion_tokens=request.output_tokens,
total_tokens=request.input_tokens + request.output_tokens,
),
)
usage: Final = _usage_for_estimate(request)
mock_response: Final = ModelResponse(model=resolved_model, usage=usage)
# Create a logging object to capture cost breakdown
litellm_logging_obj: Final = LiteLLMLoggingObj(
@ -580,92 +631,73 @@ async def estimate_cost(
function_id="cost-estimate",
)
# Use completion_cost which handles all the logic including margins/discounts
try:
cost_per_request: Final = completion_cost(
completion_response=mock_response,
model=resolved_model,
custom_llm_provider=resolved_provider,
custom_cost_per_token=resolved.custom_cost_per_token,
litellm_logging_obj=litellm_logging_obj,
)
except Exception as e:
raise HTTPException(
status_code=404,
detail={
"error": f"Could not calculate cost for model '{request.model}' (resolved to '{resolved_model}'): {e}"
},
)
# Pinning one moment keeps an off-peak window that opens mid-quote from pricing the totals on
# one side of it and the reported rates on the other.
with pinned_billing_time(current_billing_time()):
# Use completion_cost which handles all the logic including margins/discounts
try:
cost_per_request: Final = completion_cost(
completion_response=mock_response,
model=resolved_model,
custom_llm_provider=resolved_provider,
custom_cost_per_token=resolved.custom_cost_per_token,
litellm_logging_obj=litellm_logging_obj,
)
except Exception as e: # noqa: BLE001 # completion_cost raises a bare Exception for an unpriceable model
raise HTTPException(
status_code=404,
detail={
"error": f"Could not calculate cost for model '{request.model}' (resolved to '{resolved_model}'): {e}"
},
)
# Get cost breakdown from the logging object
cost_breakdown: Final = litellm_logging_obj.cost_breakdown
# The rates come back from the pricing call itself rather than a second lookup, so they are the
# ones the cost lines above billed at even when completion_cost infers a provider this endpoint
# never resolved (an unrouted "xai/grok-4" prices on xai's inclusive tier thresholds; a lookup
# here without that provider would report the sub-200k rate for a line billed above it).
rates: Final = litellm_logging_obj.billed_token_rates
per_request: Final = _cost_lines(cost_per_request, litellm_logging_obj.cost_breakdown)
daily: Final = per_request.times(request.num_requests_per_day)
monthly: Final = per_request.times(request.num_requests_per_month)
input_cost: Final = cost_breakdown.get("input_cost", 0.0) if cost_breakdown else 0.0
output_cost: Final = cost_breakdown.get("output_cost", 0.0) if cost_breakdown else 0.0
margin_cost: Final = cost_breakdown.get("margin_total_amount", 0.0) if cost_breakdown else 0.0
model_info: Final = _lookup_model_info(resolved_model)
mapped_input_price: Final = model_info.get("input_cost_per_token") if model_info is not None else None
mapped_output_price: Final = model_info.get("output_cost_per_token") if model_info is not None else None
model_info: Final = _lookup_model_info(resolved_model, resolved_provider)
mapped_provider: Final = model_info.get("litellm_provider") if model_info is not None else None
input_cost_per_token: Final = (
resolved.custom_cost_per_token["input_cost_per_token"]
if resolved.custom_cost_per_token is not None
else mapped_input_price
)
output_cost_per_token: Final = (
resolved.custom_cost_per_token["output_cost_per_token"]
if resolved.custom_cost_per_token is not None
else mapped_output_price
)
custom_llm_provider: Final = mapped_provider if mapped_provider is not None else resolved_provider
# Calculate daily and monthly costs
(
daily_cost,
daily_input_cost,
daily_output_cost,
daily_margin_cost,
) = _calculate_period_costs(
num_requests=request.num_requests_per_day,
cost_per_request=cost_per_request,
input_cost=input_cost,
output_cost=output_cost,
margin_cost=margin_cost,
)
(
monthly_cost,
monthly_input_cost,
monthly_output_cost,
monthly_margin_cost,
) = _calculate_period_costs(
num_requests=request.num_requests_per_month,
cost_per_request=cost_per_request,
input_cost=input_cost,
output_cost=output_cost,
margin_cost=margin_cost,
)
return CostEstimateResponse(
model=request.model,
input_tokens=request.input_tokens,
output_tokens=request.output_tokens,
cache_read_input_tokens=request.cache_read_input_tokens,
cache_creation_input_tokens=request.cache_creation_input_tokens,
reasoning_tokens=request.reasoning_tokens,
num_requests_per_day=request.num_requests_per_day,
num_requests_per_month=request.num_requests_per_month,
cost_per_request=cost_per_request,
input_cost_per_request=input_cost,
output_cost_per_request=output_cost,
margin_cost_per_request=margin_cost,
daily_cost=daily_cost,
daily_input_cost=daily_input_cost,
daily_output_cost=daily_output_cost,
daily_margin_cost=daily_margin_cost,
monthly_cost=monthly_cost,
monthly_input_cost=monthly_input_cost,
monthly_output_cost=monthly_output_cost,
monthly_margin_cost=monthly_margin_cost,
input_cost_per_token=input_cost_per_token,
output_cost_per_token=output_cost_per_token,
cost_per_request=per_request.total_cost,
input_cost_per_request=per_request.input_cost,
output_cost_per_request=per_request.output_cost,
margin_cost_per_request=per_request.margin_cost,
cache_read_cost_per_request=per_request.cache_read_cost,
cache_creation_cost_per_request=per_request.cache_creation_cost,
reasoning_cost_per_request=per_request.reasoning_cost,
daily_cost=daily.total_cost if daily is not None else None,
daily_input_cost=daily.input_cost if daily is not None else None,
daily_output_cost=daily.output_cost if daily is not None else None,
daily_margin_cost=daily.margin_cost if daily is not None else None,
daily_cache_read_cost=daily.cache_read_cost if daily is not None else None,
daily_cache_creation_cost=daily.cache_creation_cost if daily is not None else None,
daily_reasoning_cost=daily.reasoning_cost if daily is not None else None,
monthly_cost=monthly.total_cost if monthly is not None else None,
monthly_input_cost=monthly.input_cost if monthly is not None else None,
monthly_output_cost=monthly.output_cost if monthly is not None else None,
monthly_margin_cost=monthly.margin_cost if monthly is not None else None,
monthly_cache_read_cost=monthly.cache_read_cost if monthly is not None else None,
monthly_cache_creation_cost=monthly.cache_creation_cost if monthly is not None else None,
monthly_reasoning_cost=monthly.reasoning_cost if monthly is not None else None,
input_cost_per_token=rates.input_cost_per_token if rates is not None else None,
output_cost_per_token=rates.output_cost_per_token if rates is not None else None,
cache_read_input_token_cost=rates.cache_read_input_token_cost if rates is not None else None,
cache_creation_input_token_cost=rates.cache_creation_input_token_cost if rates is not None else None,
output_cost_per_reasoning_token=rates.output_cost_per_reasoning_token if rates is not None else None,
provider=custom_llm_provider,
)

4
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -2673,6 +2673,8 @@ if MCP_AVAILABLE:
"""
Updates the MCP Server in the db.
Partial update: a field left out of the payload keeps its stored value, and a field sent as null is cleared.
Parameters:
- payload: UpdateMCPServerRequest - Required. The updated mcp server data.
```
@ -3098,6 +3100,8 @@ if MCP_AVAILABLE:
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
litellm_changed_by: str | None = Header(None),
):
"""Partial update: a field left out keeps its stored value, and a field sent as null is cleared, except
``toolset_name`` and ``tools``, which a toolset always has; empty the tool selection with an explicit []."""
prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy")
if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role:
raise HTTPException(

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

@ -5601,7 +5601,7 @@ async def team_model_add(
updated_team: Final = await _team_db(prisma_client).update(
where={"team_id": data.team_id},
data={"updated_at": datetime.now(timezone.utc)},
include={"object_permission": True},
include={"litellm_model_table": True, "object_permission": True},
)
if updated_team is None:
raise HTTPException(
@ -5688,7 +5688,7 @@ async def team_model_delete(
updated_team: Final = await _team_db(prisma_client).update(
where={"team_id": data.team_id},
data={"models": updated_models},
include={"object_permission": True},
include={"litellm_model_table": True, "object_permission": True},
)
if updated_team is None:
raise HTTPException(

23
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -22,7 +22,6 @@ from html import escape
from types import MappingProxyType
from typing import (
TYPE_CHECKING,
Annotated,
Any,
Final,
Literal,
@ -42,7 +41,7 @@ if TYPE_CHECKING:
import jwt
from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status
from fastapi.responses import RedirectResponse
from pydantic import BaseModel, BeforeValidator, ConfigDict, TypeAdapter, ValidationError
from pydantic import BaseModel, TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
@ -95,6 +94,7 @@ from litellm.proxy.auth.auth_utils import (
)
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.auth.team_grants import TeamModelAliasTable
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.admin_ui_utils import (
admin_ui_disabled,
@ -209,31 +209,14 @@ def _team_detail_db(repo: TeamRepository) -> "TableActions[_TeamDetailRow]":
return repo.table
_MODEL_ALIASES_ADAPTER: Final = TypeAdapter(dict[str, str])
_SSO_TOKEN_CLAIMS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def _decode_model_aliases(value: object) -> object:
"""``/team/new`` stores team model aliases as a JSON-encoded string in the Json column."""
if not isinstance(value, str):
return value
try:
return _MODEL_ALIASES_ADAPTER.validate_json(value)
except ValidationError:
return None
class _TeamModelAliasTable(BaseModel):
model_config = ConfigDict(protected_namespaces=())
model_aliases: Annotated[Mapping[str, str] | None, BeforeValidator(_decode_model_aliases)] = None
class _TeamRowGrants(BaseModel):
team_id: str
team_alias: str | None = None
models: tuple[str, ...] = ()
litellm_model_table: _TeamModelAliasTable | None = None
litellm_model_table: TeamModelAliasTable | None = None
class CliSsoTeamDetail(BaseModel):

408
litellm/proxy/proxy_server.py
View file

@ -17,6 +17,7 @@ import time
import traceback
import warnings
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Collection, Mapping, MutableMapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from types import MappingProxyType, UnionType
from typing import (
@ -131,6 +132,7 @@ from litellm.router_utils.auto_router_tuning_baseline import (
snapshot_tuning_baselines,
tuning_limit_violation,
)
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.utils import (
ModelResponse,
ModelResponseStream,
@ -138,11 +140,7 @@ from litellm.types.utils import (
TextCompletionResponse,
TokenCountResponse,
)
from litellm.utils import (
_invalidate_model_cost_lowercase_map,
load_credentials_from_list,
reapply_runtime_model_cost_registrations,
)
from litellm.utils import load_credentials_from_list
if TYPE_CHECKING:
from aiohttp import ClientSession
@ -426,6 +424,7 @@ from litellm.proxy.db.exception_handler import (
)
from litellm.proxy.db.gateway_request_tracking import (
GatewayRequestAccumulator,
GatewayRequestRedisBuffer,
flush_gateway_requests,
)
from litellm.proxy.db.proxy_worker_heartbeat import (
@ -2357,6 +2356,17 @@ open_telemetry_logger: OpenTelemetry | None = None
gateway_request_accumulator: Final = GatewayRequestAccumulator()
### INITIALIZE GLOBAL LOGGING OBJECT ###
proxy_logging_obj: ProxyLogging = ProxyLogging(user_api_key_cache=user_api_key_cache, premium_user=premium_user)
def _gateway_request_redis_buffer() -> GatewayRequestRedisBuffer | None:
"""Shares the spend writer's transaction-buffer Redis and pod lock when use_redis_transaction_buffer is on."""
writer: Final = proxy_logging_obj.db_spend_update_writer
redis_cache: Final = writer.redis_update_buffer.redis_cache
if redis_cache is None or not writer.redis_update_buffer._should_commit_spend_updates_to_redis():
return None
return GatewayRequestRedisBuffer(redis_cache=redis_cache, pod_lock_manager=writer.pod_lock_manager)
### REDIS QUEUE ###
async_result: Final = None
celery_app_conn: Final = None
@ -2707,6 +2717,12 @@ async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float)
return fallback_spend, False
@dataclass(frozen=True, slots=True)
class _PendingSpendIncrement:
counter_key: str
increment: float
async def increment_spend_counters(
token: str | None,
team_id: str | None,
@ -2741,7 +2757,7 @@ async def increment_spend_counters(
cost: Final[float] = response_cost
async def _key_scope(key_token: str) -> None:
async def _key_scope(key_token: str) -> tuple[_PendingSpendIncrement | BaseException, ...]:
# key_token arrives pre-hashed from metadata["user_api_key"] (auth flow
# hashes raw "sk-..." keys before they reach the callback). The
# startswith("sk-") check is a safety net matching update_cache —
@ -2752,30 +2768,29 @@ async def increment_spend_counters(
hash_token(token=key_token) if isinstance(key_token, str) and key_token.startswith("sk-") else key_token
)
key_counter_key: Final = f"spend:key:{hashed_token}"
if key_counter_key not in reserved_counter_keys:
await _init_and_increment_spend_counter(
counter_key=key_counter_key,
source_cache_key=hashed_token,
increment=cost,
key_pending: Final[tuple[_PendingSpendIncrement, ...]] = (
()
if key_counter_key in reserved_counter_keys
else (
await _prepare_spend_counter_increment(
counter_key=key_counter_key,
source_cache_key=hashed_token,
increment=cost,
),
)
key_obj: Final[object] = await user_api_key_cache.async_get_cache(key=hashed_token)
if key_obj is None:
return
key_budget_limits = getattr(key_obj, "budget_limits", None) or (
key_obj.get("budget_limits") if isinstance(key_obj, dict) else None
)
if isinstance(key_budget_limits, str):
key_budget_limits = json.loads(key_budget_limits)
if not isinstance(key_budget_limits, list):
return
for window in key_budget_limits:
duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration
key_window_reset_at = window.get("reset_at") if isinstance(window, dict) else window.reset_at
key_window_counter = f"spend:key:{hashed_token}:window:{duration}"
async def _key_window_increment(window: object) -> _PendingSpendIncrement | None:
duration = (
window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None)
)
key_window_reset_at = (
window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None)
)
key_window_counter: Final = f"spend:key:{hashed_token}:window:{duration}"
key_window_start = get_budget_window_start(window)
if key_window_counter not in reserved_counter_keys:
await _init_and_increment_window_spend_counter(
pending_window: Final = (
await _prepare_window_spend_counter_increment(
counter_key=key_window_counter,
entity_type="Key",
entity_id=hashed_token,
@ -2783,6 +2798,9 @@ async def increment_spend_counters(
window_start=key_window_start,
increment=cost,
)
if key_window_counter not in reserved_counter_keys
else None
)
await _enqueue_window_spend_row_update(
entity_type=Litellm_EntityType.KEY,
entity_id=hashed_token,
@ -2792,33 +2810,48 @@ async def increment_spend_counters(
increment=cost,
request_started_at=request_started_at,
)
return pending_window
async def _team_scope(scope_team_id: str) -> None:
team_counter_key: Final = f"spend:team:{scope_team_id}"
if team_counter_key not in reserved_counter_keys:
await _init_and_increment_spend_counter(
counter_key=team_counter_key,
source_cache_key=f"team_id:{scope_team_id}",
increment=cost,
)
team_obj: Final[object] = await user_api_key_cache.async_get_cache(key=f"team_id:{scope_team_id}")
if team_obj is None:
return
team_budget_limits = getattr(team_obj, "budget_limits", None) or (
team_obj.get("budget_limits") if isinstance(team_obj, dict) else None
key_obj: Final[object] = await user_api_key_cache.async_get_cache(key=hashed_token)
if key_obj is None:
return key_pending
key_budget_limits = getattr(key_obj, "budget_limits", None) or (
key_obj.get("budget_limits") if isinstance(key_obj, dict) else None
)
if isinstance(team_budget_limits, str):
team_budget_limits = json.loads(team_budget_limits)
if not isinstance(team_budget_limits, list):
return
for window in team_budget_limits:
duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration
team_window_reset_at = window.get("reset_at") if isinstance(window, dict) else window.reset_at
team_window_counter = f"spend:team:{scope_team_id}:window:{duration}"
if isinstance(key_budget_limits, str):
key_budget_limits = json.loads(key_budget_limits)
if not isinstance(key_budget_limits, list):
return key_pending
window_pending: Final = await asyncio.gather(
*(_key_window_increment(window) for window in key_budget_limits), return_exceptions=True
)
return key_pending + tuple(item for item in window_pending if item is not None)
async def _team_scope(scope_team_id: str) -> tuple[_PendingSpendIncrement | BaseException, ...]:
team_counter_key: Final = f"spend:team:{scope_team_id}"
team_pending: Final[tuple[_PendingSpendIncrement, ...]] = (
()
if team_counter_key in reserved_counter_keys
else (
await _prepare_spend_counter_increment(
counter_key=team_counter_key,
source_cache_key=f"team_id:{scope_team_id}",
increment=cost,
),
)
)
async def _team_window_increment(window: object) -> _PendingSpendIncrement | None:
duration = (
window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None)
)
team_window_reset_at = (
window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None)
)
team_window_counter: Final = f"spend:team:{scope_team_id}:window:{duration}"
team_window_start = get_budget_window_start(window)
if team_window_counter not in reserved_counter_keys:
await _init_and_increment_window_spend_counter(
pending_window: Final = (
await _prepare_window_spend_counter_increment(
counter_key=team_window_counter,
entity_type="Team",
entity_id=scope_team_id,
@ -2826,6 +2859,9 @@ async def increment_spend_counters(
window_start=team_window_start,
increment=cost,
)
if team_window_counter not in reserved_counter_keys
else None
)
await _enqueue_window_spend_row_update(
entity_type=Litellm_EntityType.TEAM,
entity_id=scope_team_id,
@ -2835,25 +2871,47 @@ async def increment_spend_counters(
increment=cost,
request_started_at=request_started_at,
)
return pending_window
async def _team_member_scope(scope_user_id: str, scope_team_id: str) -> None:
team_obj: Final[object] = await user_api_key_cache.async_get_cache(key=f"team_id:{scope_team_id}")
if team_obj is None:
return team_pending
team_budget_limits = getattr(team_obj, "budget_limits", None) or (
team_obj.get("budget_limits") if isinstance(team_obj, dict) else None
)
if isinstance(team_budget_limits, str):
team_budget_limits = json.loads(team_budget_limits)
if not isinstance(team_budget_limits, list):
return team_pending
window_pending: Final = await asyncio.gather(
*(_team_window_increment(window) for window in team_budget_limits), return_exceptions=True
)
return team_pending + tuple(item for item in window_pending if item is not None)
async def _team_member_scope(
scope_user_id: str, scope_team_id: str
) -> tuple[_PendingSpendIncrement | BaseException, ...]:
team_member_counter_key: Final = f"spend:team_member:{scope_user_id}:{scope_team_id}"
if team_member_counter_key in reserved_counter_keys:
return
await _init_and_increment_spend_counter(
counter_key=team_member_counter_key,
source_cache_key=f"team_membership:{scope_user_id}:{scope_team_id}",
increment=cost,
return ()
return (
await _prepare_spend_counter_increment(
counter_key=team_member_counter_key,
source_cache_key=f"team_membership:{scope_user_id}:{scope_team_id}",
increment=cost,
),
)
async def _user_scope(scope_user_id: str) -> None:
async def _user_scope(scope_user_id: str) -> tuple[_PendingSpendIncrement | BaseException, ...]:
user_counter_key: Final = f"spend:user:{scope_user_id}"
if user_counter_key in reserved_counter_keys:
return
await _init_and_increment_spend_counter(
counter_key=user_counter_key,
source_cache_key=scope_user_id,
increment=cost,
return ()
return (
await _prepare_spend_counter_increment(
counter_key=user_counter_key,
source_cache_key=scope_user_id,
increment=cost,
),
)
scope_coros: Final = tuple(
@ -2863,7 +2921,7 @@ async def increment_spend_counters(
_team_scope(team_id) if team_id is not None else None,
_team_member_scope(user_id, team_id) if user_id is not None and team_id is not None else None,
_user_scope(user_id) if user_id is not None else None,
_increment_end_user_and_tag_spend_counters(
_prepare_end_user_and_tag_spend_increments(
end_user_id=end_user_id,
tags=tags,
response_cost=cost,
@ -2871,14 +2929,14 @@ async def increment_spend_counters(
)
if end_user_id is not None or tags is not None
else None,
_increment_model_access_group_spend_counters(
_prepare_model_access_group_spend_increments(
model_access_groups=model_access_groups,
response_cost=cost,
reserved_counter_keys=reserved_counter_keys,
)
if model_access_groups
else None,
_increment_org_spend_counter(
_prepare_org_spend_increment(
org_id=org_id,
response_cost=cost,
reserved_counter_keys=reserved_counter_keys,
@ -2893,7 +2951,20 @@ async def increment_spend_counters(
# as orphaned tasks that race the caller's reservation-counter invalidation;
# all scopes settle, then the first error propagates as before.
scope_results: Final = await asyncio.gather(*scope_coros, return_exceptions=True)
scope_errors: Final = [r for r in scope_results if isinstance(r, BaseException)]
scope_errors: Final = tuple(
item
for scope in scope_results
for item in (scope if isinstance(scope, tuple) else (scope,))
if isinstance(item, BaseException)
)
pending: Final = tuple(
item
for scope in scope_results
if not isinstance(scope, BaseException)
for item in scope
if not isinstance(item, BaseException)
)
await _apply_spend_counter_increments(pending=pending)
if scope_errors:
raise scope_errors[0]
@ -2936,41 +3007,49 @@ async def _reconcile_budget_reservation_for_counter_update(
return reserved_counter_keys
async def _increment_end_user_and_tag_spend_counters(
async def _prepare_end_user_and_tag_spend_increments(
end_user_id: str | None,
tags: list[str] | None,
response_cost: float,
reserved_counter_keys: set[str],
) -> None:
if end_user_id is not None:
await _init_and_increment_unreserved_spend_counter(
counter_key=f"spend:end_user:{end_user_id}",
source_cache_key=end_user_cache_key(end_user_id),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
if tags is None:
return
seen_tags: Final[set[str]] = set()
for tag_name in tags:
if not tag_name or not isinstance(tag_name, str) or tag_name in seen_tags:
continue
seen_tags.add(tag_name)
await _init_and_increment_unreserved_spend_counter(
counter_key=f"spend:tag:{tag_name}",
source_cache_key=tag_cache_key(tag_name),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
) -> tuple[_PendingSpendIncrement | BaseException, ...]:
unique_tags: Final = (
tuple(dict.fromkeys(tag for tag in tags if tag and isinstance(tag, str))) if tags is not None else ()
)
results: Final = await asyncio.gather(
*(
coro
for coro in (
_prepare_unreserved_spend_counter_increment(
counter_key=f"spend:end_user:{end_user_id}",
source_cache_key=end_user_cache_key(end_user_id),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
if end_user_id is not None
else None,
*(
_prepare_unreserved_spend_counter_increment(
counter_key=f"spend:tag:{tag_name}",
source_cache_key=tag_cache_key(tag_name),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
for tag_name in unique_tags
),
)
if coro is not None
),
return_exceptions=True,
)
return tuple(item for item in results if item is not None)
async def _increment_model_access_group_spend_counters(
async def _prepare_model_access_group_spend_increments(
model_access_groups: Sequence[object],
response_cost: float,
reserved_counter_keys: set[str],
) -> None:
) -> tuple[_PendingSpendIncrement | BaseException, ...]:
"""Charge the model access groups that authorized this request.
Without this the counter auth reads is written only by the reservation path, so
@ -2984,55 +3063,63 @@ async def _increment_model_access_group_spend_counters(
unique_groups: Final = tuple(
dict.fromkeys(group for group in model_access_groups if group and isinstance(group, str))
)
for group in unique_groups:
await _init_and_increment_unreserved_spend_counter(
counter_key=model_access_group_spend_counter_key(group),
source_cache_key=model_access_group_cache_key(group),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
results: Final = await asyncio.gather(
*(
_prepare_unreserved_spend_counter_increment(
counter_key=model_access_group_spend_counter_key(group),
source_cache_key=model_access_group_cache_key(group),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
for group in unique_groups
),
return_exceptions=True,
)
return tuple(item for item in results if item is not None)
async def _increment_org_spend_counter(
async def _prepare_org_spend_increment(
org_id: str | None,
response_cost: float,
reserved_counter_keys: set[str],
) -> None:
) -> tuple[_PendingSpendIncrement, ...]:
if org_id is None:
return
return ()
await _init_and_increment_unreserved_spend_counter(
pending: Final = await _prepare_unreserved_spend_counter_increment(
counter_key=f"spend:org:{org_id}",
source_cache_key=[f"org_id:{org_id}:with_budget", f"org_id:{org_id}"],
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
return (pending,) if pending is not None else ()
async def _init_and_increment_unreserved_spend_counter(
async def _prepare_unreserved_spend_counter_increment(
counter_key: str,
source_cache_key: str | list[str],
increment: float,
reserved_counter_keys: set[str],
) -> None:
) -> _PendingSpendIncrement | None:
if counter_key in reserved_counter_keys:
return
return None
await _init_and_increment_spend_counter(
return await _prepare_spend_counter_increment(
counter_key=counter_key,
source_cache_key=source_cache_key,
increment=increment,
)
async def _init_and_increment_spend_counter(
async def _prepare_spend_counter_increment(
counter_key: str,
source_cache_key: str | list[str],
increment: float,
):
) -> _PendingSpendIncrement:
"""
Initialize counter from the authoritative DB spend value if not yet
set, then atomically increment in both in-memory and Redis.
set, then return the pending increment for the caller to apply in one
pipelined Redis call.
On first access per pod:
1. Check spend_counter_cache (in-memory -> Redis via DualCache)
@ -3044,13 +3131,13 @@ async def _init_and_increment_spend_counter(
the counter as absent and seed it. Using increment means the worst case
is over-counting (conservative, blocks slightly early) rather than
under-counting (would allow overspend).
4. Increment atomically (both in-memory + Redis)
4. Increment is returned for the caller to apply via pipeline
"""
await _ensure_spend_counter_initialized(
counter_key=counter_key,
source_cache_key=source_cache_key,
)
await _increment_spend_counter_cache(counter_key=counter_key, increment=increment)
return _PendingSpendIncrement(counter_key=counter_key, increment=increment)
async def _enqueue_window_spend_row_update(
@ -3102,20 +3189,20 @@ async def _enqueue_window_spend_row_update(
)
async def _init_and_increment_window_spend_counter(
async def _prepare_window_spend_counter_increment(
counter_key: str,
entity_type: str,
entity_id: str,
window_duration: str | None,
window_start: datetime | None,
increment: float,
):
) -> _PendingSpendIncrement | None:
if window_start is None:
verbose_proxy_logger.warning(
"Skipping spend counter increment for invalid budget window %s",
counter_key,
)
return
return None
initialized: Final = await _ensure_window_spend_counter_initialized(
counter_key=counter_key,
@ -3125,8 +3212,8 @@ async def _init_and_increment_window_spend_counter(
window_start=window_start,
)
if initialized is False:
return
await _increment_spend_counter_cache(counter_key=counter_key, increment=increment)
return None
return _PendingSpendIncrement(counter_key=counter_key, increment=increment)
async def _ensure_spend_counter_initialized(
@ -3259,6 +3346,32 @@ async def _invalidate_spend_counter(counter_key: str):
)
async def _apply_spend_counter_increments(pending: Sequence[_PendingSpendIncrement]) -> None:
if not pending:
return
redis_cache: Final = spend_counter_cache.redis_cache
if redis_cache is None:
for item in pending:
await spend_counter_cache.async_increment_cache(
key=item.counter_key,
value=item.increment,
refresh_ttl=True,
)
return
ttl: Final = redis_cache.get_ttl()
increment_list: Final = [ # mutable-ok: async_increment_pipeline signature requires list[RedisPipelineIncrementOperation]
RedisPipelineIncrementOperation(key=item.counter_key, increment_value=item.increment, ttl=ttl)
for item in pending
]
try:
results: Final = await redis_cache.async_increment_pipeline(increment_list=increment_list)
except Exception:
await asyncio.gather(*(_invalidate_spend_counter(counter_key=item.counter_key) for item in pending))
raise
for item, current_value in zip(pending, results or ()):
spend_counter_cache.in_memory_cache.set_cache(key=item.counter_key, value=current_value)
async def update_cache(
token: str | None,
user_id: str | None,
@ -4436,20 +4549,9 @@ def resolve_classifier_plugin(
def _swap_in_model_cost_map(new_model_cost_map: dict) -> int:
"""Adopt a freshly fetched cost map into this process's litellm state, return the model count"""
litellm.model_cost = new_model_cost_map
# Invalidate case-insensitive lookup map since model_cost was replaced
_invalidate_model_cost_lowercase_map()
# Repopulate provider model sets (e.g. litellm.anthropic_models) so that
# wildcard patterns like "anthropic/*" include any newly added models.
litellm.add_known_models(model_cost_map=new_model_cost_map)
# Counted before the re-apply below, which writes into this same dict, so the
# number reported describes the fetched price data alone.
fetched_model_count: Final = len(new_model_cost_map) if new_model_cost_map else 0
# The swap discards everything registered at runtime (deployment model_info,
# register_model overrides), so put it back on top of the fresh catalog.
reapply_runtime_model_cost_registrations()
return fetched_model_count
from litellm.litellm_core_utils.get_model_cost_map import adopt_model_cost_map
return adopt_model_cost_map(new_model_cost_map)
def should_load_db_object(object_type: str | SupportedDBObjectType) -> bool:
@ -9543,7 +9645,7 @@ class ProxyStartupEvent:
flush_gateway_requests,
"interval",
seconds=batch_writing_interval,
args=(prisma_client, gateway_request_accumulator),
args=(prisma_client, gateway_request_accumulator, _gateway_request_redis_buffer()),
id="update_gateway_requests_job",
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
@ -17750,6 +17852,7 @@ async def reload_model_cost_map(
# Immediately reload the model cost map in the current pod
from litellm.litellm_core_utils.get_model_cost_map import (
ModelCostMapReloadUnavailable,
get_model_cost_map_provenance,
refetch_model_cost_map,
)
@ -17763,6 +17866,7 @@ async def reload_model_cost_map(
models_count = _swap_in_model_cost_map(reload_result.model_cost_map)
current_time = utc_now()
proxy_config.model_cost_map_loaded_at = current_time
provenance: Final = get_model_cost_map_provenance()
# Publish a new revision so every other pod reloads on its next poll; this pod has
# already served it, so adopt it here rather than reloading again a tick later
@ -17777,6 +17881,7 @@ async def reload_model_cost_map(
"status": "success",
"models_count": models_count,
"timestamp": current_time.isoformat(),
**provenance,
}
except HTTPException:
raise
@ -17897,12 +18002,17 @@ async def get_model_cost_map_reload_status(
try:
global prisma_client
from litellm.litellm_core_utils.get_model_cost_map import (
get_model_cost_map_provenance,
)
provenance: Final = get_model_cost_map_provenance()
if prisma_client is None:
verbose_proxy_logger.info("No database connection, returning not scheduled")
return reload_schedule_status(None)
return {**reload_schedule_status(None), **provenance}
return reload_schedule_status(await read_reload_schedule(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME))
schedule: Final = await read_reload_schedule(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME)
return {**reload_schedule_status(schedule), **provenance}
except Exception as e:
verbose_proxy_logger.exception("Failed to get model cost map reload status: %s", e)
raise HTTPException(
@ -17930,6 +18040,9 @@ async def get_model_cost_map_source(
- url: the remote URL that was attempted (null when env-forced local)
- is_env_forced: true if LITELLM_LOCAL_MODEL_COST_MAP=True forced local usage
- fallback_reason: human-readable reason why remote failed (null on success)
- loaded_at: when this pod last loaded the map
- source_revision: git blob id of the loaded file, what git rev-parse <commit>:<path> prints for it
- etag: the ETag of the remote fetch (null for the bundled backup)
- model_count: number of models in the currently loaded cost map
"""
# Read-only source info — admin viewers can read.
@ -18484,6 +18597,31 @@ async def _stream_mcp_asgi_response(handle_fn, scope: dict, receive) -> "Streami
########################################################
@app.api_route(
"/mcp/proxy",
methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"], # mutable-ok: FastAPI route methods
)
async def proxy_mcp_route(request: Request) -> Response:
"""Serve the fixed three-tool MCP proxy surface."""
from litellm.proxy._experimental.mcp_server.mcp_context import ( # pyright: ignore[reportPrivateUsage] # route-owned mode
_mcp_proxy_mode, # pyright: ignore[reportPrivateUsage] # route-owned mode
)
from litellm.proxy._experimental.mcp_server.server import handle_streamable_http_mcp
from litellm.proxy._experimental.mcp_server.utils import is_mcp_available
if not is_mcp_available():
raise HTTPException(status_code=404, detail="Not Found")
token: Final = _mcp_proxy_mode.set(True)
try:
scope: Final = dict(request.scope) # mutable-ok: ASGI scope rewrite
scope["_original_path"] = scope.get("path", "")
scope["path"] = BASE_MCP_ROUTE
return await _stream_mcp_asgi_response(handle_streamable_http_mcp, scope, request.receive)
finally:
_mcp_proxy_mode.reset(token)
@app.api_route(
BASE_MCP_ROUTE,
methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"],

43
litellm/proxy/public_endpoints/autorouter_presets.json
View file

@ -10,7 +10,12 @@
"REASONING": ["claude-opus-5"]
},
"tier_model_configs": {
"REASONING": [{ "model_name": "claude-opus-5", "litellm_params": { "reasoning_effort": "high" } }]
"REASONING": [
{
"model_name": "claude-opus-5",
"litellm_params": { "reasoning_effort": "high" }
}
]
},
"classifier_type": "heuristic_v2",
"escalation_keywords": ["LITELLM ESCALATE"],
@ -23,16 +28,21 @@
},
"anthropic_family": {
"label": "Anthropic Family",
"description": "Routes across the Claude model family: Haiku for simple queries, Sonnet for medium, Opus for complex, Opus at high thinking for reasoning.",
"description": "Routes across the Claude model family: Haiku for simple queries, Sonnet for medium, Opus for complex, Fable 5.1 at high thinking for reasoning.",
"complexity_router_config": {
"tiers": {
"SIMPLE": ["claude-haiku-4-5"],
"MEDIUM": ["claude-sonnet-5"],
"COMPLEX": ["claude-opus-5"],
"REASONING": ["claude-opus-5"]
"REASONING": ["claude-fable-5-1"]
},
"tier_model_configs": {
"REASONING": [{ "model_name": "claude-opus-5", "litellm_params": { "reasoning_effort": "high" } }]
"REASONING": [
{
"model_name": "claude-fable-5-1",
"litellm_params": { "reasoning_effort": "high" }
}
]
},
"classifier_type": "heuristic",
"escalation_keywords": ["LITELLM ESCALATE"],
@ -73,8 +83,18 @@
"REASONING": ["claude-opus-5"]
},
"tier_model_configs": {
"MEDIUM": [{ "model_name": "muse-spark-1.2", "litellm_params": { "reasoning_effort": "xhigh" } }],
"COMPLEX": [{ "model_name": "kimi-k3", "litellm_params": { "reasoning_effort": "max" } }]
"MEDIUM": [
{
"model_name": "muse-spark-1.2",
"litellm_params": { "reasoning_effort": "xhigh" }
}
],
"COMPLEX": [
{
"model_name": "kimi-k3",
"litellm_params": { "reasoning_effort": "max" }
}
]
},
"classifier_type": "llm",
"classifier_llm_config": {
@ -93,16 +113,21 @@
},
"openai_family": {
"label": "OpenAI Family",
"description": "Routes across the GPT model family: Luna for simple queries, Terra for medium, Sol for complex, Sol at xhigh thinking for reasoning.",
"description": "Routes across the GPT model family: Luna for simple queries, Terra for medium, Sol for complex, Astra at xhigh thinking for reasoning.",
"complexity_router_config": {
"tiers": {
"SIMPLE": ["gpt-5.6-luna"],
"MEDIUM": ["gpt-5.6-terra"],
"COMPLEX": ["gpt-5.6-sol"],
"REASONING": ["gpt-5.6-sol"]
"REASONING": ["gpt-6-astra"]
},
"tier_model_configs": {
"REASONING": [{ "model_name": "gpt-5.6-sol", "litellm_params": { "reasoning_effort": "xhigh" } }]
"REASONING": [
{
"model_name": "gpt-6-astra",
"litellm_params": { "reasoning_effort": "xhigh" }
}
]
},
"classifier_type": "heuristic",
"escalation_keywords": ["LITELLM ESCALATE"],

32
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -173,6 +173,29 @@ def _raise_counter_budget_exceeded(
)
_UNBILLED_ROUTES: Final[frozenset[str]] = frozenset(
{
"/models",
"/v1/models",
"/utils/token_counter",
"/responses/input_tokens",
"/v1/responses/input_tokens",
"/openai/v1/responses/input_tokens",
}
)
_TOKEN_COUNTING_SEGMENTS: Final[frozenset[str]] = frozenset({"count_tokens", "count-tokens"})
_TOKEN_COUNTING_ACTION: Final = "countTokens"
def _is_token_counting_route(route: str) -> bool:
resource, _, action = route.rsplit("/", 1)[-1].partition(":")
return resource in _TOKEN_COUNTING_SEGMENTS or action == _TOKEN_COUNTING_ACTION
def _is_unbilled_route(route: str) -> bool:
return route in _UNBILLED_ROUTES or _is_token_counting_route(route)
async def reserve_budget_for_request(
request_body: dict,
route: str,
@ -190,14 +213,7 @@ async def reserve_budget_for_request(
) -> dict | None:
if valid_token is None or not RouteChecks.is_llm_api_route(route=route):
return None
if route in {
"/models",
"/v1/models",
"/utils/token_counter",
"/responses/input_tokens",
"/v1/responses/input_tokens",
"/openai/v1/responses/input_tokens",
}:
if _is_unbilled_route(route):
return None
if get_model_from_request(request_body, route, llm_router=llm_router) is None:
return None

183
litellm/proxy/spend_tracking/key_metadata_recovery.py
View file

@ -1,5 +1,7 @@
import asyncio
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Set as AbstractSet
from datetime import datetime, timedelta
from types import MappingProxyType
from typing import Final, TypeVar
@ -7,6 +9,13 @@ from pydantic import BaseModel, TypeAdapter
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import (
SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS,
SPEND_LOG_KEY_METADATA_CACHE_TTL,
SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL,
SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS,
)
from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash
from litellm.proxy.utils import PrismaClient
from litellm.repositories.user_repository import UserRepository
@ -27,6 +36,33 @@ WHERE encode(sha256(convert_to(token, 'UTF8')), 'hex') = ANY($1::text[])
ORDER BY token, deleted_at DESC
"""
_SPEND_LOG_ALIAS_SQL: Final = """
SELECT api_key AS digest,
MIN(key_alias) AS first_alias,
MAX(key_alias) AS last_alias,
MIN(team_id) AS first_team,
MAX(team_id) AS last_team,
MIN(user_id) AS first_owner,
MAX(user_id) AS last_owner
FROM (
SELECT api_key,
NULLIF(metadata->>'user_api_key_alias', '') AS key_alias,
COALESCE(NULLIF(team_id, ''), NULLIF(metadata->>'user_api_key_team_id', '')) AS team_id,
COALESCE(NULLIF("user", ''), NULLIF(metadata->>'user_api_key_user_id', '')) AS user_id
FROM "LiteLLM_SpendLogs"
WHERE api_key = ANY($1::text[])
AND "startTime" >= $2::timestamp
AND "startTime" < $3::timestamp
) named
WHERE COALESCE(key_alias, user_id, team_id) IS NOT NULL
GROUP BY api_key
"""
_SPEND_LOG_STATEMENT_TIMEOUT_SQL: Final = f"SET LOCAL statement_timeout = {SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS}"
_SPEND_LOG_TRANSACTION_TIMEOUT: Final = timedelta(milliseconds=2 * SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS)
_HASHED_JWT_PREFIX: Final = "hashed-jwt-"
class KeyMetadataDict(TypedDict, total=False):
key_alias: ReadOnly[str | None]
@ -42,7 +78,35 @@ class _TokenDigestRow(BaseModel):
user_id: str | None = None
def _unanimous(first: str | None, last: str | None) -> str | None:
return first if first == last else None
class _SpendLogDigestRow(BaseModel):
digest: str
first_alias: str | None = None
last_alias: str | None = None
first_team: str | None = None
last_team: str | None = None
first_owner: str | None = None
last_owner: str | None = None
def metadata(self) -> KeyMetadataDict:
return KeyMetadataDict(
key_alias=_unanimous(self.first_alias, self.last_alias),
team_id=_unanimous(self.first_team, self.last_team),
user_id=_unanimous(self.first_owner, self.last_owner),
)
_TOKEN_DIGEST_ROWS: Final = TypeAdapter(tuple[_TokenDigestRow, ...])
_SPEND_LOG_DIGEST_ROWS: Final = TypeAdapter(tuple[_SpendLogDigestRow, ...])
_CACHED_KEY_METADATA: Final = TypeAdapter(KeyMetadataDict)
_SPEND_LOG_METADATA_CACHE: Final = InMemoryCache(
max_size_in_memory=SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS,
default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL,
)
_SPEND_LOG_QUERY_LOCK: Final = asyncio.Lock()
_EMPTY_KEY_METADATA: Final[Mapping[str, KeyMetadataDict]] = MappingProxyType({})
_EMPTY_EMAILS: Final[Mapping[str, str]] = MappingProxyType({})
@ -138,14 +202,6 @@ async def recover_double_hashed_key_metadata(
prisma_client: PrismaClient,
missing_keys: AbstractSet[str],
) -> Mapping[str, KeyMetadataDict]:
"""
Recover key_alias/team_id/user_id for DailyUserSpend.api_key values that
were double-hashed by the v1.99 spend-log provenance gate.
Those rows store hash(VerificationToken.token) instead of the token, so the
exact join misses. Postgres hashes the token column itself, one pass over
active keys and one over deleted keys, so no key row crosses the wire.
"""
sha_missing: Final = frozenset(key for key in missing_keys if is_valid_sha256_hash(key))
if not sha_missing:
return _EMPTY_KEY_METADATA
@ -168,6 +224,117 @@ async def recover_double_hashed_key_metadata(
return MappingProxyType({**from_active, **from_deleted})
def _is_spend_log_digest(key: str) -> bool:
return is_valid_sha256_hash(key.removeprefix(_HASHED_JWT_PREFIX))
def _spend_log_cache_key(digest: str, window: tuple[datetime, datetime]) -> str:
start, end = window
return f"spend_log_key_metadata:{digest}:{start.isoformat()}:{end.isoformat()}"
def _cached_spend_log_metadata(
cache: InMemoryCache,
digests: AbstractSet[str],
window: tuple[datetime, datetime],
) -> Mapping[str, KeyMetadataDict]:
return MappingProxyType(
{
digest: _CACHED_KEY_METADATA.validate_python(cached)
for digest in digests
for cached in (cache.get_cache(_spend_log_cache_key(digest, window)),)
if cached is not None
}
)
async def _spend_log_rows_within_the_statement_timeout(
prisma_client: PrismaClient,
digests: AbstractSet[str],
window: tuple[datetime, datetime],
) -> Sequence[Mapping[str, object]]:
start, end = window
async with prisma_client.db.tx(timeout=_SPEND_LOG_TRANSACTION_TIMEOUT) as transaction:
await transaction.execute_raw(_SPEND_LOG_STATEMENT_TIMEOUT_SQL)
return await transaction.query_raw(_SPEND_LOG_ALIAS_SQL, sorted(digests), start, end)
async def _query_spend_log_metadata(
prisma_client: PrismaClient,
digests: AbstractSet[str],
window: tuple[datetime, datetime],
) -> Mapping[str, KeyMetadataDict] | None:
rows: Final = await _db_or_empty(
lambda: _spend_log_rows_within_the_statement_timeout(prisma_client, digests, window),
"Failed spend-log alias recovery for %d missing keys: %s",
len(digests),
)
if rows is None:
return None
return MappingProxyType(
{
row.digest: meta
for row in _SPEND_LOG_DIGEST_ROWS.validate_python(rows)
for meta in (row.metadata(),)
if row.digest in digests and any(meta.values())
}
)
def _remember_spend_log_metadata(
cache: InMemoryCache, digest: str, window: tuple[datetime, datetime], meta: KeyMetadataDict | None
) -> None:
key: Final = _spend_log_cache_key(digest, window)
if meta is not None:
cache.set_cache(key, meta)
return
missed_before: Final = f"{key}:missed-before"
if cache.get_cache(missed_before) is not None:
cache.set_cache(key, KeyMetadataDict())
return
cache.set_cache(key, KeyMetadataDict(), ttl=SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL)
cache.set_cache(missed_before, True)
async def _spend_log_metadata_one_query_at_a_time(
prisma_client: PrismaClient,
cache: InMemoryCache,
lock: asyncio.Lock,
digests: AbstractSet[str],
window: tuple[datetime, datetime],
) -> Mapping[str, KeyMetadataDict]:
async with lock:
settled: Final = _cached_spend_log_metadata(cache, digests, window)
pending: Final = digests - frozenset(settled)
fresh: Final = (
await _query_spend_log_metadata(prisma_client, pending, window) if pending else _EMPTY_KEY_METADATA
)
found: Final = fresh if fresh is not None else _EMPTY_KEY_METADATA
for digest in pending:
_remember_spend_log_metadata(cache, digest, window, found.get(digest))
return MappingProxyType({**settled, **found})
async def recover_key_metadata_from_spend_logs(
prisma_client: PrismaClient,
missing_keys: AbstractSet[str],
window: tuple[datetime, datetime],
cache: InMemoryCache = _SPEND_LOG_METADATA_CACHE,
lock: asyncio.Lock = _SPEND_LOG_QUERY_LOCK,
) -> Mapping[str, KeyMetadataDict]:
digests: Final = frozenset(key for key in missing_keys if _is_spend_log_digest(key))
if not digests:
return _EMPTY_KEY_METADATA
cached: Final = _cached_spend_log_metadata(cache, digests, window)
uncached: Final = digests - frozenset(cached)
settled: Final = (
await _spend_log_metadata_one_query_at_a_time(prisma_client, cache, lock, uncached, window)
if uncached
else _EMPTY_KEY_METADATA
)
return MappingProxyType({digest: meta for digest, meta in (*cached.items(), *settled.items()) if meta})
def _row_with_recovered_fields(
row: Mapping[str, object],
recovered: Mapping[str, KeyMetadataDict],

43
litellm/responses/utils.py
View file

@ -543,6 +543,49 @@ class ResponsesAPIRequestUtils:
return request_input
@staticmethod
def strip_encrypted_reasoning_from_input(request_input: object) -> None:
"""Drop reasoning items the routed deployment cannot decrypt, keeping their readable summary.
Mutates ``request_input`` in place: the router's fallback snapshot shares this
list object, so a rebound list would replay the stripped items on the fallback hop.
"""
if not isinstance(request_input, list):
return
items: Final = cast(list[object], request_input) # cast-ok: untyped client json
stripped: Final = tuple(ResponsesAPIRequestUtils._without_encrypted_reasoning(item) for item in items)
items[:] = (item for item in stripped if item is not None) # rebind-ok: list shared with fallback snapshot
@staticmethod
def _without_encrypted_reasoning(item: object) -> object | None:
if not isinstance(item, dict):
return item
reasoning: Final = cast(Mapping[str, object], item) # cast-ok: untyped client json
if reasoning.get("type") != "reasoning" or not reasoning.get("encrypted_content"):
return reasoning
readable: Final = any(
ResponsesAPIRequestUtils._has_readable_text(reasoning.get(key)) for key in ("summary", "content")
)
if not readable:
return None
kept: Final[dict[str, object]] = { # mutable-ok: request item rebuilt without the undecryptable keys
key: value for key, value in reasoning.items() if key not in ("encrypted_content", "id")
}
return kept
@staticmethod
def _has_readable_text(value: object) -> bool:
"""A reasoning item's ``summary``/``content`` carries readable text: a non-empty string, or a
list holding at least one block with a non-empty ``text`` field (summary_text / output_text)."""
if isinstance(value, str):
return bool(value.strip())
if isinstance(value, list):
return any(
isinstance(block, dict) and bool(cast(Mapping[str, object], block).get("text")) # cast-ok: untyped json
for block in value
)
return False
@staticmethod
def _build_responses_api_response_id(
custom_llm_provider: str | None,

81
litellm/router.py
View file

@ -1319,6 +1319,43 @@ class Router:
if isinstance(litellm.input_callback, list):
litellm.input_callback = [c for c in litellm.input_callback if id(c) not in selector_ids]
def _apply_updated_routing_strategy_args(self) -> None:
"""
Re-link the default group's selector to the current `routing_strategy_args`.
Selectors freeze their `RoutingArgs` at construction, so a runtime args
update would otherwise keep serving the boot-time values until restart.
Latency/usage state survives the rebuild: it lives in the shared router
cache, not on the selector.
"""
strategy: Final = self._normalize_strategy(self.routing_strategy)
if strategy == "lar1":
from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy
apply_lar1_routing_strategy(self, self.routing_strategy_args)
return
attr: Final = self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.get(strategy or "")
current: Final = getattr(self, attr, None) if attr is not None else None
if attr is None or current is None:
return
try:
rebuilt: Final = self._build_strategy_selector(
strategy=strategy or "",
routing_strategy_args=self.routing_strategy_args,
)
except (TypeError, ValidationError):
verbose_router_logger.exception(
"Invalid routing_strategy_args %s for '%s'; keeping the previous ones",
self.routing_strategy_args,
strategy,
)
return
self._unregister_router_selectors((current,))
setattr(self, attr, rebuilt)
def routing_strategy_init(self, routing_strategy: RoutingStrategy | str, routing_strategy_args: dict):
verbose_router_logger.info("Routing strategy: %s", routing_strategy)
self._validate_routing_strategy(routing_strategy)
@ -11115,6 +11152,40 @@ class Router:
return ids
def get_candidate_model_ids_for_route(self, model: str, team_id: str | None = None) -> frozenset[str]:
"""
Deployment ids that could serve ``model`` for ``team_id``, unioned across the paths
the router resolves a route through: ``model_group_alias``, a routing group, the
``model_name`` and team indexes, and wildcard pattern routes. Read-only and
side-effect-free, unlike ``_common_checks_available_deployment`` which also applies
fallbacks and can raise. Lets a pre-call check tell a genuine cross-group route from
same-group unavailability without re-deriving that precedence at the call site, and
without leaking deployment ids into request kwargs bound for the provider.
"""
resolved: Final = self._get_model_from_alias(model=model) or model
routing_group_members: Final = self._get_routing_group_deployments(model=resolved, team_id=team_id)
if routing_group_members is not None:
return self._deployment_ids(routing_group_members)
if resolved in self.model_names:
return self._deployment_ids(self._get_all_deployments(model_name=resolved, team_id=team_id))
team_router: Final = self.team_pattern_routers.get(team_id) if team_id is not None else None
return self._deployment_ids(
(
*self._get_all_deployments(model_name=resolved, team_id=team_id),
*(self.pattern_router.route(resolved) or ()),
*((team_router.route(resolved) or ()) if team_router is not None else ()),
)
)
@staticmethod
def _deployment_ids(deployments: Sequence[Mapping[str, object]]) -> frozenset[str]:
return frozenset(
str(model_info["id"])
for deployment in deployments
for model_info in (deployment.get("model_info"),)
if isinstance(model_info, Mapping) and model_info.get("id") is not None
)
def has_model_id(self, candidate_id: str) -> bool:
"""
O(1) membership check for a deployment ID without allocating large lists.
@ -11832,7 +11903,7 @@ class Router:
_existing_router_settings: Final = self.get_settings()
rebuild_routing_groups = False
relink_lar1_from_args = False
routing_args_updated = False
for var in kwargs:
if var in RUNTIME_UPDATABLE_ROUTER_SETTINGS:
if var in _int_settings:
@ -11871,15 +11942,13 @@ class Router:
)
rebuild_routing_groups = True
elif var == "routing_strategy_args":
relink_lar1_from_args = True
routing_args_updated = True
setattr(self, var, value)
else:
verbose_router_logger.debug("Setting %s is not allowed", var)
if relink_lar1_from_args and self._normalize_strategy(self.routing_strategy) == "lar1":
from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy
apply_lar1_routing_strategy(self, self.routing_strategy_args)
if routing_args_updated:
self._apply_updated_routing_strategy_args()
if rebuild_routing_groups:
self._init_routing_groups(self._routing_groups_input)

5
litellm/router_strategy/complexity_router/classification_rubrics.py
View file

@ -108,6 +108,11 @@ _CALIBRATION_EXAMPLES: Final[Mapping[ClassificationRubric, str]] = MappingProxyT
BUSINESS_TIER_CRITERIA: Final[Mapping[ComplexityTier, str]] = MappingProxyType(
{
ComplexityTier.NON_REASONING: (
"operational requests whose whole job is to pass information along or put it in a requested "
"shape: relaying or reformatting tool or system output, acknowledging a completed action, or "
"extracting a stated field. Use it only when no judgment about the content is asked for."
),
ComplexityTier.SIMPLE: (
"greetings, chitchat, or lookups of a fact, policy, price, or date with a short known answer. "
"Never for analysis, strategy, or non-trivial work, even if the request is only one sentence."

41
litellm/router_strategy/complexity_router/complexity_router.py
View file

@ -118,8 +118,20 @@ def _tier_name(tier: ComplexityTier | str) -> str:
return tier.value if isinstance(tier, ComplexityTier) else tier
def _built_in_tier_or_none(tier_name: str) -> ComplexityTier | None:
"""The built-in tier a `tiers` key names, or None when the key is an operator-defined name."""
return ComplexityTier.__members__.get(tier_name)
_CLASSIFICATION_TIER_CRITERIA: Final[Mapping[ComplexityTier, str]] = MappingProxyType(
{
ComplexityTier.NON_REASONING: (
"operational requests whose whole job is to pass information along or put it in a "
"requested shape: relaying or reformatting tool output, acknowledging a completed action, "
"or extracting a stated value. Use it only when no judgment about the content is asked for; "
"the moment the request is to summarize, compare, explain, debug, or decide, it belongs "
"in a higher tier however short it is."
),
ComplexityTier.SIMPLE: (
"greetings, chitchat, or factual lookups with a short known answer. Do not use this tier for "
"unsolved problems, proofs, deep theory, multi-step analysis, or non-trivial code, even if the "
@ -1244,7 +1256,7 @@ class ComplexityRouter(CustomLogger):
"""
if self.config.has_custom_tiers:
return tuple(dict.fromkeys(model for models in self._tier_pools().values() for model in models))
for tier in reversed(TIER_SEVERITY_ORDER):
for tier in reversed(self.config.active_tier_severity_order()):
models = self.config.tiers.get(tier.value)
if models:
return tuple(models) if isinstance(models, list) else (models,)
@ -1890,7 +1902,11 @@ class ComplexityRouter(CustomLogger):
default_model: Final = self.config.default_model
pools: Final = self._tier_pools()
tier: Final = next(
(candidate for candidate in TIER_SEVERITY_ORDER if default_model in pools.get(candidate.value, ())),
(
candidate
for candidate in self.config.active_tier_severity_order()
if default_model in pools.get(candidate.value, ())
),
ComplexityTier.MEDIUM,
)
return ClassificationOutcome(
@ -2279,7 +2295,8 @@ class ComplexityRouter(CustomLogger):
return self._fitting_tier_fallback(classified_tier, fit_filter)
request_type: Final = classify_prompt(user_message)
classified_idx: Final = TIER_SEVERITY_ORDER.index(classified_tier)
severity_order: Final = self.config.active_tier_severity_order()
classified_idx: Final = severity_order.index(classified_tier)
pools: Final = self._tier_pools()
classified_candidates: Final = _allowed(tuple(pools.get(_tier_name(classified_tier), ())), fit_filter)
cold_start_candidates: Final = tuple(
@ -2343,9 +2360,7 @@ class ComplexityRouter(CustomLogger):
distance = 0
else:
model_tiers = self._model_tiers.get(model, (classified_tier,))
distance = min(
abs(TIER_SEVERITY_ORDER.index(model_tier) - classified_idx) for model_tier in model_tiers
)
distance = min(abs(severity_order.index(model_tier) - classified_idx) for model_tier in model_tiers)
score = quality_weight * quality_sample + cost_weight * cost_score - penalty_weight * distance
candidate_scores.append(
{
@ -2660,10 +2675,15 @@ class ComplexityRouter(CustomLogger):
def _tier_for_model(self, model: str) -> ComplexityTier | None:
"""Return the most-severe configured tier whose pool contains this model."""
pools: Final = self._tier_pools()
matched: Final = tuple(ComplexityTier(tier_name) for tier_name, models in pools.items() if model in models)
order: Final = self.config.active_tier_severity_order()
matched: Final = tuple(
tier
for tier_name, models in pools.items()
if model in models and (tier := _built_in_tier_or_none(tier_name)) is not None and tier in order
)
if not matched:
return None
return max(matched, key=TIER_SEVERITY_ORDER.index)
return max(matched, key=order.index)
def _escalate_tier(self, tier: ComplexityTier | str) -> ComplexityTier | str:
"""Bump a tier one step up to the next-higher configured tier.
@ -2678,9 +2698,10 @@ class ComplexityRouter(CustomLogger):
if self.config.has_custom_tiers:
return tier
configured: Final = frozenset(self.config.tiers)
current_index: Final = TIER_SEVERITY_ORDER.index(tier)
order: Final = self.config.active_tier_severity_order()
current_index: Final = order.index(tier)
higher_tiers: Final = tuple(
candidate for candidate in TIER_SEVERITY_ORDER[current_index + 1 :] if candidate.value in configured
candidate for candidate in order[current_index + 1 :] if candidate.value in configured
)
return higher_tiers[0] if higher_tiers else tier

74
litellm/router_strategy/complexity_router/config.py
View file

@ -29,6 +29,7 @@ from .tier_predictor import TrainedTierArtifact
class ComplexityTier(str, Enum):
"""Complexity tiers for routing decisions."""
NON_REASONING = "NON_REASONING"
SIMPLE = "SIMPLE"
MEDIUM = "MEDIUM"
COMPLEX = "COMPLEX"
@ -62,6 +63,16 @@ TIER_SEVERITY_ORDER: Final[tuple[ComplexityTier, ...]] = (
ComplexityTier.REASONING,
)
NON_REASONING_TIER_SEVERITY_ORDER: Final[tuple[ComplexityTier, ...]] = (
ComplexityTier.NON_REASONING,
*TIER_SEVERITY_ORDER,
)
def tier_severity_order(non_reasoning_enabled: bool) -> tuple[ComplexityTier, ...]:
return NON_REASONING_TIER_SEVERITY_ORDER if non_reasoning_enabled else TIER_SEVERITY_ORDER
DEFAULT_TIER_DISTANCE_PENALTY: Final[float] = 0.5
DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE: Final[int] = 3
@ -142,6 +153,9 @@ def normalize_classification_examples(value: str | None) -> str | None:
return _normalize_operator_section(value, "classification_examples", MAX_CLASSIFICATION_EXAMPLES_CHARS)
_BUILT_IN_TIER_NAMES: Final[str] = ", ".join(ComplexityTier.__members__)
class TierDefinition(BaseModel):
"""An operator-defined tier: the name the LLM classifier must return and its rubric description."""
@ -152,7 +166,7 @@ class TierDefinition(BaseModel):
default=None,
description=(
"What belongs in this tier; rendered as this tier's bullet in the classifier rubric. "
"Required unless the name is a built-in tier (SIMPLE/MEDIUM/COMPLEX/REASONING), which "
f"Required unless the name is a built-in tier ({_BUILT_IN_TIER_NAMES}), which "
"inherits the built-in criteria when omitted"
),
)
@ -174,7 +188,7 @@ class TierDefinition(BaseModel):
if description is None and name.upper() not in ComplexityTier.__members__:
raise ValueError(
f"tier_definitions entry {name!r} must have a description: only the built-in tiers "
"(SIMPLE, MEDIUM, COMPLEX, REASONING) carry one the rubric can inherit"
f"({_BUILT_IN_TIER_NAMES}) carry one the rubric can inherit"
)
rendered_on_one_line: Final = (name, description or "")
if any("\n" in part or "\r" in part for part in rendered_on_one_line):
@ -711,6 +725,20 @@ class ComplexityRouterConfig(BaseModel):
default_factory=dict,
)
enable_non_reasoning_tier: bool = Field(
default=False,
description=(
"Add NON_REASONING as a fifth built-in tier below SIMPLE, for operational agent traffic "
"that relays or reformats information rather than reasoning about it. Off by default: "
"turning it on adds a rung to this router's ladder, a bullet to the LLM classifier's "
"rubric, and a value the classifier may return, all of which move tier decisions and "
"spend on an already-deployed router. Requires an LLM classifier or a custom classifier "
"plugin, since the heuristic scorers cannot produce the tier, and a model in `tiers` "
"under the NON_REASONING key. Escalation still walks up from it, and it is never the "
"savings baseline or a `heuristic_v2` prediction."
),
)
tier_definitions: tuple[TierDefinition, ...] | None = Field(
default=None,
description=(
@ -1514,11 +1542,15 @@ class ComplexityRouterConfig(BaseModel):
which still makes it a dependency on every one of those requests."""
return self.classifier_type in LLM_CLASSIFIER_TYPES
def active_tier_severity_order(self) -> tuple[ComplexityTier, ...]:
"""This router's built-in ladder, ascending; not meaningful for a custom tier set."""
return tier_severity_order(self.enable_non_reasoning_tier)
def tier_names(self) -> tuple[str, ...]:
"""The active tier names: the defined names, or the built-in set in severity order."""
if self.tier_definitions is not None:
return tuple(definition.name for definition in self.tier_definitions)
return tuple(tier.value for tier in TIER_SEVERITY_ORDER)
return tuple(tier.value for tier in self.active_tier_severity_order())
def classifier_wire_labels(self) -> tuple[str, ...]:
"""The tier names the classifier is told to emit: defined names, or the display labels."""
@ -1610,6 +1642,36 @@ class ComplexityRouterConfig(BaseModel):
if present
)
@model_validator(mode="after")
def _validate_non_reasoning_tier(self) -> "ComplexityRouterConfig":
"""Require a classifier that can emit the opt-in tier and a model to route it to."""
non_reasoning_key: Final = ComplexityTier.NON_REASONING.value
if not self.enable_non_reasoning_tier:
if not self.has_custom_tiers and non_reasoning_key in self.tiers:
raise ValueError(
f"tiers names {non_reasoning_key} but enable_non_reasoning_tier is False, so no request "
"can route there; set enable_non_reasoning_tier: true or drop the tier"
)
return self
if self.has_custom_tiers:
raise ValueError(
"enable_non_reasoning_tier cannot be combined with tier_definitions: a custom tier set "
f"replaces the built-in ladder, so name a tier {non_reasoning_key} in tier_definitions instead"
)
if self.classifier_type not in ("llm", "custom"):
raise ValueError(
f"enable_non_reasoning_tier requires classifier_type 'llm' or 'custom', got "
f"{self.classifier_type!r}: the heuristic scorers only produce the four tiers from SIMPLE up, "
f"so nothing would ever classify as {non_reasoning_key}"
)
if not self.tiers.get(non_reasoning_key):
raise ValueError(
f"enable_non_reasoning_tier requires tiers to map {non_reasoning_key} to at least one model: "
"the tier exists to send operational traffic somewhere cheaper, and an unconfigured tier "
"would fall through to the default model"
)
return self
@model_validator(mode="after")
def _validate_tier_definitions(self) -> "ComplexityRouterConfig":
if self.tier_definitions is None:
@ -1632,7 +1694,7 @@ class ComplexityRouterConfig(BaseModel):
if self.classifier_type in ("heuristic", "heuristic_v2", "heuristic_first", "hybrid"):
raise ValueError(
"tier_definitions requires classifier_type 'llm' or 'custom': the heuristic scorer only "
"produces the four built-in tiers, as does heuristic_v2"
"produces the built-in tiers from SIMPLE up, as does heuristic_v2"
)
conflicts: Final = self._tier_definition_conflicts()
if conflicts:
@ -1786,7 +1848,7 @@ class ComplexityRouterConfig(BaseModel):
def labeled_tiers(self) -> tuple[tuple[ComplexityTier, str], ...]:
"""Every tier paired with its display name, in ascending severity order."""
return tuple((tier, self.tier_label(tier)) for tier in TIER_SEVERITY_ORDER)
return tuple((tier, self.tier_label(tier)) for tier in self.active_tier_severity_order())
def tier_for_label(self, label: str) -> ComplexityTier | None:
"""Resolve a display name back to its tier, case-insensitively, then canonical names."""
@ -1794,7 +1856,7 @@ class ComplexityRouterConfig(BaseModel):
labeled: Final = self.labeled_tiers()
return next(
(tier for tier, tier_label in labeled if tier_label.casefold() == folded),
next((tier for tier in TIER_SEVERITY_ORDER if tier.value.casefold() == folded), None),
next((tier for tier, _ in labeled if tier.value.casefold() == folded), None),
)

21
litellm/router_strategy/lowest_latency.py
View file

@ -3,8 +3,11 @@
import random
from collections.abc import Sequence
from datetime import datetime, timedelta
from math import ceil
from typing import TYPE_CHECKING, Any, Final
from pydantic import Field
import litellm
from litellm import ModelResponse, token_counter, verbose_logger
from litellm.caching.caching import DualCache
@ -24,6 +27,7 @@ class RoutingArgs(LiteLLMPydanticObjectBase):
ttl: float = 1 * 60 * 60 # 1 hour
lowest_latency_buffer: float = 0
max_latency_list_size: int = 10
ttft_percentile: float | None = Field(default=None, gt=0, le=1)
def _average_latency(samples: Sequence[float]) -> float:
@ -32,6 +36,12 @@ def _average_latency(samples: Sequence[float]) -> float:
return sum(samples) / len(samples)
def _percentile_latency(samples: Sequence[float], percentile: float) -> float:
values: Final = sorted(samples)
index: Final = ceil(len(values) * percentile) - 1
return values[index]
def _ttft_seconds(elapsed: timedelta | float) -> float:
if isinstance(elapsed, timedelta):
return elapsed.total_seconds()
@ -427,14 +437,17 @@ class LowestLatencyLoggingHandler(CustomLogger):
item_rpm = item_map.get(precise_minute, {}).get("rpm", 0)
item_tpm = item_map.get(precise_minute, {}).get("tpm", 0)
# get average latency or average ttft (depending on streaming/non-streaming)
use_ttft = (
request_kwargs is not None
and request_kwargs.get("stream", None) is not None
and request_kwargs["stream"] is True
and len(item_ttft_latency) > 0
)
average_latency = _average_latency(item_ttft_latency if use_ttft else item_latency)
selected_latency = (
_percentile_latency(item_ttft_latency, self.routing_args.ttft_percentile)
if use_ttft and self.routing_args.ttft_percentile is not None
else _average_latency(item_ttft_latency if use_ttft else item_latency)
)
# -------------- #
# Debugging Logic
@ -443,7 +456,7 @@ class LowestLatencyLoggingHandler(CustomLogger):
# this helps a user to debug why the router picked a specfic deployment #
_deployment_api_base = _deployment.get("litellm_params", {}).get("api_base", "")
if _deployment_api_base is not None:
_latency_per_deployment[_deployment_api_base] = average_latency
_latency_per_deployment[_deployment_api_base] = selected_latency
# -------------- #
# End of Debugging Logic
# -------------- #
@ -453,7 +466,7 @@ class LowestLatencyLoggingHandler(CustomLogger):
): # if user passed in tpm / rpm in the model_list
continue
else:
potential_deployments.append((_deployment, average_latency))
potential_deployments.append((_deployment, selected_latency))
if len(potential_deployments) == 0:
return None

73
litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py
View file

@ -37,13 +37,13 @@ Safe to enable globally:
"""
import time
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final, Optional, Protocol, cast
import httpx
from litellm._logging import verbose_router_logger
from litellm.exceptions import (
BadRequestError,
RateLimitError,
ServiceUnavailableError,
)
@ -158,6 +158,23 @@ class EncryptedContentAffinityCheck(CustomLogger):
return deployment
return None
@staticmethod
def _request_team_id(request_kwargs: Mapping[str, object]) -> str | None:
containers: Final = (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata"))
team_ids: Final = (c.get("user_api_key_team_id") for c in containers if isinstance(c, Mapping))
return next((tid for tid in team_ids if isinstance(tid, str)), None)
def _routed_group_candidate_model_ids(self, request_kwargs: Mapping[str, object], model: str) -> frozenset[str]:
"""
Deployment ids that could serve this turn's routed ``model``, as the router
resolves a route (model_group_alias / routing group / model_name / team /
pattern). Delegates to the router so the full precedence is not re-derived here
and no deployment ids are written into request kwargs bound for the provider.
"""
if self.router is None:
return frozenset()
return self.router.get_candidate_model_ids_for_route(model=model, team_id=self._request_team_id(request_kwargs))
@staticmethod
def _encryption_boundary_key(
litellm_params: object,
@ -225,10 +242,14 @@ class EncryptedContentAffinityCheck(CustomLogger):
"""
If the request ``input`` contains litellm-encoded item IDs, decode the
embedded ``model_id`` and pin the request to that deployment. Raises
``RateLimitError`` / ``ServiceUnavailableError`` / ``BadRequestError``
when the originating deployment is unavailable and no encryption-boundary
peer exists, rather than dispatching a doomed request to a non-peer
deployment. The 429/503 split mirrors the originating cooldown's status:
``RateLimitError`` / ``ServiceUnavailableError`` when the originating
deployment is a member of the routed model group but currently unavailable
and no encryption-boundary peer exists, rather than dispatching a doomed
request to a non-peer deployment. When the origin is not a member of the
routed group (an auto-router tier change, a model switch with no peer, a
removed deployment, or an unknown/forged marker), the encrypted reasoning is
stripped and the request dispatches with its readable history instead. The
429/503 split mirrors the originating cooldown's status:
a 429-induced cooldown surfaces as 429 (with ``Retry-After`` set to the
remaining cooldown window) so OpenAI-compatible clients back off and
retry after the deployment is eligible again.
@ -285,12 +306,34 @@ class EncryptedContentAffinityCheck(CustomLogger):
request_kwargs["_encrypted_content_affinity_pinned"] = True
return boundary_matches
# Dispatching to a non-peer would guarantee an upstream
# `invalid_encrypted_content` 400, so fail fast with a clearer error.
# The origin cannot serve this turn's routed group and no peer shares the boundary, so its
# encrypted reasoning can never decrypt here. Strip it, keep the readable history, and dispatch
# to the routed group instead of failing. Membership is tested by deployment id against the set
# the router actually resolved for this route, not by model-group name, so an alias, a
# provider-qualified spelling, a team-public name, or a pattern route of the same group is not
# mistaken for a tier change. An unknown origin (a removed deployment, or a forged marker) is
# treated the same as a cross-group one, which also denies an authenticated caller a
# deployment-id existence oracle: a real cross-group id and a nonexistent id both strip and
# dispatch rather than returning distinguishable responses. Only a genuine same-group member
# that is currently unavailable falls through to the fail-fast, preserving the cooldown contract.
routed_group_model_ids: Final = (
self._routed_group_candidate_model_ids(request_kwargs, model) if originating is not None else frozenset()
)
if str(model_id) not in routed_group_model_ids:
verbose_router_logger.debug(
"EncryptedContentAffinityCheck: model_id=%s is not a candidate for the routed group %s; "
"forwarding without its encrypted reasoning",
model_id,
model,
)
ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input)
return typed_healthy_deployments
# The origin is a member of the routed group but currently unavailable (cooled down); fail fast
# rather than dispatching to a non-peer, which would guarantee an upstream 400.
raise await self._unavailable_origin_error(
model=model,
model_id=model_id,
originating=originating,
parent_otel_span=parent_otel_span,
)
@ -298,25 +341,11 @@ class EncryptedContentAffinityCheck(CustomLogger):
self,
model: str,
model_id: str,
originating: Deployment | None,
parent_otel_span: Span | None,
) -> Exception:
# Public error messages intentionally omit the originating ``model_id`` so
# an authenticated caller forging encrypted-content markers cannot use the
# error surface to enumerate which deployment IDs exist on this router.
if originating is None:
return BadRequestError(
message=(
"The deployment that produced this encrypted_content is no "
"longer configured on this router, and no deployment on the "
"same encryption boundary is available. Re-issue the request "
"without the stale encrypted_content items, or restore the "
"originating deployment."
),
model=model,
llm_provider="",
)
cooldown: Final = await self._get_origin_cooldown(model_id=model_id, parent_otel_span=parent_otel_span)
if cooldown is not None and str(cooldown.get("status_code")) == "429":

2
litellm/utils.py
View file

@ -3079,7 +3079,7 @@ def register_model(
# Convert stringified numbers to appropriate numeric types
loaded_model_cost = model_cost
elif isinstance(model_cost, str):
loaded_model_cost = litellm.get_model_cost_map(url=model_cost)
loaded_model_cost = litellm.get_model_cost_map(url=model_cost, max_attempts=1)
if persist_across_reloads:
_registrations: Final[Mapping[str, Mapping[str, object]]] = loaded_model_cost

136
model_prices_and_context_window.json
View file

@ -3626,6 +3626,79 @@
"supports_xhigh_reasoning_effort": true,
"supports_minimal_reasoning_effort": false
},
"azure_ai/gpt-chat-latest": {
"cache_read_input_token_cost": 5e-07,
"deprecation_date": "2026-12-02",
"input_cost_per_token": 5e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 3e-05,
"source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"azure_ai/codex-mini": {
"cache_read_input_token_cost": 3.75e-07,
"deprecation_date": "2026-11-15",
"input_cost_per_token": 1.5e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "responses",
"output_cost_per_token": 6e-06,
"source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/",
"supported_endpoints": [
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true
},
"azure_ai/whisper": {
"deprecation_date": "2026-12-15",
"input_cost_per_second": 0.0001,
"litellm_provider": "azure_ai",
"mode": "audio_transcription",
"output_cost_per_second": 0.0001,
"source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/"
},
"azure_ai/gpt-5.5-2026-04-23": {
"cache_read_input_token_cost": 5e-07,
"cache_read_input_token_cost_above_272k_tokens": 1e-06,
@ -4029,13 +4102,29 @@
"supports_minimal_reasoning_effort": false
},
"azure_ai/model_router": {
"deprecation_date": "2027-05-20",
"input_cost_per_token": 1.4e-07,
"output_cost_per_token": 0,
"litellm_provider": "azure_ai",
"max_input_tokens": 200000,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/aoai/",
"comment": "Flat cost of $0.14 per M input tokens for Azure AI Foundry Model Router infrastructure. Use pattern: azure_ai/model_router/<deployment-name> where deployment-name is your Azure deployment (e.g., azure-model-router)"
},
"azure_ai/model-router": {
"deprecation_date": "2027-05-20",
"input_cost_per_token": 1.4e-07,
"output_cost_per_token": 0,
"litellm_provider": "azure_ai",
"max_input_tokens": 200000,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/aoai/",
"comment": "Catalog-name twin of azure_ai/model_router: the flat $0.14 per M input tokens is the router's own fee, the routed model is priced on top of it"
},
"azure/eu/gpt-4o-2024-08-06": {
"deprecation_date": "2027-04-14",
"cache_read_input_token_cost": 1.375e-06,
@ -10347,6 +10436,18 @@
"/v1/ocr"
]
},
"azure_ai/cohere-command-a": {
"input_cost_per_token": 2.5e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 131072,
"max_output_tokens": 8182,
"max_tokens": 8182,
"mode": "chat",
"output_cost_per_token": 1e-05,
"source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/cohere/",
"supports_function_calling": true,
"supports_tool_choice": true
},
"azure_ai/doc-intelligence/prebuilt-read": {
"litellm_provider": "azure_ai",
"ocr_cost_per_page": 0.0015,
@ -10698,6 +10799,41 @@
"supports_vision": true,
"supports_web_search": true
},
"azure_ai/grok-4-20-reasoning": {
"cache_read_input_token_cost": 1.25e-06,
"deprecation_date": "2027-04-06",
"input_cost_per_token": 1.25e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 262000,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/grok/",
"supports_function_calling": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true,
"supports_reasoning": true
},
"azure_ai/grok-4-20-non-reasoning": {
"cache_read_input_token_cost": 1.25e-06,
"deprecation_date": "2027-04-06",
"input_cost_per_token": 1.25e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 262000,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 2.5e-06,
"source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/grok/",
"supports_function_calling": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"azure_ai/grok-4-fast-non-reasoning": {
"deprecation_date": "2026-05-01",
"input_cost_per_token": 2e-07,

2
osv-scanner.toml
View file

@ -1,6 +1,6 @@
[[IgnoredVulns]]
id = "GHSA-w8v5-vhqr-4h9v"
ignoreUntil = 2026-09-09
ignoreUntil = 2026-10-01
reason = "diskcache has no fixed release published; remove this entry once one exists"
[[IgnoredVulns]]

2
pyproject.toml
View file

@ -68,7 +68,7 @@ proxy = [
"azure-storage-blob>=12.28.0,<13.0",
"mcp>=1.28.1,<2.0",
"litellm-proxy-extras==0.4.95",
"litellm-enterprise==0.1.65",
"litellm-enterprise==0.1.66",
"RestrictedPython>=8.5,<9.0",
"rich>=13.9.4,<14.0",
"InquirerPy>=0.3.4,<1.0",

1
tests/code_coverage_tests/router_code_coverage.py
View file

@ -87,6 +87,7 @@ ignored_function_names = [
"_delete_claude_code_session_router_binding", # Tested through Redis cleanup failure in test_router.py
"_resolve_claude_code_session_router", # Tested through Claude Code session routing in test_router.py
"_get_claude_code_session_router_binding", # Tested through the two-worker session routing test in test_router.py
"_apply_updated_routing_strategy_args", # Tested via update_settings in test_lowest_latency.py (file lacks "router" in name)
]

30
tests/e2e/batches/COVERAGE.md
View file

@ -120,6 +120,36 @@ create traverse gateway -> gateway -> OpenAI (LIT-5347, PR #36240). The pin:
nested managed ids round-trip retrieve. This self-chaining only needs the proxy to
reach its own `PROXY_BASE_URL`, which holds both locally and on the e2e stage.
## Cleanup
Batch teardown cancels active batches before deleting their input files and keys.
Raw file IDs from both `model_param` and `provider_fallback` uploads use the upload
provider when deleted. Model-encoded and managed file IDs route themselves
File deletion and batch cancellation check their responses and retry transient
failures up to three times. Teardown attempts every registered cleanup before
reporting failures as test errors. Already deleted files and batches that are
terminal are safe to clean up again. Managed batch cancellation polls for up to eleven minutes
before input deletion: the ten-minute provider window plus a propagation margin.
Accepted cancellation may still report validating or in_progress while the provider
updates its state. Raw and model-encoded batches are polled until cancelling or
terminal before input deletion. OpenAI and Azure lifecycle cleanup also deletes
output and error files returned by terminal batches. Bedrock deletion uses a signed S3 DELETE
restricted to the configured storage buckets and managed file prefixes. The low-RPM
test submits with its restricted key and cleans up with the test administrator key
Managed deletion forwards the deployment's trusted bucket configuration and returns
the requested managed file ID even when stored output metadata carries a provider ID
Azure input uploads request `expires_after` anchored to `created_at` with
`seconds=1209600`, and the lifecycle tests check the returned expiry. This is a
fallback for interrupted runs: immediate deletion remains the normal cleanup.
Azure's minimum supported native expiry is 14 days, so a three-day expiry cannot
be requested through its Files API
The Azure entry in `files_settings` must use `api_version: 2025-04-01-preview`
for raw uploads to honor expiry, matching the batch deployment's API version
## Terminal state + cost write-back (cross-run marker baton)
The 24h completion window rules out submit-and-wait inside one run, so

140
tests/e2e/batches/batch_cleanup.py Normal file
View file

@ -0,0 +1,140 @@
from builtins import ExceptionGroup
from collections.abc import Callable
from itertools import count
from time import monotonic, sleep
from typing import Final, Protocol
from batch_client import BatchObject, FileDeleteResponse
from capabilities import is_managed_id
from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError
from pydantic import BaseModel
CLEANUP_DELAYS: Final = (1.0, 2.0, 4.0)
BATCH_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "expired", "cancelled"})
BATCH_PENDING_STATUSES: Final = frozenset({"validating", "in_progress", "finalizing", "cancelling"})
BATCH_CANCEL_TIMEOUT_SECONDS: Final = 660.0
BATCH_CANCEL_POLL_SECONDS: Final = 10.0
class BatchCleanupClient(Protocol):
def delete_file(self, file_id: str, *, key: str, provider: str | None = None) -> Result[FileDeleteResponse]: ...
def retrieve_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: ...
def cancel_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: ...
def cleanup_result[R: BaseModel](
action: Callable[[], Result[R]], *, wait: Callable[[float], None] = sleep
) -> Result[R]:
for delay, result in ((delay, action()) for delay in CLEANUP_DELAYS):
match result:
case NetworkError() | RateLimitedError():
wait(delay)
case UnknownApiError(status_code=code) if code in {408, 429, 500, 502, 503, 504}:
wait(delay)
case _:
return result
return action()
def _require_cleanup_success[R: BaseModel](result: Result[R], operation: str) -> R:
match result:
case Success(data=data):
return data
case UnknownApiError(status_code=code):
raise AssertionError(f"{operation} failed: HTTP {code}")
case _:
raise AssertionError(f"{operation} failed: {result.kind}")
def cleanup_file(client: BatchCleanupClient, file_id: str, *, key: str, provider: str | None = None) -> None:
result: Final = cleanup_result(lambda: client.delete_file(file_id, key=key, provider=provider))
if isinstance(result, UnknownApiError) and result.status_code == 404:
return
deleted: Final = _require_cleanup_success(result, f"Delete file {file_id}")
assert deleted.deleted is True or (
deleted.deleted is None and is_managed_id(file_id) and deleted.id == file_id and deleted.object == "file"
), f"Delete file {file_id} did not confirm deletion"
def cleanup_batch(
client: BatchCleanupClient,
batch_id: str,
*,
key: str,
provider: str | None = None,
delete_output_files: bool = False,
wait: Callable[[float], None] = sleep,
clock: Callable[[], float] = monotonic,
) -> None:
needs_terminal_state: Final = is_managed_id(batch_id)
fetched: Final = _require_cleanup_success(
cleanup_result(lambda: client.retrieve_batch(batch_id, key=key, provider=provider)),
f"Retrieve batch {batch_id} for cleanup",
)
if fetched.status in BATCH_TERMINAL_STATUSES:
if delete_output_files:
_cleanup_batch_outputs(client, fetched, key=key, provider=provider)
return
if fetched.status == "cancelling" and not needs_terminal_state:
return
result: Final = (
Success(status_code=200, data=fetched)
if fetched.status == "cancelling"
else cleanup_result(lambda: client.cancel_batch(batch_id, key=key, provider=provider))
)
conflicted: Final = isinstance(result, UnknownApiError) and result.status_code in {400, 409}
if not conflicted:
cancelled: Final = _require_cleanup_success(result, f"Cancel batch {batch_id}")
assert cancelled.status in BATCH_TERMINAL_STATUSES | BATCH_PENDING_STATUSES, (
f"Cancel batch {batch_id} left status {cancelled.status}"
)
if cancelled.status in BATCH_TERMINAL_STATUSES:
if delete_output_files:
_cleanup_batch_outputs(client, cancelled, key=key, provider=provider)
return
if cancelled.status == "cancelling" and not needs_terminal_state:
return
deadline: Final = clock() + BATCH_CANCEL_TIMEOUT_SECONDS
for current in (
_require_cleanup_success(
cleanup_result(lambda: client.retrieve_batch(batch_id, key=key, provider=provider)),
f"Retrieve batch {batch_id} after cancellation",
)
for _ in count()
):
if current.status in BATCH_TERMINAL_STATUSES:
if delete_output_files:
_cleanup_batch_outputs(client, current, key=key, provider=provider)
return
assert current.status in ({"cancelling"} if conflicted else BATCH_PENDING_STATUSES), (
f"Cancel batch {batch_id} left status {current.status}"
)
if current.status == "cancelling" and not needs_terminal_state:
return
assert clock() < deadline, (
f"Batch {batch_id} cancellation did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s"
)
wait(BATCH_CANCEL_POLL_SECONDS)
def _cleanup_batch_outputs(client: BatchCleanupClient, batch: BatchObject, *, key: str, provider: str | None) -> None:
errors: Final = tuple(
error
for file_id in dict.fromkeys((batch.output_file_id, batch.error_file_id))
if file_id is not None and file_id != batch.input_file_id
if (error := _output_cleanup_error(client, file_id, key=key, provider=provider)) is not None
)
if errors:
raise ExceptionGroup(f"Batch {batch.id} output cleanup failed", errors)
def _output_cleanup_error(
client: BatchCleanupClient, file_id: str, *, key: str, provider: str | None
) -> Exception | None:
try:
cleanup_file(client, file_id, key=key, provider=provider)
except Exception as error:
return error
return None

18
tests/e2e/batches/batch_client.py
View file

@ -13,8 +13,9 @@ co-located here because only this suite uses them.
from __future__ import annotations
from dataclasses import dataclass
from typing import Final, Literal
from pydantic import BaseModel
from pydantic import BaseModel, Field
from proxy_client import ProxyClient
from e2e_http import (
@ -27,6 +28,18 @@ from e2e_http import (
from models import LiteLLMParamsBody
UPLOAD_FILENAME = "batch_input.jsonl"
AZURE_FILE_EXPIRY_SECONDS: Final = 14 * 24 * 60 * 60
class ExpiringFileUploadForm(FileUploadForm):
expires_after_anchor: Literal["created_at"] = Field(default="created_at", alias="expires_after[anchor]")
expires_after_seconds: int = Field(default=AZURE_FILE_EXPIRY_SECONDS, alias="expires_after[seconds]")
def batch_upload_form(provider: str, *, target_model_names: str | None = None) -> FileUploadForm:
if provider == "azure":
return ExpiringFileUploadForm(target_model_names=target_model_names)
return FileUploadForm(target_model_names=target_model_names)
class FileObject(BaseModel):
@ -37,6 +50,7 @@ class FileObject(BaseModel):
bytes: int | None = None
status: str | None = None
created_at: int | None = None
expires_at: int | None = None
class FileList(BaseModel):
@ -85,7 +99,7 @@ class BatchList(BaseModel):
class FileDeleteResponse(BaseModel):
id: str
object: str | None = None
deleted: bool
deleted: bool | None = None
class BatchCreateBody(BaseModel):

4
tests/e2e/batches/capabilities.py
View file

@ -108,6 +108,10 @@ class Capability:
def id(self) -> str:
return f"{self.provider}-{self.scenario}"
@property
def file_provider(self) -> str | None:
return self.provider if self.scenario in {"model_param", "provider_fallback"} else None
@property
def jsonl_model(self) -> str:
# Always the provider deployment name. Unified routes via

10
tests/e2e/batches/conftest.py
View file

@ -13,7 +13,7 @@ the proxy config.
from __future__ import annotations
import os
from typing import Iterator
from typing import Final, Iterator
import pytest
@ -21,6 +21,7 @@ from batch_client import BatchClient, build_client
from capabilities import PROVIDERS
from e2e_config import MANAGED_FILES_OPT_IN_ENV
from e2e_http import NoBody
from lifecycle import ResourceManager
from proxy_client import ProxyClient
@ -52,6 +53,13 @@ def client(proxy: ProxyClient) -> BatchClient:
return build_client(proxy)
@pytest.fixture
def resources(client: BatchClient) -> Iterator[ResourceManager]:
manager: Final = ResourceManager(client=client.proxy, strict_cleanup=True)
yield manager
manager.teardown()
@pytest.fixture(scope="session")
def batch_deployments(client: BatchClient) -> Iterator[None]:
probe = client.proxy.probe("/health/liveliness", params=NoBody())

313
tests/e2e/batches/test_batch_cleanup.py Normal file
View file

@ -0,0 +1,313 @@
from builtins import ExceptionGroup
from collections.abc import Callable
from typing import Final
from unittest.mock import Mock, call
import pytest
from batch_cleanup import BATCH_CANCEL_TIMEOUT_SECONDS, CLEANUP_DELAYS, cleanup_batch, cleanup_file, cleanup_result
from batch_client import AZURE_FILE_EXPIRY_SECONDS, BatchObject, FileDeleteResponse, batch_upload_form
from capabilities import CAPABILITIES, Capability
from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError
from lifecycle import ResourceManager
from models import KeyGenerateBody
MANAGED_FILE_ID: Final = "bGl0ZWxsbV9wcm94eTtmaWxlLTE="
MANAGED_BATCH_ID: Final = "bGl0ZWxsbV9wcm94eTtiYXRjaC0x"
class ExpectedCalls[T]:
def __init__(self, values: tuple[T, ...]) -> None:
self.values: Final = values
self.recorder: Final = Mock()
def __call__(self, value: T) -> None:
self.recorder(value)
def assert_done(self) -> None:
assert tuple(self.recorder.call_args_list) == tuple(call(value) for value in self.values)
class CleanupClient:
def __init__(
self,
*,
calls: ExpectedCalls[str],
files: tuple[Result[FileDeleteResponse], ...] = (),
batches: tuple[Result[BatchObject], ...] = (),
cancellations: tuple[Result[BatchObject], ...] = (),
) -> None:
self.calls: Final = calls
self.file_response: Final[Callable[[], Result[FileDeleteResponse]]] = Mock(side_effect=files)
self.batch_response: Final[Callable[[], Result[BatchObject]]] = Mock(side_effect=batches)
self.cancel_response: Final[Callable[[], Result[BatchObject]]] = Mock(side_effect=cancellations)
def delete_file(self, file_id: str, *, key: str, provider: str | None = None) -> Result[FileDeleteResponse]:
self.calls(f"delete {provider} {file_id}")
return self.file_response()
def retrieve_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]:
self.calls(f"retrieve {provider} {batch_id}")
return self.batch_response()
def cancel_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]:
self.calls(f"cancel {provider} {batch_id}")
return self.cancel_response()
def generate_key(self, body: KeyGenerateBody) -> str:
return "test-key"
def delete_key(self, key: str) -> None:
self.calls(f"delete key {key}")
def delete_customers(self, user_ids: list[str]) -> None:
self.calls(f"delete customers {user_ids}")
def batch(status: str) -> Success[BatchObject]:
return Success(status_code=200, data=BatchObject(id="batch-1", status=status))
def deleted_file(*, deleted: bool = True) -> Success[FileDeleteResponse]:
return Success(status_code=200, data=FileDeleteResponse(id="file-1", deleted=deleted))
class TestFileCleanup:
def test_managed_delete_accepts_the_deleted_file_object(self) -> None:
response: Final = Success(
status_code=200, data=FileDeleteResponse.model_validate({"id": MANAGED_FILE_ID, "object": "file"})
)
client: Final = CleanupClient(calls=ExpectedCalls((f"delete None {MANAGED_FILE_ID}",)), files=(response,))
cleanup_file(client, MANAGED_FILE_ID, key="test-key")
client.calls.assert_done()
@pytest.mark.parametrize("file_id", ["file-1", MANAGED_FILE_ID])
def test_a_success_status_without_a_deletion_confirmation_is_rejected(self, file_id: str) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls((f"delete None {file_id}",)),
files=(Success(status_code=200, data=FileDeleteResponse(id=file_id)),),
)
with pytest.raises(AssertionError, match="did not confirm deletion"):
cleanup_file(client, file_id, key="test-key")
client.calls.assert_done()
@pytest.mark.parametrize("cap", CAPABILITIES, ids=[cap.id for cap in CAPABILITIES])
def test_deletes_raw_files_through_the_upload_provider(self, cap: Capability) -> None:
expected_provider: Final = cap.provider if cap.scenario in {"model_param", "provider_fallback"} else None
client: Final = CleanupClient(
calls=ExpectedCalls((f"delete {expected_provider} file-1",)), files=(deleted_file(),)
)
cleanup_file(client, "file-1", key="test-key", provider=cap.file_provider)
client.calls.assert_done()
def test_failed_delete_is_reported_after_remaining_resources_are_cleaned(self) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(("delete azure file-1", "delete key test-key")),
files=(UnknownApiError(status_code=403, body="secret response"),),
)
manager: Final = ResourceManager(client=client, strict_cleanup=True)
key: Final = manager.key()
manager.defer(lambda: cleanup_file(client, "file-1", key=key, provider="azure"))
with pytest.raises(ExceptionGroup) as caught:
manager.teardown()
client.calls.assert_done()
assert len(caught.value.exceptions) == 1
assert str(caught.value.exceptions[0]) == "Delete file file-1 failed: HTTP 403"
def test_success_response_must_confirm_deletion(self) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(("delete None file-1",)), files=(deleted_file(deleted=False),)
)
with pytest.raises(AssertionError, match="did not confirm deletion"):
cleanup_file(client, "file-1", key="test-key")
client.calls.assert_done()
def test_cleanup_is_idempotent_when_file_is_already_deleted(self) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(("delete azure file-1",)),
files=(UnknownApiError(status_code=404, body="missing"),),
)
cleanup_file(client, "file-1", key="test-key", provider="azure")
client.calls.assert_done()
def test_default_resource_cleanup_keeps_existing_best_effort_behavior(self) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(("delete None file-1", "delete key test-key")),
files=(UnknownApiError(status_code=403, body="forbidden"),),
)
manager: Final = ResourceManager(client=client)
key: Final = manager.key()
manager.defer(lambda: cleanup_file(client, "file-1", key=key))
manager.teardown()
client.calls.assert_done()
class TestCleanupRetries:
@pytest.mark.parametrize(
"failure",
[NetworkError(message="offline"), RateLimitedError(), UnknownApiError(status_code=503, body="unavailable")],
)
def test_transient_error_retries_and_returns_success(self, failure: Result[FileDeleteResponse]) -> None:
responses: Final = (failure, deleted_file())
outcomes: Final = Mock(side_effect=responses)
delays: Final = ExpectedCalls((1.0,))
result: Final[Result[FileDeleteResponse]] = cleanup_result(outcomes, wait=delays)
assert isinstance(result, Success) and result.data.deleted
delays.assert_done()
def test_persistent_error_has_bounded_retries(self) -> None:
failure: Final = UnknownApiError(status_code=503, body="unavailable")
outcomes: Final = Mock(return_value=failure)
delays: Final = ExpectedCalls(CLEANUP_DELAYS)
result: Final[Result[FileDeleteResponse]] = cleanup_result(outcomes, wait=delays)
assert result is failure
delays.assert_done()
assert outcomes.call_count == len(CLEANUP_DELAYS) + 1
def test_permanent_error_is_not_retried(self) -> None:
failure: Final = UnknownApiError(status_code=403, body="forbidden")
responses: Final = (failure, deleted_file())
outcomes: Final = Mock(side_effect=responses)
delays: Final = ExpectedCalls[float](())
assert cleanup_result(outcomes, wait=delays) is failure
delays.assert_done()
assert outcomes.call_count == 1
class TestBatchCancellation:
def test_cancelling_batch_is_polled_until_terminal_without_cancelling_again(self) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls((f"retrieve None {MANAGED_BATCH_ID}",) * 3),
batches=(batch("cancelling"), batch("cancelling"), batch("cancelled")),
)
delays: Final = ExpectedCalls((10.0,))
cleanup_batch(client, MANAGED_BATCH_ID, key="test-key", wait=delays)
client.calls.assert_done()
delays.assert_done()
def test_cancellation_timeout_is_reported_but_file_and_key_cleanup_still_run(self) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(
(
f"retrieve None {MANAGED_BATCH_ID}",
f"retrieve None {MANAGED_BATCH_ID}",
"delete None file-1",
"delete key test-key",
)
),
batches=(batch("cancelling"), batch("cancelling")),
files=(deleted_file(),),
)
times: Final = (0.0, BATCH_CANCEL_TIMEOUT_SECONDS)
ticks: Final[Callable[[], float]] = Mock(side_effect=times)
manager: Final = ResourceManager(client=client, strict_cleanup=True)
key: Final = manager.key()
manager.defer(lambda: cleanup_file(client, "file-1", key=key))
manager.defer(lambda: cleanup_batch(client, MANAGED_BATCH_ID, key=key, clock=ticks))
with pytest.raises(ExceptionGroup) as caught:
manager.teardown()
assert "cancellation did not finish" in str(caught.value.exceptions[0])
client.calls.assert_done()
@pytest.mark.parametrize("status", ["completed", "failed", "expired", "cancelled"])
def test_inactive_batch_needs_no_cancellation(self, status: str) -> None:
client: Final = CleanupClient(calls=ExpectedCalls(("retrieve None batch-1",)), batches=(batch(status),))
cleanup_batch(client, "batch-1", key="test-key")
client.calls.assert_done()
def test_active_batch_is_cancelled_through_its_provider(self) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(("retrieve azure batch-1", "cancel azure batch-1")),
batches=(batch("in_progress"), batch("cancelled")),
cancellations=(batch("cancelling"),),
)
cleanup_batch(client, "batch-1", key="test-key", provider="azure")
client.calls.assert_done()
@pytest.mark.parametrize("batch_id", ["batch-1", MANAGED_BATCH_ID])
@pytest.mark.parametrize("pending_status", ["validating", "in_progress"])
def test_accepted_cancellation_waits_through_stale_provider_status(
self, batch_id: str, pending_status: str
) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(
(
f"retrieve vertex_ai {batch_id}",
f"cancel vertex_ai {batch_id}",
f"retrieve vertex_ai {batch_id}",
f"retrieve vertex_ai {batch_id}",
f"retrieve vertex_ai {batch_id}",
"delete vertex_ai file-1",
"delete key test-key",
)
),
batches=(batch("validating"), batch(pending_status), batch(pending_status), batch("cancelled")),
cancellations=(batch(pending_status),),
files=(deleted_file(),),
)
delays: Final = ExpectedCalls((10.0, 10.0))
manager: Final = ResourceManager(client=client, strict_cleanup=True)
key: Final = manager.key()
manager.defer(lambda: cleanup_file(client, "file-1", key=key, provider="vertex_ai"))
manager.defer(lambda: cleanup_batch(client, batch_id, key=key, provider="vertex_ai", wait=delays))
manager.teardown()
client.calls.assert_done()
delays.assert_done()
@pytest.mark.parametrize("output_delete_fails", [False, True])
def test_batch_that_completed_before_cleanup_deletes_output_and_error_files(
self, output_delete_fails: bool
) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(("retrieve openai batch-1", "delete openai file-output", "delete openai file-error")),
batches=(
Success(
status_code=200,
data=BatchObject(
id="batch-1",
status="completed",
input_file_id="file-input",
output_file_id="file-output",
error_file_id="file-error",
),
),
),
files=(
UnknownApiError(status_code=403, body="forbidden") if output_delete_fails else deleted_file(),
deleted_file(),
),
)
if output_delete_fails:
with pytest.raises(ExceptionGroup, match="output cleanup failed"):
cleanup_batch(client, "batch-1", key="test-key", provider="openai", delete_output_files=True)
else:
cleanup_batch(client, "batch-1", key="test-key", provider="openai", delete_output_files=True)
client.calls.assert_done()
@pytest.mark.parametrize("status", ["completed", "in_progress"])
def test_cancellation_conflict_is_accepted_only_when_batch_became_inactive(self, status: str) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(("retrieve None batch-1", "cancel None batch-1", "retrieve None batch-1")),
batches=(batch("in_progress"), batch(status)),
cancellations=(UnknownApiError(status_code=409, body="conflict"),),
)
if status == "completed":
cleanup_batch(client, "batch-1", key="test-key")
else:
with pytest.raises(AssertionError, match="Cancel batch batch-1 left status in_progress"):
cleanup_batch(client, "batch-1", key="test-key")
client.calls.assert_done()
class TestAzureFileExpiry:
def test_azure_form_serializes_native_expiry_for_the_proxy(self) -> None:
form: Final = batch_upload_form("azure", target_model_names="azure-test")
assert form.model_dump(by_alias=True, exclude_none=True) == {
"purpose": "batch",
"target_model_names": "azure-test",
"expires_after[anchor]": "created_at",
"expires_after[seconds]": AZURE_FILE_EXPIRY_SECONDS,
}
@pytest.mark.parametrize("provider", ["openai", "vertex_ai", "bedrock"])
def test_other_providers_keep_their_existing_upload_fields(self, provider: str) -> None:
assert batch_upload_form(provider).model_dump(by_alias=True, exclude_none=True) == {"purpose": "batch"}

88
tests/e2e/batches/test_batches_e2e.py
View file

@ -21,14 +21,16 @@ import os
import re
import time
from datetime import datetime, timedelta, timezone
from typing import Callable
import pytest
from pydantic import BaseModel
from e2e_config import PROXY_BASE_URL, unique_marker
from e2e_config import MASTER_KEY, PROXY_BASE_URL, unique_marker
from batch_cleanup import cleanup_batch, cleanup_file
from batch_client import (
AZURE_FILE_EXPIRY_SECONDS,
batch_upload_form,
UPLOAD_FILENAME,
BatchClient,
BatchCreateBody,
@ -155,19 +157,19 @@ def upload_for_scenario(
if cap.scenario == "encoded":
return client.upload_file(
content=content,
form=FileUploadForm(purpose="batch"),
form=batch_upload_form(cap.provider),
model=cap.model,
key=key,
)
if cap.scenario == "unified":
return client.upload_file(
content=content,
form=FileUploadForm(purpose="batch", target_model_names=cap.model),
form=batch_upload_form(cap.provider, target_model_names=cap.model),
key=key,
)
return client.upload_file(
content=content,
form=FileUploadForm(purpose="batch"),
form=batch_upload_form(cap.provider),
key=key,
provider=cap.provider,
)
@ -188,20 +190,11 @@ def create_for_scenario(
def op_provider(cap: Capability) -> str | None:
"""provider_fallback ids are raw, so retrieve/cancel/list/delete need the provider
"""provider_fallback batch ids are raw, so retrieve/cancel/list need the provider
hint; the other scenarios encode it into the id and route automatically."""
return cap.provider if cap.scenario == "provider_fallback" else None
def quietly(action: Callable[[], object]) -> Callable[[], None]:
"""Adapt a value-returning call into a best-effort cleanup the teardown can run."""
def run() -> None:
action()
return run
def assert_file_object(file: FileObject, *, provider: str) -> None:
assert file.object == "file", f"file.object={file.object!r}"
assert file.purpose == "batch", f"file.purpose={file.purpose!r}"
@ -209,6 +202,10 @@ def assert_file_object(file: FileObject, *, provider: str) -> None:
if provider != "bedrock":
assert file.bytes > 0, f"file.bytes={file.bytes!r}"
assert file.status, "file.status missing"
if provider == "azure":
assert file.expires_at is not None, "Azure batch input has no automatic expiry"
assert file.created_at is not None
assert file.expires_at - file.created_at == AZURE_FILE_EXPIRY_SECONDS
assert (
file.created_at is not None and file.created_at > 0
), "file.created_at missing"
@ -249,7 +246,7 @@ def test_batch_lifecycle(
file = unwrap(upload_for_scenario(client, cap, render_jsonl(cap.jsonl_model), key))
resources.defer(
quietly(lambda: client.delete_file(file.id, key=key, provider=provider))
lambda: cleanup_file(client, file.id, key=key, provider=cap.file_provider)
)
assert_file_object(file, provider=cap.provider)
assert matches_id_shape(
@ -260,7 +257,9 @@ def test_batch_lifecycle(
require_successful_call(created)
batch = BatchObject.model_validate_json(created.body)
resources.defer(
quietly(lambda: client.cancel_batch(batch.id, key=key, provider=provider))
lambda: cleanup_batch(
client, batch.id, key=key, provider=provider, delete_output_files=cap.provider in {"openai", "azure"}
)
)
assert batch.id, f"create returned no batch id (body={created.body[:200]})"
@ -339,7 +338,7 @@ def test_batch_key_model_access_denied(
denied_upload = client.upload_file(
content=render_jsonl(AZURE_BATCH_MODEL),
form=FileUploadForm(purpose="batch"),
form=batch_upload_form("azure"),
model=AZURE_BATCH_MODEL,
key=key,
)
@ -356,7 +355,7 @@ def test_batch_key_model_access_denied(
)
).id
resources.defer(
quietly(lambda: client.delete_file(raw_file, key=key, provider="openai"))
lambda: cleanup_file(client, raw_file, key=key, provider="openai")
)
denied_create = client.create_batch(
@ -383,6 +382,7 @@ def test_file_upload_and_delete_outputs(
key=key,
)
)
resources.defer(lambda: cleanup_file(client, file.id, key=key))
assert_file_object(file, provider="openai")
deleted = unwrap(client.delete_file(file.id, key=key))
@ -458,12 +458,12 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row(
key=key,
)
)
resources.defer(quietly(lambda: client.delete_file(file.id, key=key)))
resources.defer(lambda: cleanup_file(client, file.id, key=key))
created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key)
require_successful_call(created)
batch = BatchObject.model_validate_json(created.body)
resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key)))
resources.defer(lambda: cleanup_batch(client, batch.id, key=key))
_ = client.proxy.poll_logs_for_key(key, min_rows=1)
@ -517,7 +517,7 @@ class TestBatchFileContent:
key=key,
)
)
resources.defer(quietly(lambda: client.delete_file(file.id, key=key)))
resources.defer(lambda: cleanup_file(client, file.id, key=key))
assert file.id
downloaded = client.proxy.transport.download(
@ -559,11 +559,11 @@ class TestBatchFileContent:
file = unwrap(
client.upload_file(
content=payload,
form=FileUploadForm(purpose="batch", target_model_names=provider.model),
form=batch_upload_form(provider.name, target_model_names=provider.model),
key=key,
)
)
resources.defer(quietly(lambda: client.delete_file(file.id, key=key)))
resources.defer(lambda: cleanup_file(client, file.id, key=key))
assert_file_object(file, provider=provider.name)
assert is_managed_id(file.id), (
f"{provider.name}: unified upload must return a managed file id, got {file.id!r}"
@ -626,7 +626,7 @@ class TestOpenAIFiles:
)
)
resources.defer(
quietly(lambda: client.delete_file(file.id, key=key, provider="openai"))
lambda: cleanup_file(client, file.id, key=key, provider="openai")
)
listed = unwrap(client.list_files(key=key))
@ -690,7 +690,7 @@ class TestOpenAIFiles:
key=key,
)
)
resources.defer(quietly(lambda: client.delete_file(file.id, key=key)))
resources.defer(lambda: cleanup_file(client, file.id, key=key))
fetched = unwrap(client.retrieve_file(file.id, key=key))
assert fetched.id == file.id, "retrieve must echo the uploaded file id"
@ -760,7 +760,7 @@ class TestBatchRateLimitErrorMapping:
key=key,
)
)
resources.defer(quietly(lambda: client.delete_file(file.id, key=key)))
resources.defer(lambda: cleanup_file(client, file.id, key=key))
created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key)
@ -803,7 +803,7 @@ class TestBatchEnqueuedTokenLimit:
"""
def _upload_batch_file(
self, client: BatchClient, resources: ResourceManager, key: str
self, client: BatchClient, resources: ResourceManager, key: str, *, cleanup_key: str | None = None
) -> FileObject:
file = unwrap(
client.upload_file(
@ -813,7 +813,7 @@ class TestBatchEnqueuedTokenLimit:
key=key,
)
)
resources.defer(quietly(lambda: client.delete_file(file.id, key=key)))
resources.defer(lambda: cleanup_file(client, file.id, key=cleanup_key or key))
return file
def _generate_enqueued_key(
@ -850,7 +850,7 @@ class TestBatchEnqueuedTokenLimit:
marker="rpm",
rpm_limit=BATCH_RL_RPM_LIMIT,
)
file = self._upload_batch_file(client, resources, key)
file = self._upload_batch_file(client, resources, key, cleanup_key=MASTER_KEY)
created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key)
@ -861,7 +861,7 @@ class TestBatchEnqueuedTokenLimit:
)
require_successful_call(created)
batch = BatchObject.model_validate_json(created.body)
resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key)))
resources.defer(lambda: cleanup_batch(client, batch.id, key=MASTER_KEY, delete_output_files=True))
@pytest.mark.covers(
"quota_management.ratelimit.batch_enqueued_tokens.blocks_when_exhausted",
@ -904,7 +904,7 @@ class TestBatchEnqueuedTokenLimit:
first = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key)
require_successful_call(first)
first_batch = BatchObject.model_validate_json(first.body)
resources.defer(quietly(lambda: client.cancel_batch(first_batch.id, key=key)))
resources.defer(lambda: cleanup_batch(client, first_batch.id, key=key))
blocked = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key)
assert blocked.status_code == 429, (
@ -928,7 +928,7 @@ class TestBatchEnqueuedTokenLimit:
)
require_successful_call(retried)
retry_batch = BatchObject.model_validate_json(retried.body)
resources.defer(quietly(lambda: client.cancel_batch(retry_batch.id, key=key)))
resources.defer(lambda: cleanup_batch(client, retry_batch.id, key=key))
ASSUME_ROLE_RAW_MODEL = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0"
@ -984,13 +984,13 @@ class TestBedrockBatchAssumeRole:
key=key,
)
)
resources.defer(quietly(lambda: client.delete_file(file.id, key=key)))
resources.defer(lambda: cleanup_file(client, file.id, key=key))
assert_file_object(file, provider="bedrock")
created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key)
require_successful_call(created)
batch = BatchObject.model_validate_json(created.body)
resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key)))
resources.defer(lambda: cleanup_batch(client, batch.id, key=key))
assert batch.id, f"assume-role create returned no batch id: {created.body[:200]}"
assert is_managed_id(batch.id), (
@ -1044,7 +1044,7 @@ class TestGeminiFiles:
key=key,
)
)
resources.defer(quietly(lambda: client.delete_file(file.id, key=key)))
resources.defer(lambda: cleanup_file(client, file.id, key=key))
assert_file_object(file, provider="gemini")
assert file.id, "gemini file upload returned no id"
@ -1099,13 +1099,13 @@ class TestHostedVllmBatch:
key=key,
)
)
resources.defer(quietly(lambda: client.delete_file(file.id, key=key)))
resources.defer(lambda: cleanup_file(client, file.id, key=key))
assert_file_object(file, provider="hosted_vllm")
created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key)
require_successful_call(created)
batch = BatchObject.model_validate_json(created.body)
resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key)))
resources.defer(lambda: cleanup_batch(client, batch.id, key=key))
assert batch.id, f"hosted_vllm create returned no batch id: {created.body[:200]}"
assert batch.status in CREATED_BATCH_STATUSES, (
@ -1192,7 +1192,7 @@ class TestBatchFailurePaths:
key=key,
)
)
resources.defer(quietly(lambda: client.delete_file(file.id, key=key)))
resources.defer(lambda: cleanup_file(client, file.id, key=key))
created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key)
require_successful_call(created)
@ -1243,12 +1243,12 @@ class TestBatchFailurePaths:
file = unwrap(
client.upload_file(
content=render_jsonl(AZURE_BATCH_RAW_MODEL),
form=FileUploadForm(purpose="batch"),
form=batch_upload_form("azure"),
model=AZURE_BATCH_MODEL,
key=key,
)
)
resources.defer(quietly(lambda: client.delete_file(file.id, key=key)))
resources.defer(lambda: cleanup_file(client, file.id, key=key))
assert decoded_model_from_id(file.id) == AZURE_BATCH_MODEL, (
f"upload did not encode the azure deployment into the file id: {file.id!r}"
)
@ -1258,7 +1258,7 @@ class TestBatchFailurePaths:
)
require_successful_call(created)
batch = BatchObject.model_validate_json(created.body)
resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key)))
resources.defer(lambda: cleanup_batch(client, batch.id, key=key))
assert decoded_model_from_id(batch.id) == AZURE_BATCH_MODEL, (
"create with a foreign encoded file id must route by the file's embedded model, "
@ -1307,7 +1307,7 @@ class TestBatchSecondHop:
key=key,
)
)
resources.defer(quietly(lambda: client.delete_file(file.id, key=key)))
resources.defer(lambda: cleanup_file(client, file.id, key=key))
assert is_managed_id(file.id), (
f"second-hop unified upload must return a managed file id, got {file.id!r}"
)
@ -1315,7 +1315,7 @@ class TestBatchSecondHop:
created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key)
require_successful_call(created)
batch = BatchObject.model_validate_json(created.body)
resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key)))
resources.defer(lambda: cleanup_batch(client, batch.id, key=key))
assert is_managed_id(batch.id), (
f"second-hop create must return a managed batch id, got {batch.id!r}"

3
tests/e2e/batches/test_managed_files_enforcement_e2e.py
View file

@ -21,6 +21,7 @@ from typing import Iterator
import pytest
from batch_client import BatchClient, FileObject
from batch_cleanup import cleanup_file
from capabilities import batch_model_name, is_managed_id, openai_batch_params
from e2e_config import unique_marker
from e2e_http import FileUploadForm, Result, UnknownApiError, unwrap
@ -108,7 +109,7 @@ def test_cross_user_managed_id_denied_owner_allowed(
key=owner_key,
)
)
resources.defer(lambda: client.delete_file(uploaded.id, key=owner_key))
resources.defer(lambda: cleanup_file(client, uploaded.id, key=owner_key))
assert is_managed_id(uploaded.id), f"expected a managed unified file id, got {uploaded.id}"
denied = client.retrieve_file(uploaded.id, key=other_key)

8
tests/e2e/coverage_registry/mcp.yaml
View file

@ -119,3 +119,11 @@
assertions: [succeeds]
source: "server.py:1089"
rationale: Smoke; rarely used; same auth model as tools
- id: mcp.list_tools.api_key.toolset_scoped
module: mcp
tier: P0
operation: list_tools
auth_family: api_key
assertions: [toolset_scoped]
source: "user_api_key_auth_mcp.py:2137"
rationale: "A key granted a toolset lists exactly the toolset's tools: the rest of the server's catalog stays hidden and every stored name resolves"

10
tests/e2e/coverage_registry/mgmt.yaml
View file

@ -76,3 +76,13 @@
- {id: mgmt.credential_migration.check.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4252", rationale: "Encryption migration (smoke)"}
- {id: mgmt.credential.new.serves_request, module: mgmt, tier: P1, surface: api, assertions: [serves_request], source: "credential_endpoints/endpoints.py:42", rationale: "Stored credential resolves into a deployment and serves a live /messages request"}
- {id: mgmt.model.test_connection.happy_path, module: mgmt, tier: P0, surface: api, assertions: [happy_path], source: "_health_endpoints.py:1785", rationale: "Test Connection for a responses-mode Bedrock Mantle deployment reaches the live provider and reports success; this exact shape 500ed on an acompletion partial before v1.91.0", fail_before_fix: proven}
- {id: mgmt.mcp_server.new.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1577", rationale: "Every field of an admin-created MCP server reads back verbatim, by id and in the list, on every replica"}
- {id: mgmt.mcp_server.list.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1112", rationale: "The MCP page grid lists a created server with the same field values its detail view reports"}
- {id: mgmt.mcp_server.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "mcp_management_endpoints.py:2665", rationale: "A dashboard edit of one field leaves the others intact and is visible on every replica after one save; edits that took several saves to stick were a customer defect"}
- {id: mgmt.mcp_server.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "mcp_management_endpoints.py:2665", rationale: "An explicit null clears the stored field (absent keeps, null clears)"}
- {id: mgmt.mcp_server.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:2139", rationale: "A deleted server is gone by id and from the list on every replica"}
- {id: mgmt.mcp_toolset.new.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3009", rationale: "Toolset tools read back under the exact server_id and tool_name written; a toolset stored under one name and read under another granted nothing"}
- {id: mgmt.mcp_toolset.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "mcp_management_endpoints.py:3098", rationale: "Editing the description leaves the tools and name intact"}
- {id: mgmt.mcp_toolset.update.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3098", rationale: "Narrowing the tools to one entry reads back exactly that entry"}
- {id: mgmt.mcp_toolset.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "mcp_management_endpoints.py:3098", fail_before_fix: proven, rationale: "An explicit null clears the stored description; the update used to drop null and keep the old value"}
- {id: mgmt.mcp_toolset.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3149", rationale: "A deleted toolset is gone by id and from the list on every replica"}

64
tests/e2e/e2e_http.py
View file

@ -49,6 +49,11 @@ class AnthropicHeaders(AuthHeaders):
anthropic_version: str = Field(default="2023-06-01", alias="anthropic-version")
class PartialBody(BaseModel):
"""A body for a partial-update route (absent = keep, null = clear): a field left
unset is omitted from the wire, and a field set to None is sent as JSON null."""
class NoBody(BaseModel):
"""Empty body/query for routes that take none."""
@ -252,6 +257,13 @@ def assert_auth_denied(result: StreamingResponse, context: str) -> None:
f"{context}: expected 401/403, got {result.status_code}: {result.body[:300]}"
)
def wire_body(json: BaseModel) -> dict[str, object]:
if isinstance(json, PartialBody):
return json.model_dump(by_alias=True, exclude_unset=True)
return json.model_dump(by_alias=True, exclude_none=True)
def _headers(headers: BaseModel) -> dict[str, str]:
dumped: dict[str, object] = headers.model_dump(by_alias=True, exclude_none=True)
return {key: str(value) for key, value in dumped.items()}
@ -307,9 +319,26 @@ def request_with_retry[T: RetryableResponse](
return issue()
def _classify[R: BaseModel](
resp: requests.Response, response_type: type[R]
) -> Result[R]:
class ClassifiableResponse(Protocol):
"""What classifying an outcome reads off a response. requests.Response satisfies
it, and so does a fake, so the classification rules are testable on their own."""
@property
def status_code(self) -> int: ...
@property
def ok(self) -> bool: ...
@property
def text(self) -> str: ...
@property
def content(self) -> bytes: ...
def json(self) -> object: ...
def classify[R: BaseModel](resp: ClassifiableResponse, response_type: type[R]) -> Result[R]:
if resp.status_code == 401:
return UnauthorizedError(body=resp.text)
if resp.status_code == 429:
@ -317,7 +346,8 @@ def _classify[R: BaseModel](
if not resp.ok:
return UnknownApiError(status_code=resp.status_code, body=resp.text)
try:
return Success(status_code=resp.status_code, data=response_type.model_validate(resp.json()))
payload: Final[object] = resp.json() if resp.content else {}
return Success(status_code=resp.status_code, data=response_type.model_validate(payload))
except Exception as exc: # noqa: BLE001 - any parse/validation failure is a value
return ValidationError(message=str(exc))
@ -335,13 +365,13 @@ def post[R: BaseModel](
lambda: requests.post(
str(url),
headers=_headers(headers),
json=json.model_dump(by_alias=True, exclude_none=True),
json=wire_body(json),
timeout=timeout,
)
)
except requests.RequestException as exc:
return NetworkError(message=str(exc))
return _classify(resp, response_type)
return classify(resp, response_type)
def get[R: BaseModel](
@ -363,7 +393,7 @@ def get[R: BaseModel](
)
except requests.RequestException as exc:
return NetworkError(message=str(exc))
return _classify(resp, response_type)
return classify(resp, response_type)
def get_external[R: BaseModel](
@ -383,7 +413,7 @@ def get_external[R: BaseModel](
)
except requests.RequestException as exc:
return NetworkError(message=str(exc))
return _classify(resp, response_type)
return classify(resp, response_type)
def delete[R: BaseModel](
@ -400,14 +430,14 @@ def delete[R: BaseModel](
lambda: requests.delete(
str(url),
headers=_headers(headers),
json=json.model_dump(by_alias=True, exclude_none=True),
json=wire_body(json),
params=_params(params),
timeout=timeout,
)
)
except requests.RequestException as exc:
return NetworkError(message=str(exc))
return _classify(resp, response_type)
return classify(resp, response_type)
def patch[R: BaseModel](
@ -423,13 +453,13 @@ def patch[R: BaseModel](
lambda: requests.patch(
str(url),
headers=_headers(headers),
json=json.model_dump(by_alias=True, exclude_none=True),
json=wire_body(json),
timeout=timeout,
)
)
except requests.RequestException as exc:
return NetworkError(message=str(exc))
return _classify(resp, response_type)
return classify(resp, response_type)
def put[R: BaseModel](
@ -445,13 +475,13 @@ def put[R: BaseModel](
lambda: requests.put(
str(url),
headers=_headers(headers),
json=json.model_dump(by_alias=True, exclude_none=True),
json=wire_body(json),
timeout=timeout,
)
)
except requests.RequestException as exc:
return NetworkError(message=str(exc))
return _classify(resp, response_type)
return classify(resp, response_type)
def probe(
@ -555,7 +585,7 @@ def send(
str(url),
headers=_headers(headers),
params=_params(params),
json=json.model_dump(by_alias=True, exclude_none=True),
json=wire_body(json),
stream=stream,
timeout=timeout,
)
@ -605,7 +635,7 @@ def upload[R: BaseModel](
)
except requests.RequestException as exc:
return NetworkError(message=str(exc))
return _classify(resp, response_type)
return classify(resp, response_type)
def stream_binary(
@ -623,7 +653,7 @@ def stream_binary(
resp = requests.post(
str(url),
headers=_headers(headers),
json=json.model_dump(by_alias=True, exclude_none=True),
json=wire_body(json),
stream=True,
timeout=timeout,
)

25
tests/e2e/guardrails/guardrails_client.py
View file

@ -7,7 +7,7 @@ from __future__ import annotations
import time
from collections.abc import Callable
from dataclasses import dataclass
from typing import Literal
from typing import Final, Literal
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, settle_propagation, unique_marker
from e2e_http import NoBody, Result, StreamingResponse, Success, unwrap
@ -405,6 +405,29 @@ def build_client(proxy: ProxyClient) -> GuardrailsClient:
return GuardrailsClient(proxy=proxy)
def poll_until_guardrail_applied(
call: Callable[[], StreamingResponse],
guardrail_name: str,
*,
timeout: float = POLL_TIMEOUT,
interval: float = POLL_INTERVAL,
now: Callable[[], float] = time.monotonic,
sleep: Callable[[float], None] = time.sleep,
) -> StreamingResponse:
deadline: Final = now() + timeout
if not (result := call()).ok:
return result
while (
guardrail_name
not in (name.strip() for name in result.headers.get("x-litellm-applied-guardrails", "").split(","))
and (remaining := deadline - now()) > 0
):
sleep(min(interval, remaining))
if now() >= deadline or not (result := call()).ok:
break
return result
def poll_until_blocked[R: BaseModel](call: Callable[[], Result[R]]) -> Result[R]:
"""Retry a call that a guardrail should reject until it is, returning the last result.

66
tests/e2e/guardrails/test_guardrails_client.py Normal file
View file

@ -0,0 +1,66 @@
from dataclasses import dataclass
from itertools import chain, repeat
from typing import Final
import pytest
from e2e_http import StreamingResponse
from guardrails_client import poll_until_guardrail_applied
@dataclass
class Clock:
elapsed: float = 0.0
def now(self) -> float:
return self.elapsed
def sleep(self, seconds: float) -> None:
self.elapsed += seconds
def _response(applied: str, status: int = 200) -> StreamingResponse:
return StreamingResponse(status_code=status, body="{}", headers={"x-litellm-applied-guardrails": applied})
def test_waits_for_requested_guardrail_after_an_unrelated_global_guardrail() -> None:
clock: Final = Clock()
expected: Final = _response("global-filter, tool-permission")
responses: Final = iter((_response("global-filter"), expected))
result: Final = poll_until_guardrail_applied(
lambda: next(responses), "tool-permission", timeout=5, interval=2, now=clock.now, sleep=clock.sleep
)
assert result is expected
assert clock.elapsed == 2
@pytest.mark.parametrize("applied", ("", "global-filter", "tool-permission-sibling"))
def test_missing_exact_guardrail_returns_failure_evidence_at_deadline(applied: str) -> None:
clock: Final = Clock()
missing: Final = _response(applied)
responses: Final = iter((missing, missing, missing))
result: Final = poll_until_guardrail_applied(
lambda: next(responses), "tool-permission", timeout=5, interval=2, now=clock.now, sleep=clock.sleep
)
assert result is missing
assert clock.elapsed == 5
with pytest.raises(StopIteration):
next(responses)
@pytest.mark.parametrize("status", (400, 401, 429, 500))
def test_http_failure_is_not_hidden_by_a_later_success(status: int) -> None:
clock: Final = Clock()
failed: Final = _response("", status)
responses: Final = iter(chain((failed,), repeat(_response("tool-permission"))))
result: Final = poll_until_guardrail_applied(
lambda: next(responses), "tool-permission", timeout=5, interval=2, now=clock.now, sleep=clock.sleep
)
assert result is failed
assert clock.elapsed == 0

24
tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py
View file

@ -30,6 +30,7 @@ from guardrails_client import (
ToolPermissionParamsBody,
ToolPermissionRuleBody,
poll_until_blocked,
poll_until_guardrail_applied,
)
from lifecycle import ResourceManager
from models import ChatResponse, ChatTool, ChatToolFunction
@ -84,8 +85,8 @@ def _register_tool_permission(client: GuardrailsClient, resources: ResourceManag
resources.defer(lambda: client.delete_guardrail(guardrail_id))
def _applied_guardrails(outcome: StreamingResponse) -> str:
return outcome.headers.get("x-litellm-applied-guardrails", "")
def _applied_guardrails(outcome: StreamingResponse) -> tuple[str, ...]:
return tuple(name.strip() for name in outcome.headers.get("x-litellm-applied-guardrails", "").split(","))
def _tool_call_names(response: ChatResponse) -> tuple[str, ...]:
@ -144,14 +145,17 @@ class TestToolPermissionPreCall:
name = f"e2e-toolperm-allow-{unique_marker()}"
_register_tool_permission(client, resources, name=name)
outcome = client.chat_raw(
scoped_key,
MODEL,
TOOL_PROMPT,
guardrails=[name],
max_tokens=128,
tools=[ALLOWED_TOOL],
tool_choice="required",
outcome = poll_until_guardrail_applied(
lambda: client.chat_raw(
scoped_key,
MODEL,
TOOL_PROMPT,
guardrails=[name],
max_tokens=128,
tools=[ALLOWED_TOOL],
tool_choice="required",
),
name,
)
assert outcome.ok, f"the permitted tool must be served, got {outcome.status_code}: {outcome.body[:400]}"

23
tests/e2e/lifecycle.py
View file

@ -8,8 +8,9 @@ ResourceManager; the test registers a cleanup for every resource it creates, and
the fixture's teardown releases them all even when the test body raises.
"""
from builtins import ExceptionGroup
from dataclasses import dataclass, field
from typing import Callable, List, Protocol, runtime_checkable
from typing import Callable, Final, List, Protocol, runtime_checkable
from proxy_client import ProxyClient
from models import KeyGenerateBody
@ -52,6 +53,7 @@ class ResourceManager:
"""
client: ResourceClient
strict_cleanup: bool = False
_cleanups: List[Callable[[], object]] = field(
default_factory=list
) # mutable-ok: append-only teardown registry
@ -82,8 +84,17 @@ class ResourceManager:
return customer_id
def teardown(self) -> None:
for cleanup in reversed(self._cleanups):
try:
cleanup()
except Exception:
pass # best-effort: a failed cleanup must not block the rest
failures: Final = tuple(
failure for cleanup in reversed(self._cleanups)
if (failure := _run_cleanup(cleanup)) is not None
)
if failures and self.strict_cleanup:
raise ExceptionGroup("Resource cleanup failed", failures)
def _run_cleanup(cleanup: Callable[[], object]) -> Exception | None:
try:
cleanup()
except Exception as exc:
return exc
return None

24
tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py
View file

@ -16,8 +16,10 @@ into a chat completion chunk. Two customer-visible contracts only hold on that p
from __future__ import annotations
from typing import Final, Literal
import pytest
from pydantic import BaseModel
from pydantic import BaseModel, Field
from e2e_config import unique_marker
from e2e_http import StreamingResponse
@ -51,7 +53,8 @@ class _BridgeChoice(BaseModel):
class _BridgeChunk(BaseModel):
id: str
choices: list[_BridgeChoice] = []
object: Literal["chat.completion.chunk"]
choices: list[_BridgeChoice] = Field(default_factory=list)
class _WeatherArgs(BaseModel):
@ -103,16 +106,19 @@ class TestResponsesBridgeChatCompletionsStreaming:
resources.key(),
ChatBody(
model=bridged_model,
messages=[ChatMessage(role="user", content=f"Count from 1 to 5, one number per line. {unique_marker()}")],
messages=[
ChatMessage(role="user", content=f"Count from 1 to 5, one number per line. {unique_marker()}")
],
max_tokens=64,
stream=True,
),
)
chunks = _bridge_chunks(result)
ids = {chunk.id for chunk in chunks}
chunks: Final = _bridge_chunks(result)
assert len(chunks) > 1, "the shared-id contract needs more than one streamed chunk"
ids: Final = frozenset(chunk.id for chunk in chunks)
assert len(ids) == 1, f"bridged stream used {len(ids)} different chunk ids: {sorted(ids)[:5]}"
assert ids.pop().startswith("chatcmpl-"), f"bridged chunk id is not chat-completion shaped: {chunks[0].id}"
assert chunks[0].id.strip(), "bridged stream emitted an empty chunk id"
@pytest.mark.covers(
"llm.chat_completions.openai.basic.stream.bridge_streams_sse",
@ -134,9 +140,9 @@ class TestResponsesBridgeChatCompletionsStreaming:
chunks = _bridge_chunks(result)
content = "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices)
assert content.strip(), f"bridged stream completed with no content deltas: {result.stream_events[:3]}"
assert any(
choice.finish_reason for chunk in chunks for choice in chunk.choices
), f"bridged stream never emitted a finish_reason: {result.stream_events[-3:]}"
assert any(choice.finish_reason for chunk in chunks for choice in chunk.choices), (
f"bridged stream never emitted a finish_reason: {result.stream_events[-3:]}"
)
assert result.stream_done, f"bridged stream did not terminate with [DONE]: {result.stream_events[-2:]}"
@pytest.mark.covers(

124
tests/e2e/logging/datadog_reader.py
View file

@ -12,8 +12,12 @@ empty result. External reads go through ``e2e_http``.
from __future__ import annotations
import math
import random
import time
from dataclasses import dataclass
from collections.abc import Callable, Mapping
from dataclasses import dataclass, field
from typing import Final
import pytest
from pydantic import BaseModel, ConfigDict, Field
@ -27,17 +31,33 @@ from e2e_config import (
DD_SITE,
POLL_TIMEOUT,
)
from e2e_http import URL, Headers, RateLimitedError, Success, post
from e2e_http import URL, Headers, StreamingResponse, send
#: How many rate-limited responses in a row one search tolerates before the
#: hard fail; each retry sleeps a full search interval, so this rides out a
#: burst from a concurrent consumer of the org-wide search budget.
_RATE_LIMIT_RETRIES = 5
type SearchCall = Callable[[str, float], StreamingResponse]
def _seconds(value: str | None) -> float | None:
if value is None:
return None
try:
seconds: Final = float(value)
except ValueError:
return None
return seconds if math.isfinite(seconds) and seconds >= 0 else None
def _rate_limit_delay(headers: Mapping[str, str]) -> float:
delays: Final = tuple(
delay
for name in ("x-ratelimit-reset", "retry-after")
if (delay := _seconds(headers.get(name))) is not None
)
return max(1.0, max(delays, default=DD_SEARCH_INTERVAL))
class _DdAuthHeaders(Headers):
api_key: str = Field(serialization_alias="DD-API-KEY")
app_key: str = Field(serialization_alias="DD-APPLICATION-KEY")
api_key: str = Field(serialization_alias="DD-API-KEY", repr=False)
app_key: str = Field(serialization_alias="DD-APPLICATION-KEY", repr=False)
class _SearchFilter(BaseModel):
@ -88,8 +108,12 @@ class _SearchResponse(BaseModel):
@dataclass(frozen=True, slots=True)
class DdLogsReader:
site: str
api_key: str
app_key: str
api_key: str = field(repr=False)
app_key: str = field(repr=False)
search: SearchCall | None = field(default=None, repr=False)
now: Callable[[], float] = field(default=time.monotonic, repr=False)
sleep: Callable[[float], None] = field(default=time.sleep, repr=False)
jitter: Callable[[], float] = field(default=random.random, repr=False)
def events_for_marker(self, marker: str) -> list[DdLogEvent]:
"""Every ingested event whose attributes carry the marker. DataDog
@ -108,25 +132,28 @@ class DdLogsReader:
a single event. A 429 backs off and retries - the search budget is
org-wide, so another consumer can empty it under us - while any other
failure stays a hard fail."""
for _ in range(_RATE_LIMIT_RETRIES):
result = post(
URL(f"https://api.{self.site}/api/v2/logs/events/search"),
headers=_DdAuthHeaders(api_key=self.api_key, app_key=self.app_key),
json=_SearchRequest(filter=_SearchFilter(query=query)),
response_type=_SearchResponse,
timeout=30.0,
)
match result:
case Success(data=page):
return [event.attributes for event in page.data]
case RateLimitedError(retry_after_seconds=retry_after):
time.sleep(retry_after if retry_after else DD_SEARCH_INTERVAL)
case failure:
pytest.fail(f"DataDog Logs Search API at api.{self.site} failed: {failure}")
return self._events_for_query(query, self.now() + POLL_TIMEOUT)
def _events_for_query(self, query: str, deadline: float) -> list[DdLogEvent]:
search: Final = self.search or self._search_page
while (remaining := deadline - self.now()) > 0:
if (result := search(query, min(30.0, remaining))).ok:
return [event.attributes for event in _SearchResponse.model_validate_json(result.body).data]
if result.status_code != 429:
pytest.fail(f"DataDog Logs Search API at api.{self.site} failed with HTTP {result.status_code}")
if (delay := min(_rate_limit_delay(result.headers) + self.jitter(), deadline - self.now())) > 0:
self.sleep(delay)
pytest.fail(
f"DataDog Logs Search API at api.{self.site} still rate-limited after "
f"{_RATE_LIMIT_RETRIES} retries {DD_SEARCH_INTERVAL}s apart - the org-wide "
"logs_public_search_api budget (2 requests per 10s) is exhausted by another consumer"
f"DataDog Logs Search API at api.{self.site} remained rate-limited for {POLL_TIMEOUT}s; "
"the org-wide logs_public_search_api budget is exhausted"
)
def _search_page(self, query: str, timeout: float) -> StreamingResponse:
return send(
URL(f"https://api.{self.site}/api/v2/logs/events/search"),
headers=_DdAuthHeaders(api_key=self.api_key, app_key=self.app_key),
json=_SearchRequest(filter=_SearchFilter(query=query)),
timeout=timeout,
)
def poll_events_for_marker(self, marker: str) -> list[DdLogEvent]:
@ -140,33 +167,42 @@ class DdLogsReader:
hide from the exactly-one assertion - real-DataDog jitter can surface
one call's two events tens of seconds apart. Searches pace at
DD_SEARCH_INTERVAL, not POLL_INTERVAL, to respect the search API's
request budget. At the deadline the last result is returned as-is."""
deadline = time.monotonic() + POLL_TIMEOUT
while time.monotonic() < deadline:
events = self.events_for_query(query)
request budget. Discovery, quota retries, and duplicate detection share
one POLL_TIMEOUT deadline; an incomplete settle window fails closed."""
deadline: Final = self.now() + POLL_TIMEOUT
while (remaining := deadline - self.now()) > 0:
events = self._events_for_query(query, deadline)
if events:
return self._settled_events_for_query(query, events)
time.sleep(DD_SEARCH_INTERVAL)
return self.events_for_query(query)
return self._settled_events_for_query(query, events, deadline)
if (remaining := deadline - self.now()) > 0:
self.sleep(min(DD_SEARCH_INTERVAL, remaining))
return []
def _settled_events_for_query(self, query: str, events: list[DdLogEvent]) -> list[DdLogEvent]:
def _settled_events_for_query(self, query: str, events: list[DdLogEvent], deadline: float) -> list[DdLogEvent]:
"""Re-read at every search interval until the settle window closes; a
duplicate ends the watch early because more waiting cannot clear it.
Keep the last non-empty result: a transient empty search (index lag)
must not erase events already confirmed earlier in the settle window.
A successful final search must reach the full settle window before the
shared read-back deadline; otherwise duplicate detection is incomplete.
"""
settle_deadline = time.monotonic() + DD_SETTLE_SECONDS
settle_deadline: Final = self.now() + DD_SETTLE_SECONDS
last_nonempty = events
while time.monotonic() < settle_deadline:
time.sleep(DD_SEARCH_INTERVAL)
latest = self.events_for_query(query)
if not latest:
continue
if len(events) > 1:
return events
while (remaining := deadline - self.now()) > 0:
self.sleep(min(DD_SEARCH_INTERVAL, remaining))
if self.now() >= deadline:
break
latest = self._events_for_query(query, deadline)
if len(latest) > 1:
return latest
last_nonempty = latest
return last_nonempty
if latest:
last_nonempty = latest
if self.now() >= settle_deadline:
return last_nonempty
pytest.fail(f"DataDog log delivery could not complete its duplicate-detection window within {POLL_TIMEOUT}s")
def build_dd_logs_reader() -> DdLogsReader:

223
tests/e2e/logging/test_datadog_reader.py Normal file
View file

@ -0,0 +1,223 @@
import json
from collections.abc import Iterator, Sequence
from dataclasses import dataclass
from typing import Final
import pytest
from datadog_reader import DdLogsReader
from datadog_reader import _DdAuthHeaders # pyright: ignore[reportPrivateUsage] # verifies private auth-header serialization
from e2e_config import DD_SEARCH_INTERVAL, POLL_TIMEOUT
from e2e_http import StreamingResponse
def test_failure_diagnostics_hide_credentials_without_changing_auth_headers() -> None:
api_key: Final = "test-datadog-api-secret"
app_key: Final = "test-datadog-app-secret"
reader: Final = DdLogsReader(site="datadoghq.com", api_key=api_key, app_key=app_key)
headers: Final = _DdAuthHeaders(api_key=api_key, app_key=app_key)
for value in (reader, headers):
assert api_key not in repr(value)
assert app_key not in repr(value)
assert headers.model_dump(by_alias=True) == {
"DD-API-KEY": api_key,
"DD-APPLICATION-KEY": app_key,
}
@dataclass
class Clock:
elapsed: float = 0.0
def now(self) -> float:
return self.elapsed
def sleep(self, seconds: float) -> None:
self.elapsed += seconds
@dataclass
class Search:
responses: Iterator[StreamingResponse]
calls: tuple[tuple[str, float], ...] = ()
def __call__(self, query: str, timeout: float) -> StreamingResponse:
self.calls += ((query, timeout),)
return next(self.responses)
def _page(*event_ids: str) -> StreamingResponse:
return StreamingResponse(
status_code=200,
body=json.dumps({"data": [{"attributes": {"attributes": {"id": event_id}}} for event_id in event_ids]}),
)
def _reader(responses: Sequence[StreamingResponse], clock: Clock) -> tuple[DdLogsReader, Search]:
search: Final = Search(iter(responses))
return DdLogsReader(
site="us5.datadoghq.com",
api_key="test-api-secret",
app_key="test-app-secret",
search=search,
now=clock.now,
sleep=clock.sleep,
jitter=lambda: 0.25,
), search
def test_429_honors_server_reset_and_preserves_duplicate_events() -> None:
clock: Final = Clock()
reader, search = _reader(
(StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": "6"}), _page("first", "duplicate")),
clock,
)
events: Final = reader.events_for_query("test-marker")
assert tuple(event.attributes["id"] for event in events) == ("first", "duplicate")
assert clock.elapsed == 6.25
assert search.calls == (("test-marker", 30.0), ("test-marker", 30.0))
@pytest.mark.parametrize("reset", ("", "invalid", "nan", "inf", "-1"))
def test_invalid_reset_uses_search_interval(reset: str) -> None:
clock: Final = Clock()
reader, _ = _reader(
(StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": reset}), _page()), clock
)
assert reader.events_for_query("test-marker") == []
assert clock.elapsed == DD_SEARCH_INTERVAL + 0.25
def test_zero_reset_cannot_create_a_busy_retry_loop() -> None:
clock: Final = Clock()
reader, _ = _reader(
(StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": "0"}), _page()), clock
)
assert reader.events_for_query("test-marker") == []
assert clock.elapsed == 1.25
def test_retry_after_is_not_shortened_by_an_earlier_reset() -> None:
clock: Final = Clock()
reader, _ = _reader(
(StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": "2", "retry-after": "8"}), _page()),
clock,
)
assert reader.events_for_query("test-marker") == []
assert clock.elapsed == 8.25
def test_rate_limit_wait_stops_at_deadline_without_issuing_another_request() -> None:
clock: Final = Clock()
reader, search = _reader(
(StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT * 10)}),), clock
)
with pytest.raises(pytest.fail.Exception, match="remained rate-limited"):
reader.events_for_query("test-marker")
assert clock.elapsed == POLL_TIMEOUT
assert search.calls == (("test-marker", 30.0),)
def test_late_retry_cannot_receive_a_fresh_request_timeout() -> None:
clock: Final = Clock()
reader, search = _reader(
(StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT - 5)}), _page()),
clock,
)
assert reader.events_for_query("test-marker") == []
assert search.calls == (("test-marker", 30.0), ("test-marker", 4.75))
@pytest.mark.parametrize("status", (-1, 401, 403, 500))
def test_non_quota_failures_are_not_retried_or_treated_as_empty_results(status: int) -> None:
clock: Final = Clock()
reader, search = _reader((StreamingResponse(status_code=status, body=""), _page()), clock)
with pytest.raises(pytest.fail.Exception, match=f"failed with HTTP {status}"):
reader.events_for_query("test-marker")
assert search.calls == (("test-marker", 30.0),)
assert clock.elapsed == 0
def test_polling_quota_retries_share_the_original_deadline() -> None:
clock: Final = Clock()
reader, search = _reader(
(_page(), StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT)})),
clock,
)
with pytest.raises(pytest.fail.Exception, match="remained rate-limited"):
reader.poll_events_for_query("test-marker")
assert clock.elapsed == POLL_TIMEOUT
assert len(search.calls) == 2
def test_empty_polling_does_not_start_a_final_search_after_its_deadline() -> None:
clock: Final = Clock()
attempts: Final = int(POLL_TIMEOUT / DD_SEARCH_INTERVAL)
reader, search = _reader((_page(),) * attempts, clock)
assert reader.poll_events_for_query("test-marker") == []
assert clock.elapsed == POLL_TIMEOUT
assert len(search.calls) == attempts
def test_settlement_quota_retries_keep_the_remaining_readback_budget() -> None:
clock: Final = Clock()
empty_reads: Final = int(POLL_TIMEOUT / DD_SEARCH_INTERVAL) - 2
reader, search = _reader(
(_page(),) * empty_reads
+ (_page("first"), StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT)})),
clock,
)
with pytest.raises(pytest.fail.Exception, match="remained rate-limited"):
reader.poll_events_for_query("test-marker")
assert clock.elapsed == POLL_TIMEOUT
assert search.calls[-1] == ("test-marker", DD_SEARCH_INTERVAL)
assert len(search.calls) == empty_reads + 2
def test_settlement_detects_a_duplicate_on_the_final_search() -> None:
clock: Final = Clock()
reader, _ = _reader((_page("first"), _page("first"), _page(), _page("first", "duplicate")), clock)
events: Final = reader.poll_events_for_query("test-marker")
assert tuple(event.attributes["id"] for event in events) == ("first", "duplicate")
assert clock.elapsed == 30
def test_settlement_keeps_confirmed_events_through_empty_searches() -> None:
clock: Final = Clock()
reader, _ = _reader((_page("first"), _page(), _page(), _page()), clock)
events: Final = reader.poll_events_for_query("test-marker")
assert tuple(event.attributes["id"] for event in events) == ("first",)
assert clock.elapsed == 30
def test_late_delivery_cannot_pass_without_a_complete_settle_window() -> None:
clock: Final = Clock()
empty_reads: Final = int(POLL_TIMEOUT / DD_SEARCH_INTERVAL) - 2
reader, search = _reader((_page(),) * empty_reads + (_page("first"), _page("first")), clock)
with pytest.raises(pytest.fail.Exception, match="duplicate-detection window"):
reader.poll_events_for_query("test-marker")
assert clock.elapsed == POLL_TIMEOUT
assert len(search.calls) == empty_reads + 2

35
tests/e2e/management/management_client.py
View file

@ -43,6 +43,9 @@ from models import (
KeyResetSpendBody,
KeyResetSpendResponse,
KeyUpdateBody,
McpServerCreateBody,
McpServerRow,
McpServerUpdateBody,
ModelDeleteBody,
OrgDeleteBody,
OrgInfoParams,
@ -537,6 +540,38 @@ class ManagementClient:
).root
)
def create_mcp_server(self, body: McpServerCreateBody) -> McpServerRow:
return unwrap(
self.proxy.transport.post(
"/v1/mcp/server",
headers=self.proxy.transport.master,
json=body,
response_type=McpServerRow,
)
)
def update_mcp_server(self, body: McpServerUpdateBody) -> McpServerRow:
"""PUT /v1/mcp/server, the call behind the dashboard's Save Changes: a partial
update where a field left unset keeps its stored value and None clears it."""
return unwrap(
self.proxy.transport.put(
"/v1/mcp/server",
headers=self.proxy.transport.master,
json=body,
response_type=McpServerRow,
)
)
def delete_mcp_server(self, server_id: str) -> Result[NoBody]:
"""DELETE /v1/mcp/server/{server_id}. Returns the outcome so the act phase can
unwrap it while a deferred teardown can ignore an already-deleted server."""
return self.proxy.transport.delete(
f"/v1/mcp/server/{server_id}",
headers=self.proxy.transport.master,
json=NoBody(),
response_type=NoBody,
)
def chat_status(self, key: str, model: str, content: str) -> StreamingResponse:
return self.proxy.transport.send(
"/chat/completions",

294
tests/e2e/management/test_mcp_lifecycle_e2e.py Normal file
View file

@ -0,0 +1,294 @@
"""Live e2e: the MCP server and toolset management routes' lifecycle contract.
Two customer defects sit on these routes, and each step here is the read-back that
would have caught one of them: a dashboard edit that took several saves to stick
because the read landed on a replica the write had not reached, and a toolset whose
tools were stored under one name and read back under another, so it granted
nothing. Every read-back therefore polls every replica that serves the route
(ProxyClient.read_back_everywhere) and asserts the exact values written, and both
update routes are held to the same partial-update contract: a field left out of the
payload keeps its stored value, a field sent as null is cleared. The server URL is
unreachable on purpose; only persistence is under test, never a tool call.
"""
from __future__ import annotations
from collections.abc import Callable, Mapping
from typing import Final
import pytest
from e2e_config import unique_marker
from e2e_http import unwrap
from lifecycle import ResourceManager
from management_client import ManagementClient
from models import (
McpInfo,
McpServerCreateBody,
McpServerListResponse,
McpServerRow,
McpServerUpdateBody,
ToolsetCreateBody,
ToolsetListResponse,
ToolsetRow,
ToolsetTool,
ToolsetUpdateBody,
)
pytestmark = pytest.mark.e2e
UNREACHABLE_URL: Final = "https://e2e-fake-mcp.test.local/mcp"
def _create_server(client: ManagementClient, resources: ResourceManager) -> tuple[McpServerCreateBody, str]:
name: Final = f"e2e_mcp_lifecycle_{unique_marker()}"
body: Final = McpServerCreateBody(
server_name=name,
alias=name,
url=UNREACHABLE_URL,
transport="http",
description="e2e lifecycle server",
mcp_info=McpInfo(
server_name=f"{name} (display)",
description="shown on the MCP page",
logo_url="https://e2e.test.local/logo.png",
),
)
server_id: Final = client.create_mcp_server(body).server_id
resources.defer(lambda: client.delete_mcp_server(server_id))
return body, server_id
def _assert_server_matches(row: McpServerRow, written: McpServerCreateBody, *, where: str) -> None:
stored: Final = (row.server_name, row.alias, row.url, row.transport, row.description, row.mcp_info)
expected: Final = (
written.server_name,
written.alias,
written.url,
written.transport,
written.description,
written.mcp_info,
)
assert stored == expected, f"{where}: stored {stored}, expected {expected}"
def _server_everywhere(
client: ManagementClient, server_id: str, *, settled: Callable[[McpServerRow], bool]
) -> Mapping[str, McpServerRow]:
return client.proxy.read_body_back_everywhere(f"/v1/mcp/server/{server_id}", McpServerRow, settled=settled)
def _listed_server_everywhere(client: ManagementClient, server_id: str) -> Mapping[str, McpServerRow]:
listings: Final = client.proxy.read_body_back_everywhere(
"/v1/mcp/server",
McpServerListResponse,
settled=lambda rows: any(row.server_id == server_id for row in rows.root),
)
return {replica: next(row for row in rows.root if row.server_id == server_id) for replica, rows in listings.items()}
class TestMcpServerLifecycle:
@pytest.mark.covers("mgmt.mcp_server.new.persists")
def test_create_persists_every_field_on_every_replica(
self, client: ManagementClient, resources: ResourceManager
) -> None:
body, server_id = _create_server(client, resources)
by_id: Final = _server_everywhere(client, server_id, settled=lambda row: row.server_id == server_id)
for replica, row in by_id.items():
_assert_server_matches(row, body, where=f"GET /v1/mcp/server/{server_id} on {replica}")
@pytest.mark.skip(
reason=(
"product gap: GET /v1/mcp/server builds each row from the in-memory registry, whose "
"_build_mcp_server_table sets description from mcp_info['description'], so the list "
"reports the mcp_info description while GET /v1/mcp/server/{server_id} reports the "
"stored description column. A server created with both set to different text reads "
"back with two different descriptions depending on the route"
)
)
@pytest.mark.covers("mgmt.mcp_server.list.persists")
def test_created_server_is_listed_with_every_field(
self, client: ManagementClient, resources: ResourceManager
) -> None:
body, server_id = _create_server(client, resources)
for replica, row in _listed_server_everywhere(client, server_id).items():
_assert_server_matches(row, body, where=f"GET /v1/mcp/server on {replica}")
@pytest.mark.covers("mgmt.mcp_server.update.preserves_unrelated_fields")
def test_updating_only_the_alias_keeps_every_other_field_on_every_replica(
self, client: ManagementClient, resources: ResourceManager
) -> None:
body, server_id = _create_server(client, resources)
renamed: Final = f"{body.alias}_renamed"
_ = client.update_mcp_server(McpServerUpdateBody(server_id=server_id, alias=renamed))
after_one_put: Final = _server_everywhere(client, server_id, settled=lambda row: row.alias == renamed)
for replica, row in after_one_put.items():
_assert_server_matches(
row,
body.model_copy(update={"alias": renamed}),
where=f"GET /v1/mcp/server/{server_id} on {replica} after one PUT of alias",
)
@pytest.mark.covers("mgmt.mcp_server.update.clear_persists")
def test_clearing_the_description_with_null_reads_back_null(
self, client: ManagementClient, resources: ResourceManager
) -> None:
body, server_id = _create_server(client, resources)
_ = client.update_mcp_server(McpServerUpdateBody(server_id=server_id, description=None))
cleared: Final = _server_everywhere(client, server_id, settled=lambda row: row.description is None)
for replica, row in cleared.items():
_assert_server_matches(
row,
body.model_copy(update={"description": None}),
where=f"GET /v1/mcp/server/{server_id} on {replica} after PUT description=null",
)
@pytest.mark.covers("mgmt.mcp_server.delete.persists")
def test_delete_removes_the_server_from_every_replica(
self, client: ManagementClient, resources: ResourceManager
) -> None:
_, server_id = _create_server(client, resources)
_ = unwrap(client.delete_mcp_server(server_id))
gone: Final = client.proxy.gone_everywhere(f"/v1/mcp/server/{server_id}")
assert set(gone.values()) == {404}, f"a deleted server must 404 on every replica; got {dict(gone)}"
listings: Final = client.proxy.read_body_back_everywhere(
"/v1/mcp/server",
McpServerListResponse,
settled=lambda rows: all(row.server_id != server_id for row in rows.root),
)
for replica, rows in listings.items():
assert all(row.server_id != server_id for row in rows.root), (
f"GET /v1/mcp/server on {replica} still lists the deleted server {server_id}"
)
def _create_toolset(
client: ManagementClient, resources: ResourceManager, server_id: str
) -> tuple[ToolsetCreateBody, str]:
body: Final = ToolsetCreateBody(
toolset_name=f"e2e_toolset_{unique_marker()}",
description="e2e lifecycle toolset",
tools=[
ToolsetTool(server_id=server_id, tool_name="search_datadog_logs"),
ToolsetTool(server_id=server_id, tool_name="get_datadog_metric"),
],
)
toolset_id: Final = client.proxy.create_toolset(body).toolset_id
resources.defer(lambda: client.proxy.delete_toolset(toolset_id))
return body, toolset_id
def _assert_toolset_matches(row: ToolsetRow, written: ToolsetCreateBody, *, where: str) -> None:
stored: Final = (row.toolset_name, row.description, row.tools)
expected: Final = (written.toolset_name, written.description, written.tools)
assert stored == expected, f"{where}: stored {stored}, expected {expected}"
def _toolset_everywhere(
client: ManagementClient, toolset_id: str, *, settled: Callable[[ToolsetRow], bool]
) -> Mapping[str, ToolsetRow]:
return client.proxy.read_body_back_everywhere(f"/v1/mcp/toolset/{toolset_id}", ToolsetRow, settled=settled)
class TestMcpToolsetLifecycle:
@pytest.mark.covers("mgmt.mcp_toolset.new.persists")
def test_create_persists_both_tools_under_the_exact_names_written(
self, client: ManagementClient, resources: ResourceManager
) -> None:
_, server_id = _create_server(client, resources)
body, toolset_id = _create_toolset(client, resources, server_id)
by_id: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.toolset_id == toolset_id)
for replica, row in by_id.items():
_assert_toolset_matches(row, body, where=f"GET /v1/mcp/toolset/{toolset_id} on {replica}")
listings: Final = client.proxy.read_body_back_everywhere(
"/v1/mcp/toolset",
ToolsetListResponse,
settled=lambda rows: any(row.toolset_id == toolset_id for row in rows.root),
)
for replica, rows in listings.items():
_assert_toolset_matches(
next(row for row in rows.root if row.toolset_id == toolset_id),
body,
where=f"GET /v1/mcp/toolset on {replica}",
)
@pytest.mark.covers("mgmt.mcp_toolset.update.preserves_unrelated_fields")
def test_updating_only_the_description_keeps_the_tools_and_name(
self, client: ManagementClient, resources: ResourceManager
) -> None:
_, server_id = _create_server(client, resources)
body, toolset_id = _create_toolset(client, resources, server_id)
_ = client.proxy.update_toolset(ToolsetUpdateBody(toolset_id=toolset_id, description="edited"))
edited: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.description == "edited")
for replica, row in edited.items():
_assert_toolset_matches(
row,
body.model_copy(update={"description": "edited"}),
where=f"GET /v1/mcp/toolset/{toolset_id} on {replica} after PUT of description",
)
@pytest.mark.covers("mgmt.mcp_toolset.update.persists")
def test_updating_the_tools_to_one_entry_reads_back_exactly_that_entry(
self, client: ManagementClient, resources: ResourceManager
) -> None:
_, server_id = _create_server(client, resources)
body, toolset_id = _create_toolset(client, resources, server_id)
kept: Final = body.tools[:1]
_ = client.proxy.update_toolset(ToolsetUpdateBody(toolset_id=toolset_id, tools=kept))
narrowed: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.tools == kept)
for replica, row in narrowed.items():
_assert_toolset_matches(
row,
body.model_copy(update={"tools": kept}),
where=f"GET /v1/mcp/toolset/{toolset_id} on {replica} after PUT of one tool",
)
@pytest.mark.covers("mgmt.mcp_toolset.update.clear_persists")
def test_clearing_the_description_with_null_reads_back_null(
self, client: ManagementClient, resources: ResourceManager
) -> None:
_, server_id = _create_server(client, resources)
body, toolset_id = _create_toolset(client, resources, server_id)
_ = client.proxy.update_toolset(ToolsetUpdateBody(toolset_id=toolset_id, description=None))
cleared: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.description is None)
for replica, row in cleared.items():
_assert_toolset_matches(
row,
body.model_copy(update={"description": None}),
where=f"GET /v1/mcp/toolset/{toolset_id} on {replica} after PUT description=null",
)
@pytest.mark.covers("mgmt.mcp_toolset.delete.persists")
def test_delete_removes_the_toolset_from_every_replica(
self, client: ManagementClient, resources: ResourceManager
) -> None:
_, server_id = _create_server(client, resources)
_, toolset_id = _create_toolset(client, resources, server_id)
_ = unwrap(client.proxy.delete_toolset(toolset_id))
gone: Final = client.proxy.gone_everywhere(f"/v1/mcp/toolset/{toolset_id}")
assert set(gone.values()) == {404}, f"a deleted toolset must 404 on every replica; got {dict(gone)}"
listings: Final = client.proxy.read_body_back_everywhere(
"/v1/mcp/toolset",
ToolsetListResponse,
settled=lambda rows: all(row.toolset_id != toolset_id for row in rows.root),
)
for replica, rows in listings.items():
assert all(row.toolset_id != toolset_id for row in rows.root), (
f"GET /v1/mcp/toolset on {replica} still lists the deleted toolset {toolset_id}"
)

7
tests/e2e/mcp/datadog_mcp.py
View file

@ -3,6 +3,7 @@
from __future__ import annotations
import os
from collections.abc import Sequence
from e2e_config import datadog_mcp_url, unique_marker
from lifecycle import ResourceManager
@ -35,7 +36,11 @@ def register_datadog_mcp(
resources: ResourceManager,
*,
mcp_access_groups: list[str] | None = None,
allowed_tools: Sequence[str] | None = (SEARCH_LOGS_TOOL,),
) -> str:
"""Register the core Datadog toolset with its credentials from the env. By default
the server exposes only `search_datadog_logs`; pass `allowed_tools=None` to expose
every tool the core toolset serves."""
assert_dd_mcp_creds()
name = f"e2e_dd_mcp_{unique_marker()}"
server_id = client.register_server(
@ -47,7 +52,7 @@ def register_datadog_mcp(
"DD-API-KEY": _dd_api_key(),
"DD-APPLICATION-KEY": _dd_app_key(),
},
allowed_tools=[SEARCH_LOGS_TOOL],
allowed_tools=None if allowed_tools is None else list(allowed_tools),
mcp_access_groups=mcp_access_groups,
)
resources.defer(lambda: client.delete_server(server_id))

39
tests/e2e/mcp/mcp_client.py
View file

@ -16,11 +16,11 @@ import time
from collections.abc import Mapping
from dataclasses import dataclass
from pydantic import BaseModel, ConfigDict, Field, RootModel
from pydantic import BaseModel, ConfigDict, Field
from e2e_config import settle_propagation
from e2e_http import Headers, NoBody, Result, Success, UnknownApiError, unwrap
from models import KeyGenerateBody, ObjectPermission
from models import KeyGenerateBody, McpServerListResponse, McpServerRow, ObjectPermission
from proxy_client import ProxyClient
McpToolArg = str | int | float | bool | list[str] | dict[str, str]
@ -46,16 +46,6 @@ class McpServerNewResponse(BaseModel):
server_id: str
class McpServerRow(BaseModel):
server_id: str
alias: str | None = None
url: str | None = None
class McpServersListResponse(RootModel[list[McpServerRow]]):
pass
class McpToolMcpInfo(BaseModel):
server_id: str | None = None
alias: str | None = None
@ -193,7 +183,7 @@ class McpClient:
"/v1/mcp/server",
headers=self.proxy.transport.master,
params=NoBody(),
response_type=McpServersListResponse,
response_type=McpServerListResponse,
)
).root
@ -224,11 +214,16 @@ class McpClient:
user_id: str,
mcp_servers: list[str] | None,
mcp_access_groups: list[str] | None = None,
mcp_toolsets: list[str] | None = None,
models: list[str] | None = None,
) -> str:
object_permission = (
ObjectPermission(mcp_servers=mcp_servers, mcp_access_groups=mcp_access_groups)
if mcp_servers is not None or mcp_access_groups is not None
ObjectPermission(
mcp_servers=mcp_servers,
mcp_access_groups=mcp_access_groups,
mcp_toolsets=mcp_toolsets,
)
if mcp_servers is not None or mcp_access_groups is not None or mcp_toolsets is not None
else None
)
return self.proxy.generate_key(
@ -272,6 +267,20 @@ class McpClient:
)
time.sleep(self.proxy.poll_interval)
def await_tools(self, key: str, server_id: str, *, expected: frozenset[str]) -> frozenset[str]:
"""Poll tools/list until `server_id`'s tools as `key` sees them are exactly
`expected`, and return the last listing either way, so the caller's equality
assertion names the difference. Fails at poll_timeout only when the read
itself never succeeded."""
deadline = time.monotonic() + self.proxy.poll_timeout
while True:
result = self.list_tools(key)
if isinstance(result, Success) and result.data.tool_names_for_server(server_id) == expected:
return expected
if time.monotonic() >= deadline:
return unwrap(result).tool_names_for_server(server_id)
time.sleep(self.proxy.poll_interval)
def await_call_tool(
self,
key: str,

95
tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py Normal file
View file

@ -0,0 +1,95 @@
"""Live e2e: a key granted a toolset lists exactly the toolset's tools.
An admin registers the real Datadog remote MCP server with its whole core toolset
exposed, discovers two of its tool names through a key granted the server outright,
and curates a toolset naming exactly those two. A second key is granted the server
plus that toolset, and its tools/list must come back as exactly those two names: no
more, so the rest of the server's catalog stays hidden behind the toolset, and no
fewer, so a tool stored under one name and read under another (which granted
nothing) fails here first. Requires DD_API_KEY + DD_APP_KEY (the suite's real MCP
upstream).
"""
from __future__ import annotations
from typing import Final
import pytest
from datadog_mcp import SEARCH_LOGS_TOOL, register_datadog_mcp
from e2e_config import unique_marker
from e2e_http import unwrap
from lifecycle import ResourceManager
from mcp_client import McpClient
from models import ToolsetCreateBody, ToolsetTool
pytestmark = pytest.mark.e2e
def _key(
client: McpClient,
resources: ResourceManager,
label: str,
*,
server_id: str,
toolset_id: str | None = None,
) -> str:
key: Final = client.generate_key(
user_id=f"e2e-mcp-{label}-{unique_marker()}",
mcp_servers=[server_id],
mcp_toolsets=None if toolset_id is None else [toolset_id],
)
resources.defer(lambda: client.proxy.delete_key(key))
return key
def _wire_prefix(wire_name: str, tool_name: str, catalog: frozenset[str]) -> str:
"""The prefix tools/list puts in front of one server's tool names, measured off a
tool whose own name is known rather than guessed from the alias. A toolset grants
by the tool's own name, never the wire name, and the prefix is whatever the proxy
is configured to build (the alias, or a short server id), so measuring it is the
only way to cross between the two."""
assert wire_name.endswith(tool_name), f"tools/list served {wire_name!r}, expected it to end with {tool_name!r}"
prefix: Final = wire_name[: len(wire_name) - len(tool_name)]
unprefixed: Final = frozenset(name for name in catalog if not name.startswith(prefix))
assert not unprefixed, (
f"every tool of one server shares the wire prefix {prefix!r}, so {sorted(unprefixed)} "
f"cannot be reduced to the names a toolset grants by"
)
return prefix
class TestMcpToolsetEnforcement:
@pytest.mark.covers("mcp.list_tools.api_key.toolset_scoped")
def test_key_granted_a_toolset_lists_exactly_its_tools(self, client: McpClient, resources: ResourceManager) -> None:
server_id: Final = register_datadog_mcp(client, resources, allowed_tools=None)
client.await_registered(server_id)
catalog_key: Final = _key(client, resources, "catalog", server_id=server_id)
known_wire: Final = client.await_tool(catalog_key, server_id, SEARCH_LOGS_TOOL)
catalog: Final = unwrap(client.list_tools(catalog_key)).tool_names_for_server(server_id)
assert len(catalog) > 2, (
f"the Datadog core toolset must serve more tools than the toolset names, or the "
f"restriction has nothing to hide; got {sorted(catalog)}"
)
prefix: Final = _wire_prefix(known_wire, SEARCH_LOGS_TOOL, catalog)
chosen_wire: Final = frozenset(sorted(catalog)[:2])
chosen: Final = frozenset(name.removeprefix(prefix) for name in chosen_wire)
toolset: Final = client.proxy.create_toolset(
ToolsetCreateBody(
toolset_name=f"e2e_toolset_{unique_marker()}",
description="two Datadog tools",
tools=[ToolsetTool(server_id=server_id, tool_name=name) for name in sorted(chosen)],
)
)
resources.defer(lambda: client.proxy.delete_toolset(toolset.toolset_id))
assert frozenset(tool.tool_name for tool in toolset.tools) == chosen, (
f"toolset stored {toolset.tools}, expected the two names {sorted(chosen)} verbatim"
)
scoped_key: Final = _key(client, resources, "toolset", server_id=server_id, toolset_id=toolset.toolset_id)
listed: Final = client.await_tools(scoped_key, server_id, expected=chosen_wire)
assert listed == chosen_wire, (
f"a key granted the toolset must list exactly its two tools; "
f"got {sorted(listed)}, expected {sorted(chosen_wire)}"
)

73
tests/e2e/models.py
View file

@ -10,6 +10,7 @@ from collections.abc import Sequence
from datetime import datetime
from typing import Final, Literal
from e2e_http import PartialBody
from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_serializer, model_validator
# ---------- keys ----------
@ -55,6 +56,7 @@ class KeyMetadata(BaseModel):
class ObjectPermission(BaseModel):
mcp_servers: list[str] | None = None
mcp_access_groups: list[str] | None = None
mcp_toolsets: list[str] | None = None
class KeyGenerateBody(BaseModel):
@ -77,7 +79,7 @@ class KeyGenerateBody(BaseModel):
allowed_passthrough_routes: list[str] | None = None
metadata: KeyMetadata | None = None
object_permission: ObjectPermission | None = None
router_settings: "RouterSettingsOverride | None" = None
router_settings: RouterSettingsOverride | None = None
class KeyGenerateResponse(BaseModel):
@ -516,6 +518,15 @@ class CountTokensResponse(BaseModel):
# ---------- mcp servers ----------
class McpInfo(BaseModel):
"""The `mcp_info` display block stored on an MCP server; only the fields the
lifecycle test writes and reads back."""
server_name: str | None = None
description: str | None = None
logo_url: str | None = None
class McpServerCreateBody(BaseModel):
"""POST /v1/mcp/server. For a gateway-managed OAuth server, `auth_type` is
`oauth2` and `oauth2_flow` is `authorization_code`; the upstream endpoints
@ -530,6 +541,18 @@ class McpServerCreateBody(BaseModel):
oauth2_flow: Literal["client_credentials", "authorization_code"] | None = None
authorization_url: str | None = None
token_url: str | None = None
server_name: str | None = None
description: str | None = None
mcp_info: McpInfo | None = None
class McpServerUpdateBody(PartialBody):
"""PUT /v1/mcp/server: a field left unset keeps its stored value, a field set
to None is cleared."""
server_id: str
alias: str | None = None
description: str | None = None
class McpServerInfo(BaseModel):
@ -543,6 +566,54 @@ class McpServerInfo(BaseModel):
allow_all_keys: bool | None = None
class McpServerRow(McpServerInfo):
"""A stored MCP server as the create, get, and list routes return it: the
fields the lifecycle test asserts survive the round trip."""
server_name: str | None = None
transport: str | None = None
description: str | None = None
mcp_info: McpInfo | None = None
class McpServerListResponse(RootModel[list[McpServerRow]]):
"""GET /v1/mcp/server answers with a bare array of servers."""
class ToolsetTool(BaseModel):
server_id: str
tool_name: str
class ToolsetCreateBody(BaseModel):
toolset_name: str
description: str | None = None
tools: list[ToolsetTool]
class ToolsetUpdateBody(PartialBody):
"""PUT /v1/mcp/toolset: a field left unset keeps its stored value, a field set
to None is cleared."""
toolset_id: str
description: str | None = None
tools: list[ToolsetTool] | None = None
class ToolsetRow(BaseModel):
"""A stored toolset as POST /v1/mcp/toolset, GET /v1/mcp/toolset/{toolset_id},
and each row of GET /v1/mcp/toolset return it."""
toolset_id: str
toolset_name: str
description: str | None = None
tools: list[ToolsetTool] = Field(default_factory=list)
class ToolsetListResponse(RootModel[list[ToolsetRow]]):
"""GET /v1/mcp/toolset answers with a bare array of toolsets."""
class EmbedBody(BaseModel):
model: str
input: str

216
tests/e2e/proxy_client.py
View file

@ -12,6 +12,7 @@ import time
import warnings
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from functools import reduce
from datetime import datetime
from types import MappingProxyType
from typing import Final
@ -26,6 +27,7 @@ from e2e_http import (
Result,
StreamingResponse,
Success,
UnknownApiError,
is_ok,
unwrap,
)
@ -70,6 +72,9 @@ from models import (
SpendLogsPage,
SpendLogsPageParams,
SpendLogsParams,
ToolsetCreateBody,
ToolsetRow,
ToolsetUpdateBody,
)
from e2e_config import (
CONTROL_PLANE_BASE_URL,
@ -82,7 +87,7 @@ from e2e_config import (
SLOW_PROVIDER_TIMEOUT_SECONDS,
settle_propagation,
)
from transport import HttpTransport, SplitTransport, Transport
from transport import HttpTransport, SplitTransport, Transport, is_control_plane_path
RowsPredicate = Callable[[list[SpendLogRow]], bool]
@ -235,6 +240,99 @@ def servable_timeout_message(
)
type ReplicaRead[T] = Callable[[float], T]
@dataclass(frozen=True, slots=True)
class EverywhereConverged[T]:
"""Every replica answered with something `settled` accepts, keyed by replica."""
answers: Mapping[str, T]
@dataclass(frozen=True, slots=True)
class NeverConvergedOn[T]:
"""`replica` ran out its budget without an answer `settled` accepts; `last` is
its final answer, so the failure can say what that replica still serves."""
replica: str
last: T
def _last_answer[T](
read: ReplicaRead[T],
*,
settled: Callable[[T], bool],
timeout: float,
interval: float,
request_timeout: float,
now: Callable[[], float],
sleep: Callable[[float], None],
) -> T:
"""Poll `read` until `settled` accepts its answer or `timeout` runs out, and
return the last answer either way. Each read's request timeout is clamped to
the budget left, and the final poll runs even when less than an interval
remains, so a deadline never skips the read that would have settled."""
deadline: Final = now() + timeout
answer = read(min(request_timeout, timeout))
while not settled(answer):
remaining = deadline - now()
if remaining <= 0:
return answer
sleep(min(interval, remaining))
answer = read(min(request_timeout, remaining))
return answer
def await_everywhere[T](
reads: Mapping[str, ReplicaRead[T]],
*,
settled: Callable[[T], bool],
timeout: float,
interval: float,
request_timeout: float,
now: Callable[[], float],
sleep: Callable[[float], None],
) -> EverywhereConverged[T] | NeverConvergedOn[T]:
"""`_last_answer` against every replica in turn, each with the full budget, so a
write counts as visible only once the last replica reflects it, and stop at the
first replica that never converges. Clock and sleep are injected."""
def read_replica(
outcome: EverywhereConverged[T] | NeverConvergedOn[T],
item: tuple[str, ReplicaRead[T]],
) -> EverywhereConverged[T] | NeverConvergedOn[T]:
if isinstance(outcome, NeverConvergedOn):
return outcome
replica, read = item
answer: Final = _last_answer(
read,
settled=settled,
timeout=timeout,
interval=interval,
request_timeout=request_timeout,
now=now,
sleep=sleep,
)
if not settled(answer):
return NeverConvergedOn(replica=replica, last=answer)
return EverywhereConverged(answers=MappingProxyType({**outcome.answers, replica: answer}))
initial: Final[EverywhereConverged[T] | NeverConvergedOn[T]] = EverywhereConverged(answers=MappingProxyType({}))
return reduce(read_replica, reads.items(), initial)
def _is_not_found[R: BaseModel](result: Result[R]) -> bool:
return isinstance(result, UnknownApiError) and result.status_code == 404
def _status_of[R: BaseModel](result: Result[R]) -> int:
match result:
case Success(status_code=status_code) | UnknownApiError(status_code=status_code):
return status_code
case _:
return -1
type Poller[T] = Callable[[], T]
@ -321,6 +419,7 @@ def converge_timeout_message(*, what: str, replica: str, timeout: float, last_re
class ProxyClient:
transport: Transport
replicas: Mapping[str, Transport]
control_replicas: Mapping[str, Transport]
poll_timeout: float = 120.0
poll_interval: float = 5.0
model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT
@ -569,6 +668,112 @@ class ProxyClient:
if not is_ok(result):
warnings.warn(f"delete_model({model_id!r}) failed: {result}", stacklevel=2)
# ---- replica read-back ----------------------------------------------
def replicas_for(self, path: str) -> Mapping[str, Transport]:
"""The replicas that serve `path`: every data-plane replica for an LLM route,
and for a management route the control-plane replicas, since the data-plane
replicas trim management routes and answer them 404. A monolith serves both
from every replica, so a management read-back polls all of them; a split
deployment exposes one control-plane address (there is one backend process
behind it on the stack these suites run against), so it polls that. A
control plane fronting several backends would need its own replica list to
prove each one converged, the way PROXY_REPLICA_URLS does for the gateways.
Never empty: a read-back against no replica would assert nothing and pass."""
replicas: Final = self.control_replicas if is_control_plane_path(path) else self.replicas
assert replicas, f"no replica is configured to serve {path}, so a read-back there would prove nothing"
return replicas
def read_body_back_everywhere[R: BaseModel](
self, path: str, response_type: type[R], *, settled: Callable[[R], bool]
) -> Mapping[str, R]:
"""GET `path` on every replica that serves it, polling each to poll_timeout
until `settled` accepts its body, and fail naming the first replica that
never converged. Returns each replica's settled body, keyed by replica, so
the caller can assert the rest of it."""
outcome: Final = await_everywhere(
{url: self._reader(transport, path, response_type) for url, transport in self.replicas_for(path).items()},
settled=lambda result: isinstance(result, Success) and settled(result.data),
timeout=self.poll_timeout,
interval=self.poll_interval,
request_timeout=REQUEST_TIMEOUT,
now=time.monotonic,
sleep=time.sleep,
)
match outcome:
case EverywhereConverged(answers=answers):
return MappingProxyType({url: unwrap(result) for url, result in answers.items()})
case NeverConvergedOn(replica=replica, last=last):
raise AssertionError(
f"GET {path} on {replica} never converged within {self.poll_timeout}s of the write; "
f"last read: {last}"
)
def gone_everywhere(self, path: str) -> Mapping[str, int]:
"""Poll GET `path` on every replica that serves it until each stops serving
it, and fail naming the first replica that still does at poll_timeout.
Returns each replica's final status, so the caller asserts the 404 itself."""
outcome: Final = await_everywhere(
{url: self._reader(transport, path, NoBody) for url, transport in self.replicas_for(path).items()},
settled=_is_not_found,
timeout=self.poll_timeout,
interval=self.poll_interval,
request_timeout=REQUEST_TIMEOUT,
now=time.monotonic,
sleep=time.sleep,
)
match outcome:
case EverywhereConverged(answers=answers):
return MappingProxyType({url: _status_of(result) for url, result in answers.items()})
case NeverConvergedOn(replica=replica, last=last):
raise AssertionError(
f"GET {path} on {replica} still answers {self.poll_timeout}s after the delete; last read: {last}"
)
@staticmethod
def _reader[R: BaseModel](transport: Transport, path: str, response_type: type[R]) -> ReplicaRead[Result[R]]:
return lambda request_timeout: transport.get(
path,
headers=transport.master,
params=NoBody(),
response_type=response_type,
timeout=request_timeout,
)
# ---- mcp toolsets ---------------------------------------------------
def create_toolset(self, body: ToolsetCreateBody) -> ToolsetRow:
return unwrap(
self.transport.post(
"/v1/mcp/toolset",
headers=self.transport.master,
json=body,
response_type=ToolsetRow,
)
)
def update_toolset(self, body: ToolsetUpdateBody) -> ToolsetRow:
"""PUT /v1/mcp/toolset: a partial update where a field left unset keeps its
stored value and None clears it."""
return unwrap(
self.transport.put(
"/v1/mcp/toolset",
headers=self.transport.master,
json=body,
response_type=ToolsetRow,
)
)
def delete_toolset(self, toolset_id: str) -> Result[NoBody]:
"""DELETE /v1/mcp/toolset/{toolset_id}. Returns the outcome so the act phase
can unwrap it while a deferred teardown can ignore an already-deleted row."""
return self.transport.delete(
f"/v1/mcp/toolset/{toolset_id}",
headers=self.transport.master,
json=NoBody(),
response_type=NoBody,
)
def create_credential(self, body: CredentialCreateBody) -> None:
unwrap(
self.transport.post(
@ -736,7 +941,10 @@ def build_proxy_client(
base URLs are the same for a monolithic proxy, so routing is then a no-op.
``replica_urls`` (PROXY_REPLICA_URLS) names every data-plane replica the model
barrier polls directly; it is the data-plane URL itself unless the stack
exports each gateway's own address.
exports each gateway's own address. Management read-backs poll those same
replicas when the two planes share a base URL (a monolith, where every replica
serves every route) and the control plane alone when they differ (a split
deployment, where the data-plane replicas do not serve management routes).
The endpoints are injectable for callers that resolve the proxy some other
way than ``e2e_config``'s env names (see ``claude_code/_env.py``); they must
@ -764,9 +972,13 @@ def build_proxy_client(
for url in replica_urls
}
)
control_replicas: Final = (
replicas if control_plane_base_url == base_url else MappingProxyType({control_plane_base_url: split.control})
)
return ProxyClient(
transport=split,
replicas=replicas,
control_replicas=control_replicas,
poll_timeout=POLL_TIMEOUT,
poll_interval=POLL_INTERVAL,
)

239
tests/e2e/router/test_auto_router_regressions_e2e.py
View file

@ -41,6 +41,7 @@ which stores either the registered alias or the provider-prefixed form.
import json
import os
from collections.abc import Iterator
from contextlib import ExitStack
from dataclasses import dataclass
from typing import Final
@ -120,19 +121,10 @@ class ResponsesApiResponse(BaseModel):
@dataclass(frozen=True, slots=True)
class TagSplitDeployments:
"""Scenario A mirrors the customer-shaped config from GitHub issue #36619:
plain deployment registered first, tier deployment and marker both tagged.
Scenario B flips both axes for GitHub issue #36621: marker registered first
and its tier deployment left untagged, so routing depends neither on
registration order nor on tier deployments carrying tags."""
tag_a: str
shared_a: str
tier_a: str
tag_b: str
shared_b: str
tier_b: str
class TagSplitDeployment:
tag: str
shared: str
tier: str
@dataclass(frozen=True, slots=True)
@ -173,9 +165,7 @@ def _uniform_tier_config(tier_model: str) -> dict[str, object]:
}
def _key_for(
proxy: ProxyClient, resources: ResourceManager, models: list[str], tag_filtering: bool = False
) -> str:
def _key_for(proxy: ProxyClient, resources: ResourceManager, models: list[str], tag_filtering: bool = False) -> str:
key: Final = proxy.generate_key(
KeyGenerateBody(
models=models,
@ -211,46 +201,61 @@ def _assert_served_only_by(rows: list[SpendLogRow], allowed: frozenset[str], con
)
@pytest.fixture(scope="module")
def split(proxy: ProxyClient) -> Iterator[TagSplitDeployments]:
@pytest.fixture(scope="class")
def router_stack() -> Iterator[ExitStack]:
with ExitStack() as stack:
yield stack
def _register_models(
proxy: ProxyClient, stack: ExitStack, registrations: tuple[tuple[str, LiteLLMParamsBody], ...]
) -> None:
for name, params in registrations:
stack.callback(proxy.delete_model, proxy.create_model(name, params))
def _tag_split(proxy: ProxyClient, stack: ExitStack, *, marker_first: bool) -> TagSplitDeployment:
marker: Final = unique_marker()
deployments: Final = TagSplitDeployments(
tag_a=f"e2e-split-a-{marker}",
shared_a=f"e2e-autoroute-a-{marker}",
tier_a=f"e2e-tier-a-{marker}",
tag_b=f"e2e-split-b-{marker}",
shared_b=f"e2e-autoroute-b-{marker}",
tier_b=f"e2e-tier-b-{marker}",
named: Final = TagSplitDeployment(
tag=f"e2e-split-{marker}",
shared=f"e2e-autoroute-{marker}",
tier=f"e2e-tier-{marker}",
)
anthropic_key: Final = _provider_key("ANTHROPIC_API_KEY")
marker_params_a: Final = LiteLLMParamsBody(
model="auto_router/complexity_router",
complexity_router_config=_uniform_tier_config(deployments.tier_a),
tags=[deployments.tag_a],
marker_registration: Final = (
named.shared,
LiteLLMParamsBody(
model="auto_router/complexity_router",
complexity_router_config=_uniform_tier_config(named.tier),
tags=[named.tag],
),
)
marker_params_b: Final = LiteLLMParamsBody(
model="auto_router/complexity_router",
complexity_router_config=_uniform_tier_config(deployments.tier_b),
tags=[deployments.tag_b],
tier_registration: Final = (
named.tier,
LiteLLMParamsBody(model=CHEAP_MODEL, api_key=anthropic_key, tags=None if marker_first else [named.tag]),
)
registrations: Final[tuple[tuple[str, LiteLLMParamsBody], ...]] = (
(deployments.shared_a, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=anthropic_key)),
(deployments.tier_a, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=anthropic_key, tags=[deployments.tag_a])),
(deployments.shared_a, marker_params_a),
(deployments.shared_b, marker_params_b),
(deployments.tier_b, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=anthropic_key)),
(deployments.shared_b, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=anthropic_key)),
plain_registration: Final = (named.shared, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=anthropic_key))
registrations: Final = (
(marker_registration, tier_registration, plain_registration)
if marker_first
else (plain_registration, tier_registration, marker_registration)
)
created: Final = tuple(proxy.create_model(name, params) for name, params in registrations)
try:
yield deployments
finally:
for model_id in created:
proxy.delete_model(model_id)
_register_models(proxy, stack, registrations)
return named
@pytest.fixture(scope="module")
def zero_priced_alias(proxy: ProxyClient) -> Iterator[ZeroPricedAlias]:
@pytest.fixture(scope="class")
def plain_first_split(proxy: ProxyClient, router_stack: ExitStack) -> TagSplitDeployment:
return _tag_split(proxy, router_stack, marker_first=False)
@pytest.fixture(scope="class")
def marker_first_split(proxy: ProxyClient, router_stack: ExitStack) -> TagSplitDeployment:
return _tag_split(proxy, router_stack, marker_first=True)
@pytest.fixture(scope="class")
def zero_priced_alias(proxy: ProxyClient, router_stack: ExitStack) -> ZeroPricedAlias:
marker: Final = unique_marker()
named: Final = ZeroPricedAlias(alias=f"e2e-priced-alias-{marker}", tier=f"e2e-priced-tier-{marker}")
alias_params: Final = LiteLLMParamsBody(
@ -263,16 +268,12 @@ def zero_priced_alias(proxy: ProxyClient) -> Iterator[ZeroPricedAlias]:
(named.tier, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))),
(named.alias, alias_params),
)
created: Final = tuple(proxy.create_model(name, params) for name, params in registrations)
try:
yield named
finally:
for model_id in created:
proxy.delete_model(model_id)
_register_models(proxy, router_stack, registrations)
return named
@pytest.fixture(scope="module")
def heuristic_split(proxy: ProxyClient) -> Iterator[HeuristicSplit]:
@pytest.fixture(scope="class")
def heuristic_split(proxy: ProxyClient, router_stack: ExitStack) -> HeuristicSplit:
marker: Final = unique_marker()
named: Final = HeuristicSplit(
alias=f"e2e-heuristic-router-{marker}",
@ -289,16 +290,12 @@ def heuristic_split(proxy: ProxyClient) -> Iterator[HeuristicSplit]:
(named.strong, LiteLLMParamsBody(model=STRONG_MODEL, api_key=_provider_key("OPENAI_API_KEY"))),
(named.alias, LiteLLMParamsBody(model="auto_router/complexity_router", complexity_router_config=config)),
)
created: Final = tuple(proxy.create_model(name, params) for name, params in registrations)
try:
yield named
finally:
for model_id in created:
proxy.delete_model(model_id)
_register_models(proxy, router_stack, registrations)
return named
@pytest.fixture(scope="module")
def semantic_auto_router(proxy: ProxyClient) -> Iterator[SemanticAutoRouter]:
@pytest.fixture(scope="class")
def semantic_auto_router(proxy: ProxyClient, router_stack: ExitStack) -> SemanticAutoRouter:
marker: Final = unique_marker()
named: Final = SemanticAutoRouter(
marker=f"e2e-semantic-router-{marker}",
@ -321,16 +318,12 @@ def semantic_auto_router(proxy: ProxyClient) -> Iterator[SemanticAutoRouter]:
(named.fallback, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))),
(named.marker, marker_params),
)
created: Final = tuple(proxy.create_model(name, params) for name, params in registrations)
try:
yield named
finally:
for model_id in created:
proxy.delete_model(model_id)
_register_models(proxy, router_stack, registrations)
return named
@pytest.fixture(scope="module")
def credentialed_alias(proxy: ProxyClient) -> Iterator[CredentialedAlias]:
@pytest.fixture(scope="class")
def credentialed_alias(proxy: ProxyClient, router_stack: ExitStack) -> CredentialedAlias:
marker: Final = unique_marker()
named: Final = CredentialedAlias(alias=f"e2e-cred-alias-{marker}", tier=f"e2e-cred-tier-{marker}")
alias_params: Final = LiteLLMParamsBody(
@ -342,104 +335,110 @@ def credentialed_alias(proxy: ProxyClient) -> Iterator[CredentialedAlias]:
(named.tier, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))),
(named.alias, alias_params),
)
created: Final = tuple(proxy.create_model(name, params) for name, params in registrations)
try:
yield named
finally:
for model_id in created:
proxy.delete_model(model_id)
_register_models(proxy, router_stack, registrations)
return named
class TestTagSplitRouting:
@pytest.mark.covers("reliability.routing.tagged_marker.request_tag_selects_marker")
def test_body_tagged_chat_routes_through_the_marker_to_its_tier(
self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments
self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment
) -> None:
"""Pins GitHub issue #36619: with tag filtering on, a chat request whose
body metadata tags match the tagged marker under a shared model name is
answered by the marker's tier deployment, not by the plain deployment
that was registered under the name first."""
key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True)
chat: Final = unwrap(proxy.chat(key, _hello_chat_body(split.shared_a, tags=[split.tag_a])))
key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True)
chat: Final = unwrap(proxy.chat(key, _hello_chat_body(plain_first_split.shared, tags=[plain_first_split.tag])))
assert chat.choices, "tagged chat through the shared name returned no choices"
rows: Final = proxy.poll_logs_for_key(key, min_rows=1)
_assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "body-tagged chat on the shared name")
_assert_served_only_by(rows, CHEAP_SERVED | {plain_first_split.tier}, "body-tagged chat on the shared name")
@pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment")
def test_untagged_chat_is_always_served_by_the_plain_deployment(
self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments
self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment
) -> None:
"""Pins GitHub issue #36620: untagged chat requests to the shared name
succeed on every call and are all served by the plain deployment; the
tagged marker never captures them, so no intermittent auto-router
errors and no tier hijacking."""
key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True)
key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True)
for _ in range(5):
chat = unwrap(proxy.chat(key, _hello_chat_body(split.shared_a)))
chat = unwrap(proxy.chat(key, _hello_chat_body(plain_first_split.shared)))
assert chat.choices, "untagged chat through the shared name returned no choices"
rows: Final = proxy.poll_logs_for_key(key, min_rows=5)
_assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged chat on the shared name")
_assert_served_only_by(rows, PLAIN_SERVED | {plain_first_split.shared}, "untagged chat on the shared name")
@pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment")
def test_untagged_messages_is_served_by_the_plain_deployment(
self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments
self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment
) -> None:
"""Pins GitHub issue #36620 on the /v1/messages surface: an untagged
Anthropic-native request to the shared name is served by the plain
deployment, not captured by the tagged marker."""
key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True)
answer: Final = unwrap(proxy.messages(key, _hello_messages_body(split.shared_a)))
key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True)
answer: Final = unwrap(proxy.messages(key, _hello_messages_body(plain_first_split.shared)))
assert answer.content or answer.choices, "untagged /v1/messages returned neither content nor choices"
rows: Final = proxy.poll_logs_for_key(key, min_rows=1)
_assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged /v1/messages on the shared name")
_assert_served_only_by(
rows, PLAIN_SERVED | {plain_first_split.shared}, "untagged /v1/messages on the shared name"
)
class TestUntaggedTierDeployments:
@pytest.mark.covers("reliability.routing.tagged_marker.header_tag_selects_marker")
def test_header_tagged_messages_routes_through_the_marker_to_an_untagged_tier(
self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments
self, proxy: ProxyClient, resources: ResourceManager, marker_first_split: TagSplitDeployment
) -> None:
"""Pins GitHub issue #36621: a /v1/messages request tagged only via the
x-litellm-tags header selects the tagged marker, and the rewrite still
lands on the tier deployment even though that deployment carries no
tags, because the marker consumed the routing tags."""
key: Final = _key_for(proxy, resources, [split.shared_b, split.tier_b], tag_filtering=True)
headers: Final = TaggedAnthropicHeaders(authorization=f"Bearer {key}", x_litellm_tags=split.tag_b)
key: Final = _key_for(
proxy, resources, [marker_first_split.shared, marker_first_split.tier], tag_filtering=True
)
headers: Final = TaggedAnthropicHeaders(authorization=f"Bearer {key}", x_litellm_tags=marker_first_split.tag)
answer: Final = unwrap(
proxy.transport.post(
"/v1/messages",
headers=headers,
json=_hello_messages_body(split.shared_b),
json=_hello_messages_body(marker_first_split.shared),
response_type=AnthropicMessagesResponse,
)
)
assert answer.content or answer.choices, "header-tagged /v1/messages returned neither content nor choices"
rows: Final = proxy.poll_logs_for_key(key, min_rows=1)
_assert_served_only_by(rows, CHEAP_SERVED | {split.tier_b}, "header-tagged /v1/messages on the shared name")
_assert_served_only_by(
rows, CHEAP_SERVED | {marker_first_split.tier}, "header-tagged /v1/messages on the shared name"
)
@pytest.mark.covers("reliability.routing.tagged_marker.untagged_tier_deployments_still_served")
def test_body_tagged_chat_reaches_the_untagged_tier_after_marker_rewrite(
self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments
self, proxy: ProxyClient, resources: ResourceManager, marker_first_split: TagSplitDeployment
) -> None:
"""Pins the tag-consumption half of GitHub issue #36621: after the
tagged marker rewrites the request to its tier model, the consumed
routing tags no longer constrain deployment selection, so the untagged
tier deployment serves the request instead of a strict-tag denial."""
key: Final = _key_for(proxy, resources, [split.shared_b, split.tier_b], tag_filtering=True)
chat: Final = unwrap(proxy.chat(key, _hello_chat_body(split.shared_b, tags=[split.tag_b])))
key: Final = _key_for(
proxy, resources, [marker_first_split.shared, marker_first_split.tier], tag_filtering=True
)
chat: Final = unwrap(
proxy.chat(key, _hello_chat_body(marker_first_split.shared, tags=[marker_first_split.tag]))
)
assert chat.choices, "body-tagged chat through the marker-first shared name returned no choices"
rows: Final = proxy.poll_logs_for_key(key, min_rows=1)
_assert_served_only_by(rows, CHEAP_SERVED | {split.tier_b}, "body-tagged chat with untagged tier")
_assert_served_only_by(rows, CHEAP_SERVED | {marker_first_split.tier}, "body-tagged chat with untagged tier")
@pytest.mark.covers("reliability.routing.tagged_marker.tag_semantics_stay_strict")
def test_tagged_call_straight_at_an_untagged_deployment_stays_denied(
self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments
self, proxy: ProxyClient, resources: ResourceManager, marker_first_split: TagSplitDeployment
) -> None:
"""The tag-consumption fix must not loosen strict tag semantics: a
tagged request aimed directly at an untagged deployment (no marker
involved) is still rejected with the 401 tags-configuration error."""
key: Final = _key_for(proxy, resources, [split.tier_b], tag_filtering=True)
result: Final = proxy.chat(key, _hello_chat_body(split.tier_b, tags=[split.tag_b]))
key: Final = _key_for(proxy, resources, [marker_first_split.tier], tag_filtering=True)
result: Final = proxy.chat(key, _hello_chat_body(marker_first_split.tier, tags=[marker_first_split.tag]))
assert isinstance(result, UnauthorizedError), (
f"expected the tagged direct call to an untagged deployment to be denied with 401, got {result}"
)
@ -451,37 +450,39 @@ class TestUntaggedTierDeployments:
class TestResponsesApiTagRouting:
@pytest.mark.covers("reliability.routing.tagged_marker.responses_input_routes_through_marker")
def test_header_tagged_responses_with_string_input_routes_to_the_tier(
self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments
self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment
) -> None:
"""Pins the /v1/responses surface of the tag split (GitHub issues
#36620/#36621): a /v1/responses request with string input, tagged via
the x-litellm-tags header, succeeds and routes through the tagged
marker to its tier."""
key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True)
headers: Final = TaggedAuthHeaders(authorization=f"Bearer {key}", x_litellm_tags=split.tag_a)
key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True)
headers: Final = TaggedAuthHeaders(authorization=f"Bearer {key}", x_litellm_tags=plain_first_split.tag)
body: Final = ResponsesBody(
model=split.shared_a, input=f"say hello {unique_marker()}", max_output_tokens=64
model=plain_first_split.shared, input=f"say hello {unique_marker()}", max_output_tokens=64
)
answer: Final = unwrap(
proxy.transport.post("/v1/responses", headers=headers, json=body, response_type=ResponsesApiResponse)
)
assert answer.id, "header-tagged /v1/responses returned no response id"
rows: Final = proxy.poll_logs_for_key(key, min_rows=1)
_assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "header-tagged /v1/responses string input")
_assert_served_only_by(
rows, CHEAP_SERVED | {plain_first_split.tier}, "header-tagged /v1/responses string input"
)
@pytest.mark.covers("reliability.routing.tagged_marker.responses_input_routes_through_marker")
def test_body_tagged_responses_with_list_input_routes_to_the_tier(
self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments
self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment
) -> None:
"""Pins the body-tag and list-input combination of the same split:
/v1/responses with litellm_metadata.tags and structured input items
routes through the tagged marker to its tier."""
key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True)
key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True)
body: Final = ResponsesBody(
model=split.shared_a,
model=plain_first_split.shared,
input=[ResponsesInputItem(role="user", content=f"say hello {unique_marker()}")],
max_output_tokens=64,
litellm_metadata=ResponsesTagMetadata(tags=[split.tag_a]),
litellm_metadata=ResponsesTagMetadata(tags=[plain_first_split.tag]),
)
answer: Final = unwrap(
proxy.transport.post(
@ -493,18 +494,18 @@ class TestResponsesApiTagRouting:
)
assert answer.id, "body-tagged /v1/responses returned no response id"
rows: Final = proxy.poll_logs_for_key(key, min_rows=1)
_assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "body-tagged /v1/responses list input")
_assert_served_only_by(rows, CHEAP_SERVED | {plain_first_split.tier}, "body-tagged /v1/responses list input")
@pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment")
def test_untagged_responses_is_served_by_the_plain_deployment(
self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments
self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment
) -> None:
"""Pins the untagged half of the /v1/responses tag split: an untagged
request to the shared name is served by the plain deployment, matching
the chat and messages surfaces."""
key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True)
key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True)
body: Final = ResponsesBody(
model=split.shared_a, input=f"say hello {unique_marker()}", max_output_tokens=64
model=plain_first_split.shared, input=f"say hello {unique_marker()}", max_output_tokens=64
)
answer: Final = unwrap(
proxy.transport.post(
@ -516,7 +517,9 @@ class TestResponsesApiTagRouting:
)
assert answer.id, "untagged /v1/responses returned no response id"
rows: Final = proxy.poll_logs_for_key(key, min_rows=1)
_assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged /v1/responses on the shared name")
_assert_served_only_by(
rows, PLAIN_SERVED | {plain_first_split.shared}, "untagged /v1/responses on the shared name"
)
class TestStrategyAliasPricing:
@ -551,9 +554,7 @@ class TestComplexityHeuristicScope:
while the accompanying ~2KB agent system prompt is packed with enough
reasoning and complexity keywords that scoring the combined text lands
in REASONING; only ask-only scoring keeps this on the cheap tier."""
key: Final = _key_for(
proxy, resources, [heuristic_split.alias, heuristic_split.cheap, heuristic_split.strong]
)
key: Final = _key_for(proxy, resources, [heuristic_split.alias, heuristic_split.cheap, heuristic_split.strong])
body: Final = ChatBody(
model=heuristic_split.alias,
messages=[

91
tests/e2e/test_e2e_http.py
View file

@ -13,13 +13,24 @@ monkeypatches anything.
from __future__ import annotations
from collections.abc import Callable, Iterator, Mapping, Sequence
from dataclasses import dataclass, field
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
import pytest
from e2e_http import RETRY_ATTEMPTS, TRANSIENT_STATUSES, request_with_retry, streaming_outcome
from e2e_http import (
RETRY_ATTEMPTS,
TRANSIENT_STATUSES,
NoBody,
PartialBody,
Success,
ValidationError,
classify,
request_with_retry,
streaming_outcome,
wire_body,
)
from pydantic import BaseModel, TypeAdapter
@dataclass
@ -33,10 +44,10 @@ class FakeResponse:
@dataclass
class SleepRecorder:
delays: list[float] = field(default_factory=list)
delays: tuple[float, ...] = ()
def __call__(self, seconds: float) -> None:
self.delays.append(seconds)
self.delays += (seconds,)
def _issue_from(responses: Sequence[FakeResponse]) -> Callable[[], FakeResponse]:
@ -55,7 +66,7 @@ class TestTransientRetryPolicy:
sleep = SleepRecorder()
result = request_with_retry(_issue_from(responses), sleep=sleep)
assert result is responses[0]
assert sleep.delays == []
assert sleep.delays == ()
assert responses[0].close_calls == 0
def test_429_is_never_retried(self) -> None:
@ -63,7 +74,7 @@ class TestTransientRetryPolicy:
sleep = SleepRecorder()
result = request_with_retry(_issue_from(responses), sleep=sleep)
assert result is responses[0]
assert sleep.delays == []
assert sleep.delays == ()
assert responses[0].close_calls == 0
def test_overloaded_529_retries_with_backoff_then_returns_the_success(self) -> None:
@ -71,7 +82,7 @@ class TestTransientRetryPolicy:
sleep = SleepRecorder()
result = request_with_retry(_issue_from(responses), sleep=sleep)
assert result is responses[1]
assert sleep.delays == [0.5]
assert sleep.delays == (0.5,)
assert responses[0].close_calls == 1
assert responses[1].close_calls == 0
@ -80,7 +91,7 @@ class TestTransientRetryPolicy:
sleep = SleepRecorder()
result = request_with_retry(_issue_from(responses), sleep=sleep)
assert result is responses[RETRY_ATTEMPTS - 1]
assert sleep.delays == [0.5, 1.0]
assert sleep.delays == (0.5, 1.0)
assert [r.close_calls for r in responses] == [1, 1, 0, 0]
@ -134,3 +145,65 @@ class TestStreamEventArrivals:
assert result.stream_events == []
assert result.stream_event_arrivals == []
assert result.body == "bad request"
class _ServerUpdate(PartialBody):
server_id: str
alias: str | None = None
description: str | None = None
class _ServerCreate(BaseModel):
alias: str
description: str | None = None
class TestWireBody:
"""A partial-update body must put exactly the caller's choice on the wire: an
omitted field stays off it so the route keeps the stored value, and an explicit
None goes out as JSON null so the route clears it. Plain bodies keep dropping
None, which is what every create route expects."""
def test_partial_body_omits_unset_fields_and_sends_explicit_none_as_null(self) -> None:
assert wire_body(_ServerUpdate(server_id="s1", description=None)) == {"server_id": "s1", "description": None}
assert wire_body(_ServerUpdate(server_id="s1", alias="renamed")) == {"server_id": "s1", "alias": "renamed"}
def test_plain_body_drops_none_fields(self) -> None:
assert wire_body(_ServerCreate(alias="a", description=None)) == {"alias": "a"}
_JSON: Final[TypeAdapter[object]] = TypeAdapter(object)
@dataclass
class FakeJsonResponse:
"""The `classify` view of a response: a status, the raw body bytes, and the
parse that would raise on an empty one."""
status_code: int
content: bytes
@property
def ok(self) -> bool:
return self.status_code < 400
@property
def text(self) -> str:
return self.content.decode()
def json(self) -> object:
return _JSON.validate_json(self.content)
class TestClassifyEmptyBody:
"""A delete that answers 202 with no body is a success, not a parse failure:
the MCP server and toolset delete routes both answer that way, and reading it
as a failure would hide a delete that did not happen behind one that did."""
def test_empty_2xx_body_is_a_success(self) -> None:
result: Final = classify(FakeJsonResponse(status_code=202, content=b""), NoBody)
assert isinstance(result, Success) and result.status_code == 202
def test_body_that_is_not_json_is_still_a_validation_failure(self) -> None:
result: Final = classify(FakeJsonResponse(status_code=200, content=b"<html/>"), NoBody)
assert isinstance(result, ValidationError)

95
tests/e2e/test_proxy_client.py
View file

@ -15,28 +15,35 @@ from collections.abc import Iterable, Mapping
from dataclasses import dataclass
from itertools import chain, repeat
from types import MappingProxyType
from typing import Final
from typing import Final, cast
import pytest
from e2e_config import parse_replica_urls
from e2e_http import Result, Success
from models import KeyInfo, KeyInfoResponse, ModelListEntry, ModelsListResponse
from proxy_client import (
Poller,
ConvergeOutcome,
Converged,
EverywhereConverged,
ModelsPoller,
NeverConvergedOn,
NotConverged,
NotServableOn,
Poller,
ProxyClient,
ReplicaRead,
Servable,
await_converged_everywhere,
await_everywhere,
await_servable_everywhere,
first_lagging_replica,
build_proxy_client,
converge_timeout_message,
first_lagging_replica,
)
from transport import Transport
MODEL: Final = "gpt-under-test"
_NO_TRANSPORTS: Final = cast(Transport, None)
TIMEOUT: Final = 10.0
INTERVAL: Final = 2.0
RPM_BEFORE_UPDATE: Final = 100
@ -187,3 +194,83 @@ class TestParseReplicaUrls:
def test_falls_back_to_the_data_plane_address_when_unset(self) -> None:
assert parse_replica_urls("", "http://lb") == ("http://lb",)
def _answers(answers: Iterable[str]) -> ReplicaRead[str]:
it: Final = iter(answers)
return lambda _timeout: next(it)
def _await_everywhere(reads: Mapping[str, ReplicaRead[str]]) -> EverywhereConverged[str] | NeverConvergedOn[str]:
clock: Final = FakeClock()
return await_everywhere(
reads,
settled=lambda answer: answer == "renamed",
timeout=TIMEOUT,
interval=INTERVAL,
request_timeout=5.0,
now=clock.now,
sleep=clock.sleep,
)
class TestAwaitEverywhere:
def test_waits_for_the_lagging_replica_and_returns_every_settled_answer(self) -> None:
reads: Final = {
"gateway-1": _answers(repeat("renamed")),
"gateway-2": _answers(chain(repeat("stale", 2), repeat("renamed"))),
}
outcome: Final = _await_everywhere(reads)
assert isinstance(outcome, EverywhereConverged)
assert dict(outcome.answers) == {"gateway-1": "renamed", "gateway-2": "renamed"}
def test_names_the_replica_that_never_converges_with_what_it_last_served(self) -> None:
reads: Final = {
"gateway-1": _answers(repeat("renamed")),
"gateway-2": _answers(repeat("stale")),
}
assert _await_everywhere(reads) == NeverConvergedOn(replica="gateway-2", last="stale")
def test_polls_until_the_deadline_before_giving_up(self) -> None:
lagging: Final = chain(repeat("stale", int(TIMEOUT / INTERVAL)), repeat("renamed"))
outcome: Final = _await_everywhere({"gateway-1": _answers(lagging)})
assert isinstance(outcome, EverywhereConverged), outcome
class TestReplicasFor:
def test_split_deployment_reads_management_routes_back_from_the_control_plane(self) -> None:
client: Final = build_proxy_client(
base_url="http://lb",
control_plane_base_url="http://backend",
replica_urls=("http://gateway-1", "http://gateway-2"),
)
assert set(client.replicas_for("/key/info")) == {"http://backend"}
assert set(client.replicas_for("/v1/models")) == {"http://gateway-1", "http://gateway-2"}
def test_monolith_reads_management_routes_back_from_every_replica(self) -> None:
client: Final = build_proxy_client(
base_url="http://lb",
control_plane_base_url="http://lb",
replica_urls=("http://pod-1", "http://pod-2"),
)
assert set(client.replicas_for("/key/info")) == {"http://pod-1", "http://pod-2"}
def test_mcp_admin_routes_read_back_from_every_data_plane_replica(self) -> None:
"""/v1/mcp/* is a lazily mounted feature, so a data-plane replica serves it
too and answers from its own in-memory registry. Routing it to the control
plane would leave every replica but that one unproven, and would move the
tools/list barrier in mcp_client off the plane that serves tools/list."""
client: Final = build_proxy_client(
base_url="http://lb",
control_plane_base_url="http://backend",
replica_urls=("http://gateway-1", "http://gateway-2"),
)
assert set(client.replicas_for("/v1/mcp/server/abc")) == {"http://gateway-1", "http://gateway-2"}
assert set(client.replicas_for("/v1/mcp/toolset/abc")) == {"http://gateway-1", "http://gateway-2"}
def test_a_route_no_replica_serves_is_refused_rather_than_read_back_vacuously(self) -> None:
"""A read-back over zero replicas would satisfy every predicate and assert
nothing, so asking for one fails instead of passing silently."""
client: Final = ProxyClient(transport=_NO_TRANSPORTS, replicas={}, control_replicas={})
with pytest.raises(AssertionError, match="no replica is configured"):
_ = client.replicas_for("/v1/models")

10
tests/e2e/ui/helpers/navigation.ts
View file

@ -73,3 +73,13 @@ export async function clickTeamId(page: PlaywrightPage, teamId: string): Promise
await cell.click();
await expect(page.getByText("Back to Teams")).toBeVisible({ timeout: 10_000 });
}
export async function openKeyDetail(page: PlaywrightPage, alias: string): Promise<void> {
await page.getByPlaceholder("Search by key alias or ID").fill(alias);
const row = page.getByRole("row").filter({ hasText: alias });
await expect(row, `key row "${alias}" never appeared on the Virtual Keys page`).toBeVisible({ timeout: 15_000 });
await row.getByRole("button", { name: alias }).click();
await expect(page.getByText("Back to Keys"), `key detail for "${alias}" never opened`).toBeVisible({
timeout: 15_000,
});
}

51
tests/e2e/ui/helpers/traffic.ts
View file

@ -1,4 +1,4 @@
import { APIRequestContext, expect } from "@playwright/test";
import { APIRequestContext, APIResponse, expect } from "@playwright/test";
/** Model names served by fixtures/config.yml, both backed by the mock LLM server. */
export const CHAT_MODEL_A = "fake-openai-gpt-4";
@ -15,6 +15,9 @@ export const masterKey = (): string => process.env.LITELLM_MASTER_KEY || "sk-123
export const rootPath = (): string => process.env.SERVER_ROOT_PATH ?? "";
/** Date.now() alone collides: `--repeat-each` starts its copies inside the same millisecond. */
export const uniqueSuffix = (): string => `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
interface ChatOptions {
model: string;
prompt: string;
@ -25,9 +28,8 @@ interface ChatOptions {
traceId?: string;
}
/** POST /v1/chat/completions and return the completion id (the Logs Request ID). */
export async function sendChatCompletion(request: APIRequestContext, opts: ChatOptions): Promise<string> {
const res = await request.post(`${rootPath()}/v1/chat/completions`, {
const postChatCompletion = (request: APIRequestContext, opts: ChatOptions): Promise<APIResponse> =>
request.post(`${rootPath()}/v1/chat/completions`, {
headers: {
Authorization: `Bearer ${opts.apiKey ?? masterKey()}`,
"Content-Type": "application/json",
@ -39,12 +41,26 @@ export async function sendChatCompletion(request: APIRequestContext, opts: ChatO
...(opts.traceId ? { litellm_trace_id: opts.traceId } : {}),
},
});
/** POST /v1/chat/completions and return the completion id (the Logs Request ID). */
export async function sendChatCompletion(request: APIRequestContext, opts: ChatOptions): Promise<string> {
const res = await postChatCompletion(request, opts);
expect(res.ok(), `chat completion for ${opts.model} failed (${res.status()}): ${await res.text()}`).toBe(true);
const body = await res.json();
expect(body.choices?.[0]?.message?.content).toContain(MOCK_RESPONSE_TEXT);
return body.id as string;
}
export interface ChatAttempt {
status: number;
body: string;
}
export async function attemptChatCompletion(request: APIRequestContext, opts: ChatOptions): Promise<ChatAttempt> {
const res = await postChatCompletion(request, opts);
return { status: res.status(), body: await res.text() };
}
/** `key` is the sk- value to authenticate with; `token` is its hash, which spend aggregates are keyed by. */
export async function createVirtualKey(
request: APIRequestContext,
@ -66,6 +82,33 @@ export async function createVirtualKey(
};
}
export interface KeyInfo {
key_alias: string | null;
max_budget: number | null;
budget_duration: string | null;
budget_reset_at: string | null;
blocked: boolean | null;
models: string[];
team_id: string | null;
}
export async function readKeyInfo(request: APIRequestContext, token: string): Promise<KeyInfo> {
const res = await request.get(`${rootPath()}/key/info?key=${encodeURIComponent(token)}`, {
headers: { Authorization: `Bearer ${masterKey()}` },
});
expect(res.ok(), `GET /key/info for ${token} failed (${res.status()}): ${await res.text()}`).toBe(true);
const body = await res.json();
return body.info as KeyInfo;
}
export async function deleteVirtualKey(request: APIRequestContext, token: string): Promise<void> {
const res = await request.post(`${rootPath()}/key/delete`, {
headers: { Authorization: `Bearer ${masterKey()}`, "Content-Type": "application/json" },
data: { keys: [token] },
});
expect(res.ok(), `key delete for ${token} failed (${res.status()}): ${await res.text()}`).toBe(true);
}
/** Spend logs are flushed on a timer, so an assertion straight after a completion races the writer. */
export async function waitForSpendLog(
request: APIRequestContext,

208
tests/e2e/ui/tests/internal-user/internalUserKeyScope.spec.ts Normal file
View file

@ -0,0 +1,208 @@
import { test, expect, type APIRequestContext } from "@playwright/test";
import { Page } from "../../fixtures/pages";
import {
dismissFeedbackPopup,
navigateToPage,
openKeyDetail,
} from "../../helpers/navigation";
import {
CHAT_MODEL_A,
CHAT_MODEL_B,
MOCK_RESPONSE_TEXT,
attemptChatCompletion,
createVirtualKey,
deleteVirtualKey,
masterKey,
readKeyInfo,
rootPath,
uniqueSuffix,
} from "../../helpers/traffic";
const MEMBER_PASSWORD = "E2e-Team-Member-Pass-1!";
interface CreatedTeam {
readonly team_id: string;
}
function assertCreatedTeam(body: unknown): asserts body is CreatedTeam {
expect(body, "/team/new returned no team_id").toMatchObject({
team_id: expect.any(String),
});
}
async function postAsMaster(
request: APIRequestContext,
path: string,
data: Record<string, unknown>,
): Promise<unknown> {
const res = await request.post(`${rootPath()}${path}`, {
headers: {
Authorization: `Bearer ${masterKey()}`,
"Content-Type": "application/json",
},
data,
});
expect(
res.ok(),
`POST ${path} failed (${res.status()}): ${await res.text()}`,
).toBe(true);
return res.json();
}
test.describe("Internal User - own team key model scope", () => {
test.use({ storageState: { cookies: [], origins: [] } });
test("a team member narrows their own key's models and the proxy enforces it", async ({
page,
request,
}) => {
const suffix = uniqueSuffix();
const email = `team-member-${suffix}@test.local`;
const userId = `e2e-key-scope-user-${suffix}`;
const alias = `e2e-key-scope-${suffix}`;
const team = await postAsMaster(request, "/team/new", {
team_alias: `E2E Key Scope ${suffix}`,
models: [CHAT_MODEL_A, CHAT_MODEL_B],
team_member_permissions: ["/key/generate", "/key/update", "/key/info"],
});
assertCreatedTeam(team);
const teamId = team.team_id;
try {
await postAsMaster(request, "/user/new", {
user_id: userId,
user_email: email,
user_role: "internal_user",
auto_create_key: false,
});
await postAsMaster(request, "/user/update", {
user_id: userId,
password: MEMBER_PASSWORD,
});
await postAsMaster(request, "/team/member_add", {
team_id: teamId,
member: { role: "user", user_id: userId },
});
const created = await createVirtualKey(request, {
key_alias: alias,
team_id: teamId,
user_id: userId,
models: [],
});
try {
await page.goto("/ui/login");
await page.getByPlaceholder("Enter your username").fill(email);
await page
.getByPlaceholder("Enter your password")
.fill(MEMBER_PASSWORD);
await page.getByRole("button", { name: "Login", exact: true }).click();
await expect(
page.locator("a", { hasText: "Virtual Keys" }),
`${email} never reached the dashboard`,
).toBeVisible({ timeout: 30_000 });
await dismissFeedbackPopup(page);
await navigateToPage(page, Page.ApiKeys);
await openKeyDetail(page, alias);
await page.getByRole("tab", { name: "Settings" }).click();
await page.getByRole("button", { name: "Edit Settings" }).click();
await page.getByRole("combobox", { name: "Select models" }).click();
await expect(
page.getByRole("option", { name: CHAT_MODEL_A, exact: true }),
`the Models dropdown does not offer ${CHAT_MODEL_A} to a team member`,
).toBeVisible({ timeout: 15_000 });
await expect(
page.getByRole("option", { name: CHAT_MODEL_B, exact: true }),
`the Models dropdown does not offer ${CHAT_MODEL_B} to a team member`,
).toBeVisible();
await page
.getByRole("option", { name: CHAT_MODEL_A, exact: true })
.click();
await page.keyboard.press("Escape");
const updated = page.waitForResponse(
(res) =>
res.url().includes("/key/update") &&
res.request().method() === "POST",
);
await page.getByRole("button", { name: "Save Changes" }).click();
const updateStatus = (await updated).status();
expect(
updateStatus,
"a team member's own-key edit was refused",
).toBeGreaterThanOrEqual(200);
expect(
updateStatus,
"a team member's own-key edit was refused",
).toBeLessThan(300);
await expect(
page.getByText("Key updated successfully").first(),
).toBeVisible({ timeout: 15_000 });
await expect
.poll(
async () => (await readKeyInfo(request, created.token)).models,
{
message: `the narrowed model scope never reached /key/info for ${alias}`,
timeout: 20_000,
},
)
.toEqual([CHAT_MODEL_A]);
await expect
.poll(
async () =>
await attemptChatCompletion(request, {
model: CHAT_MODEL_B,
prompt: `out of scope ${suffix}`,
apiKey: created.key,
}),
{
message: `${CHAT_MODEL_B} was still served after the key was narrowed to ${CHAT_MODEL_A}`,
timeout: 30_000,
},
)
.toMatchObject({
status: 403,
body: expect.stringContaining(CHAT_MODEL_B),
});
const inScope = await attemptChatCompletion(request, {
model: CHAT_MODEL_A,
prompt: `in scope ${suffix}`,
apiKey: created.key,
});
expect(
inScope,
`${CHAT_MODEL_A} is no longer served by the narrowed key`,
).toMatchObject({
status: 200,
body: expect.stringContaining(MOCK_RESPONSE_TEXT),
});
} finally {
await deleteVirtualKey(request, created.token);
}
} finally {
await request.post(`${rootPath()}/user/delete`, {
headers: {
Authorization: `Bearer ${masterKey()}`,
"Content-Type": "application/json",
},
data: { user_ids: [userId] },
});
await request.post(`${rootPath()}/team/delete`, {
headers: {
Authorization: `Bearer ${masterKey()}`,
"Content-Type": "application/json",
},
data: { team_ids: [teamId] },
});
}
});
});

196
tests/e2e/ui/tests/internal-user/modelsByTeam.spec.ts Normal file
View file

@ -0,0 +1,196 @@
import {
test as base,
expect,
type Locator,
type Page as PlaywrightPage,
} from "@playwright/test";
import {
E2E_TEAM_CRUD_ALIAS,
E2E_TEAM_ORG_ALIAS,
INTERNAL_USER_STORAGE_PATH,
} from "../../constants";
import { Page } from "../../fixtures/pages";
import { navigateToPage } from "../../helpers/navigation";
import { readBack } from "../../helpers/roundTrip";
import { CHAT_MODEL_A, CHAT_MODEL_B, masterKey } from "../../helpers/traffic";
const MOCK_LLM_BASE = `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`;
const CURRENT_TEAM_VIEW = "Current Team Models";
const ALL_MODELS_VIEW = "All Available Models";
const PERSONAL_TEAM = "Personal";
const teamSelector = (page: PlaywrightPage): Locator =>
page.getByRole("combobox", { name: "Current team", exact: true });
const viewSelector = (page: PlaywrightPage): Locator =>
page.getByRole("combobox", { name: "View", exact: true });
async function chooseOption(
page: PlaywrightPage,
selector: Locator,
optionName: string,
): Promise<void> {
await selector.click();
const option = page.getByRole("option", { name: optionName, exact: true });
await expect(option, `option ${optionName} is offered`).toBeVisible({
timeout: 10_000,
});
await option.click();
await expect(
selector,
`${optionName} is the selection the control now reports`,
).toContainText(optionName, {
timeout: 10_000,
});
}
async function deleteDeployment(
page: PlaywrightPage,
id: string,
): Promise<void> {
const post = () =>
page.request.post("/model/delete", {
headers: { Authorization: `Bearer ${masterKey()}` },
data: { id },
});
const deleted = await post().catch(() => post());
expect(
deleted.ok(),
`cleanup: /model/delete ${id} returned ${deleted.status()}`,
).toBe(true);
}
function modelRow(page: PlaywrightPage, modelName: string): Locator {
return page.getByRole("row").filter({ hasText: modelName });
}
async function isRegistered(
page: PlaywrightPage,
modelName: string,
): Promise<boolean> {
const body = await readBack<{ data: { model_name?: string }[] }>(
page,
"/v2/model/info",
);
return body.data.some((row) => row.model_name === modelName);
}
const uniqueSuffix = (): string =>
`${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
const test = base.extend<{ ungrantedModelName: string }>({
ungrantedModelName: async ({ page }, use) => {
const ungrantedModelName = `e2e-ungranted-${uniqueSuffix()}`;
const created = await page.request.post("/model/new", {
headers: { Authorization: `Bearer ${masterKey()}` },
data: {
model_name: ungrantedModelName,
litellm_params: {
model: `openai/${ungrantedModelName}`,
api_base: MOCK_LLM_BASE,
api_key: "fake-key",
},
model_info: {},
},
});
expect(
created.ok(),
`/model/new failed: ${created.status()} ${await created.text()}`,
).toBe(true);
const ungrantedModelId = (await created.json()).model_info?.id;
expect(ungrantedModelId, "model id from /model/new").toBeTruthy();
try {
await expect
.poll(async () => await isRegistered(page, ungrantedModelName), {
message: `deployment ${ungrantedModelName} never appeared in /v2/model/info after create`,
timeout: 60_000,
})
.toBe(true);
await use(ungrantedModelName);
} finally {
await deleteDeployment(page, ungrantedModelId);
}
},
});
test.describe("Models and Endpoints for an internal user", () => {
test.use({ storageState: INTERNAL_USER_STORAGE_PATH });
test("shows an internal user exactly the models of the team they select", async ({
page,
ungrantedModelName,
}) => {
await navigateToPage(page, Page.Models);
await expect(
page.getByRole("tab", { name: "Your Models" }),
"an internal user lands on their own models tab, not an admin-only view",
).toBeVisible({ timeout: 15_000 });
await expect(
viewSelector(page),
"the models table opens scoped to the selected team",
).toContainText(CURRENT_TEAM_VIEW, { timeout: 15_000 });
await expect(
modelRow(page, ungrantedModelName),
`the personal view lists ${ungrantedModelName}, so it is on the proxy and reachable from this page`,
).toHaveCount(1, { timeout: 30_000 });
await chooseOption(page, teamSelector(page), E2E_TEAM_CRUD_ALIAS);
await expect(
modelRow(page, CHAT_MODEL_A),
`${E2E_TEAM_CRUD_ALIAS} lists ${CHAT_MODEL_A}`,
).toHaveCount(1, {
timeout: 15_000,
});
await expect(
modelRow(page, CHAT_MODEL_B),
`${E2E_TEAM_CRUD_ALIAS} lists ${CHAT_MODEL_B}`,
).toHaveCount(1, {
timeout: 15_000,
});
await expect(
modelRow(page, ungrantedModelName),
`${ungrantedModelName} is on the proxy but not granted to ${E2E_TEAM_CRUD_ALIAS}, so it must not be listed`,
).toHaveCount(0);
await chooseOption(page, teamSelector(page), E2E_TEAM_ORG_ALIAS);
await expect(
modelRow(page, CHAT_MODEL_A),
`${E2E_TEAM_ORG_ALIAS} lists ${CHAT_MODEL_A}`,
).toHaveCount(1, {
timeout: 15_000,
});
await expect(
page.getByTestId("pagination-range"),
`${E2E_TEAM_ORG_ALIAS} lists the one model it grants and nothing else`,
).toHaveText("Showing 1-1 of 1", { timeout: 15_000 });
await expect(
modelRow(page, CHAT_MODEL_B),
`${CHAT_MODEL_B} belongs to another team and must not leak into ${E2E_TEAM_ORG_ALIAS}`,
).toHaveCount(0);
await expect(
modelRow(page, ungrantedModelName),
`${ungrantedModelName} is granted to no team and must not leak into ${E2E_TEAM_ORG_ALIAS}`,
).toHaveCount(0);
await chooseOption(page, viewSelector(page), ALL_MODELS_VIEW);
await expect(
modelRow(page, CHAT_MODEL_A),
`switching to ${ALL_MODELS_VIEW} leaves the table populated rather than blanking it`,
).toHaveCount(1, { timeout: 15_000 });
await page.reload();
await expect(
teamSelector(page),
"the team selection is not persisted across a reload, so the table returns to the personal view",
).toContainText(PERSONAL_TEAM, { timeout: 15_000 });
await expect(
viewSelector(page),
"the view selection is not persisted across a reload either",
).toContainText(CURRENT_TEAM_VIEW, { timeout: 15_000 });
await expect(
modelRow(page, ungrantedModelName),
"the personal view still renders models after a reload rather than coming back empty",
).toHaveCount(1, { timeout: 30_000 });
});
});

252
tests/e2e/ui/tests/modelsPage/editLitellmParams.spec.ts Normal file
View file

@ -0,0 +1,252 @@
import {
test as base,
expect,
type Page as PlaywrightPage,
} from "@playwright/test";
import { ADMIN_STORAGE_PATH } from "../../constants";
import { Page } from "../../fixtures/pages";
import { navigateToPage } from "../../helpers/navigation";
import { captureRequestBody, readBack } from "../../helpers/roundTrip";
import { masterKey, sendChatCompletion } from "../../helpers/traffic";
const MOCK_LLM_BASE = `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`;
const CUSTOM_PARAM = "extra_headers";
const CUSTOM_PARAM_VALUE = { "X-E2E-Edit-Probe": "one" };
type StoredParams = Record<string, unknown>;
async function readStoredParams(
page: PlaywrightPage,
modelId: string,
): Promise<StoredParams> {
const body = await readBack<{ data: { litellm_params: StoredParams }[] }>(
page,
`/model/info?litellm_model_id=${modelId}`,
);
return body.data[0]?.litellm_params ?? {};
}
function paramsEditor(page: PlaywrightPage) {
return page.getByPlaceholder('"rpm": 100');
}
async function editParams(
page: PlaywrightPage,
mutate: (params: StoredParams) => StoredParams,
): Promise<void> {
await page.getByRole("button", { name: "Edit Settings" }).click();
const editor = paramsEditor(page);
await expect(
editor,
"the LiteLLM Params editor is reachable on every visit to the edit form",
).toBeVisible({
timeout: 15_000,
});
const shown = JSON.parse(await editor.inputValue()) as StoredParams;
await editor.fill(JSON.stringify(mutate(shown), null, 2));
}
async function deleteDeployment(
page: PlaywrightPage,
id: string,
): Promise<void> {
const post = () =>
page.request.post("/model/delete", {
headers: { Authorization: `Bearer ${masterKey()}` },
data: { id },
});
const deleted = await post().catch(() => post());
expect(
deleted.ok(),
`cleanup: /model/delete ${id} returned ${deleted.status()}`,
).toBe(true);
}
const uniqueSuffix = (): string =>
`${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
const test = base.extend<{
deployment: { readonly modelName: string; readonly createdModelId: string };
}>({
deployment: async ({ page, request }, use) => {
const modelName = `e2e-edit-params-${uniqueSuffix()}`;
const created = await page.request.post("/model/new", {
headers: { Authorization: `Bearer ${masterKey()}` },
data: {
model_name: modelName,
litellm_params: {
model: `openai/${modelName}`,
api_base: MOCK_LLM_BASE,
api_key: "fake-key",
},
model_info: {},
},
});
expect(
created.ok(),
`/model/new failed: ${created.status()} ${await created.text()}`,
).toBe(true);
const createdModelId = (await created.json()).model_info?.id;
expect(createdModelId, "model id from /model/new").toBeTruthy();
try {
await expect
.poll(
async () => {
try {
await sendChatCompletion(request, {
model: modelName,
prompt: `warmup ${modelName}`,
});
return true;
} catch {
return false;
}
},
{
message: `deployment ${modelName} never became routable after /model/new`,
timeout: 60_000,
},
)
.toBe(true);
await use({ modelName, createdModelId });
} finally {
await deleteDeployment(page, createdModelId);
}
},
});
test.describe("Edit LiteLLM Params on a deployment", () => {
test.use({ storageState: ADMIN_STORAGE_PATH });
test("params added on a deployment can be re-edited, and the deployment keeps serving", async ({
page,
request,
deployment: { modelName, createdModelId },
}) => {
await navigateToPage(page, Page.Models);
const modelIdCell = page.getByTestId(`model-id-${createdModelId}`);
await expect(
modelIdCell,
`the Models table lists ${modelName}`,
).toBeVisible({ timeout: 15_000 });
await modelIdCell.click();
await expect(page.getByText("Back to Models").first()).toBeVisible({
timeout: 15_000,
});
await editParams(page, (params) => ({
...params,
temperature: 0.2,
[CUSTOM_PARAM]: CUSTOM_PARAM_VALUE,
}));
const firstSave = await captureRequestBody(
page,
{ method: "PATCH", urlIncludes: `/model/${createdModelId}/update` },
async () => {
await page.getByRole("button", { name: "Save Changes" }).click();
},
);
expect(
firstSave.litellm_params?.temperature,
"the added temperature goes on the wire",
).toBe(0.2);
expect(
firstSave.litellm_params?.[CUSTOM_PARAM],
`the added ${CUSTOM_PARAM} goes on the wire`,
).toEqual(CUSTOM_PARAM_VALUE);
expect(
firstSave.litellm_params?.model,
"a params edit does not rewrite the upstream model",
).toBe(`openai/${modelName}`);
expect(
firstSave.litellm_params?.api_base,
"a params edit does not rewrite the api base",
).toBe(MOCK_LLM_BASE);
expect(
firstSave.litellm_params,
"the credential is never re-sent, so a masked placeholder cannot overwrite the stored key",
).not.toHaveProperty("api_key");
await expect
.poll(
async () => (await readStoredParams(page, createdModelId)).temperature,
{
message: "the added temperature never reached the stored deployment",
timeout: 20_000,
},
)
.toBe(0.2);
const afterFirstSave = await readStoredParams(page, createdModelId);
expect(
afterFirstSave[CUSTOM_PARAM],
`the added ${CUSTOM_PARAM} reached the stored deployment`,
).toEqual(CUSTOM_PARAM_VALUE);
expect(
afterFirstSave.model,
"the stored upstream model survived the edit",
).toBe(`openai/${modelName}`);
expect(
afterFirstSave.api_base,
"the stored api base survived the edit",
).toBe(MOCK_LLM_BASE);
await editParams(page, (params) => ({
...Object.fromEntries(
Object.entries(params).filter(([key]) => key !== CUSTOM_PARAM),
),
temperature: 0.7,
}));
const secondSave = await captureRequestBody(
page,
{ method: "PATCH", urlIncludes: `/model/${createdModelId}/update` },
async () => {
await page.getByRole("button", { name: "Save Changes" }).click();
},
);
expect(
secondSave.litellm_params?.temperature,
"a param set by an earlier save can be edited again",
).toBe(0.7);
expect(
secondSave.litellm_params,
`dropping ${CUSTOM_PARAM} from the editor drops it from the request the UI sends`,
).not.toHaveProperty(CUSTOM_PARAM);
expect(
secondSave.litellm_params?.model,
"a second params edit still leaves the upstream model alone",
).toBe(`openai/${modelName}`);
expect(
secondSave.litellm_params?.api_base,
"a second params edit still leaves the api base alone",
).toBe(MOCK_LLM_BASE);
expect(
secondSave.litellm_params,
"the credential is still never re-sent",
).not.toHaveProperty("api_key");
await expect
.poll(
async () => (await readStoredParams(page, createdModelId)).temperature,
{
message:
"the re-edited temperature never reached the stored deployment",
timeout: 20_000,
},
)
.toBe(0.7);
await page.reload();
await expect(
page
.getByRole("tabpanel", { name: "Overview" })
.getByText('"temperature": 0.7'),
"reopening the deployment renders the re-edited value, not the one from the first save",
).toBeVisible({ timeout: 20_000 });
await sendChatCompletion(request, {
model: modelName,
prompt: `still serving ${modelName}`,
});
});
});

245
tests/e2e/ui/tests/modelsPage/modelHealthStatus.spec.ts Normal file
View file

@ -0,0 +1,245 @@
import {
test as base,
expect,
type Locator,
type Page as PlaywrightPage,
} from "@playwright/test";
import { ADMIN_STORAGE_PATH } from "../../constants";
import { Page } from "../../fixtures/pages";
import { navigateToPage } from "../../helpers/navigation";
import { readBack } from "../../helpers/roundTrip";
import { masterKey } from "../../helpers/traffic";
const MOCK_LLM_BASE = `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`;
const UNREACHABLE_BASE = "http://127.0.0.1:9/v1";
async function isRegistered(
page: PlaywrightPage,
modelName: string,
): Promise<boolean> {
const body = await readBack<{ data: { model_name?: string }[] }>(
page,
"/v2/model/info",
);
return body.data.some((row) => row.model_name === modelName);
}
function healthRow(page: PlaywrightPage, modelName: string): Locator {
return page.getByRole("row").filter({ hasText: modelName });
}
function pageOf(label: string): { current: number; total: number } {
const [current, total] = label
.replace("Page ", "")
.split(" of ")
.map((part) => Number(part.trim()));
return { current, total };
}
async function locateHealthRow(
page: PlaywrightPage,
modelName: string,
): Promise<Locator> {
const pageLabel = page.getByTestId("pagination-page");
await expect(
pageLabel,
"the health table reports which page it is showing",
).toBeVisible({ timeout: 20_000 });
const deadline = Date.now() + 60_000;
while (Date.now() < deadline) {
const row = healthRow(page, modelName);
const onThisPage = await row
.first()
.waitFor({ state: "visible", timeout: 3_000 })
.then(() => true)
.catch(() => false);
if (onThisPage) return row;
const { current, total } = pageOf(await pageLabel.innerText());
const goTo = current < total ? current + 1 : 1;
if (total === 1) continue;
await page
.getByRole("button", {
name: current < total ? "Go to next page" : "Go to first page",
})
.click();
await expect(pageLabel).toContainText(`Page ${goTo} of`, {
timeout: 15_000,
});
}
return healthRow(page, modelName);
}
async function openHealthTab(page: PlaywrightPage): Promise<void> {
await page.getByRole("tab", { name: "Health Status" }).click();
await expect(
page.getByRole("heading", { name: "Model Health Status" }),
).toBeVisible({ timeout: 15_000 });
}
async function expectStatus(
page: PlaywrightPage,
modelName: string,
status: string,
): Promise<void> {
const row = await locateHealthRow(page, modelName);
await expect(row, `${modelName} has one row in the health table`).toHaveCount(
1,
{ timeout: 20_000 },
);
await expect(
row.getByText(status, { exact: true }),
`the Health Status cell for ${modelName} reads ${status}`,
).toHaveCount(1, { timeout: 60_000 });
}
async function deleteDeployment(
page: PlaywrightPage,
id: string,
): Promise<void> {
const post = () =>
page.request.post("/model/delete", {
headers: { Authorization: `Bearer ${masterKey()}` },
data: { id },
});
const deleted = await post().catch(() => post());
expect(
deleted.ok(),
`cleanup: /model/delete ${id} returned ${deleted.status()}`,
).toBe(true);
}
const uniqueSuffix = (): string =>
`${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
async function withDeployment(
page: PlaywrightPage,
prefix: string,
apiBase: string,
use: (name: string) => Promise<void>,
): Promise<void> {
const name = `${prefix}-${uniqueSuffix()}`;
const created = await page.request.post("/model/new", {
headers: { Authorization: `Bearer ${masterKey()}` },
data: {
model_name: name,
litellm_params: {
model: `openai/${name}`,
api_base: apiBase,
api_key: "fake-key",
},
model_info: {},
},
});
expect(
created.ok(),
`/model/new for ${name} failed: ${created.status()} ${await created.text()}`,
).toBe(true);
const id = (await created.json()).model_info?.id;
expect(id, `model id from /model/new for ${name}`).toBeTruthy();
try {
await expect
.poll(() => isRegistered(page, name), {
message: `deployment ${name} never appeared in /v2/model/info after create`,
timeout: 60_000,
})
.toBe(true);
await use(name);
} finally {
await deleteDeployment(page, id);
}
}
const test = base.extend<{ reachableName: string; unreachableName: string }>({
reachableName: async ({ page }, use) => {
await withDeployment(page, "e2e-health-up", MOCK_LLM_BASE, use);
},
unreachableName: async ({ page }, use) => {
await withDeployment(page, "e2e-health-down", UNREACHABLE_BASE, use);
},
});
test.describe("Model health status", () => {
test.use({ storageState: ADMIN_STORAGE_PATH });
test("Run Health Check reports a reachable deployment healthy and an unreachable one unhealthy", async ({
page,
reachableName,
unreachableName,
}) => {
await navigateToPage(page, Page.Models);
await openHealthTab(page);
for (const name of [reachableName, unreachableName]) {
const row = await locateHealthRow(page, name);
await expect(row, `${name} has one row in the health table`).toHaveCount(
1,
{ timeout: 20_000 },
);
await row
.getByRole("button", { name: "Run Health Check", exact: true })
.click();
}
await expectStatus(page, reachableName, "healthy");
await expect(
healthRow(page, reachableName).getByText("unhealthy", { exact: true }),
"a reachable deployment is never reported unhealthy",
).toHaveCount(0);
await expectStatus(page, unreachableName, "unhealthy");
const successDetail = (
await locateHealthRow(page, reachableName)
).getByRole("button", {
name: "View response details",
});
await expect(
successDetail,
`${reachableName} offers its health check response for inspection`,
).toBeVisible({ timeout: 60_000 });
await successDetail.click();
const successDialog = page.getByRole("dialog");
await expect(
successDialog.getByRole("heading", {
name: `Health Check Response - ${reachableName}`,
}),
"the healthy deployment's detail opens its own response dialog",
).toBeVisible({ timeout: 10_000 });
await successDialog.getByRole("button", { name: "Close" }).last().click();
await expect(successDialog).toBeHidden({ timeout: 10_000 });
const errorDetail = (
await locateHealthRow(page, unreachableName)
).getByRole("button", {
name: "View full error details",
});
await expect(
errorDetail,
`${unreachableName} offers its health check error for inspection`,
).toBeVisible({ timeout: 60_000 });
await errorDetail.click();
const errorDialog = page.getByRole("dialog");
await expect(
errorDialog.getByRole("heading", {
name: `Health Check Error - ${unreachableName}`,
}),
"the unreachable deployment's detail opens its own error dialog",
).toBeVisible({ timeout: 10_000 });
await expect(
errorDialog,
"the error dialog carries the upstream connection failure, not a generic message",
).toContainText(/connection error/i, { timeout: 10_000 });
await expect(
errorDialog,
"the error dialog names the endpoint that could not be reached",
).toContainText(UNREACHABLE_BASE);
await errorDialog.getByRole("button", { name: "Close" }).last().click();
await expect(errorDialog).toBeHidden({ timeout: 10_000 });
await page.reload();
await openHealthTab(page);
await expectStatus(page, reachableName, "healthy");
await expectStatus(page, unreachableName, "unhealthy");
});
});

112
tests/e2e/ui/tests/proxy-admin/keyBlocking.spec.ts Normal file
View file

@ -0,0 +1,112 @@
import { test as base, expect } from "@playwright/test";
import { ADMIN_STORAGE_PATH } from "../../constants";
import { Page } from "../../fixtures/pages";
import { dismissFeedbackPopup, navigateToPage, openKeyDetail } from "../../helpers/navigation";
import {
CHAT_MODEL_A,
MOCK_RESPONSE_TEXT,
attemptChatCompletion,
createVirtualKey,
deleteVirtualKey,
readKeyInfo,
sendChatCompletion,
uniqueSuffix,
} from "../../helpers/traffic";
interface ScopedKey {
alias: string;
token: string;
apiKey: string;
}
const test = base.extend<{ scopedKey: ScopedKey }>({
scopedKey: async ({ page }, use) => {
const alias = `e2e-block-key-${uniqueSuffix()}`;
const created = await createVirtualKey(page.request, {
key_alias: alias,
models: [CHAT_MODEL_A],
});
await use({ alias, token: created.token, apiKey: created.key });
await deleteVirtualKey(page.request, created.token);
},
});
test.describe("Proxy Admin - Key blocking", () => {
test.use({ storageState: ADMIN_STORAGE_PATH });
test("blocking a key stops it serving and unblocking restores it", async ({ page, scopedKey }) => {
const { alias, token, apiKey } = scopedKey;
await sendChatCompletion(page.request, {
model: CHAT_MODEL_A,
prompt: `pre-block ${alias}`,
apiKey,
});
await navigateToPage(page, Page.ApiKeys);
await dismissFeedbackPopup(page);
await openKeyDetail(page, alias);
await page.getByRole("button", { name: "More key actions" }).click();
await page.getByRole("menuitem", { name: "Block Key" }).click();
const blockDialog = page.getByRole("dialog", { name: "Block Key" });
await expect(blockDialog, "the Block Key confirmation never opened").toBeVisible({ timeout: 10_000 });
await blockDialog.getByRole("button", { name: "Block", exact: true }).click();
await expect
.poll(async () => (await readKeyInfo(page.request, token)).blocked, {
message: "the key never came back blocked from /key/info",
timeout: 20_000,
})
.toBe(true);
await expect
.poll(
async () =>
await attemptChatCompletion(page.request, {
model: CHAT_MODEL_A,
prompt: "blocked",
apiKey,
}),
{
message: "a blocked key was still served by /v1/chat/completions",
timeout: 30_000,
},
)
.toMatchObject({ status: 401, body: expect.stringContaining("blocked") });
await page.reload();
await expect(
page.getByText("Blocked", { exact: true }),
"the reloaded key detail does not show the key as blocked",
).toBeVisible({ timeout: 15_000 });
await page.getByRole("button", { name: "More key actions" }).click();
await page.getByRole("menuitem", { name: "Unblock Key" }).click();
const unblockDialog = page.getByRole("dialog", { name: "Unblock Key" });
await expect(unblockDialog, "the Unblock Key confirmation never opened").toBeVisible({ timeout: 10_000 });
await unblockDialog.getByRole("button", { name: "Unblock", exact: true }).click();
await expect
.poll(async () => (await readKeyInfo(page.request, token)).blocked, {
message: "the key never came back unblocked from /key/info",
timeout: 20_000,
})
.toBe(false);
await expect
.poll(
async () =>
await attemptChatCompletion(page.request, {
model: CHAT_MODEL_A,
prompt: "unblocked",
apiKey,
}),
{
message: "an unblocked key is still refused by /v1/chat/completions",
timeout: 30_000,
},
)
.toMatchObject({ status: 200, body: expect.stringContaining(MOCK_RESPONSE_TEXT) });
});
});

101
tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts Normal file
View file

@ -0,0 +1,101 @@
import { test as base, expect } from "@playwright/test";
import { ADMIN_STORAGE_PATH, E2E_TEAM_CRUD_ID } from "../../constants";
import { Page } from "../../fixtures/pages";
import { dismissFeedbackPopup, navigateToPage, openKeyDetail } from "../../helpers/navigation";
import { captureRequestBody } from "../../helpers/roundTrip";
import { CHAT_MODEL_A, createVirtualKey, deleteVirtualKey, readKeyInfo, uniqueSuffix } from "../../helpers/traffic";
interface ScopedKey {
alias: string;
token: string;
}
const test = base.extend<{ scopedKey: ScopedKey }>({
scopedKey: async ({ page }, use) => {
const alias = `e2e-budget-window-${uniqueSuffix()}`;
const created = await createVirtualKey(page.request, {
key_alias: alias,
team_id: E2E_TEAM_CRUD_ID,
models: [CHAT_MODEL_A],
});
await use({ alias, token: created.token });
await deleteVirtualKey(page.request, created.token);
},
});
test.describe("Proxy Admin - Key budget window", () => {
test.use({ storageState: ADMIN_STORAGE_PATH });
test("a monthly spend cap survives a reload, and clearing the window keeps the cap", async ({ page, scopedKey }) => {
const { alias, token } = scopedKey;
const before = await readKeyInfo(page.request, token);
expect(before.max_budget, "a freshly generated key starts with no budget").toBeNull();
await navigateToPage(page, Page.ApiKeys);
await dismissFeedbackPopup(page);
await openKeyDetail(page, alias);
await page.getByRole("tab", { name: "Settings" }).click();
await page.getByRole("button", { name: "Edit Settings" }).click();
await page.getByRole("spinbutton", { name: "Max Budget (USD)" }).fill("12.5");
await page.getByLabel("Reset Budget", { exact: true }).click();
await page.getByRole("option", { name: "monthly", exact: true }).click();
await page.getByRole("button", { name: "Save Changes" }).click();
await expect
.poll(async () => (await readKeyInfo(page.request, token)).max_budget, {
message: "the $12.50 cap never reached /key/info",
timeout: 20_000,
})
.toBe(12.5);
await expect
.poll(async () => (await readKeyInfo(page.request, token)).budget_duration, {
message: "the monthly reset window never reached /key/info",
timeout: 20_000,
})
.toBe("30d");
const capped = await readKeyInfo(page.request, token);
const resetAt = new Date(capped.budget_reset_at ?? "");
expect(Number.isNaN(resetAt.getTime()), "a monthly window left the key with no budget_reset_at").toBe(false);
expect(resetAt.getTime(), "budget_reset_at was set in the past").toBeGreaterThan(Date.now());
expect(resetAt.getUTCDate(), "a monthly window resets on the 1st, a daily one would not").toBe(1);
await page.reload();
await expect(
page.getByRole("paragraph").filter({ hasText: "of $12.50" }),
"the reloaded key detail does not render the $12.50 cap",
).toBeVisible({ timeout: 15_000 });
await page.getByRole("tab", { name: "Settings" }).click();
await expect(
page.getByTestId("budget-reset-value"),
"the reloaded key detail does not name the 30d reset window",
).toHaveText(/Every 30d/, { timeout: 15_000 });
await page.getByRole("button", { name: "Edit Settings" }).click();
await page.getByLabel("Reset Budget", { exact: true }).click();
await page.getByRole("option", { name: "Never resets", exact: true }).click();
const cleared = await captureRequestBody(page, { method: "POST", urlIncludes: "/key/update" }, async () => {
await page.getByRole("button", { name: "Save Changes" }).click();
});
expect(cleared).toHaveProperty("budget_duration");
expect(cleared.budget_duration, "clearing the window must send budget_duration: null explicitly").toBeNull();
await expect
.poll(async () => (await readKeyInfo(page.request, token)).budget_duration, {
message: "the reset window was never cleared on /key/info",
timeout: 20_000,
})
.toBeNull();
const after = await readKeyInfo(page.request, token);
expect(after.budget_reset_at, "clearing the reset window left a stale next-reset timestamp").toBeNull();
expect(after.max_budget, "clearing the reset window also wiped the spend cap").toBe(12.5);
expect(after.models, "editing the budget left the key's models untouched").toEqual(before.models);
expect(after.team_id, "editing the budget left the key's team untouched").toEqual(before.team_id);
});
});

15
tests/mcp_tests/mcp_server.py
View file

@ -1,10 +1,12 @@
# math_server.py
import argparse
import os
from typing import Final
from mcp.server.fastmcp import FastMCP
from mcp.server.fastmcp import Context, FastMCP
mcp = FastMCP("Math")
ADD_OFFSET: Final = int(os.getenv("MCP_ADD_OFFSET", "0"))
def _parse_args() -> argparse.Namespace:
@ -31,7 +33,7 @@ def _parse_args() -> argparse.Namespace:
@mcp.tool()
def add(a: int, b: int) -> int:
"""Add two numbers"""
return a + b
return a + b + ADD_OFFSET
@mcp.tool()
@ -40,6 +42,15 @@ def multiply(a: int, b: int) -> int:
return a * b
@mcp.tool()
def request_headers(ctx: Context) -> dict[str, str]:
request: Final = ctx.request_context.request
return {
"authorization": request.headers.get("authorization", "") if request is not None else "",
"x-request-tag": request.headers.get("x-request-tag", "") if request is not None else "",
}
def main() -> None:
args = _parse_args()
transport = (args.transport or "stdio").lower()

3
tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml
View file

@ -23,3 +23,6 @@ mcp_servers:
transport: http
url: http://127.0.0.1:0/mcp
allow_all_keys: true
math_restricted:
transport: http
url: http://127.0.0.1:0/mcp

564
tests/mcp_tests/test_proxy_mcp_e2e.py
View file

@ -1,26 +1,39 @@
import asyncio
import json
import os
import queue
import socket
import subprocess
import sys
import tempfile
import threading
import time
import typing
from contextlib import asynccontextmanager, contextmanager
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
import httpx
import pytest
import uvicorn
import yaml
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client
from mcp.types import CallToolResult
from starlette.requests import Request
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._experimental.mcp_server.tool_search import handle_mcp_proxy_tool
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, ProxyException, UserAPIKeyAuth
from litellm.proxy.proxy_server import (
app as proxy_app,
)
from litellm.proxy.proxy_server import (
cleanup_router_config_variables,
initialize,
)
CONFIG_TEMPLATE_PATH = Path("tests/mcp_tests/test_configs/test_config_mcp_e2e.yaml")
MCP_SERVER_SCRIPT = Path("tests/mcp_tests/mcp_server.py")
PROJECT_ROOT = Path(__file__).resolve().parents[2]
@ -45,28 +58,49 @@ def _clear_proxy_database_env() -> typing.Iterator[None]:
mp.undo()
def _initialize_proxy(config_path: str) -> None:
async def _initialize_proxy(config_path: str) -> None:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
cleanup_router_config_variables()
asyncio.run(initialize(config=config_path, debug=True))
await initialize(config=config_path, debug=True)
for server_id, upstream in tuple(global_mcp_server_manager.registry.items()):
if upstream.server_name != "math_restricted":
continue
global_mcp_server_manager.registry[server_id] = upstream.model_copy(
update={"tool_name_to_display_name": {"add": "Add Numbers"}}
)
@dataclass(frozen=True)
class ProxyRig:
url: str
config_path: str
loop: asyncio.AbstractEventLoop
def _start_proxy_server(
config_path: str,
) -> tuple[str, uvicorn.Server, threading.Thread, socket.socket]:
_initialize_proxy(config_path)
) -> tuple[ProxyRig, uvicorn.Server, threading.Thread, socket.socket]:
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sock.bind(("127.0.0.1", 0))
host, port = sock.getsockname()
config = uvicorn.Config(proxy_app, host=host, port=port, log_level="warning")
config = uvicorn.Config(proxy_app, host=host, port=port, log_level="warning", lifespan="off")
server = uvicorn.Server(config)
loop = asyncio.new_event_loop()
async def _serve() -> None:
from litellm.proxy._experimental.mcp_server import server as mcp_server
await _initialize_proxy(config_path)
async with proxy_app.router.lifespan_context(proxy_app), mcp_server.lifespan(proxy_app):
await server.serve(sockets=[sock])
def _run() -> None:
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
loop.run_until_complete(server.serve(sockets=[sock]))
with asyncio.Runner(loop_factory=lambda: loop) as runner:
runner.run(_serve())
thread = threading.Thread(target=_run, daemon=True)
thread.start()
@ -79,79 +113,93 @@ def _start_proxy_server(
raise TimeoutError("Proxy server did not start in time")
time.sleep(0.05)
return f"http://{host}:{port}", server, thread, sock
return ProxyRig(f"http://{host}:{port}", config_path, loop), server, thread, sock
@pytest.fixture(scope="session")
def math_streamable_http_server() -> str:
@contextmanager
def _math_http_server(offset: int) -> typing.Iterator[str]:
host = "127.0.0.1"
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.bind((host, 0))
_, port = sock.getsockname()
cmd = [
sys.executable,
str(MCP_SERVER_SCRIPT),
"--transport",
"http",
"--host",
host,
"--port",
str(port),
]
env = os.environ.copy()
server_process = subprocess.Popen(
cmd,
cwd=str(PROJECT_ROOT),
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
start_time = time.time()
while True:
if server_process.poll() is not None:
stdout, stderr = server_process.communicate()
raise RuntimeError(
f"Streamable HTTP MCP server exited early.\nSTDOUT: {stdout.decode()}\nSTDERR: {stderr.decode()}"
)
with tempfile.TemporaryFile() as server_log:
process = subprocess.Popen(
[sys.executable, str(MCP_SERVER_SCRIPT), "--transport", "http", "--host", host, "--port", str(port)],
cwd=str(PROJECT_ROOT),
stdout=server_log,
stderr=subprocess.STDOUT,
env={**os.environ, "MCP_ADD_OFFSET": str(offset)},
)
try:
with socket.create_connection((host, port), timeout=0.1):
break
except OSError:
if time.time() - start_time > PROXY_START_TIMEOUT:
server_process.terminate()
raise TimeoutError("Streamable HTTP MCP server did not start in time")
time.sleep(0.05)
yield f"http://{host}:{port}"
server_process.terminate()
try:
server_process.wait(timeout=5)
except subprocess.TimeoutExpired:
server_process.kill()
start_time = time.monotonic()
while True:
if process.poll() is not None:
server_log.seek(0)
raise RuntimeError(f"MCP upstream exited early: {server_log.read().decode()}")
try:
with socket.create_connection((host, port), timeout=0.1):
break
except OSError:
if time.monotonic() - start_time > PROXY_START_TIMEOUT:
raise TimeoutError("Streamable HTTP MCP server did not start in time")
time.sleep(0.05)
yield f"http://{host}:{port}"
finally:
process.terminate()
try:
process.wait(timeout=5)
except subprocess.TimeoutExpired:
process.kill()
process.wait(timeout=5)
@pytest.fixture(scope="session")
def proxy_server_url(
tmp_path_factory: pytest.TempPathFactory, math_streamable_http_server: str
def math_streamable_http_server() -> typing.Iterator[str]:
with _math_http_server(100) as url:
yield url
@pytest.fixture(scope="session")
def math_restricted_server() -> typing.Iterator[str]:
with _math_http_server(200) as url:
yield url
@pytest.fixture(scope="session")
def _proxy_server(
tmp_path_factory: pytest.TempPathFactory,
math_streamable_http_server: str,
math_restricted_server: str,
):
config_dir = tmp_path_factory.mktemp("mcp_e2e")
config_path = config_dir / "config.yaml"
config = yaml.safe_load(CONFIG_TEMPLATE_PATH.read_text())
config["mcp_servers"]["math_streamable_http"][
"url"
] = f"{math_streamable_http_server}/mcp"
config["mcp_servers"]["math_stdio"]["command"] = sys.executable
config["mcp_servers"]["math_streamable_http"]["url"] = f"{math_streamable_http_server}/mcp"
config["mcp_servers"]["math_restricted"]["url"] = f"{math_restricted_server}/mcp"
config["general_settings"]["custom_auth"] = f"{__name__}.authorize_proxy_key"
config["litellm_settings"]["callbacks"] = [f"{__name__}.proxy_call_recorder"]
config["mcp_servers"]["math_restricted"]["mcp_info"] = {"mcp_server_cost_info": {"default_cost_per_query": 0.25}}
config_path.write_text(yaml.safe_dump(config))
server_url, server, thread, sock = _start_proxy_server(str(config_path))
rig, server, thread, sock = _start_proxy_server(str(config_path))
yield server_url
try:
yield rig
finally:
server.should_exit = True
thread.join(timeout=10)
sock.close()
assert not thread.is_alive(), "Proxy did not shut down"
server.should_exit = True
thread.join(timeout=10)
sock.close()
@pytest.fixture
def proxy_server_url(_proxy_server: ProxyRig, setup_and_teardown: None) -> str:
asyncio.run_coroutine_threadsafe(_initialize_proxy(_proxy_server.config_path), _proxy_server.loop).result(
timeout=30
)
return _proxy_server.url
class TestProxyMcpSimpleConnections:
@ -177,9 +225,7 @@ class TestProxyMcpSimpleConnections:
assert text == "7"
@pytest.mark.asyncio
async def test_proxy_mcp_streamable_http_roundtrip(
self, proxy_server_url: str
) -> None:
async def test_proxy_mcp_streamable_http_roundtrip(self, proxy_server_url: str) -> None:
async with asyncio.timeout(20):
async with streamablehttp_client(
url=f"{proxy_server_url}/mcp",
@ -197,12 +243,10 @@ class TestProxyMcpSimpleConnections:
assert result.content
first_content = result.content[0]
text = getattr(first_content, "text", None)
assert text == "11"
assert text == "111"
@pytest.mark.asyncio
async def test_proxy_mcp_lists_all_servers_without_header(
self, proxy_server_url: str
) -> None:
async def test_proxy_mcp_lists_all_servers_without_header(self, proxy_server_url: str) -> None:
async with asyncio.timeout(20):
async with streamablehttp_client(
url=f"{proxy_server_url}/mcp",
@ -220,22 +264,16 @@ class TestProxyMcpSimpleConnections:
}
assert expected_tool_names <= tool_names
async def _call_and_get_text(
tool_name: str, *, a: int, b: int
) -> str | None:
result = await session.call_tool(
tool_name, arguments={"a": a, "b": b}
)
async def _call_and_get_text(tool_name: str, *, a: int, b: int) -> str | None:
result = await session.call_tool(tool_name, arguments={"a": a, "b": b})
assert result.content
first_content = result.content[0]
return getattr(first_content, "text", None)
stdio_result = await _call_and_get_text("math_stdio-add", a=2, b=3)
streamable_result = await _call_and_get_text(
"math_streamable_http-add", a=4, b=5
)
streamable_result = await _call_and_get_text("math_streamable_http-add", a=4, b=5)
assert stdio_result == "5"
assert streamable_result == "9"
assert streamable_result == "109"
class TestProxyMcpStatelessBehavior:
@ -254,9 +292,7 @@ class TestProxyMcpStatelessBehavior:
"""
@pytest.mark.asyncio
async def test_independent_clients_no_shared_session(
self, proxy_server_url: str
) -> None:
async def test_independent_clients_no_shared_session(self, proxy_server_url: str) -> None:
"""Two independent clients connect and operate without sharing session state."""
async with asyncio.timeout(30):
# --- Client A: connect, initialize, call tool ---
@ -269,9 +305,7 @@ class TestProxyMcpStatelessBehavior:
) as (read_a, write_a, _get_sid_a):
async with ClientSession(read_a, write_a) as session_a:
await session_a.initialize()
result_a = await session_a.call_tool(
"add", arguments={"a": 10, "b": 20}
)
result_a = await session_a.call_tool("add", arguments={"a": 10, "b": 20})
assert result_a.content
text_a = getattr(result_a.content[0], "text", None)
assert text_a == "30"
@ -293,9 +327,359 @@ class TestProxyMcpStatelessBehavior:
await session_b.initialize()
tools = await session_b.list_tools()
assert any(t.name.endswith("add") for t in tools.tools)
result_b = await session_b.call_tool(
"add", arguments={"a": 100, "b": 200}
)
result_b = await session_b.call_tool("add", arguments={"a": 100, "b": 200})
assert result_b.content
text_b = getattr(result_b.content[0], "text", None)
assert text_b == "300"
PROXY_MODE_TOOLS = frozenset({"search_tools", "get_tool_schema", "call_tool"})
def _payload(result: typing.Any) -> typing.Any:
assert result.content, f"empty tool result: {result}"
return json.loads(result.content[0].text)
def _proxy_session(proxy_server_url: str, **extra_headers: str):
return streamablehttp_client(
url=f"{proxy_server_url}/mcp/proxy",
headers={"Authorization": PROXY_AUTHORIZATION_HEADER, **extra_headers},
)
class TestProxyMcpSchemaDiscoveryMode:
"""Drive /mcp/proxy over the real streamable-HTTP transport with the MCP SDK client:
the fixed three-tool surface, opaque-id discovery, schema-validated execution against
two upstreams that expose the same tool name, and the operations the surface refuses."""
@pytest.mark.asyncio
async def test_initialize_and_list_expose_only_discovery_tools(self, proxy_server_url: str) -> None:
async with asyncio.timeout(20):
async with _proxy_session(proxy_server_url) as (read, write, _sid):
async with ClientSession(read, write) as session:
init = await session.initialize()
assert init.capabilities.tools is not None
assert init.capabilities.prompts is None
assert init.capabilities.resources is None
listed = await session.list_tools()
assert {tool.name for tool in listed.tools} == PROXY_MODE_TOOLS
@pytest.mark.asyncio
async def test_search_schema_and_call_round_trip_keeps_server_identity(self, proxy_server_url: str) -> None:
async with asyncio.timeout(30):
async with _proxy_session(proxy_server_url) as (read, write, _sid):
async with ClientSession(read, write) as session:
await session.initialize()
hits = _payload(await session.call_tool("search_tools", arguments={"query": "add"}))
by_name = {hit["name"]: hit for hit in hits}
assert {"math_stdio-add", "math_streamable_http-add"} <= set(by_name)
assert all("inputSchema" not in hit for hit in hits)
assert by_name["math_stdio-add"]["tool_id"] != by_name["math_streamable_http-add"]["tool_id"]
schema = _payload(
await session.call_tool(
"get_tool_schema", arguments={"tool_id": by_name["math_stdio-add"]["tool_id"]}
)
)
assert schema["name"] == "math_stdio-add"
assert set(schema["inputSchema"]["required"]) == {"a", "b"}
assert schema["outputSchema"]["properties"]["result"]["type"] == "integer"
stdio = await session.call_tool(
"call_tool",
arguments={"tool_id": by_name["math_stdio-add"]["tool_id"], "arguments": {"a": 3, "b": 4}},
)
http = await session.call_tool(
"call_tool",
arguments={
"tool_id": by_name["math_streamable_http-add"]["tool_id"],
"arguments": {"a": 5, "b": 6},
},
)
assert stdio.isError is False and stdio.content[0].text == "7"
assert http.isError is False and http.content[0].text == "111"
@pytest.mark.asyncio
async def test_server_scope_header_narrows_discovery(self, proxy_server_url: str) -> None:
async with asyncio.timeout(20):
async with _proxy_session(proxy_server_url, **{"x-mcp-servers": "math_streamable_http"}) as (
read,
write,
_sid,
):
async with ClientSession(read, write) as session:
await session.initialize()
hits = _payload(await session.call_tool("search_tools", arguments={"query": "add"}))
assert {hit["name"] for hit in hits} == {"math_streamable_http-add"}
@pytest.mark.asyncio
async def test_rejections_never_reach_upstream(self, proxy_server_url: str) -> None:
from mcp.shared.exceptions import McpError
from mcp.types import METHOD_NOT_FOUND
async with asyncio.timeout(30):
async with _proxy_session(proxy_server_url) as (read, write, _sid):
async with ClientSession(read, write) as session:
await session.initialize()
hits = _payload(await session.call_tool("search_tools", arguments={"query": "add"}))
tool_id = next(hit["tool_id"] for hit in hits if hit["name"] == "math_stdio-add")
bad_args = await session.call_tool(
"call_tool", arguments={"tool_id": tool_id, "arguments": {"a": "three", "b": 4}}
)
assert bad_args.isError is True and "Invalid arguments" in bad_args.content[0].text
stale = await session.call_tool("get_tool_schema", arguments={"tool_id": "0" * 32})
assert stale.isError is True and "unauthorized tool_id" in stale.content[0].text
for not_an_object in ("wrong", False):
refused_args = await session.call_tool(
"call_tool", arguments={"tool_id": tool_id, "arguments": not_an_object}
)
assert refused_args.isError is True and "object" in refused_args.content[0].text
direct = await session.call_tool("math_stdio-add", arguments={"a": 1, "b": 2})
assert direct.isError is True and "unavailable on /mcp/proxy" in direct.content[0].text
for operation in (session.list_prompts, session.list_resources):
with pytest.raises(McpError) as refused:
await operation()
assert refused.value.error.code == METHOD_NOT_FOUND
async def authorize_proxy_key(request: Request, api_key: str) -> UserAPIKeyAuth:
permissions = {
"sk-1234": LiteLLM_ObjectPermissionTable(object_permission_id="open", mcp_servers=["math_stdio"]),
"sk-restricted": LiteLLM_ObjectPermissionTable(
object_permission_id="restricted", mcp_servers=["math_restricted"]
),
"sk-none": LiteLLM_ObjectPermissionTable(object_permission_id="none", mcp_servers=["no-mcp-servers"]),
"sk-add-only": LiteLLM_ObjectPermissionTable(
object_permission_id="add-only", mcp_servers=["math_stdio"], mcp_tool_permissions={"math_stdio": ["add"]}
),
}
permission = permissions.get(api_key)
if permission is None:
raise ProxyException(message="Unknown test key", type="authentication_error", param=None, code=401)
return UserAPIKeyAuth(api_key=api_key, user_id=api_key, object_permission=permission)
class ProxyCallRecorder(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.events: queue.Queue[str] = queue.Queue()
self.failures: queue.Queue[str] = queue.Queue()
async def async_log_success_event(
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
payload = kwargs.get("standard_logging_object")
if isinstance(payload, dict) and payload.get("call_type") == "call_mcp_tool":
self.events.put(json.dumps(payload, default=str))
async def async_log_failure_event(
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
payload = kwargs.get("standard_logging_object")
if isinstance(payload, dict) and payload.get("call_type") == "call_mcp_tool":
self.failures.put(json.dumps(payload, default=str))
proxy_call_recorder = ProxyCallRecorder()
@asynccontextmanager
async def _scoped_session(url: str, key: str = "sk-1234", **headers: str) -> typing.AsyncIterator[ClientSession]:
async with asyncio.timeout(30):
async with _proxy_session(url, Authorization=f"Bearer {key}", **headers) as (read, write, _sid):
async with ClientSession(read, write) as session:
await session.initialize()
yield session
async def _search(session: ClientSession, query: str) -> dict[str, str]:
result = await session.call_tool("search_tools", arguments={"query": query})
assert result.isError is False, result
return {hit["name"]: hit["tool_id"] for hit in _payload(result)}
async def _call(session: ClientSession, tool_id: str, a: int = 3, b: int = 4) -> CallToolResult:
return await session.call_tool("call_tool", arguments={"tool_id": tool_id, "arguments": {"a": a, "b": b}})
def _assert_unauthorized(result: CallToolResult) -> None:
assert result.isError is True
assert result.content[0].text == "Unknown or unauthorized tool_id"
class TestProxyMcpAuthorizationScope:
@pytest.mark.asyncio
async def test_server_grant_bounds_search_and_blocks_foreign_ids(self, proxy_server_url: str) -> None:
async with _scoped_session(proxy_server_url, "sk-restricted") as granted:
restricted_id = (await _search(granted, "add"))["math_restricted-add"]
assert (await _call(granted, restricted_id)).content[0].text == "207"
async with _scoped_session(proxy_server_url) as ungranted:
assert set(await _search(ungranted, "add")) == {"math_stdio-add", "math_streamable_http-add"}
_assert_unauthorized(await ungranted.call_tool("get_tool_schema", {"tool_id": restricted_id}))
_assert_unauthorized(await _call(ungranted, restricted_id))
@pytest.mark.asyncio
async def test_no_mcp_servers_sentinel_hides_every_tool(self, proxy_server_url: str) -> None:
async with _scoped_session(proxy_server_url) as granted:
tool_id = (await _search(granted, "add"))["math_stdio-add"]
async with _scoped_session(proxy_server_url, "sk-none") as session:
assert await _search(session, "add") == {}
_assert_unauthorized(await session.call_tool("get_tool_schema", {"tool_id": tool_id}))
_assert_unauthorized(await _call(session, tool_id))
@pytest.mark.asyncio
async def test_tool_grant_hides_ungranted_tools_and_blocks_their_ids(self, proxy_server_url: str) -> None:
async with _scoped_session(proxy_server_url) as granted:
multiply_id = (await _search(granted, "multiply"))["math_stdio-multiply"]
async with _scoped_session(proxy_server_url, "sk-add-only", **{"x-mcp-servers": "math_stdio"}) as session:
ids = await _search(session, "add multiply request_headers")
assert set(ids) == {"math_stdio-add"}
assert (await _call(session, ids["math_stdio-add"])).content[0].text == "7"
_assert_unauthorized(await session.call_tool("get_tool_schema", {"tool_id": multiply_id}))
_assert_unauthorized(await _call(session, multiply_id))
@pytest.mark.asyncio
async def test_same_named_tools_keep_distinct_ids_and_reach_their_own_upstream(self, proxy_server_url: str) -> None:
async with _scoped_session(proxy_server_url, "sk-restricted") as session:
ids = await _search(session, "add")
assert set(ids) == {"math_stdio-add", "math_streamable_http-add", "math_restricted-add"}
assert len(set(ids.values())) == 3
assert all(len(tool_id) == 32 for tool_id in ids.values())
for name, expected in (
("math_stdio-add", "7"),
("math_streamable_http-add", "107"),
("math_restricted-add", "207"),
):
schema = _payload(await session.call_tool("get_tool_schema", {"tool_id": ids[name]}))
assert schema["name"] == name
assert schema["tool_id"] == ids[name]
result = await _call(session, ids[name])
assert result.isError is False
assert result.content[0].text == expected
@pytest.mark.asyncio
async def test_server_scope_header_narrows_grants_and_blocks_out_of_scope_ids(self, proxy_server_url: str) -> None:
async with _scoped_session(proxy_server_url, "sk-restricted") as unscoped:
other_id = (await _search(unscoped, "add"))["math_stdio-add"]
async with _scoped_session(
proxy_server_url, "sk-restricted", **{"x-mcp-servers": "math_restricted"}
) as session:
ids = await _search(session, "add")
assert set(ids) == {"math_restricted-add"}
_assert_unauthorized(await session.call_tool("get_tool_schema", {"tool_id": other_id}))
_assert_unauthorized(await _call(session, other_id))
assert (await _call(session, ids["math_restricted-add"])).content[0].text == "207"
@pytest.mark.asyncio
@pytest.mark.parametrize("key", [None, "sk-invalid"])
async def test_missing_or_invalid_key_cannot_initialize(self, proxy_server_url: str, key: str | None) -> None:
async with httpx.AsyncClient() as client:
response = await client.post(
f"{proxy_server_url}/mcp/proxy",
headers={
"Accept": "application/json, text/event-stream",
**({"Authorization": f"Bearer {key}"} if key else {}),
},
json={
"jsonrpc": "2.0",
"id": 1,
"method": "initialize",
"params": {
"protocolVersion": "2025-03-26",
"capabilities": {},
"clientInfo": {"name": "auth-test", "version": "1"},
},
},
)
assert response.status_code == 401, response.text
@pytest.mark.asyncio
async def test_server_headers_are_forwarded_only_to_the_named_upstream(self, proxy_server_url: str) -> None:
for tag in ("first-request", "second-request"):
async with _scoped_session(
proxy_server_url,
"sk-restricted",
**{
"x-mcp-math_restricted-authorization": f"Bearer {tag}",
"x-mcp-math_restricted-x-request-tag": tag,
},
) as session:
ids = await _search(session, "request_headers")
for name, expected in (
("math_restricted", {"authorization": f"Bearer {tag}", "x-request-tag": tag}),
("math_streamable_http", {"authorization": "", "x-request-tag": ""}),
):
result = await session.call_tool(
"call_tool", {"tool_id": ids[f"{name}-request_headers"], "arguments": {}}
)
assert result.isError is False
assert _payload(result) == expected
@pytest.mark.asyncio
async def test_proxy_call_emits_spend_log(self, proxy_server_url: str) -> None:
async with _scoped_session(proxy_server_url, "sk-restricted") as session:
tool_id = (await _search(session, "add"))["math_restricted-add"]
result = await _call(session, tool_id, 123, 456)
assert result.isError is False and result.content[0].text == "779"
async with asyncio.timeout(10):
while True:
payload = json.loads(await asyncio.to_thread(proxy_call_recorder.events.get, True, 5))
if payload.get("metadata", {}).get("mcp_tool_call_metadata", {}).get("arguments") == {
"a": 123,
"b": 456,
}:
break
assert payload["call_type"] == "call_mcp_tool"
assert payload["response_cost"] == 0.25
assert payload["status"] == "success"
assert payload["metadata"]["mcp_tool_call_metadata"]["mcp_server_name"] == "math_restricted"
assert payload["metadata"]["mcp_tool_call_metadata"]["name"] == "add"
assert payload["metadata"]["mcp_tool_call_metadata"]["namespaced_tool_name"] == "math_restricted/add"
@pytest.mark.asyncio
async def test_proxy_scope_exception_returns_iserror_and_emits_failure_log(self, proxy_server_url: str) -> None:
async with _scoped_session(
proxy_server_url,
"sk-none",
**{"x-mcp-servers": "math_restricted", "x-litellm-call-id": "proxy-scope-denial"},
) as session:
result = await session.call_tool("call_tool", {"tool_id": "denied-scope", "arguments": {}})
assert result.isError is True
assert result.content[0].text == (
"Error: The key is not allowed to access the requested MCP servers: math_restricted"
)
async with asyncio.timeout(10):
while True:
payload = json.loads(await asyncio.to_thread(proxy_call_recorder.failures.get, True, 5))
if payload["id"] == "proxy-scope-denial":
break
assert payload["call_type"] == "call_mcp_tool"
assert payload["status"] == "failure"
assert payload["response_cost"] == 0
assert "math_restricted" in payload["error_str"]
@pytest.mark.parametrize("arguments", ["wrong", False, None, [], 0])
def test_handler_rejects_non_object_arguments(
self, proxy_server_url: str, _proxy_server: ProxyRig, arguments: object
) -> None:
async def check() -> None:
auth = UserAPIKeyAuth(
object_permission=LiteLLM_ObjectPermissionTable(
object_permission_id="validation", mcp_servers=["math_stdio"]
)
)
hits = _payload(await handle_mcp_proxy_tool("search_tools", {"query": "add"}, auth))
tool_id = next(hit["tool_id"] for hit in hits if hit["name"] == "math_stdio-add")
result = await handle_mcp_proxy_tool("call_tool", {"tool_id": tool_id, "arguments": arguments}, auth)
assert result.isError is True
assert result.content[0].text == "arguments must be an object"
asyncio.run_coroutine_threadsafe(check(), _proxy_server.loop).result(timeout=30)

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