Merge remote-tracking branch 'origin/main' into HEAD

This commit is contained in:
kerry 2026-09-16 16:18:47 +00:00
commit d398c9199c
110 changed files with 8974 additions and 1459 deletions

View file

@ -9,7 +9,7 @@ commands:
parameters:
category:
type: enum
enum: ["backend", "client"]
enum: ["backend", "client", "provider-harness"]
default: "backend"
steps:
- run:
@ -2918,19 +2918,30 @@ jobs:
provider_replay_harness:
docker:
- *python312_image
- image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f
working_directory: ~/project
resource_class: medium
environment:
E2E_CACHE_TEST_REDIS_URL: redis://127.0.0.1:6379/0
E2E_PROVIDER_CACHE: "0"
E2E_FIXTURE_MODE: live
steps:
- checkout
- skip_if_unrelated_changes:
category: provider-harness
- setup_litellm_test_deps
- wait_for_service:
url: tcp://localhost:6379
- run:
name: Test provider replay harness
name: Test provider capture and replay harness
command: |
mkdir -p test-results/provider-replay-harness
uv run --no-sync pytest -q --noconftest -o addopts= -o pythonpath=tests/e2e -p no:rerunfailures \
--junitxml=test-results/provider-replay-harness/junit.xml \
tests/e2e/test_provider_edge.py tests/e2e/test_fixture_bundle.py \
tests/e2e/test_fixture_canonical.py tests/e2e/test_fixture_mode.py \
tests/code_coverage_tests/test_provider_replay_harness.py
tests/code_coverage_tests/test_provider_replay_harness.py \
tests/code_coverage_tests/test_provider_cache.py
- store_test_results:
path: test-results/provider-replay-harness

View file

@ -1,13 +1,19 @@
#!/usr/bin/env bash
set -uo pipefail
category="${1:?usage: classify_changes.sh <backend|client|ui>}"
category="${1:?usage: classify_changes.sh <backend|client|ui|provider-harness>}"
has_client=false
has_backend=false
has_ci=false
has_provider_harness=false
while IFS= read -r file || [ -n "$file" ]; do
[ -n "$file" ] || continue
case "$file" in
tests/e2e/*/*.py) : ;;
tests/e2e/*.py | tests/code_coverage_tests/test_provider_cache.py | tests/code_coverage_tests/test_provider_replay_harness.py | tests/test_litellm/test_circleci_path_filter.py | .circleci/* | pyproject.toml | uv.lock)
has_provider_harness=true ;;
esac
case "$file" in
ui/* | tests/e2e/ui/*) has_client=true ;;
docs/* | *.md | *.mdx) : ;;
@ -17,6 +23,9 @@ while IFS= read -r file || [ -n "$file" ]; do
done
case "$category" in
provider-harness)
[ "$has_provider_harness" = true ] && echo run || echo skip
;;
backend)
[ "$has_backend" = true ] && echo run || echo skip
;;

View file

@ -1,7 +1,7 @@
#!/usr/bin/env bash
set -uo pipefail
category="${1:?usage: path_filter.sh <backend|client>}"
category="${1:?usage: path_filter.sh <backend|client|provider-harness>}"
here="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
run_full() {
@ -36,5 +36,5 @@ if [ "$decision" = run ]; then
run_full "$category-relevant changes detected"
fi
echo "path-filter[$category]: only unrelated (docs/client) changes detected; halting job as successful"
echo "path-filter[$category]: only unrelated changes detected; halting job as successful"
circleci-agent step halt

View file

@ -26,6 +26,7 @@ on:
- ui/Dockerfile
- ui/nginx.conf
- .github/workflows/image-scan.yml
- .grype.yaml
schedule:
- cron: "41 6 * * *"
workflow_dispatch:
@ -93,6 +94,7 @@ jobs:
GRYPE_MATCH_PYTHON_USING_CPES: "true"
run: |
"$RUNNER_TEMP/grype" litellm-image-scan:${{ github.sha }} \
--config .grype.yaml \
--only-fixed \
--fail-on high \
--output table

13
.grype.yaml Normal file
View file

@ -0,0 +1,13 @@
# Wolfi's security database names zlib 1.3.3-r0 as the fix for CVE-2026-85091,
# but the newest zlib published to the Wolfi apk repo is 1.3.2-r7, so every
# wolfi-base digest reports it and no `apk upgrade` can clear it.
# Drop this once Wolfi ships zlib >= 1.3.3-r0; expected by 2026-10-15.
ignore:
- vulnerability: CVE-2026-85091
package:
name: zlib
type: apk
- vulnerability: GHSA-g5fp-32jq-cfw2
package:
name: zlib
type: apk

View file

@ -0,0 +1,23 @@
-- AlterTable
ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "total_response_time_ms" BIGINT NOT NULL DEFAULT 0;
ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "timed_requests" BIGINT NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "total_response_time_ms" BIGINT NOT NULL DEFAULT 0;
ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "timed_requests" BIGINT NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "total_response_time_ms" BIGINT NOT NULL DEFAULT 0;
ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "timed_requests" BIGINT NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "total_response_time_ms" BIGINT NOT NULL DEFAULT 0;
ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "timed_requests" BIGINT NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "total_response_time_ms" BIGINT NOT NULL DEFAULT 0;
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "timed_requests" BIGINT NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "total_response_time_ms" BIGINT NOT NULL DEFAULT 0;
ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "timed_requests" BIGINT NOT NULL DEFAULT 0;

View file

@ -801,6 +801,8 @@ model LiteLLM_DailyUserSpend {
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
total_response_time_ms BigInt @default(0)
timed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@ -837,6 +839,8 @@ model LiteLLM_DailyOrganizationSpend {
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
total_response_time_ms BigInt @default(0)
timed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@ -873,6 +877,8 @@ model LiteLLM_DailyEndUserSpend {
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
total_response_time_ms BigInt @default(0)
timed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@ -908,6 +914,8 @@ model LiteLLM_DailyAgentSpend {
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
total_response_time_ms BigInt @default(0)
timed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@ -943,6 +951,8 @@ model LiteLLM_DailyTeamSpend {
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
total_response_time_ms BigInt @default(0)
timed_requests BigInt @default(0)
ptu_flat_cost Float @default(0.0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@ -981,6 +991,8 @@ model LiteLLM_DailyTagSpend {
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
total_response_time_ms BigInt @default(0)
timed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt

View file

@ -1167,7 +1167,9 @@ class ModelResponseIterator:
# (matches OpenAI behavior and non-streaming Anthropic implementation)
if self.converted_response_format_tool:
finish_reason = "stop"
usage: Final = self._handle_usage(anthropic_usage_chunk=message_delta["usage"])
usage: Final = (
self._handle_usage(anthropic_usage_chunk=message_delta["usage"]) if "usage" in message_delta else None
)
container: Final = message_delta["delta"].get("container")
return finish_reason, usage, container

View file

@ -18,7 +18,9 @@ if TYPE_CHECKING:
import litellm
def cost_per_token(model: str, usage: "Usage", service_tier: str | None = None) -> tuple[float, float]:
def cost_per_token(
model: str, usage: "Usage", service_tier: str | None = None, model_info: "ModelInfo | None" = None
) -> tuple[float, float]:
"""
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
@ -27,6 +29,7 @@ def cost_per_token(model: str, usage: "Usage", service_tier: str | None = None)
- usage: LiteLLM Usage block, containing anthropic caching information
- service_tier: the service tier the request was served at (e.g. "priority"),
read from the Anthropic response usage and used to select tier-specific pricing
- model_info: effective deployment prices, when they override public rates
Returns:
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
@ -36,16 +39,23 @@ def cost_per_token(model: str, usage: "Usage", service_tier: str | None = None)
usage=usage,
custom_llm_provider="anthropic",
service_tier=service_tier,
model_info=model_info,
)
# Apply provider_specific_entry multipliers for geo/speed routing
try:
model_info: Final = litellm.get_model_info(model=model, custom_llm_provider="anthropic")
provider_specific_entry: Final[dict] = model_info.get("provider_specific_entry") or {}
effective_info: Final = (
model_info
if model_info is not None
else litellm.get_model_info(model=model, custom_llm_provider="anthropic")
)
provider_specific_entry: Final = effective_info.get("provider_specific_entry")
geo_multiplier: Final = get_provider_specific_geo_multiplier(model_info=model_info, usage=usage)
geo_multiplier: Final = get_provider_specific_geo_multiplier(model_info=effective_info, usage=usage)
speed_multiplier: Final = (
provider_specific_entry.get("fast", 1.0) if getattr(usage, "speed", None) == "fast" else 1.0
provider_specific_entry.get("fast", 1.0)
if provider_specific_entry and getattr(usage, "speed", None) == "fast"
else 1.0
)
if speed_multiplier != 1.0:

View file

@ -376,10 +376,9 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
usage_dict: UsageDelta = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(
chunk.usage
)
merged_chunk["usage"] = usage_dict
if self.applied_edits and "context_management" not in merged_chunk:
merged_chunk["context_management"] = ContextManagementResponse(applied_edits=list(self.applied_edits))
return self._augment_message_delta_usage(merged_chunk)
return self._augment_message_delta_usage({**merged_chunk, "usage": usage_dict})
def _handle_choiceless_chunk(self, chunk: "ModelResponseStream") -> bool:
"""Consume an OpenAI-compatible chunk that carries no ``choices``.
@ -448,8 +447,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
}
iterations.append(message_iteration)
augmented_usage["iterations"] = iterations
augmented["usage"] = augmented_usage
return augmented
return {**augmented, "usage": augmented_usage}
def _next_compaction_event(self) -> dict[str, object] | None:
"""Return the next compaction content-block SSE event, or ``None``.

View file

@ -2,9 +2,11 @@
For calculating cost of fireworks ai serverless inference models.
"""
import math
from datetime import datetime
from typing import Final
from typing import (
Final,
cast, # noqa: TID251 # the fallback entry is a dict copy of a ReadOnly TypedDict; no cast-free way to retype it
)
from litellm.constants import (
FIREWORKS_AI_4_B,
@ -12,12 +14,10 @@ from litellm.constants import (
FIREWORKS_AI_56_B_MOE,
FIREWORKS_AI_176_B_MOE,
)
from litellm.litellm_core_utils.llm_cost_calc.utils import TokenRates, apply_off_peak_pricing
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
from litellm.types.utils import ModelInfo, Usage
from litellm.utils import get_model_info
NO_CACHE_READ_RATE: Final = float("nan")
# Extract the number of billion parameters from the model name
# only used for together_computer LLMs
@ -67,6 +67,28 @@ def _resolve_model_info(model: str) -> ModelInfo:
return get_model_info(model=base_model, custom_llm_provider="fireworks_ai")
def _with_cache_read_fallback(model_info: ModelInfo) -> ModelInfo:
"""Entries without a cache-read rate keep the previous calculator's input-rate fallback for cached
reads (LIT-7845 tracks the documented discount); the shared map is never mutated, so a copy carries it."""
input_rate: Final = model_info.get("input_cost_per_token")
if model_info.get("cache_read_input_token_cost") is not None or input_rate is None:
return model_info
off_peak: Final = model_info.get("off_peak_pricing")
if off_peak is None or "cache_read_input_token_cost" in off_peak:
return cast(ModelInfo, {**model_info, "cache_read_input_token_cost": input_rate})
return cast(
ModelInfo,
{
**model_info,
"cache_read_input_token_cost": input_rate,
"off_peak_pricing": {
**off_peak,
"cache_read_input_token_cost": off_peak.get("input_cost_per_token", input_rate),
},
},
)
def cost_per_token(model: str, usage: Usage, current_time: datetime | None = None) -> tuple[float, float]:
"""
Calculates the cost per token for a given model, prompt tokens, and completion tokens,
@ -80,29 +102,11 @@ def cost_per_token(model: str, usage: Usage, current_time: datetime | None = Non
Returns:
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
"""
model_info: Final = _resolve_model_info(model)
standard_cache_read_rate: Final = model_info.get("cache_read_input_token_cost")
rates: Final = apply_off_peak_pricing(
model_info,
current_time,
TokenRates(
input_rate=model_info["input_cost_per_token"] or 0.0,
output_rate=model_info["output_cost_per_token"] or 0.0,
cache_read_rate=standard_cache_read_rate if standard_cache_read_rate is not None else NO_CACHE_READ_RATE,
cache_creation_rate=0.0,
reasoning_rate=None,
),
model_info: Final = _with_cache_read_fallback(_resolve_model_info(model))
return generic_cost_per_token(
model=model,
usage=usage,
custom_llm_provider="fireworks_ai",
model_info=model_info,
current_time=current_time,
)
cache_read_rate: Final[float] = rates.input_rate if math.isnan(rates.cache_read_rate) else rates.cache_read_rate
prompt_tokens_details: Final = usage.prompt_tokens_details
cached_tokens: Final[int] = (
prompt_tokens_details.cached_tokens
if prompt_tokens_details is not None and prompt_tokens_details.cached_tokens is not None
else 0
)
non_cached_prompt_tokens: Final[int] = max(usage.prompt_tokens - cached_tokens, 0)
prompt_cost: Final[float] = non_cached_prompt_tokens * rates.input_rate + cached_tokens * cache_read_rate
completion_cost: Final[float] = usage.completion_tokens * rates.output_rate
return prompt_cost, completion_cost

View file

@ -124,6 +124,12 @@ def _unsupported_reasoning_effort(reasoning_effort: str) -> UnsupportedParamsErr
)
def _served_model_name(model_version: object) -> str | None:
if not isinstance(model_version, str) or not model_version:
return None
return model_version.split("@", 1)[0]
class VertexAIBaseConfig:
def get_mapped_special_auth_params(self) -> dict:
"""
@ -1951,6 +1957,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
def _check_prompt_level_content_filter(
processed_chunk: GenerateContentResponseBody,
response_id: str | None,
model: str | None = None,
) -> Optional["ModelResponseStream"]:
"""
Check if prompt is blocked due to content filtering at the prompt level.
@ -1990,7 +1997,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
enhancements=None,
)
model_response: Final = ModelResponseStream(choices=[choice], id=response_id)
model_response: Final = ModelResponseStream(choices=[choice], id=response_id, model=model)
return model_response
return None
@ -2434,7 +2441,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
completion_response = GenerateContentResponseBody(**completion_response)
## GET MODEL ##
model_response.model = model
served: Final = _served_model_name(completion_response.get("modelVersion"))
model_response.model = served if served is not None else model
## CHECK IF RESPONSE FLAGGED
if "promptFeedback" in completion_response and "blockReason" in completion_response["promptFeedback"]:
@ -3264,12 +3272,18 @@ class ModelResponseIterator:
processed_chunk: Final = GenerateContentResponseBody(**chunk)
response_id: Final = processed_chunk.get("responseId")
model_response = ModelResponseStream(choices=[], id=response_id)
served: Final = _served_model_name(processed_chunk.get("modelVersion"))
model_response = ModelResponseStream(
choices=[],
id=response_id,
model=served,
)
# Check if prompt is blocked due to content filtering
blocked_response: Final = VertexGeminiConfig._check_prompt_level_content_filter(
processed_chunk=processed_chunk,
response_id=response_id,
model=served,
)
if blocked_response is not None:
model_response = blocked_response

View file

@ -49,7 +49,6 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
Inherits from OpenAIResponsesAPIConfig since XAI's Responses API is largely
compatible with OpenAI's, with a few differences:
- Does not support the 'instructions' parameter
- Requires code_interpreter tools to have 'container' field removed
- Recommends store=false when sending images
@ -60,20 +59,6 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.XAI
def get_supported_openai_params(self, model: str) -> list:
"""
Get supported parameters for XAI Responses API.
XAI supports most OpenAI Responses API params except 'instructions'.
"""
supported_params: Final = super().get_supported_openai_params(model)
# Remove 'instructions' as it's not supported by XAI
if "instructions" in supported_params:
supported_params.remove("instructions")
return supported_params
def _transform_web_search_tool(self, tool: Mapping[str, object]) -> Mapping[str, object]:
"""
Transform web_search tool to XAI format.
@ -158,19 +143,13 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
Map parameters for XAI Responses API.
Handles XAI-specific transformations:
1. Drops 'instructions' parameter (not supported)
2. Transforms code_interpreter tools to remove 'container' field
3. Transforms web_search tools to XAI format (removes search_context_size, adds filters)
4. Transforms x_search tools to XAI format
5. Sets store=false when images are detected (recommended by XAI)
1. Transforms code_interpreter tools to remove 'container' field
2. Transforms web_search tools to XAI format (removes search_context_size, adds filters)
3. Transforms x_search tools to XAI format
4. Sets store=false when images are detected (recommended by XAI)
"""
params: Final = dict(response_api_optional_params)
# Drop instructions parameter (not supported by XAI)
if "instructions" in params:
verbose_logger.debug("XAI Responses API does not support 'instructions' parameter. Dropping it.")
params.pop("instructions")
if "metadata" in params:
verbose_logger.debug("XAI Responses API does not support 'metadata' parameter. Dropping it.")
params.pop("metadata")

View file

@ -1,6 +1,8 @@
"""Bridge token flow: litellm identity resolution and the DCR-bridge oauth_delegate mint/refresh pipeline."""
import math
import os
import secrets
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Final, Literal
@ -12,6 +14,9 @@ from typing_extensions import assert_never
from litellm._logging import verbose_logger
from litellm.proxy._experimental.mcp_server.oauth_utils import TOKEN_NO_CACHE_HEADERS
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
_V2_GCM_PREFIX, # pyright: ignore[reportPrivateUsage] # reuse the encrypted credential's format discriminator
)
from litellm.types.mcp_server.mcp_server_manager import MCPServer
if TYPE_CHECKING:
@ -24,6 +29,7 @@ if TYPE_CHECKING:
UpstreamTokenGrant,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.handle_jwt import JWTIdentity
def _litellm_key_from_request(request: Request) -> str | None:
@ -48,6 +54,64 @@ def _litellm_key_from_request(request: Request) -> str | None:
return None
async def oauth_authorization_uses_gateway_credential(request: Request) -> bool:
"""Classify credentials for browser authorize; candidates still require full authorization."""
from litellm.proxy.auth.handle_jwt import JWTHandler # noqa: PLC0415 # proxy import cycle
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # startup owns the active auth configuration
jwt_handler,
master_key,
user_custom_auth,
)
if "x-litellm-api-key" in request.headers:
return True
token: Final = _litellm_key_from_request(request)
if token is None:
return "authorization" in request.headers
if token.startswith("sk-") or (master_key and secrets.compare_digest(token.encode(), master_key.encode())):
return True
if user_custom_auth is not None or jwt_handler.litellm_jwtauth.oidc_userinfo_enabled:
return True
if not JWTHandler.is_jwt(token):
return await _opaque_bearer_is_gateway_credential(token)
claims: Final = JWTHandler.get_unverified_claims(token)
issuer: Final = claims.get("iss") if claims is not None else None
global_issuer: Final = os.getenv("JWT_ISSUER")
# An unscoped global validator can accept issuers absent from the configured issuer list.
if not isinstance(issuer, str) or not issuer or not global_issuer:
return True
return issuer == global_issuer or any(
issuer == configured.issuer for configured in jwt_handler.litellm_jwtauth.issuers or ()
)
async def _opaque_bearer_is_gateway_credential(token: str) -> bool:
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import (
is_envelope, # noqa: PLC0415 # envelope imports bridge types
is_refresh_envelope,
)
from litellm.proxy._types import hash_token # noqa: PLC0415 # proxy import cycle
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken # noqa: PLC0415 # proxy import cycle
from litellm.proxy.auth.resolvers.exceptions import KeyNotFoundError # noqa: PLC0415 # proxy import cycle
from litellm.proxy.auth.resolvers.store import IdentityStore # noqa: PLC0415 # proxy import cycle
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # startup owns the identity store dependencies
prisma_client,
user_api_key_cache,
)
if is_envelope(token) or is_refresh_envelope(token) or token.startswith(_V2_GCM_PREFIX):
return True
try:
if ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(token) is not None:
return True
await IdentityStore(prisma_client, user_api_key_cache).resolve(hashed_token=hash_token(token))
except KeyNotFoundError:
return False
except Exception as exc: # noqa: BLE001 # an identity lookup fault must not permit cookie fallback
verbose_logger.debug("OAuth bearer ownership could not be checked (%s)", type(exc).__name__)
return True
def _key_is_active(key_obj: "UserAPIKeyAuth") -> bool:
"""``True`` when the presented key is neither blocked nor past its expiry.
@ -243,6 +307,10 @@ async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResol
return "no_active_key"
if user_object is None:
return "no_active_key"
return _active_user_record(user_object)
def _active_user_record(user_object: "LiteLLM_UserTable") -> "LiteLLM_UserTable | Literal['no_active_key']":
if isinstance(user_object.metadata, dict) and user_object.metadata.get("scim_active") is False:
return "no_active_key"
return user_object
@ -301,15 +369,137 @@ async def _revalidate_active_subject(identity: "EnvelopeIdentity") -> "_KeyResol
async def _extract_user_id_from_request(request: Request) -> str | None:
"""The litellm ``user_id`` for the token request, so a per-user token is stored under the same
identity the egress later reads it by. Storage is best-effort, so every non-resolved outcome
(including a transient DB outage) collapses to ``None`` here and the caller simply skips the store;
the bridge mint, which must status those outcomes differently, consumes
:func:`_resolve_active_litellm_key` directly."""
resolved: Final = await _resolve_active_litellm_key(request)
if not isinstance(resolved, _ResolvedKey):
"""Resolve the caller for identity binding without granting credential-write permission."""
from litellm.proxy.auth.handle_jwt import JWTIdentity # noqa: PLC0415 # proxy import cycle
resolved: Final = await _resolve_request_auth(request)
if isinstance(resolved, JWTIdentity):
return resolved.user_id
return _active_key_user_id(resolved) if resolved is not None else None
async def authorize_oauth_credential_request(request: Request, server_id: str) -> str | None:
from litellm.proxy._types import UserAPIKeyAuth # noqa: PLC0415 # proxy import cycle
resolved: Final = await _resolve_request_auth(request, f"/v1/mcp/server/{server_id}/oauth-user-credential")
if not isinstance(resolved, UserAPIKeyAuth) or not _active_key_user_id(resolved):
return None
if not await can_store_oauth_credential(request, resolved, server_id):
return None
return resolved.user_id
async def _resolve_request_auth(
request: Request, write_route: str | None = None
) -> "UserAPIKeyAuth | JWTIdentity | None":
from litellm.proxy.auth.handle_jwt import JWTHandler # noqa: PLC0415 # proxy import cycle
token: Final = _litellm_key_from_request(request)
if token is not None and JWTHandler.is_jwt(token):
return await _resolve_jwt_auth(request, token, write_route)
resolved: Final = await _resolve_active_litellm_key(request)
return resolved.key if isinstance(resolved, _ResolvedKey) else None
async def can_store_oauth_credential(request: Request, auth: "UserAPIKeyAuth", server_id: str) -> bool:
"""Apply the same write policy to request credentials and verified signed-callback users."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # registry imports auth helpers
global_mcp_server_manager,
)
from litellm.proxy._experimental.mcp_server.ui_session_utils import (
can_access_mcp_server, # noqa: PLC0415 # proxy import cycle
)
from litellm.proxy.auth.route_checks import RouteChecks # noqa: PLC0415 # proxy import cycle
from litellm.proxy.auth.user_api_key_auth import ( # noqa: PLC0415 # proxy import cycle
_run_centralized_common_checks, # pyright: ignore[reportPrivateUsage] # reuse admission policy for the credential-write action
)
write_route: Final = f"/v1/mcp/server/{server_id}/oauth-user-credential"
try:
RouteChecks.is_virtual_key_allowed_to_call_route(route=write_route, valid_token=auth, request=request)
await _run_centralized_common_checks(
user_api_key_auth_obj=auth,
request=request,
request_data={},
route=write_route,
)
return await can_access_mcp_server(auth, server_id, global_mcp_server_manager.get_allowed_mcp_servers)
except Exception as exc: # noqa: BLE001 # authorization failure must never write credentials
verbose_logger.debug("OAuth credential write not authorized (%s)", type(exc).__name__)
return False
async def _resolve_jwt_auth(
request: Request,
token: str,
write_route: str | None,
) -> "UserAPIKeyAuth | JWTIdentity | None":
from litellm.proxy._types import UserAPIKeyAuth # noqa: PLC0415 # proxy import cycle
from litellm.proxy.auth.handle_jwt import JWTAuthManager # noqa: PLC0415 # proxy import cycle
from litellm.proxy.auth.user_api_key_auth import ( # noqa: PLC0415 # proxy import cycle
_resolve_jwt_to_virtual_key, # pyright: ignore[reportPrivateUsage] # reuse admission mapping policy without provisioning a new key
)
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # proxy globals initialized at startup
general_settings,
jwt_handler,
premium_user,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
if general_settings.get("enable_jwt_auth") is not True or premium_user is not True or prisma_client is None:
return None
try:
if jwt_handler.litellm_jwtauth.is_virtual_key_mapping_configured():
claims: Final = await jwt_handler.auth_jwt(token=token)
validate: Final = jwt_handler.litellm_jwtauth.custom_validate
if validate is not None and not validate(claims):
return None
mapped: Final = await _resolve_jwt_to_virtual_key(
jwt_claims=claims,
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
if isinstance(mapped, UserAPIKeyAuth):
return None if await _key_owner_scim_deactivated(mapped) or not _active_key_user_id(mapped) else mapped
if mapped is not None:
return None
if write_route is None:
identity: Final = await JWTAuthManager.resolve_identity(
api_key=token,
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
if identity.user_object is not None and isinstance(_active_user_record(identity.user_object), str):
return None
return identity
authorized: Final = await JWTAuthManager.authorize_jwt(
api_key=token,
jwt_handler=jwt_handler,
request_data={},
general_settings=general_settings,
route=write_route,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
request_headers=dict(request.headers),
request_method=request.method,
)
resolved_user: Final = authorized["user_object"]
if resolved_user is not None and isinstance(_active_user_record(resolved_user), str):
return None
return JWTAuthManager.user_api_key_auth_from_result(authorized)
except Exception as exc: # noqa: BLE001 # public OAuth exchange stays available; unvalidated identities never write credentials
verbose_logger.debug("OAuth JWT identity could not be validated (%s)", type(exc).__name__)
return None
return _active_key_user_id(resolved.key)
_UpstreamGrantRejection = Literal["no_access_token", "expired_lifetime"]

View file

@ -32,6 +32,9 @@ from litellm.proxy._experimental.mcp_server.bridge_token_flow import (
_prepare_bridge_mint,
_prepare_bridge_refresh,
_reload_active_user_by_id,
authorize_oauth_credential_request,
can_store_oauth_credential,
oauth_authorization_uses_gateway_credential,
)
from litellm.proxy._experimental.mcp_server.faults import (
CallerRejected,
@ -836,16 +839,30 @@ async def _user_can_reach_mcp_server(user_id: str, server_id: str) -> bool:
return server_id in await global_mcp_server_manager.get_allowed_mcp_servers(admitted)
async def _bridge_authorize_access_denial(
litellm_user_id: str,
async def _resolve_oauth_authorization_user(
request: Request,
mcp_server: MCPServer,
redirect_uri: str,
state: str,
) -> RedirectResponse | None:
"""The denial redirect for a signed-in user who cannot reach the target server, or None to proceed."""
if await _user_can_reach_mcp_server(litellm_user_id, mcp_server.server_id):
return None
return _bridge_access_denied_redirect(redirect_uri, state, mcp_server)
enforce_binding: bool,
) -> str | RedirectResponse:
"""Resolve the authorization subject without replacing denied credentials with cookie grants."""
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( # noqa: PLC0415 # proxy import cycle
_user_id_from_session_cookie,
)
use_gateway_credential: Final = enforce_binding and await oauth_authorization_uses_gateway_credential(request)
request_user_id: Final = (
await authorize_oauth_credential_request(request, mcp_server.server_id) if use_gateway_credential else None
)
if use_gateway_credential and request_user_id is None:
return _bridge_access_denied_redirect(redirect_uri, state, mcp_server)
user_id: Final = request_user_id or _user_id_from_session_cookie(request)
if user_id is None:
return _redirect_to_litellm_login(request)
if not await _user_can_reach_mcp_server(user_id, mcp_server.server_id):
return _bridge_access_denied_redirect(redirect_uri, state, mcp_server)
return user_id
async def authorize_with_server(
@ -911,23 +928,12 @@ async def authorize_with_server(
# Seal the authenticated caller into state so the token exchange cannot select another credential owner.
litellm_user_id: str | None = None
if enforce_binding or (resolved_server.is_dcr_bridge and resolved_server.is_oauth_delegate):
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( # noqa: PLC0415 # inline import avoids a module-load circular import
_user_id_from_session_cookie,
subject: Final = await _resolve_oauth_authorization_user(
request, resolved_server, redirect_uri, state, enforce_binding
)
litellm_user_id = (
await _extract_user_id_from_request(request) if enforce_binding else None
) or _user_id_from_session_cookie(request)
if litellm_user_id is None:
return _redirect_to_litellm_login(request)
denial: Final = await _bridge_authorize_access_denial(
litellm_user_id=litellm_user_id,
mcp_server=resolved_server,
redirect_uri=redirect_uri,
state=state,
)
if denial is not None:
return denial
if isinstance(subject, RedirectResponse):
return subject
litellm_user_id = subject
oauth_nonce: Final = secrets.token_urlsafe(32) if enforce_binding else None
encoded_state: Final = encode_state_with_base_url(
@ -1218,12 +1224,32 @@ async def exchange_token_with_server(
user_id: Final = resolved_user_id
if user_id:
try:
await _store_per_user_token_server_side(
server=resolved_server,
user_id=user_id,
token_response=token_response,
identity_binding_proof=binding_proof,
# Identity binding above must retain the verified caller even when a write is
# denied. Authorize persistence separately, immediately before its side effect.
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler
# A sealed code delegates a verified user for this authorized server. Raw
# request credentials retain their own JWT/key restrictions during resolution.
can_store: Final = (
await can_store_oauth_credential(
request, await MCPRequestHandler.reload_admitted_user(user_id), resolved_server.server_id
)
if bridge_identity is not None
else await authorize_oauth_credential_request(request, resolved_server.server_id) == user_id
)
if can_store:
await _store_per_user_token_server_side(
server=resolved_server,
user_id=user_id,
token_response=token_response,
identity_binding_proof=binding_proof,
)
else:
verbose_logger.warning(
"OAuth credential storage not authorized for user=%s server=%s",
user_id,
resolved_server.server_id,
)
except Exception as exc:
verbose_logger.warning(
"exchange_token_with_server: server-side storage failed for user=%s server=%s: %s",
@ -1236,8 +1262,9 @@ async def exchange_token_with_server(
"exchange_token_with_server: could not resolve a LiteLLM user_id for the request, "
"so the per-user token for server=%s was NOT stored. The authorization_code egress "
"requires the stored token, so the client will be challenged with 401 on reconnect. "
"Ensure the request carries a valid LiteLLM key (x-litellm-api-key or Authorization), "
"or store it via POST /mcp/server/{id}/oauth-user-credential.",
"Ensure the request carries a valid LiteLLM key or enabled JWT identity "
"(x-litellm-api-key or Authorization), "
"or store it via POST /v1/mcp/server/{id}/oauth-user-credential.",
resolved_server.server_id,
)

View file

@ -2,6 +2,7 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable
from typing import Final
from fastapi import HTTPException
@ -137,3 +138,15 @@ async def build_effective_auth_contexts(
if admitted_context is None:
return team_contexts
return [*team_contexts, admitted_context]
async def can_access_mcp_server(
user_api_key_auth: UserAPIKeyAuth,
server_id: str,
allowed_servers: Callable[[UserAPIKeyAuth], Awaitable[list[str]]],
) -> bool:
"""Resolve server access through the same credential contexts as MCP management."""
for context in await build_effective_auth_contexts(user_api_key_auth):
if server_id in await allowed_servers(context):
return True
return False

View file

@ -3125,6 +3125,11 @@
"title": "Total Prompt Tokens",
"type": "integer"
},
"total_response_time_ms": {
"default": 0,
"title": "Total Response Time Ms",
"type": "integer"
},
"total_spend": {
"default": 0.0,
"title": "Total Spend",
@ -3135,6 +3140,11 @@
"title": "Total Successful Requests",
"type": "integer"
},
"total_timed_requests": {
"default": 0,
"title": "Total Timed Requests",
"type": "integer"
},
"total_tokens": {
"default": 0,
"title": "Total Tokens",
@ -3643,6 +3653,16 @@
"title": "Successful Requests",
"type": "integer"
},
"timed_requests": {
"default": 0,
"title": "Timed Requests",
"type": "integer"
},
"total_response_time_ms": {
"default": 0,
"title": "Total Response Time Ms",
"type": "integer"
},
"total_tokens": {
"default": 0,
"title": "Total Tokens",

View file

@ -5220,6 +5220,8 @@ class BaseDailySpendTransaction(TypedDict):
api_requests: int
successful_requests: int
failed_requests: int
total_response_time_ms: NotRequired[int] # writable-ok: the rollup queue accumulates into this key in place
timed_requests: NotRequired[int] # writable-ok: the rollup queue accumulates into this key in place
class DailyTeamSpendTransaction(BaseDailySpendTransaction):

View file

@ -15,6 +15,7 @@ import os
import re
import time
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Final, Literal, NoReturn, Protocol, TypeVar, cast
import httpx
@ -54,7 +55,7 @@ from litellm.proxy._types import (
from litellm.proxy.auth.auth_checks import can_team_access_model
from litellm.proxy.auth.resolvers.grants import GrantResolver, UserLookup, canonical_user_id
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.auth.team_grants import team_model_aliases
from litellm.proxy.auth.team_grants import team_grants, team_model_aliases
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
get_management_object_ttl,
@ -62,6 +63,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.user_repository import UserRepository
from litellm.types.agents import AgentResponse
from litellm.types.proxy.auth.auth_checks import UserNotFoundError
from .auth_checks import (
_allowed_routes_check,
@ -128,6 +130,19 @@ class _UserInfoResponse(Protocol):
def json(self) -> dict[str, object]: ...
@dataclass(frozen=True, slots=True)
class JWTIdentity:
user_id: str | None
user_object: LiteLLM_UserTable | None
agent_id: str | None
@dataclass(frozen=True, slots=True)
class _JWTProvisioning:
user_id_upsert: bool
team_id_upsert: bool
class AgentLookup(Protocol):
"""The registered-agent lookups a JWT agent claim is matched against."""
@ -1471,6 +1486,7 @@ class JWTAuthManager:
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging,
team_id_upsert: bool | None = None,
) -> tuple[str | None, LiteLLM_TeamTable | None]:
"""Find and validate specific team ID from team_id_jwt_field or team_alias_jwt_field"""
individual_team_id = jwt_handler.get_team_id(token=jwt_valid_token, default_value=None)
@ -1498,7 +1514,9 @@ class JWTAuthManager:
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert,
team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert
if team_id_upsert is None
else team_id_upsert,
)
return individual_team_id, team_object
except HTTPException as e:
@ -1726,6 +1744,7 @@ class JWTAuthManager:
proxy_logging_obj: ProxyLogging,
route: str,
org_alias: str | None = None,
user_id_upsert: bool | None = None,
) -> tuple[
LiteLLM_UserTable | None,
LiteLLM_OrganizationTable | None,
@ -1789,7 +1808,11 @@ class JWTAuthManager:
user_id=user_id,
user_email=user_email,
sso_user_id=user_id,
upsert=jwt_handler.is_upsert_user_id(valid_user_email=valid_user_email),
upsert=(
jwt_handler.is_upsert_user_id(valid_user_email=valid_user_email)
if user_id_upsert is None
else user_id_upsert
),
),
team_id=team_id,
)
@ -2010,6 +2033,7 @@ class JWTAuthManager:
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging,
team_id_upsert: bool | None = None,
) -> None:
"""Attach team context from x-litellm-team-id to an admin result.
@ -2027,7 +2051,7 @@ class JWTAuthManager:
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert,
team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert if team_id_upsert is None else team_id_upsert,
)
except Exception as e:
# Fall back to pre-PR admin behavior: honor the admin's
@ -2262,57 +2286,136 @@ class JWTAuthManager:
request_headers: dict | None = None,
request_method: str | None = None,
) -> JWTAuthBuilderResult:
"""Main authentication and authorization builder"""
# Check if OIDC UserInfo endpoint is enabled, but fall back to standard
# JWT auth if the token itself is a well-formed JWT (3-part structure).
if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled and not jwt_handler.is_jwt(token=api_key):
verbose_proxy_logger.debug("OIDC UserInfo is enabled. Fetching user info from UserInfo endpoint.")
# Use the access token to fetch user info from OIDC UserInfo endpoint
jwt_valid_token: dict = await jwt_handler.get_oidc_userinfo(token=api_key)
else:
# Default behavior: decode and validate the JWT token
jwt_valid_token = await jwt_handler.auth_jwt(token=api_key)
# Check custom validate
if jwt_handler.litellm_jwtauth.custom_validate:
if not jwt_handler.litellm_jwtauth.custom_validate(jwt_valid_token):
raise HTTPException(
status_code=403,
detail="Invalid JWT token",
)
# Check RBAC
rbac_role: Final = jwt_handler.get_rbac_role(token=jwt_valid_token)
await JWTAuthManager.check_rbac_role(
jwt_handler,
jwt_valid_token,
general_settings,
request_data,
route,
rbac_role,
return await JWTAuthManager.authorize_jwt(
api_key=api_key,
jwt_handler=jwt_handler,
request_data=request_data,
general_settings=general_settings,
route=route,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
request_headers=request_headers,
request_method=request_method,
provisioning=_JWTProvisioning(
user_id_upsert=jwt_handler.litellm_jwtauth.user_id_upsert,
team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert,
),
)
@staticmethod
async def authenticate_jwt(api_key: str, jwt_handler: JWTHandler) -> dict[str, object]:
claims: Final = (
await jwt_handler.get_oidc_userinfo(token=api_key)
if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled and not jwt_handler.is_jwt(token=api_key)
else await jwt_handler.auth_jwt(token=api_key)
)
validate: Final = jwt_handler.litellm_jwtauth.custom_validate
if validate is not None and not validate(claims):
raise HTTPException(status_code=403, detail="Invalid JWT token")
return claims
@staticmethod
async def resolve_identity(
api_key: str,
jwt_handler: JWTHandler,
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging,
) -> JWTIdentity:
claims: Final = await JWTAuthManager.authenticate_jwt(api_key, jwt_handler)
return await JWTAuthManager._resolve_claim_identity(
claims, jwt_handler, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj
)
@staticmethod
async def _resolve_claim_identity(
claims: dict[str, object],
jwt_handler: JWTHandler,
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging,
) -> JWTIdentity:
claim_user_id, user_email, valid_user_email = await JWTAuthManager.get_user_info(jwt_handler, claims)
user_id: Final = (
jwt_handler.get_object_id(token=claims, default_value=None) or claim_user_id
if jwt_handler.get_rbac_role(token=claims) == LitellmUserRoles.INTERNAL_USER
else claim_user_id
)
agent_id: Final = JWTAuthManager.resolve_agent_id(jwt_handler, claims, jwt_handler.agent_lookup)
is_admin: Final = jwt_handler.is_admin(scopes=jwt_handler.get_scopes(token=claims))
try:
user, _, _, _, canonical_id = await JWTAuthManager.get_objects(
user_id=user_id,
user_email=user_email,
org_id=None,
end_user_id=None,
team_id=None,
valid_user_email=valid_user_email,
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
route="",
user_id_upsert=False,
)
except UserNotFoundError:
if not is_admin:
raise
return JWTIdentity(user_id=user_id, user_object=None, agent_id=agent_id)
return JWTIdentity(user_id=user_id if is_admin else canonical_id, user_object=user, agent_id=agent_id)
@staticmethod
async def authorize_jwt(
api_key: str,
jwt_handler: JWTHandler,
request_data: dict[str, object],
general_settings: dict[str, object],
route: str,
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging,
request_headers: dict[str, str] | None = None,
request_method: str | None = None,
provisioning: _JWTProvisioning | None = None,
) -> JWTAuthBuilderResult:
"""Resolve and authorize JWT context; only normal admission supplies provisioning."""
handler: Final = jwt_handler
jwt_valid_token: Final = await JWTAuthManager.authenticate_jwt(api_key, handler)
team_id_upsert: Final = provisioning.team_id_upsert if provisioning is not None else False
model: Final = request_data.get("model")
requested_model: Final = model if isinstance(model, str) else None
# Check RBAC
rbac_role: Final = handler.get_rbac_role(token=jwt_valid_token)
await JWTAuthManager.check_rbac_role(handler, jwt_valid_token, general_settings, request_data, route, rbac_role)
# Check Scope Based Access
scopes: Final = jwt_handler.get_scopes(token=jwt_valid_token)
if jwt_handler.litellm_jwtauth.enforce_scope_based_access and jwt_handler.litellm_jwtauth.scope_mappings:
scopes: Final = handler.get_scopes(token=jwt_valid_token)
if handler.litellm_jwtauth.enforce_scope_based_access and handler.litellm_jwtauth.scope_mappings:
JWTAuthManager.check_scope_based_access(
scope_mappings=jwt_handler.litellm_jwtauth.scope_mappings,
scope_mappings=handler.litellm_jwtauth.scope_mappings,
scopes=scopes,
request_data=request_data,
general_settings=general_settings,
)
object_id = jwt_handler.get_object_id(token=jwt_valid_token, default_value=None)
object_id = handler.get_object_id(token=jwt_valid_token, default_value=None)
# Get basic user info
user_id, user_email, valid_user_email = await JWTAuthManager.get_user_info(jwt_handler, jwt_valid_token)
user_id, user_email, valid_user_email = await JWTAuthManager.get_user_info(handler, jwt_valid_token)
# Get IDs
org_id: Final = jwt_handler.get_org_id(token=jwt_valid_token, default_value=None)
end_user_id: Final = jwt_handler.get_end_user_id(token=jwt_valid_token, default_value=None)
org_id: Final = handler.get_org_id(token=jwt_valid_token, default_value=None)
end_user_id: Final = handler.get_end_user_id(token=jwt_valid_token, default_value=None)
team_id: str | None = None
team_object: LiteLLM_TeamTable | None = None
object_id = jwt_handler.get_object_id(token=jwt_valid_token, default_value=None)
object_id = handler.get_object_id(token=jwt_valid_token, default_value=None)
if rbac_role and object_id:
if rbac_role == LitellmUserRoles.TEAM:
@ -2321,14 +2424,14 @@ class JWTAuthManager:
user_id = object_id
agent_id: Final = JWTAuthManager.resolve_agent_id(
jwt_handler=jwt_handler,
jwt_handler=handler,
jwt_valid_token=jwt_valid_token,
agent_registry=jwt_handler.agent_lookup,
agent_registry=handler.agent_lookup,
)
# Check admin access
admin_result: Final = await JWTAuthManager.check_admin_access(
jwt_handler,
handler,
scopes,
route,
user_id,
@ -2343,18 +2446,24 @@ class JWTAuthManager:
admin_result=admin_result,
route=route,
request_headers=request_headers,
jwt_handler=jwt_handler,
jwt_handler=handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
team_id_upsert=team_id_upsert,
)
if provisioning is None:
identity: Final = await JWTAuthManager._resolve_claim_identity(
jwt_valid_token, handler, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj
)
return {**admin_result, "user_object": identity.user_object}
return admin_result
# Get team with model access
## Check if team_id is specified via x-litellm-team-id header
all_team_ids: Final = JWTAuthManager.get_all_team_ids(jwt_handler, jwt_valid_token)
specific_team_id: Final = jwt_handler.get_team_id(token=jwt_valid_token, default_value=None)
all_team_ids: Final = JWTAuthManager.get_all_team_ids(handler, jwt_valid_token)
specific_team_id: Final = handler.get_team_id(token=jwt_valid_token, default_value=None)
# The DB fallback only applies when the token carries no team identity at
# all. `get_all_jwt_team_ids` ignores `team_id_default` so a configured
@ -2364,9 +2473,9 @@ class JWTAuthManager:
# the RBAC team-role path (which already set `team_id`); otherwise a
# provisional x-litellm-team-id header could override an RBAC-asserted team.
db_team_fallback: Final = (
jwt_handler.litellm_jwtauth.fallback_to_db_teams
and not jwt_handler.get_all_jwt_team_ids(token=jwt_valid_token)
and not jwt_handler.get_team_alias(token=jwt_valid_token, default_value=None)
handler.litellm_jwtauth.fallback_to_db_teams
and not handler.get_all_jwt_team_ids(token=jwt_valid_token)
and not handler.get_team_alias(token=jwt_valid_token, default_value=None)
and team_id is None
)
if specific_team_id and not db_team_fallback:
@ -2391,7 +2500,7 @@ class JWTAuthManager:
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
team_id_upsert=(jwt_handler.litellm_jwtauth.team_id_upsert and not db_team_fallback),
team_id_upsert=(team_id_upsert and not db_team_fallback),
)
except HTTPException:
if not db_team_fallback:
@ -2403,22 +2512,23 @@ class JWTAuthManager:
team_id,
team_object,
) = await JWTAuthManager.find_and_validate_specific_team_id(
jwt_handler,
handler,
jwt_valid_token,
prisma_client,
user_api_key_cache,
parent_otel_span,
proxy_logging_obj,
team_id_upsert=team_id_upsert,
)
if not team_object and not team_id:
## CHECK USER GROUP ACCESS
team_id, team_object = await JWTAuthManager.find_team_with_model_access(
team_ids=all_team_ids,
requested_model=request_data.get("model"),
requested_model=requested_model,
route=route,
request_method=request_method,
jwt_handler=jwt_handler,
jwt_handler=handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
@ -2442,7 +2552,7 @@ class JWTAuthManager:
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert,
team_id_upsert=team_id_upsert,
)
if team_id and not JWTAuthManager._team_has_passthrough_route_access(
@ -2453,7 +2563,7 @@ class JWTAuthManager:
JWTAuthManager._raise_team_passthrough_route_denial(route=route)
# Extract alias fields for resolution (if configured)
org_alias: Final = jwt_handler.get_org_alias(token=jwt_valid_token, default_value=None)
org_alias: Final = handler.get_org_alias(token=jwt_valid_token, default_value=None)
# get_objects returns effective_user_id for downstream spend attribution (GH #26789).
(
@ -2469,25 +2579,27 @@ class JWTAuthManager:
end_user_id=end_user_id,
team_id=team_id,
valid_user_email=valid_user_email,
jwt_handler=jwt_handler,
jwt_handler=handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
route=route,
org_alias=org_alias,
user_id_upsert=provisioning.user_id_upsert if provisioning is not None else False,
)
# Derive org_id from org_object if resolved by alias
resolved_org_id: Final = org_object.organization_id if org_object else org_id
await JWTAuthManager.sync_user_role_and_teams(
jwt_handler=jwt_handler,
jwt_valid_token=jwt_valid_token,
user_object=user_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
if provisioning is not None:
await JWTAuthManager.sync_user_role_and_teams(
jwt_handler=handler,
jwt_valid_token=jwt_valid_token,
user_object=user_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
# If JWT did not resolve team_id, attempt a team fallback.
if team_id is None and db_team_fallback:
@ -2498,11 +2610,11 @@ class JWTAuthManager:
) = await JWTAuthManager._resolve_db_team_fallback(
user_object=user_object,
user_id=user_id,
requested_model=request_data.get("model"),
requested_model=requested_model,
route=route,
jwt_handler=jwt_handler,
enforce_team_based_model_access=jwt_handler.litellm_jwtauth.enforce_team_based_model_access,
team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert,
jwt_handler=handler,
enforce_team_based_model_access=handler.litellm_jwtauth.enforce_team_based_model_access,
team_id_upsert=team_id_upsert,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
@ -2530,7 +2642,7 @@ class JWTAuthManager:
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert,
team_id_upsert=team_id_upsert,
)
elif db_team_fallback and team_id == header_team_id:
JWTAuthManager._validate_header_team_in_db_membership(
@ -2540,7 +2652,7 @@ class JWTAuthManager:
if not JWTAuthManager._is_team_route_allowed(
route=route,
request_method=request_method,
jwt_handler=jwt_handler,
jwt_handler=handler,
):
raise HTTPException(
status_code=403,
@ -2550,16 +2662,17 @@ class JWTAuthManager:
)
## MAP USER TO TEAMS
await JWTAuthManager.map_user_to_teams(
user_object=user_object,
team_object=team_object,
)
if provisioning is not None:
await JWTAuthManager.map_user_to_teams(
user_object=user_object,
team_object=team_object,
)
# Validate that a valid rbac id is returned for spend tracking
JWTAuthManager.validate_object_id(
user_id=user_id,
team_id=team_id,
enforce_rbac=general_settings.get("enforce_rbac", False),
enforce_rbac=bool(general_settings.get("enforce_rbac", False)),
is_proxy_admin=False,
)
@ -2582,3 +2695,38 @@ class JWTAuthManager:
jwt_claims=jwt_valid_token,
agent_id=agent_id,
)
@staticmethod
def user_api_key_auth_from_result(
result: JWTAuthBuilderResult,
parent_otel_span: Span | None = None,
) -> UserAPIKeyAuth:
"""Keep JWT identity and permission attribution identical across consumers."""
user: Final = result["user_object"]
admin: Final = result["is_proxy_admin"]
return UserAPIKeyAuth(
api_key=None,
user_role=(
LitellmUserRoles.PROXY_ADMIN
if admin
else LitellmUserRoles(user.user_role)
if user is not None and user.user_role is not None
else LitellmUserRoles.INTERNAL_USER
),
user_id=result["user_id"],
user_email=result["user_email"],
team_id=result["team_id"],
org_id=result["org_id"],
end_user_id=result["end_user_id"],
parent_otel_span=parent_otel_span,
jwt_claims=result["jwt_claims"],
agent_id=result.get("agent_id"),
user_tpm_limit=user.tpm_limit if user is not None and not admin else None,
user_rpm_limit=user.rpm_limit if user is not None and not admin else None,
user_model_max_budget=user.model_max_budget if user is not None and not admin else None,
**team_grants(
team_object=result["team_object"],
team_membership=result.get("team_membership"),
user_id=result["user_id"],
),
)

View file

@ -1669,13 +1669,11 @@ async def _user_api_key_auth_builder(
is_proxy_admin: Final = result["is_proxy_admin"]
team_id: Final = result["team_id"]
team_object: Final = result["team_object"]
user_id: Final = result["user_id"]
user_email: Final = result["user_email"]
user_object: Final = result["user_object"]
end_user_id = result["end_user_id"]
org_id: Final = result["org_id"]
team_membership: Final[LiteLLM_TeamMembership | None] = result.get("team_membership", None)
jwt_claims = result.get("jwt_claims", None)
agent_id: Final[str | None] = result.get("agent_id")
@ -1693,40 +1691,9 @@ async def _user_api_key_auth_builder(
value=_JWT_PROXY_ADMIN_SENTINEL,
ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl,
)
return UserAPIKeyAuth(
api_key=None,
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id=user_id,
user_email=user_email,
team_id=team_id,
org_id=org_id,
end_user_id=end_user_id,
parent_otel_span=parent_otel_span,
jwt_claims=jwt_claims,
agent_id=agent_id,
**team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id),
)
return JWTAuthManager.user_api_key_auth_from_result(result, parent_otel_span)
valid_token = UserAPIKeyAuth(
api_key=None,
team_id=team_id,
user_role=(
LitellmUserRoles(user_object.user_role)
if user_object is not None and user_object.user_role is not None
else LitellmUserRoles.INTERNAL_USER
),
user_id=user_id,
user_email=user_email,
org_id=org_id,
parent_otel_span=parent_otel_span,
end_user_id=end_user_id,
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),
jwt_claims=jwt_claims,
agent_id=agent_id,
**team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id),
)
valid_token = JWTAuthManager.user_api_key_auth_from_result(result, parent_otel_span)
# AUTO_REGISTER deferred from _resolve_jwt_to_virtual_key.
# JWT policy (RBAC, scope, custom_validate, email-domain)

View file

@ -57,6 +57,8 @@ _COUNTER_COLUMNS: Final = (
"cache_read_input_tokens",
"cache_creation_input_tokens",
"compression_saved_tokens",
"total_response_time_ms",
"timed_requests",
)
_SPEND_COLUMNS: Final = (
"spend",

View file

@ -137,6 +137,19 @@ class _SpendTransactionManager(Protocol):
async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ...
def _timed_request_duration_ms(
payload: dict | SpendLogsPayload,
request_status: Literal["success", "failure"],
is_internal_call: bool,
) -> int | None:
if is_internal_call or request_status != "success":
return None
duration_ms: Final = payload.get("request_duration_ms")
if not isinstance(duration_ms, int) or duration_ms < 0:
return None
return duration_ms
def _spend_update_tx(prisma_client: PrismaClient) -> _SpendTransactionManager:
tx: Final[_SpendTransactionManager] = prisma_client.db.tx(timeout=timedelta(seconds=60))
return tx
@ -2232,6 +2245,7 @@ class DBSpendUpdateWriter:
recorded_autorouter_savings=_metadata.get("autorouter_savings"),
billed_at=payload.get("endTime"),
)
timed_duration_ms: Final = _timed_request_duration_ms(payload, request_status, is_internal_call)
daily_transaction: Final = BaseDailySpendTransaction(
date=date,
@ -2259,6 +2273,8 @@ class DBSpendUpdateWriter:
prompt_caching_savings_spend=savings_spend.prompt_caching,
gateway_injected_caching_savings_spend=savings_spend.gateway_injected_caching,
autorouter_savings_spend=0.0 if is_internal_call else savings_spend.autorouter,
total_response_time_ms=timed_duration_ms or 0,
timed_requests=0 if timed_duration_ms is None else 1,
)
return daily_transaction
except Exception as e:

View file

@ -142,6 +142,14 @@ class DailySpendUpdateQueue(BaseUpdateQueue):
payload.get("autorouter_savings_spend", 0) or 0
) + daily_transaction.get("autorouter_savings_spend", 0)
daily_transaction["total_response_time_ms"] = (
payload.get("total_response_time_ms", 0) or 0
) + daily_transaction.get("total_response_time_ms", 0)
daily_transaction["timed_requests"] = (
payload.get("timed_requests", 0) or 0
) + daily_transaction.get("timed_requests", 0)
else:
aggregated_daily_spend_update_transactions[_key] = deepcopy(payload)
return aggregated_daily_spend_update_transactions

View file

@ -1,14 +1,17 @@
"""LLM-as-a-Judge guardrail: uses an LLM to score responses against weighted criteria."""
"""LLM-as-a-Judge guardrail: uses an LLM to score requests or responses against weighted criteria."""
from collections.abc import Callable, Sequence
from collections.abc import Callable, Mapping, Sequence
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Generic, Literal, Optional, TypeVar
from fastapi import HTTPException
from pydantic import BaseModel, ConfigDict, ValidationError
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.llm_judge import (
default_router_provider,
@ -16,8 +19,9 @@ from litellm.litellm_core_utils.llm_judge import (
judge_acompletion,
parse_json_verdict,
)
from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations
from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus
from litellm.litellm_core_utils.prompt_templates.common_utils import get_last_user_message
from litellm.types.guardrails import GuardrailEventHooks, Mode, SupportedGuardrailIntegrations
from litellm.types.utils import LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN, GenericGuardrailAPIInputs, GuardrailStatus
if TYPE_CHECKING:
from litellm import Router
@ -26,18 +30,65 @@ if TYPE_CHECKING:
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import StandardLoggingEvalInformation
JUDGE_SYSTEM_PROMPT = """You are a quality judge. Evaluate the assistant's response against the criteria provided.
For each criterion, assign a score from 0 to 100 and provide concise reasoning.
JudgeInputType = Literal["request", "response"]
JudgeEventHook = GuardrailEventHooks | list[GuardrailEventHooks] | Mode
JudgeModeParam = str | list[str] | Mode | GuardrailEventHooks | list[GuardrailEventHooks] | None
_JUDGE_SYSTEM_PROMPT_TEMPLATE: Final = """You are a quality judge. Evaluate the {subject} against the criteria provided.
{focus}For each criterion, assign a score from 0 to 100 and provide concise reasoning.
Return ONLY valid JSON in this exact format:
{
{{
"verdicts": [
{"criterion_name": "<name>", "score": <0-100>, "reasoning": "<one sentence>", "passed": <true|false>, "weight": <weight>}
{{"criterion_name": "<name>", "score": <0-100>, "reasoning": "<one sentence>", "passed": <true|false>, "weight": <weight>}}
],
"overall_score": <weighted average 0-100>
}"""
}}"""
JUDGE_SYSTEM_PROMPTS: Final[MappingProxyType[JudgeInputType, str]] = MappingProxyType(
{
"request": _JUDGE_SYSTEM_PROMPT_TEMPLATE.format(
subject="request",
focus="Judge the most recent user turn; treat earlier turns in the conversation only as context.\n",
),
"response": _JUDGE_SYSTEM_PROMPT_TEMPLATE.format(subject="assistant's response", focus=""),
}
)
_JUDGE_SUBJECT_LABELS: Final[MappingProxyType[JudgeInputType, str]] = MappingProxyType(
{"request": "Latest request turn to evaluate", "response": "Assistant response to evaluate"}
)
_LIFECYCLE_HOOKS: Final[MappingProxyType[JudgeInputType, tuple[GuardrailEventHooks, ...]]] = MappingProxyType(
{
"request": (GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call, GuardrailEventHooks.logging_only),
"response": (GuardrailEventHooks.post_call, GuardrailEventHooks.logging_only),
}
)
_VALID_ON_FAILURE: Final = frozenset({"block", "log"})
_JUDGE_CALL_METADATA: Final = MappingProxyType(
{INTERNAL_CALL_ORIGIN_METADATA_KEY: LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN}
)
class _LoggedCallParams(BaseModel):
model_config = ConfigDict(frozen=True)
metadata: Mapping[str, object] | None = None
def _is_logged_judge_call(data: Mapping[str, object], event_type: GuardrailEventHooks) -> bool:
"""logging_only is the only event whose ``data`` is the SDK's model_call_details rather than the client body."""
if event_type is not GuardrailEventHooks.logging_only:
return False
try:
params: Final = _LoggedCallParams.model_validate(data.get("litellm_params") or {})
except ValidationError:
return False
return (params.metadata or {}).get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN
_default_router_provider: Final = default_router_provider
_parse_judge_verdict: Final = parse_json_verdict
_extract_text_from_content: Final = extract_text_from_content
@ -86,10 +137,29 @@ def _get_litellm_param(
return default
def _coerce_event_hook(mode: JudgeModeParam) -> JudgeEventHook:
if mode is None:
return GuardrailEventHooks.post_call
if isinstance(mode, Mode):
return mode
if isinstance(mode, list):
return [GuardrailEventHooks(hook) for hook in mode]
return GuardrailEventHooks(mode)
def _text_under_review(inputs: GenericGuardrailAPIInputs, input_type: JudgeInputType) -> str:
all_text: Final = "\n".join(inputs.get("texts") or [])
if input_type == "response":
return all_text
latest_user_turn: Final = get_last_user_message(inputs.get("structured_messages") or [])
return latest_user_turn if latest_user_turn is not None else all_text
def _build_judge_prompt(
criteria: Sequence[JudgeCriterion],
messages: Sequence[JudgeMessage],
response_text: str,
text_under_review: str,
input_type: JudgeInputType = "response",
) -> str:
criteria_block: Final = "\n".join(
f"- {c.get('name', '')} (weight {c.get('weight', 0)}%): {c.get('description', '')}" for c in criteria
@ -99,15 +169,16 @@ def _build_judge_prompt(
for m in messages
if m.get("content") is not None
)
conversation_block: Final = f"Conversation:\n{conversation}\n\n" if conversation or input_type == "response" else ""
return (
f"Criteria to evaluate:\n{criteria_block}\n\n"
f"Conversation:\n{conversation}\n\n"
f"Assistant response to evaluate:\n{response_text}"
f"{conversation_block}"
f"{_JUDGE_SUBJECT_LABELS[input_type]}:\n{text_under_review}"
)
class LLMAsAJudgeGuardrail(CustomGuardrail):
"""Post-call guardrail that judges response quality via an LLM."""
"""Guardrail that judges request (pre_call/during_call) or response (post_call) quality via an LLM."""
def __init__(
self,
@ -116,22 +187,15 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
criteria: Sequence[JudgeCriterion],
overall_threshold: float = 80.0,
on_failure: Literal["block", "log"] = "block",
event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None,
event_hook: JudgeModeParam = None,
default_on: bool = False,
router_provider: "Callable[[], Router | None] | None" = None,
**kwargs: Any,
) -> None:
_event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None
if event_hook is not None:
if isinstance(event_hook, list):
_event_hook = [GuardrailEventHooks(h) if isinstance(h, str) else h for h in event_hook]
else:
_event_hook = GuardrailEventHooks(event_hook) if isinstance(event_hook, str) else event_hook
super().__init__(
guardrail_name=guardrail_name,
supported_event_hooks=list(self.get_supported_event_hooks()),
event_hook=_event_hook or GuardrailEventHooks.post_call,
event_hook=_coerce_event_hook(event_hook),
default_on=default_on,
**kwargs,
)
@ -143,18 +207,24 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
@classmethod
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
return [GuardrailEventHooks.post_call]
return [GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call, GuardrailEventHooks.post_call]
def should_run_guardrail(self, data: Mapping[str, object], event_type: GuardrailEventHooks) -> bool:
if _is_logged_judge_call(data, event_type):
return False
return super().should_run_guardrail(data, event_type)
async def _run_judge(
self,
messages: Sequence[JudgeMessage],
response_text: str,
text_under_review: str,
input_type: JudgeInputType = "response",
) -> dict[str, object]:
judge_messages: Final[list[AllMessageValues]] = [
{"role": "system", "content": JUDGE_SYSTEM_PROMPT},
{"role": "system", "content": JUDGE_SYSTEM_PROMPTS[input_type]},
{
"role": "user",
"content": _build_judge_prompt(self.criteria, messages, response_text),
"content": _build_judge_prompt(self.criteria, messages, text_under_review, input_type),
},
]
response: Final = await judge_acompletion(
@ -163,6 +233,7 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
judge_messages,
response_format={"type": "json_object"},
temperature=0,
metadata=dict(_JUDGE_CALL_METADATA),
)
raw: Final = response.choices[0].message.content or "{}"
return _parse_judge_verdict(raw)
@ -174,13 +245,8 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> GenericGuardrailAPIInputs:
# Only evaluate post-call (response text). Fail open on pre-call.
if input_type != "response":
return inputs
texts: Final = inputs.get("texts") or []
response_text: Final = " ".join(texts)
if not response_text:
text_under_review: Final = _text_under_review(inputs, input_type)
if not text_under_review:
return inputs
start_time: Final = datetime.now()
@ -188,10 +254,12 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
judge_result: dict[str, object] = {}
try:
messages: Final[Sequence[JudgeMessage]] = request_data.get("messages") or []
messages: Final[Sequence[JudgeMessage]] = (
inputs.get("structured_messages") or request_data.get("messages") or []
)
try:
judge_result = await self._run_judge(messages, response_text)
judge_result = await self._run_judge(messages, text_under_review, input_type)
except Exception as judge_err:
verbose_logger.warning(
"llm_as_a_judge guardrail: judge call failed, failing open. Error: %s", judge_err
@ -230,7 +298,7 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
raise HTTPException(
status_code=422,
detail={
"error": "LLM judge rejected response: score below threshold",
"error": f"LLM judge rejected {input_type}: score below threshold",
"overall_score": overall_score,
"threshold": self.overall_threshold,
"verdicts": judge_result.get("verdicts", []),
@ -252,9 +320,13 @@ class LLMAsAJudgeGuardrail(CustomGuardrail):
guardrail_status=status,
start_time=start_time.timestamp(),
end_time=datetime.now().timestamp(),
event_type=GuardrailEventHooks.post_call,
event_type=self._event_type_for(input_type),
)
def _event_type_for(self, input_type: JudgeInputType) -> GuardrailEventHooks | None:
configured: Final = tuple(hook for hook in _LIFECYCLE_HOOKS[input_type] if self._event_hook_is_event_type(hook))
return configured[0] if len(configured) == 1 else None
def initialize_guardrail(
litellm_params: "LitellmParams",
@ -282,10 +354,7 @@ def initialize_guardrail(
overall_threshold: Final = float(_get_litellm_param(litellm_params, guardrail, "overall_threshold", 80.0))
mode: Final[str | None] = _get_litellm_param(litellm_params, guardrail, "mode", None)
event_hook: GuardrailEventHooks | None = None
if isinstance(mode, str) and mode in {e.value for e in GuardrailEventHooks}:
event_hook = GuardrailEventHooks(mode)
mode: Final[JudgeModeParam] = _get_litellm_param(litellm_params, guardrail, "mode", None)
instance: Final = LLMAsAJudgeGuardrail(
guardrail_name=guardrail_name,
@ -293,7 +362,7 @@ def initialize_guardrail(
criteria=criteria,
overall_threshold=overall_threshold,
on_failure=on_failure,
event_hook=event_hook,
event_hook=mode,
default_on=bool(_get_litellm_param(litellm_params, guardrail, "default_on", False)),
)
litellm.logging_callback_manager.add_litellm_callback(instance)

View file

@ -1,4 +1,7 @@
import json
import os
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Any, Final
from urllib.parse import urlparse
@ -19,20 +22,26 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.guardrails.guardrail_hooks.base import (
GuardrailConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.singulr import (
AssistantMessage,
SingulrGuardrailPayload,
SingulrGuardrailRequest,
SingulrGuardrailResponse,
SingulrMcpGuardrailPayload,
ToolCall,
ToolCallFunction,
)
from litellm.types.utils import GenericGuardrailAPIInputs
from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs
_DEFAULT_API_BASE: Final = "http://localhost:8003"
_GUARD_ENDPOINT: Final = "/api/v1/ai-gateway/litellm"
_GUARD_ENDPOINT: Final = "/api/v1/ai-gateway/litellm-v2"
_DEFAULT_TIMEOUT: Final = 30.0
_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({})
_MCP_MODEL_PREFIX: Final = "MCP:"
class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object):
@ -51,8 +60,8 @@ class SingulrGuardrail(CustomGuardrail):
**kwargs: Unpack[_CustomGuardrailOptions],
) -> None:
self.singulr_api_key = singulr_api_key or os.environ.get("SINGULR_API_KEY")
self.singulr_api_base = (singulr_api_base or os.environ.get("SINGULR_API_BASE") or _DEFAULT_API_BASE).rstrip(
"/"
self.singulr_api_base = (
(singulr_api_base or os.environ.get("SINGULR_API_BASE") or _DEFAULT_API_BASE).strip().rstrip("/")
)
parsed: Final = urlparse(self.singulr_api_base)
if parsed.scheme == "http" and parsed.hostname not in (
@ -85,6 +94,9 @@ class SingulrGuardrail(CustomGuardrail):
kwargs["supported_event_hooks"] = [
GuardrailEventHooks.pre_call,
GuardrailEventHooks.post_call,
GuardrailEventHooks.logging_only,
GuardrailEventHooks.pre_mcp_call,
GuardrailEventHooks.post_mcp_call,
]
super().__init__(**kwargs)
@ -97,52 +109,70 @@ class SingulrGuardrail(CustomGuardrail):
return SingulrGuardrailConfigModel
def _build_payload(
self,
request_data: dict[str, Any],
inputs: GenericGuardrailAPIInputs,
input_type: str,
) -> dict[str, object]:
if not request_data:
texts: Final = inputs.get("texts", [])
payload = SingulrGuardrailPayload(
input_type=input_type,
is_playground_request=True,
playground_text=texts[0] if texts else None,
@staticmethod
def _metadata_containers(request_data: Mapping[str, Any]) -> tuple[Mapping[str, Any], ...]:
litellm_params: Final = request_data.get("litellm_params") or _EMPTY_MAPPING
return tuple(
container
for container in (
request_data.get("litellm_metadata"),
request_data.get("metadata"),
litellm_params.get("litellm_metadata") if litellm_params else None,
litellm_params.get("metadata") if litellm_params else None,
)
else:
response: Final = request_data.get("response")
singulr_req_object: Final = SingulrGuardrailRequest(
model=request_data.get("model"),
messages=request_data.get("messages"),
tools=request_data.get("tools"),
model_response=response.model_dump(mode="json") if input_type == "response" and response else None,
litellm_metadata=request_data.get("litellm_metadata"),
)
payload = SingulrGuardrailPayload(
litellm_call_id=request_data.get("litellm_call_id"),
request_data=singulr_req_object,
input_type=input_type,
)
return payload.model_dump(mode="json")
def _build_headers(self) -> dict[str, str]:
return dict(
(header, value)
for header, value in (
("Content-Type", "application/json"),
("X-Singulr-Gateway-Token", self.singulr_api_key),
(
"X-Singulr-Enforcement-Entity-Id",
self.singulr_application_id or "",
),
("X-Singulr-Guardrail-Id", self.singulr_guardrail_id or ""),
)
if value
if container
)
@classmethod
def _resolve_metadata_value(cls, request_data: Mapping[str, Any], key: str) -> str | None:
for container in cls._metadata_containers(request_data=request_data):
value = container.get(key)
if value:
return value
return None
@classmethod
def _resolve_user_role_from_request_data(cls, request_data: Mapping[str, Any]) -> str | None:
for container in cls._metadata_containers(request_data=request_data):
auth = container.get("user_api_key_auth")
if isinstance(auth, UserAPIKeyAuth) and auth.user_role:
return auth.user_role.value
return None
@classmethod
def _build_metadata(cls, request_data: Mapping[str, Any]) -> Mapping[str, str] | None:
fields: Final = (
"user_api_key_alias",
"user_api_key_user_id",
"user_api_key_user_email",
"user_api_key_org_id",
"user_api_key_org_alias",
"user_api_key_team_id",
"user_api_key_team_alias",
)
resolved: Final = (
*((field, cls._resolve_metadata_value(request_data=request_data, key=field)) for field in fields),
("user_api_key_user_role", cls._resolve_user_role_from_request_data(request_data=request_data)),
)
if not any(value for _, value in resolved):
return None
return {key: value for key, value in resolved if value} # mutable-ok: short-lived JSON payload dict
@staticmethod
def _build_user_message(text: str) -> Mapping[str, Any]:
return {"role": "user", "content": text} # mutable-ok: short-lived JSON payload dict
def _build_headers(self) -> Mapping[str, str]:
all_headers: Final = MappingProxyType(
{
"Content-Type": "application/json",
"X-Singulr-Gateway-Token": self.singulr_api_key,
"X-Singulr-Enforcement-Entity-Id": self.singulr_application_id,
"X-Singulr-Guardrail-Id": self.singulr_guardrail_id,
}
)
return MappingProxyType({header: value for header, value in all_headers.items() if value})
async def _call_api(self, payload: dict[str, object]) -> SingulrGuardrailResponse | None:
endpoint: Final = f"{self.singulr_api_base}{_GUARD_ENDPOINT}"
verbose_proxy_logger.debug("Singulr: %s", endpoint)
@ -168,7 +198,7 @@ class SingulrGuardrail(CustomGuardrail):
if self.block_on_error:
raise GuardrailRaisedException(
guardrail_name=self.guardrail_name,
message=(f"Singulr API returned HTTP {exc.response.status_code}: {exc.response.text}"),
message=f"Singulr API returned HTTP {exc.response.status_code}: {exc.response.text}",
) from exc
return None
@ -190,33 +220,218 @@ class SingulrGuardrail(CustomGuardrail):
) from exc
return None
@log_guardrail_information
async def apply_guardrail(
async def _apply_guardrail_on_request(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: str,
logging_obj: "LiteLLMLoggingObj | None" = None,
texts: Sequence[str],
structured_messages: Sequence[Any],
request_data: Mapping[str, Any],
) -> GenericGuardrailAPIInputs:
payload: Final = self._build_payload(request_data, inputs, input_type)
if not payload:
return inputs
result: Final = await self._call_api(payload)
if result is None:
return inputs
verbose_proxy_logger.debug(
"Singulr: should_block=%s blocking_due_to=%s",
result.should_block,
result.blocking_due_to,
messages: Final = (
tuple(structured_messages)
if structured_messages
else tuple(self._build_user_message(text) for text in texts)
)
if result.should_block:
images: Final = inputs.get("images")
tools: Final = inputs.get("tools")
if not messages and not images and not tools:
verbose_proxy_logger.debug("Singulr: No messages, images, or tools to check after filtering")
return inputs
metadata: Final = self._build_metadata(request_data=request_data)
singulr_req_obj = SingulrGuardrailPayload(
correlation_id=request_data.get("litellm_call_id"),
model_name=inputs.get("model"),
guardrail_scope="request",
messages=messages,
images=images,
tools=tools,
metadata=metadata,
)
payload = singulr_req_obj.model_dump(mode="json")
guardrail_resp = await self._call_api(payload)
if guardrail_resp is None:
return inputs
if guardrail_resp.should_block:
raise GuardrailRaisedException(
guardrail_name=self.guardrail_name,
message=f"Blocked by Singulr: {result.blocking_due_to or 'unknown'}",
status_code=400,
message=f"Blocked by Singulr, Blocking due to {guardrail_resp.blocking_due_to or 'unknown'}",
blocked_content=True,
)
return inputs
@staticmethod
def _mcp_tool_name(request_data: Mapping[str, Any]) -> str | None:
return request_data.get("mcp_tool_name") or request_data.get("name")
@staticmethod
def _mcp_arguments(request_data: Mapping[str, Any]) -> object:
arguments: Final = request_data.get("mcp_arguments")
return arguments if arguments is not None else request_data.get("arguments")
@staticmethod
def _is_mcp_call(request_data: Mapping[str, Any], logging_obj: LiteLLMLoggingObj | None) -> bool:
call_type: Final = logging_obj.call_type if logging_obj is not None else request_data.get("call_type")
if call_type is not None:
return call_type == CallTypes.call_mcp_tool.value
model: Final = request_data.get("model")
return "mcp_tool_name" in request_data or (isinstance(model, str) and model.startswith(_MCP_MODEL_PREFIX))
async def _apply_guardrail_on_mcp_request(self, request_data: Mapping[str, Any]) -> None:
metadata: Final = self._build_metadata(request_data=request_data)
singulr_mcp_obj = SingulrMcpGuardrailPayload(
guardrail_scope="mcp_request",
tool_name=self._mcp_tool_name(request_data),
tool_arguments=self._mcp_arguments(request_data),
mcp_server_name=request_data.get("mcp_server_name"),
metadata=metadata,
)
payload = singulr_mcp_obj.model_dump(mode="json")
guardrail_resp = await self._call_api(payload)
if guardrail_resp is None:
return
if guardrail_resp.should_block:
raise GuardrailRaisedException(
guardrail_name=self.guardrail_name,
status_code=400,
message=f"Blocked by Singulr, Blocking due to {guardrail_resp.blocking_due_to or 'unknown'}",
blocked_content=True,
)
async def _apply_guardrail_on_mcp_response(
self, inputs: GenericGuardrailAPIInputs, texts: Sequence[str], request_data: Mapping[str, Any]
) -> GenericGuardrailAPIInputs:
if not texts:
return inputs
metadata: Final = self._build_metadata(request_data=request_data)
singulr_mcp_obj = SingulrMcpGuardrailPayload(
model_name=request_data.get("model"),
guardrail_scope="mcp_response",
tool_result=texts,
metadata=metadata,
)
payload = singulr_mcp_obj.model_dump(mode="json")
guardrail_resp = await self._call_api(payload)
if guardrail_resp is None:
return inputs
if guardrail_resp.should_block:
raise GuardrailRaisedException(
guardrail_name=self.guardrail_name,
status_code=400,
message=f"Blocked by Singulr, Blocking due to {guardrail_resp.blocking_due_to or 'unknown'}",
blocked_content=True,
)
return inputs
@staticmethod
def _build_tool_call(tool_call: Mapping[str, Any]) -> "ToolCall | None":
tool_call_id: Final = tool_call.get("id")
fun: Final = tool_call.get("function")
if not tool_call_id or not fun:
return None
func_name: Final = fun.get("name")
args: Final = fun.get("arguments")
if not func_name or args is None:
return None
call_type: Final = tool_call.get("type")
return ToolCall(
id=tool_call_id,
type=call_type if isinstance(call_type, str) and call_type else "function",
function=ToolCallFunction(
name=func_name,
arguments=args if isinstance(args, str) else json.dumps(args, default=str),
),
)
async def _apply_guardrail_on_response(
self, inputs: GenericGuardrailAPIInputs, texts: Sequence[str], request_data: Mapping[str, Any]
) -> GenericGuardrailAPIInputs:
combined_texts: Final = "\n".join(texts) if texts else None
tool_calls: Final = inputs.get("tool_calls", ())
tool_calls_res: Final = tuple(
tool_call_res
for tool_call_res in (self._build_tool_call(tool_call) for tool_call in tool_calls)
if tool_call_res is not None
)
assistant_message: Final = AssistantMessage(
role="assistant",
content=combined_texts,
tool_calls=tool_calls_res,
)
metadata: Final = self._build_metadata(request_data=request_data)
singulr_resp_obj = SingulrGuardrailPayload(
correlation_id=request_data.get("litellm_call_id"),
guardrail_scope="response",
model_name=request_data.get("model"),
messages=request_data.get("messages"),
images=inputs.get("images"),
response=assistant_message,
metadata=metadata,
)
payload = singulr_resp_obj.model_dump(mode="json")
guardrail_resp = await self._call_api(payload)
if guardrail_resp is None:
return inputs
if guardrail_resp.should_block:
raise GuardrailRaisedException(
guardrail_name=self.guardrail_name,
status_code=400,
message=f"Blocked by Singulr, Blocking due to {guardrail_resp.blocking_due_to or 'unknown'}",
blocked_content=True,
)
return inputs
@log_guardrail_information
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict, # mutable-ok: required by CustomGuardrail.apply_guardrail override signature
input_type: str,
logging_obj: "LiteLLMLoggingObj | None" = None,
) -> GenericGuardrailAPIInputs:
texts: Final = inputs.get("texts", ())
structured_messages: Final = inputs.get("structured_messages", ())
verbose_proxy_logger.debug(
"Singulr Guardrail: apply_guardrail called with input_type=%s, texts=%d, structured_messages=%d",
input_type,
len(texts),
len(structured_messages),
)
is_mcp_call: Final = self._is_mcp_call(request_data, logging_obj)
if input_type == "request":
if is_mcp_call:
await self._apply_guardrail_on_mcp_request(request_data=request_data)
return inputs
return await self._apply_guardrail_on_request(
inputs=inputs, texts=texts, structured_messages=structured_messages, request_data=request_data
)
elif input_type == "response":
if is_mcp_call:
return await self._apply_guardrail_on_mcp_response(
inputs=inputs, texts=texts, request_data=request_data
)
return await self._apply_guardrail_on_response(inputs=inputs, texts=texts, request_data=request_data)
return inputs

View file

@ -114,6 +114,12 @@ class DailySpendRecord(Protocol):
@property
def failed_requests(self) -> int: ...
@property
def total_response_time_ms(self) -> int: ...
@property
def timed_requests(self) -> int: ...
class _KeyMetadataDict(TypedDict, total=False):
key_alias: ReadOnly[str | None]
@ -162,6 +168,8 @@ class _GroupingSetsRow(SimpleNamespace):
api_requests: int | None
successful_requests: int | None
failed_requests: int | None
total_response_time_ms: int | None
timed_requests: int | None
class _EntityRollupRow(_GroupingSetsRow):
@ -217,6 +225,8 @@ def update_metrics(existing_metrics: SpendMetrics, record: DailySpendRecord) ->
existing_metrics.api_requests += record.api_requests or 0
existing_metrics.successful_requests += record.successful_requests or 0
existing_metrics.failed_requests += record.failed_requests or 0
existing_metrics.total_response_time_ms += record.total_response_time_ms or 0
existing_metrics.timed_requests += record.timed_requests or 0
return existing_metrics
@ -767,7 +777,9 @@ def _build_aggregated_sql_query(
SUM(autorouter_savings_spend)::float AS autorouter_savings_spend,
SUM(api_requests)::bigint AS api_requests,
SUM(successful_requests)::bigint AS successful_requests,
SUM(failed_requests)::bigint AS failed_requests
SUM(failed_requests)::bigint AS failed_requests,
SUM(total_response_time_ms)::bigint AS total_response_time_ms,
SUM(timed_requests)::bigint AS timed_requests
FROM "{pg_table}"
WHERE {where_clause}
GROUP BY GROUPING SETS (
@ -846,7 +858,9 @@ def _build_entity_rollup_sql_query(
SUM(autorouter_savings_spend)::float AS autorouter_savings_spend,
SUM(api_requests)::bigint AS api_requests,
SUM(successful_requests)::bigint AS successful_requests,
SUM(failed_requests)::bigint AS failed_requests
SUM(failed_requests)::bigint AS failed_requests,
SUM(total_response_time_ms)::bigint AS total_response_time_ms,
SUM(timed_requests)::bigint AS timed_requests
FROM "{pg_table}"
WHERE {where_clause}
GROUP BY GROUPING SETS (
@ -985,6 +999,8 @@ def _record_to_spend_metrics(record: _GroupingSetsRow) -> SpendMetrics:
api_requests=record.api_requests or 0,
successful_requests=record.successful_requests or 0,
failed_requests=record.failed_requests or 0,
total_response_time_ms=record.total_response_time_ms or 0,
timed_requests=record.timed_requests or 0,
)
@ -1246,6 +1262,8 @@ async def get_daily_activity(
total_prompt_caching_savings_spend=metadata_metrics.prompt_caching_savings_spend,
total_gateway_injected_caching_savings_spend=metadata_metrics.gateway_injected_caching_savings_spend,
total_autorouter_savings_spend=metadata_metrics.autorouter_savings_spend,
total_response_time_ms=metadata_metrics.total_response_time_ms,
total_timed_requests=metadata_metrics.timed_requests,
page=page,
total_pages=-(-total_count // page_size), # Ceiling division
has_more=(page * page_size) < total_count,
@ -1423,6 +1441,8 @@ async def get_daily_activity_aggregated(
"totals"
].gateway_injected_caching_savings_spend,
total_autorouter_savings_spend=aggregated["totals"].autorouter_savings_spend,
total_response_time_ms=aggregated["totals"].total_response_time_ms,
total_timed_requests=aggregated["totals"].timed_requests,
page=1,
total_pages=1,
has_more=False,

View file

@ -170,6 +170,7 @@ if MCP_AVAILABLE:
from litellm.proxy._experimental.mcp_server.ui_session_utils import (
admitted_user_context,
build_effective_auth_contexts,
can_access_mcp_server,
is_ui_session_credential,
)
from litellm.proxy._types import (
@ -2483,10 +2484,11 @@ if MCP_AVAILABLE:
)
return server
allowed_server_ids: Final[set[str]] = set()
for auth_context in await build_effective_auth_contexts(user_api_key_dict):
allowed_server_ids.update(await global_mcp_server_manager.get_allowed_mcp_servers(auth_context))
if server is None or server.server_id not in allowed_server_ids:
if server is None or not await can_access_mcp_server(
user_api_key_dict,
server.server_id,
global_mcp_server_manager.get_allowed_mcp_servers,
):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={

View file

@ -801,6 +801,8 @@ model LiteLLM_DailyUserSpend {
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
total_response_time_ms BigInt @default(0)
timed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@ -837,6 +839,8 @@ model LiteLLM_DailyOrganizationSpend {
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
total_response_time_ms BigInt @default(0)
timed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@ -873,6 +877,8 @@ model LiteLLM_DailyEndUserSpend {
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
total_response_time_ms BigInt @default(0)
timed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@ -908,6 +914,8 @@ model LiteLLM_DailyAgentSpend {
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
total_response_time_ms BigInt @default(0)
timed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@ -943,6 +951,8 @@ model LiteLLM_DailyTeamSpend {
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
total_response_time_ms BigInt @default(0)
timed_requests BigInt @default(0)
ptu_flat_cost Float @default(0.0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@ -981,6 +991,8 @@ model LiteLLM_DailyTagSpend {
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
total_response_time_ms BigInt @default(0)
timed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt

View file

@ -171,15 +171,25 @@ def _cost_of_usage(
) -> float | None:
"""What ``usage`` costs on ``model``, or ``None`` when the model has no pricing."""
try:
prompt_cost, completion_cost = generic_cost_per_token(
model=model.model,
usage=usage,
custom_llm_provider=model.provider,
service_tier=basis.service_tier,
data_residency=basis.data_residency,
model_info=model_info,
vertex_location=basis.vertex_location,
)
if model.provider == "anthropic":
from litellm.llms.anthropic.cost_calculation import cost_per_token
prompt_cost, completion_cost = cost_per_token(
model=model.model,
usage=usage,
service_tier=basis.service_tier,
model_info=model_info,
)
else:
prompt_cost, completion_cost = generic_cost_per_token(
model=model.model,
usage=usage,
custom_llm_provider=model.provider,
service_tier=basis.service_tier,
data_residency=basis.data_residency,
model_info=model_info,
vertex_location=basis.vertex_location,
)
except Exception as e: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models; degrade to zero savings
verbose_proxy_logger.debug(
"savings: cannot price usage for provider=%s model=%s (%s)", model.provider, model.model, e
@ -198,11 +208,6 @@ def _cache_token_split(usage: Usage) -> tuple[int, int]:
return int(read), int(created)
_CACHE_SPLIT_FIELDS: Final = frozenset(
("cached_tokens", "cache_creation_tokens", "cache_write_tokens", "cache_creation_token_details", "text_tokens")
)
def _baseline_cache_rate_keys(baseline_info: ModelInfo | None) -> tuple[bool, bool]:
"""Whether the baseline model has a ``(cache read, cache write)`` rate of its own.
@ -274,19 +279,22 @@ def _baseline_usage(usage: Usage, conversation_continuing: bool, baseline_info:
(getattr(details, field, 0) or 0) for field in ("audio_tokens", "image_tokens", "video_tokens")
)
return Usage(
prompt_tokens=usage.prompt_tokens,
completion_tokens=usage.completion_tokens,
total_tokens=usage.total_tokens,
completion_tokens_details=usage.completion_tokens_details,
prompt_tokens_details=PromptTokensDetailsWrapper(
**details.model_dump(exclude=_CACHE_SPLIT_FIELDS),
cached_tokens=reads,
cache_creation_tokens=writes,
cache_write_tokens=writes,
cache_creation_token_details=details.cache_creation_token_details if writes else None,
# Whatever no longer sits in a cache bucket is plain input on the baseline.
text_tokens=max(usage.prompt_tokens - reads - writes - other_modalities, 0),
),
**{
**usage.model_dump(),
# Rebuild through Usage so private fallback counts agree with the public buckets.
"cache_read_input_tokens": reads,
"cache_creation_input_tokens": writes,
"prompt_tokens_details": PromptTokensDetailsWrapper(
**{
**details.model_dump(),
"cached_tokens": reads,
"cache_creation_tokens": writes,
"cache_write_tokens": writes,
"cache_creation_token_details": details.cache_creation_token_details if writes else None,
"text_tokens": max(usage.prompt_tokens - reads - writes - other_modalities, 0),
}
),
},
)

View file

@ -2669,34 +2669,15 @@ class ProxyLogging:
user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict(user_api_key_dict)
else:
user_api_key_auth_dict = user_api_key_dict
# Add task to list for parallel execution
if (
"apply_guardrail" in type(callback).__dict__
and not callback.use_native_lifecycle_hooks
and user_api_key_dict is not None
and not getattr(callback, "use_native_during_call_hook", False)
):
data["guardrail_to_apply"] = callback
guardrail_task = self._run_guardrail_with_metrics(
callback,
unified_guardrail.async_moderation_hook(
user_api_key_dict=user_api_key_dict,
data=data,
call_type=call_type,
),
"during_call",
guardrail_tasks.append(
self._run_during_call_guardrail(
callback=callback,
data=data,
user_api_key_dict=user_api_key_dict,
user_api_key_auth_dict=user_api_key_auth_dict,
call_type=call_type,
)
else:
guardrail_task = self._run_guardrail_with_metrics(
callback,
callback.async_moderation_hook(
data=data,
user_api_key_dict=user_api_key_auth_dict,
call_type=call_type,
),
"during_call",
)
guardrail_tasks.append(guardrail_task)
)
# Step 2: Run all guardrail tasks in parallel
if guardrail_tasks:
@ -2708,6 +2689,41 @@ class ProxyLogging:
return data
async def _run_during_call_guardrail(
self,
callback: CustomGuardrail,
data: dict[str, object], # mutable-ok: request payload dict, guardrail_to_apply is written in place
user_api_key_dict: UserAPIKeyAuth | None,
user_api_key_auth_dict: UserAPIKeyAuth | dict[str, object] | None,
call_type: CallTypesLiteral,
) -> None:
if (
"apply_guardrail" in type(callback).__dict__
and not callback.use_native_lifecycle_hooks
and user_api_key_dict is not None
and not callback.use_native_during_call_hook
):
data["guardrail_to_apply"] = callback
await self._run_guardrail_with_metrics(
callback,
unified_guardrail.async_moderation_hook(
user_api_key_dict=user_api_key_dict,
data=data,
call_type=call_type,
),
"during_call",
)
return
await self._run_guardrail_with_metrics(
callback,
callback.async_moderation_hook(
data=data,
user_api_key_dict=user_api_key_auth_dict,
call_type=call_type,
),
"during_call",
)
async def failed_tracking_alert(
self,
error_message: str,

View file

@ -32,6 +32,9 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
)
from litellm.litellm_core_utils.thread_pool_executor import executor
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig,
)
from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils
from litellm.types.integrations.custom_logger import converted_stream_requested
from litellm.types.llms.openai import (
@ -257,6 +260,7 @@ class BaseResponsesAPIStreamingIterator:
self._failure_handled = False # Track if failure handler has been called
self._yielded_first_chunk = False
self._generated_content = ""
self._generated_tool_arguments = ""
self._completed_response_cached = False
self._completed_response_logged = False
self._completed_response_cache_hit: bool | None = None
@ -352,6 +356,10 @@ class BaseResponsesAPIStreamingIterator:
_delta: Final = getattr(openai_responses_api_chunk, "delta", None)
if isinstance(_delta, str):
self._generated_content += _delta
elif _event_type in _TOOL_ARGUMENTS_DELTA_EVENTS:
_args_delta: Final = getattr(openai_responses_api_chunk, "delta", None)
if isinstance(_args_delta, str):
self._generated_tool_arguments += _args_delta
_stream_model_id: Final = _model_id_from_metadata(self.litellm_metadata)
if _event_type in (
ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
@ -419,14 +427,41 @@ class BaseResponsesAPIStreamingIterator:
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED,
):
self.completed_response = openai_responses_api_chunk
_stamp_responses_usage_cost(getattr(openai_responses_api_chunk, "response", None), self.logging_obj)
_response_obj: Final[object] = getattr(openai_responses_api_chunk, "response", None)
_estimate_wanted: Final[bool] = _chunk_type in (
openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
)
_billed_response: Final[ResponsesAPIResponse | None] = _billed_terminal_response(
_response_obj,
(
lambda: (
_estimate_usage_safely(
self.model or "",
self.request_data.get("input"),
self.request_data,
self._generated_content + self._generated_tool_arguments,
)
if _estimate_wanted
else None
)
),
)
_terminal_chunk: Final = (
openai_responses_api_chunk
if _billed_response is None or _billed_response is _response_obj
else openai_responses_api_chunk.model_copy(update={"response": _billed_response})
)
self.completed_response = _terminal_chunk
_stamp_responses_usage_cost(_billed_response, self.logging_obj)
if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED:
self._handle_logging_failed_response()
else:
self._handle_logging_completed_response()
return _terminal_chunk
return openai_responses_api_chunk
return None
@ -655,7 +690,9 @@ class BaseResponsesAPIStreamingIterator:
if cache is None:
return
cached_response: Final = response_obj.model_dump_json()
cached_response: Final = _dump_json_safely(response_obj)
if cached_response is None:
return
if is_async:
from litellm.caching.caching_handler import create_cache_write_task
@ -1301,6 +1338,31 @@ def _add_text_like_part_events(
)
def _billed_terminal_response(
response_obj: object, estimate: Callable[[], ResponseAPIUsage | None] | None
) -> ResponsesAPIResponse | None:
if isinstance(response_obj, ResponsesAPIResponse):
return (
response_obj
if response_obj.usage is not None or estimate is None
else response_obj.model_copy(update={"usage": estimate()})
)
if not isinstance(response_obj, dict):
return None
usage: Final[object] = response_obj.get("usage") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # a model_constructed terminal event leaves response as an untyped dict
return ResponsesAPIResponse.model_construct(
**{**response_obj, "usage": usage if usage is not None or estimate is None else estimate()} # pyright: ignore[reportUnknownArgumentType, reportArgumentType] # same untyped dict spread
)
def _dump_json_safely(response: BaseModel) -> str | None:
try:
return response.model_dump_json()
except Exception as exc:
verbose_logger.debug("could not serialize completed response for cache: %s", exc)
return None
def _logging_copy(event: object) -> object:
"""Hand logging callbacks a copy, so their usage rewrite (Responses shape to chat shape) never
reaches the event the caller is iterating. The round trip through ``model_dump`` sidesteps the
@ -1332,6 +1394,56 @@ def _usage_as_model(usage: object) -> ResponseAPIUsage | None:
return None
_TOOL_ARGUMENTS_DELTA_EVENTS: Final = frozenset(
{
ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA,
ResponsesAPIStreamEvents.CUSTOM_TOOL_CALL_INPUT_DELTA,
ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA,
}
)
def _estimate_usage_from_text(
model: str,
request_input: object,
responses_api_request: Mapping[str, object],
generated_text: str,
) -> ResponseAPIUsage:
messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped
input=request_input, # pyright: ignore[reportArgumentType] # the raw Responses API input is a str or ResponseInputParam list, matching the helper's declared union
responses_api_request=dict(responses_api_request),
)
input_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped
model=model, messages=messages
)
output_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped
model=model, text=generated_text, count_response_tokens=True
)
return ResponseAPIUsage(
input_tokens=input_tokens,
output_tokens=output_tokens,
total_tokens=input_tokens + output_tokens,
)
def _estimate_usage_safely(
model: str,
request_input: object,
responses_api_request: Mapping[str, object],
generated_text: str,
) -> ResponseAPIUsage | None:
try:
return _estimate_usage_from_text(
model=model,
request_input=request_input,
responses_api_request=responses_api_request,
generated_text=generated_text,
)
except Exception as e:
verbose_logger.debug("Could not estimate usage from stream text, billing $0: %s", e)
return None
def _stamp_responses_usage_cost(
response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None
) -> None:

View file

@ -425,12 +425,34 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream])
_NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType({})
_SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object])
_SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str])
def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]:
return tuple(item for item in value if isinstance(item, str)) if isinstance(value, tuple) else ()
def _silent_experiment_targets(silent_model: object) -> tuple[str, ...]:
if silent_model is None:
return ()
try:
targets: Final = _SILENT_MODEL_ADAPTER.validate_python(silent_model)
except ValidationError:
verbose_router_logger.warning(
"silent_model must be a model name or a list of model names, got %r; skipping shadow traffic",
silent_model,
)
return ()
return (targets,) if isinstance(targets, str) else tuple(targets)
def _silent_experiment_kwargs_snapshot(kwargs: Mapping[str, object]) -> Mapping[str, object]:
metadata: Final = kwargs.get("metadata")
if not isinstance(metadata, Mapping):
return MappingProxyType({**kwargs})
return MappingProxyType({**kwargs, "metadata": dict(metadata)})
def _with_router_resolved_session_model(session: object, model_name: str) -> Mapping[str, Mapping[str, object]]:
"""
Realtime client-secret requests carry the model inside ``session`` as well, and the caller's copy of it still
@ -2455,18 +2477,17 @@ class Router:
)
silent_model: Final = litellm_params.pop("silent_model", None)
if silent_model is not None:
for silent_target in _silent_experiment_targets(silent_model):
# Mirroring traffic to a secondary model
# Use threading.Thread (not ThreadPoolExecutor) - executor.submit()
# requires pickling args, which fails when kwargs contain unpicklable
# objects (e.g. _thread.RLock from OTEL spans, loggers) in deployment.
thread: Final = threading.Thread(
threading.Thread(
target=self._silent_experiment_completion,
args=(silent_model, messages),
kwargs=kwargs,
args=(silent_target, messages),
kwargs=_silent_experiment_kwargs_snapshot(kwargs),
daemon=True,
)
thread.start()
).start()
kwargs.setdefault("messages", messages)
self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
@ -2567,9 +2588,6 @@ class Router:
silent_kwargs["metadata"]["is_silent_experiment"] = True
# Force stream=False so the response is fully consumed and callbacks fire
silent_kwargs["stream"] = False
# Pop logging objects and call IDs to ensure a fresh logging context
# This prevents collisions in the Proxy's database (spend_logs)
silent_kwargs.pop("litellm_call_id", None)
@ -2579,6 +2597,23 @@ class Router:
return silent_kwargs
async def _run_silent_experiment(
self, silent_model: str, messages: Sequence[Mapping[str, str]], silent_kwargs: Mapping[str, object]
) -> None:
remaining_kwargs: Final = MappingProxyType(
{key: value for key, value in silent_kwargs.items() if key != "stream"}
)
response: Final = await self.acompletion(
model=silent_model,
messages=cast(list[AllMessageValues], messages),
stream=bool(silent_kwargs.get("stream", False)),
**remaining_kwargs,
)
if not isinstance(response, CustomStreamWrapper):
return
async for _ in response:
pass
def _silent_experiment_completion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs):
"""
Run a silent experiment in the background (thread).
@ -2604,11 +2639,7 @@ class Router:
try:
async def _run_silent_completion():
await self.acompletion(
model=silent_model,
messages=cast(list[AllMessageValues], messages),
**silent_kwargs,
)
await self._run_silent_experiment(silent_model, messages, silent_kwargs)
# Drain any fire-and-forget tasks (e.g. alerting hooks)
# scheduled via asyncio.create_task during acompletion.
pending: Final = asyncio.all_tasks()
@ -3500,11 +3531,7 @@ class Router:
silent_kwargs["metadata"]["model_group"] = silent_model
# Trigger the silent request
await self.acompletion(
model=silent_model,
messages=cast(list[AllMessageValues], messages),
**silent_kwargs,
)
await self._run_silent_experiment(silent_model, messages, silent_kwargs)
except Exception as e:
verbose_router_logger.error("Silent experiment failed for model %s: %s", silent_model, e)
@ -3563,14 +3590,14 @@ class Router:
)
silent_model: Final = litellm_params.pop("silent_model", None)
if silent_model is not None:
for silent_target in _silent_experiment_targets(silent_model):
# Mirroring traffic to a secondary model
# This is a silent experiment, so we don't want to block the primary request
asyncio.create_task(
self._silent_experiment_acompletion(
silent_model=silent_model,
silent_model=silent_target,
messages=messages, # Use messages instead of *args
**kwargs,
**_silent_experiment_kwargs_snapshot(kwargs),
)
)

View file

@ -586,7 +586,7 @@ class MessageBlockDelta(TypedDict):
type: Literal["message_delta"]
delta: MessageDelta
usage: UsageDelta
usage: NotRequired[ReadOnly[UsageDelta]]
context_management: NotRequired[ContextManagementResponse]

View file

@ -1,24 +1,53 @@
from typing import Any
from collections.abc import Mapping, Sequence
from typing import Literal
from pydantic import BaseModel, Field
from .base import GuardrailConfigModel
class SingulrGuardrailRequest(BaseModel):
model: str | None = None
messages: list[dict[str, Any]] | None = None
tools: list[dict[str, Any]] | None = None
model_response: dict[str, Any] | None = None
litellm_metadata: dict[str, Any] | None = None
class ContentBlock(BaseModel):
type: str | None = None
text: str | None = None
class ToolCallFunction(BaseModel):
name: str
arguments: str
class ToolCall(BaseModel):
id: str
type: str = "function"
function: ToolCallFunction
class AssistantMessage(BaseModel):
role: Literal["assistant"] = "assistant"
content: str | Sequence[ContentBlock] | None = None
tool_calls: Sequence[ToolCall] | None = None
class SingulrGuardrailPayload(BaseModel):
litellm_call_id: str | None = None
request_data: SingulrGuardrailRequest | None = None
input_type: str
is_playground_request: bool | None = None
playground_text: str | None = None
correlation_id: str | None = None
model_name: str | None = None
model_provider_name: str | None = None
guardrail_scope: str | None = None
messages: Sequence[Mapping[str, object]] | None = None
images: Sequence[str] | None = None
tools: Sequence[Mapping[str, object]] | None = None
response: AssistantMessage | None = None
metadata: Mapping[str, str] | None = None
class SingulrMcpGuardrailPayload(BaseModel):
model_name: str | None = None
guardrail_scope: str | None = None
tool_name: str | None = None
tool_arguments: object = None
mcp_server_name: str | None = None
tool_result: Sequence[str] | None = None
metadata: Mapping[str, str] | None = None
class SingulrGuardrailResponse(BaseModel):

View file

@ -32,6 +32,8 @@ class SpendMetrics(BaseModel):
successful_requests: int = Field(default=0)
failed_requests: int = Field(default=0)
api_requests: int = Field(default=0)
total_response_time_ms: int = Field(default=0)
timed_requests: int = Field(default=0)
class MetricBase(BaseModel):
@ -93,6 +95,8 @@ class DailySpendMetadata(BaseModel):
total_prompt_caching_savings_spend: float = Field(default=0.0)
total_gateway_injected_caching_savings_spend: float = Field(default=0.0)
total_autorouter_savings_spend: float = Field(default=0.0)
total_response_time_ms: int = Field(default=0)
total_timed_requests: int = Field(default=0)
page: int = Field(default=1)
total_pages: int = Field(default=1)
has_more: bool = Field(default=False)
@ -125,6 +129,8 @@ class LiteLLM_DailyUserSpend(BaseModel):
api_requests: int = 0
successful_requests: int = 0
failed_requests: int = 0
total_response_time_ms: int = 0
timed_requests: int = 0
class GroupedData(TypedDict):

View file

@ -2957,6 +2957,7 @@ InternalCallOrigin = Literal[
"autorouter_classifier",
"shadow_eval_router",
"shadow_eval_judge",
"llm_as_a_judge_guardrail",
"background_response_cost_poll",
]
"""Which internal litellm feature originated a billed sub-call, so a spend log row
@ -2965,6 +2966,7 @@ records that it is not traffic the caller sent."""
AUTOROUTER_CLASSIFIER_CALL_ORIGIN: Final[InternalCallOrigin] = "autorouter_classifier"
SHADOW_EVAL_ROUTER_CALL_ORIGIN: Final[InternalCallOrigin] = "shadow_eval_router"
SHADOW_EVAL_JUDGE_CALL_ORIGIN: Final[InternalCallOrigin] = "shadow_eval_judge"
LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN: Final[InternalCallOrigin] = "llm_as_a_judge_guardrail"
BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN: Final[InternalCallOrigin] = "background_response_cost_poll"

View file

@ -801,6 +801,8 @@ model LiteLLM_DailyUserSpend {
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
total_response_time_ms BigInt @default(0)
timed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@ -837,6 +839,8 @@ model LiteLLM_DailyOrganizationSpend {
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
total_response_time_ms BigInt @default(0)
timed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@ -873,6 +877,8 @@ model LiteLLM_DailyEndUserSpend {
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
total_response_time_ms BigInt @default(0)
timed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@ -908,6 +914,8 @@ model LiteLLM_DailyAgentSpend {
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
total_response_time_ms BigInt @default(0)
timed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@ -943,6 +951,8 @@ model LiteLLM_DailyTeamSpend {
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
total_response_time_ms BigInt @default(0)
timed_requests BigInt @default(0)
ptu_flat_cost Float @default(0.0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@ -981,6 +991,8 @@ model LiteLLM_DailyTagSpend {
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
total_response_time_ms BigInt @default(0)
timed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt

View file

@ -0,0 +1,472 @@
from __future__ import annotations
import os
import shutil
import socket
import subprocess
import threading
import time
import uuid
from collections.abc import Generator
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from dataclasses import dataclass, replace
from http.client import HTTPConnection
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Final
from urllib.parse import urlsplit
import pytest
from e2e_http import NetworkError, PreparedForward, RawResponse, StreamChunk, StreamHead, forward, prepare_forward
from models import LiteLLMParamsBody
from provider_cache import CacheEdge, CacheHit, CaptureLease, exact_key, successful_response
from provider_cache_redis import PUBLISH, RedisCommands, RedisResponseStore, configured_cache, redis_store
from provider_cache_routing import LIVE_PROVIDER_REQUIRED, route_cache_model
from provider_edge import configured_cache_backend, start_provider_edge
from redis.exceptions import ConnectionError as RedisConnectionError
SECRET: Final = b"synthetic-cache-hmac-key-for-tests"
BODY: Final = b'{"model":"test","messages":[{"role":"user","content":"hello"}]}'
SUCCESS: Final = b'{"id":"provider-fixed-id","choices":[{"message":{"content":"hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}'
HEADERS: Final = {"content-type": "application/json", "authorization": "Bearer synthetic-account-one"}
class Provider(ThreadingHTTPServer):
hits: tuple[tuple[str, bytes], ...] = ()
response: bytes = SUCCESS
status: int = 200
delay: float = 0
stream: bool = False
truncated: bool = False
cookie: str = ""
class Handler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
def do_POST(self) -> None:
server: Final = self.server
assert isinstance(server, Provider)
body: Final = self.rfile.read(int(self.headers.get("content-length", "0")))
server.hits += ((self.path, body),)
time.sleep(server.delay)
self.send_response(server.status)
if server.stream:
self.send_header("content-type", "text/event-stream")
self.send_header("transfer-encoding", "chunked")
self.end_headers()
self.wfile.write(b"%x\r\n%s\r\n" % (len(server.response), server.response))
if server.truncated:
self.close_connection = True
return
self.wfile.write(b"0\r\n\r\n")
return
self.send_header("content-type", "application/json")
self.send_header("content-length", str(len(server.response)))
if server.cookie:
self.send_header("set-cookie", server.cookie)
self.end_headers()
self.wfile.write(server.response)
def log_message(self, format: str, *args: object) -> None:
pass
@pytest.fixture
def provider() -> Generator[Provider, None, None]:
server: Final = Provider(("127.0.0.1", 0), Handler)
thread: Final = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
try:
yield server
finally:
server.shutdown()
server.server_close()
thread.join(timeout=5)
@pytest.fixture(scope="module")
def redis_url(tmp_path_factory: pytest.TempPathFactory) -> Generator[str, None, None]:
configured: Final = os.environ.get("E2E_CACHE_TEST_REDIS_URL")
if configured:
yield configured
return
binary: Final = shutil.which("redis-server")
assert binary is not None, "Set E2E_CACHE_TEST_REDIS_URL or install Redis for cache integration checks"
root: Final = tmp_path_factory.mktemp("provider-cache-redis")
with socket.socket() as probe:
probe.bind(("127.0.0.1", 0))
port: Final = probe.getsockname()[1]
with (root / "redis.log").open("wb") as log:
process: Final = subprocess.Popen(
[binary, "--bind", "127.0.0.1", "--port", str(port), "--save", "", "--appendonly", "no", "--dir", str(root)],
stdout=log, stderr=subprocess.STDOUT,
)
try:
deadline: Final = time.monotonic() + 5
while True:
try:
with socket.create_connection(("127.0.0.1", port), timeout=0.1):
break
except OSError:
assert process.poll() is None and time.monotonic() < deadline
time.sleep(0.02)
yield f"redis://127.0.0.1:{port}/0"
finally:
process.terminate()
process.wait(timeout=5)
@pytest.fixture
def store(redis_url: str) -> RedisResponseStore:
return redis_store(redis_url, "test-" + uuid.uuid4().hex)
@contextmanager
def edge(cache: CacheEdge, provider: Provider) -> Generator[str, None, None]:
upstream: Final = f"http://127.0.0.1:{provider.server_port}"
running: Final = start_provider_edge(cache, mounts={"openai": upstream})
try:
yield running.edge.api_base("openai") + "/v1/chat/completions"
finally:
running.shutdown()
def call(url: str, body: bytes = BODY, headers: dict[str, str] = HEADERS) -> RawResponse:
result: Final = forward("POST", url, headers=headers, body=body, timeout=5)
assert isinstance(result, RawResponse), result
return result
def test_success_is_reusable_across_fresh_edges(store: RedisResponseStore, provider: Provider) -> None:
with edge(CacheEdge(store, SECRET), provider) as url:
assert call(url).body == SUCCESS
assert call(url).body == SUCCESS
with edge(CacheEdge(store, SECRET), provider) as other:
assert call(other).body == SUCCESS
assert len(provider.hits) == 1
@pytest.mark.parametrize("body", [BODY + b" ", BODY.replace(b"hello", b"Hello"), BODY.replace(b"test", b"test2")])
def test_any_body_change_calls_live(store: RedisResponseStore, provider: Provider, body: bytes) -> None:
with edge(CacheEdge(store, SECRET), provider) as url:
call(url)
call(url, body)
call(url, body)
assert len(provider.hits) == 2
@pytest.mark.parametrize("name,value", [("authorization", "Bearer another-account"), ("x-request-id", "one"), ("anthropic-version", "new")])
def test_changed_header_cannot_reuse(store: RedisResponseStore, provider: Provider, name: str, value: str) -> None:
with edge(CacheEdge(store, SECRET), provider) as url:
call(url)
call(url, headers=HEADERS | {name: value})
call(url + "?x=1")
assert len(provider.hits) == 3
@pytest.mark.parametrize("status,response", [(429, b'{"error":"rate limited"}'), (500, b'failed'), (200, b'{"error":"bad"}'), (200, b'not json')])
def test_failed_provider_responses_never_enter_cache(store: RedisResponseStore, provider: Provider, status: int, response: bytes) -> None:
provider.status = status
provider.response = response
with edge(CacheEdge(store, SECRET), provider) as url:
assert call(url).status_code == status
assert call(url).body == response
assert len(provider.hits) == 2
def test_cookie_setting_success_is_reused_without_the_cookie(store: RedisResponseStore, provider: Provider) -> None:
provider.cookie = "__cf_bm=synthetic-bot-management; Path=/; HttpOnly; Secure"
with edge(CacheEdge(store, SECRET), provider) as url:
replies: Final = tuple(call(url) for _ in range(2))
assert len(provider.hits) == 1
assert all(reply.body == SUCCESS and "set-cookie" not in reply.headers for reply in replies)
def test_expiry_does_not_slide(store: RedisResponseStore, provider: Provider) -> None:
short: Final = replace(store, lifetime_ms=250)
with edge(CacheEdge(short, SECRET), provider) as url:
call(url)
call(url)
time.sleep(0.3)
call(url)
call(url)
assert len(provider.hits) == 2
def test_concurrent_requests_publish_atomically(store: RedisResponseStore, provider: Provider) -> None:
provider.delay = 0.15
with edge(CacheEdge(store, SECRET), provider) as url:
with ThreadPoolExecutor(max_workers=5) as executor:
replies: Final = tuple(executor.map(lambda _: call(url).body, range(5)))
assert replies == (SUCCESS,) * 5
assert len(provider.hits) == 1
@pytest.mark.parametrize("age_past_expiry_ms", [0, 1])
def test_expired_response_is_rejected_without_physical_eviction(
store: RedisResponseStore, age_past_expiry_ms: int,
) -> None:
response_key: Final = store.keys("expired")[0]
retained: Final = store.client.eval(
"""
local clock = redis.call('TIME')
local expires = clock[1] * 1000 + math.floor(clock[2] / 1000) - tonumber(ARGV[1])
redis.call('HSET', KEYS[1], 'captured', expires - 86400000, 'expires', expires, 'payload', 'old-response')
return redis.call('PTTL', KEYS[1])
""",
1, response_key, age_past_expiry_ms,
)
assert retained == -1
replacement: Final = store.lookup("expired")
assert isinstance(replacement, CaptureLease)
assert replacement.expires_at_ms - replacement.captured_at_ms == 86_400_000
assert store.publish("expired", replacement, b"fresh-response")
hit: Final = store.lookup("expired")
assert isinstance(hit, CacheHit) and hit.payload == b"fresh-response"
@pytest.mark.parametrize("truncated", [False, True])
def test_stream_completion_controls_publication(store: RedisResponseStore, provider: Provider, truncated: bool) -> None:
provider.stream = True
provider.truncated = truncated
provider.response = b'data: {"choices":[{"index":0,"delta":{"content":"hello"},"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n'
with edge(CacheEdge(store, SECRET), provider) as url:
for _ in range(2):
result: Final = forward("POST", url, headers=HEADERS, body=BODY, timeout=5)
if truncated:
assert isinstance(result, NetworkError)
else:
assert isinstance(result, RawResponse) and result.body == provider.response
assert len(provider.hits) == (2 if truncated else 1)
def test_store_outage_preserves_provider_success(provider: Provider) -> None:
with socket.socket() as probe:
probe.bind(("127.0.0.1", 0))
port: Final = probe.getsockname()[1]
unavailable: Final = redis_store(f"redis://127.0.0.1:{port}/0", "unavailable")
with edge(CacheEdge(unavailable, SECRET), provider) as url:
assert call(url).body == SUCCESS
assert call(url).body == SUCCESS
assert len(provider.hits) == 2
def test_old_lease_cannot_overwrite_new_owner(store: RedisResponseStore) -> None:
short: Final = replace(store, lease_ms=50)
old: Final = short.lookup("key")
assert isinstance(old, CaptureLease)
time.sleep(0.08)
current: Final = short.lookup("key")
assert isinstance(current, CaptureLease)
assert not short.publish("key", old, b"old")
assert short.publish("key", current, b"new")
hit: Final = short.lookup("key")
assert isinstance(hit, CacheHit) and hit.payload == b"new"
def test_identity_preserves_values_and_never_contains_credentials() -> None:
variants: Final = (b'{}', b'{"a":null}', b'{"a":false}', b'{"a":0}', b'{"a":0.0}', b'{"a":"0"}', b' { }', None, b'')
keys: Final = tuple(exact_key(SECRET, "POST", "https://example.invalid/v1/chat/completions", HEADERS, body) for body in variants)
assert len(set(keys)) == len(variants)
assert all(len(key) == 64 and "synthetic-account" not in key for key in keys)
@pytest.mark.parametrize("payload", [b"corrupt response", '{"response":"{}","signature":"é"}'.encode()])
def test_corrupt_entry_is_replaced_by_same_successful_request(store: RedisResponseStore, provider: Provider, payload: bytes) -> None:
upstream: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions"
prepared: Final = prepare_forward("POST", upstream, HEADERS, BODY)
assert isinstance(prepared, PreparedForward)
key: Final = exact_key(SECRET, "POST", upstream, prepared.headers, BODY)
lease: Final = store.lookup(key)
assert isinstance(lease, CaptureLease)
assert store.publish(key, lease, payload)
cache: Final = CacheEdge(store, SECRET)
for _ in range(2):
head = cache.forward("POST", upstream, HEADERS, BODY, 5)
assert isinstance(head, StreamHead)
assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS
assert len(provider.hits) == 1
assert dict(cache.counters.counts) == {
"corrupt": 1, "misses": 1, "upstream_attempts": 1, "writes": 1, "hits": 1,
}
@pytest.mark.parametrize("payload", [
b'data: {}\n\ndata: [DONE]\n\n',
b'data: {"choices":[{"index":0,"delta":{}}]}\n\ndata: [DONE]\n\n',
b'data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\ndata: [DONE]',
b'data: {"error":{"message":"failed"}}\n\ndata: [DONE]\n\n',
])
def test_malformed_success_stream_is_never_cached(store: RedisResponseStore, provider: Provider, payload: bytes) -> None:
provider.stream = True
provider.response = payload
with edge(CacheEdge(store, SECRET), provider) as url:
assert call(url).body == payload
assert call(url).body == payload
assert len(provider.hits) == 2
def test_anthropic_stream_requires_start_finish_and_stop() -> None:
start: Final = b'data: {"type":"message_start","message":{}}\n\n'
finish: Final = b'data: {"type":"message_delta","delta":{"stop_reason":"end_turn"}}\n\n'
stop: Final = b'data: {"type":"message_stop"}\n\n'
url: Final = "https://example.invalid/v1/messages"
headers: Final = {"content-type": "text/event-stream"}
assert successful_response(url, 200, headers, start + finish + stop)
assert not successful_response(url, 200, headers, start + stop)
assert not successful_response(url, 200, headers, finish + stop)
assert not successful_response(url, 200, headers, start + finish)
@pytest.mark.parametrize("provider,suffix", [("openai", "/v1"), ("anthropic", "")])
def test_normal_registration_routes_supported_providers(provider: str, suffix: str) -> None:
params: Final = LiteLLMParamsBody(model=f"{provider}/test", api_key="os.environ/SYNTHETIC_KEY", timeout=12)
routed: Final = route_cache_model(params, lambda mount: f"http://edge.invalid/{mount}", enabled=True)
assert routed.api_base == f"http://edge.invalid/{provider}{suffix}"
assert routed.model_dump(exclude={"api_base"}) == params.model_dump(exclude={"api_base"})
assert params.api_base is None
@pytest.mark.parametrize("params", [
LiteLLMParamsBody(model="bedrock/test"),
LiteLLMParamsBody(model="azure/test"),
LiteLLMParamsBody(model="openai/test", api_base="https://custom.invalid/v1"),
LiteLLMParamsBody(model="openai/test", api_base=""),
LiteLLMParamsBody(model="openai/test", litellm_credential_name="named-credential"),
LiteLLMParamsBody(model="openai/test", mock_response="synthetic"),
])
def test_registration_preserves_unsupported_or_explicit_routes(params: LiteLLMParamsBody) -> None:
def unexpected_edge(mount: str) -> str:
pytest.fail(f"should not start edge for {mount}")
assert route_cache_model(params, unexpected_edge, enabled=True) is params
def test_rollback_and_live_only_policy_keep_direct_provider_route() -> None:
params: Final = LiteLLMParamsBody(model="openai/test")
assert route_cache_model(params, lambda _: "http://edge.invalid", enabled=False) is params
assert route_cache_model(params, lambda _: "http://edge.invalid", enabled=True, mode="realtime") is params
token: Final = LIVE_PROVIDER_REQUIRED.set(True)
try:
assert route_cache_model(params, lambda _: "http://edge.invalid", enabled=True) is params
finally:
LIVE_PROVIDER_REQUIRED.reset(token)
assert route_cache_model(params, lambda _: "http://edge.invalid", enabled=True).api_base == "http://edge.invalid/v1"
@dataclass(frozen=True)
class PublishOutage:
client: RedisCommands
def eval(self, script: str, numkeys: int, *args: str | bytes | int) -> object:
if script == PUBLISH:
raise RedisConnectionError("synthetic publication outage")
return self.client.eval(script, numkeys, *args)
def test_write_outage_preserves_success_without_hidden_retry(store: RedisResponseStore, provider: Provider) -> None:
unavailable: Final = replace(store, client=PublishOutage(store.client))
cache: Final = CacheEdge(unavailable, SECRET)
with edge(cache, provider) as url:
assert call(url).body == SUCCESS
assert call(url).body == SUCCESS
assert len(provider.hits) == 2
assert dict(cache.counters.counts)["write_failures"] == 2
with edge(CacheEdge(store, SECRET), provider) as url:
assert call(url).body == SUCCESS
assert call(url).body == SUCCESS
assert len(provider.hits) == 3
def test_connection_failure_releases_capture_lease(store: RedisResponseStore) -> None:
with socket.socket() as unavailable:
unavailable.bind(("127.0.0.1", 0))
url: Final = f"http://127.0.0.1:{unavailable.getsockname()[1]}/v1/chat/completions"
cache: Final = CacheEdge(store, SECRET)
assert isinstance(cache.forward("POST", url, HEADERS, BODY, 0.2), NetworkError)
prepared: Final = prepare_forward("POST", url, HEADERS, BODY)
assert isinstance(prepared, PreparedForward)
key: Final = exact_key(SECRET, "POST", url, prepared.headers, BODY)
slot: Final = store.lookup(key)
assert isinstance(slot, CaptureLease)
assert store.release(key, slot)
assert dict(cache.counters.counts)["rejected"] == 1
def test_close_before_first_chunk_releases_lease(store: RedisResponseStore, provider: Provider) -> None:
url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions"
cache: Final = CacheEdge(store, SECRET)
head: Final = cache.forward("POST", url, HEADERS, BODY, 5)
assert isinstance(head, StreamHead)
head.steps.close()
prepared: Final = prepare_forward("POST", url, HEADERS, BODY)
assert isinstance(prepared, PreparedForward)
key: Final = exact_key(SECRET, "POST", url, prepared.headers, BODY)
slot: Final = store.lookup(key)
assert isinstance(slot, CaptureLease)
assert store.release(key, slot)
def test_effective_account_change_cannot_reuse_cache(
store: RedisResponseStore, provider: Provider, monkeypatch: pytest.MonkeyPatch, tmp_path,
) -> None:
url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions"
cache: Final = CacheEdge(store, SECRET)
for account in ("account-a", "account-b", "account-b"):
netrc = tmp_path / account
netrc.write_text(f"machine 127.0.0.1 login {account} password synthetic\n")
monkeypatch.setenv("NETRC", str(netrc))
head = cache.forward("POST", url, HEADERS, BODY, 5)
assert isinstance(head, StreamHead)
assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS
assert len(provider.hits) == 2
assert dict(cache.counters.counts)["hits"] == 1
def test_enabled_environment_reuses_store_across_fresh_backends(
redis_url: str, provider: Provider, monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("E2E_PROVIDER_CACHE", "1")
monkeypatch.setenv("E2E_PROVIDER_CACHE_REDIS_URL", redis_url)
monkeypatch.setenv("E2E_PROVIDER_CACHE_HMAC_KEY", SECRET.decode())
monkeypatch.setenv("E2E_PROVIDER_CACHE_NAMESPACE", "environment-" + uuid.uuid4().hex)
configured_cache.cache_clear()
try:
for _ in range(2):
backend = configured_cache_backend()
assert isinstance(backend, CacheEdge)
with edge(backend, provider) as url:
assert call(url).body == SUCCESS
configured_cache.cache_clear()
assert len(provider.hits) == 1
monkeypatch.setenv("E2E_PROVIDER_CACHE", "0")
assert configured_cache_backend() is None
finally:
configured_cache.cache_clear()
@pytest.mark.parametrize("known_mount", (True, False))
def test_duplicate_headers_bypass_cache_and_count_live_calls(
store: RedisResponseStore, provider: Provider, known_mount: bool,
) -> None:
cache: Final = CacheEdge(store, SECRET)
with edge(cache, provider) as url:
parsed: Final = urlsplit(url)
for _ in range(2):
connection = HTTPConnection(str(parsed.hostname), parsed.port, timeout=5)
try:
connection.putrequest("POST", parsed.path if known_mount else "/unknown/v1/chat/completions")
connection.putheader("content-length", str(len(BODY)))
connection.putheader("content-type", "application/json")
connection.putheader("x-duplicate", "first")
connection.putheader("x-duplicate", "second")
connection.endheaders(BODY)
response = connection.getresponse()
assert response.status == (200 if known_mount else 404)
payload = response.read()
assert payload == SUCCESS if known_mount else b"unknown provider mount" in payload
finally:
connection.close()
assert len(provider.hits) == (2 if known_mount else 0)
assert dict(cache.counters.counts)["duplicate_header_bypass"] == 2
assert dict(cache.counters.counts).get("upstream_attempts", 0) == (2 if known_mount else 0)

View file

@ -0,0 +1,33 @@
# Shared provider-response cache
`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live
The edge caches complete successful POST responses for `/v1/chat/completions` and `/v1/messages`, including streams. Unsupported endpoints pass through. It matches the method, original URL, effective outbound headers (including authentication and HTTP-library defaults), body presence and exact body bytes using a full keyed digest. It sends the same prepared request used for matching. No prompts, random markers, JSON values or credentials are normalized away. Provider `Set-Cookie` headers are dropped before validation and never recorded: the edge already withholds them from the proxy, and OpenAI responses always carry Cloudflare bot-management cookies
An eligible miss calls the provider. A complete successful response is stored immediately even if a later test assertion fails. Provider errors, malformed responses, truncated streams and cancelled captures are not stored. Cache reads, writes and lease failures fall through to normal provider behavior; they introduce no provider retry. An already-started response cannot be restarted after a delivery failure
Recordings are shared across workers and builds through dedicated Redis, separate from the candidate's own cache. They expire 86,400 seconds after capture starts, based on Redis time. Reads never extend expiry. There is no scheduled recapture: the next miss calls the provider again. Bounded coordination reduces duplicate concurrent calls, but slow or failed captures may lead to extra live calls after the wait expires
## Configuration
The trusted runner receives:
- `E2E_PROVIDER_CACHE`: `1` to enable, `0` to use the normal live path
- `E2E_PROVIDER_CACHE_REDIS_URL`: authenticated dedicated Redis URL
- `E2E_PROVIDER_CACHE_HMAC_KEY`: dedicated secret containing at least 32 bytes
- `E2E_PROVIDER_CACHE_NAMESPACE`: shared environment namespace, independent of build and candidate revision
- `E2E_PROVIDER_CACHE_METRICS_DIR`: optional per-process counter artifact directory
Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits
Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. Ordinary assertion failures still fail E2E. The shared cache does not modify provider response IDs or make the proxy aware of replay
## Recorded response semantics
Replay preserves the original response ID, usage and end-to-end headers. The proxy can therefore deduplicate repeated provider IDs when storing spend-log rows, just as it does when a live upstream returns the same ID twice. One spend-log row per invocation is not guaranteed for identical recorded responses. Existing spend reconciliation requests use distinct prompt markers and retain their distinct-ID and row-count assertions; accounting tests are not automatically excluded from caching
Provider remaining-quota headers describe the captured response. Metrics derived from them are historical on a cache hit, not a measurement of current provider capacity. Gateway-generated API-key quota headers are a separate contract. A test of fresh provider quota or timing must use the live-provider policy; replay can still exercise how the proxy processes the recorded headers
## Qualification
`tests/code_coverage_tests/test_provider_cache.py` exercises local HTTP providers and disposable real Redis. CI runs these checks with the existing provider-edge and replay harness tests. These component checks do not establish Buildkite deployment, full-suite cross-build reuse or a genuine 24-hour expiry observation; those require separate runtime evidence

View file

@ -22,7 +22,6 @@ from typing import Final
import pytest
import requests
from e2e_config import (
CONTROL_PLANE_BASE_URL,
FIXTURE_DIR,
@ -41,6 +40,7 @@ from idp import Identity, Keycloak, keycloak_from_env
from junit_properties import attach_result_properties
from lifecycle import ProxyClientProvider, ResourceManager
from models import TeamNewBody, UserNewBody, UserNewResponse
from provider_cache_routing import LIVE_PROVIDER_REQUIRED
from provider_edge import replay_leftover_error
from proxy_client import ProxyClient, build_proxy_client
@ -85,6 +85,7 @@ def jwt_identity(idp: Keycloak, resources: ResourceManager, proxy: ProxyClient)
def pytest_configure(config: pytest.Config) -> None:
config.addinivalue_line("markers", "provider_live: requires actual provider timing, limits or state; bypass shared cache")
config.addinivalue_line(
"markers",
"e2e: live test that requires a running proxy and real provider keys",
@ -192,11 +193,13 @@ def _proxy_fail_reason() -> str | None:
return None
@pytest.hookimpl(tryfirst=True)
def pytest_runtest_setup(item: pytest.Item) -> None:
"""Hard-fail `e2e`-marked tests unless a proxy answers its liveness probe.
Unmarked tests (unit coverage of the harness) don't touch the proxy, so they
run even when none is up. Never skip for a missing proxy. Replay mode needs
the proxy too: only provider-bound traffic replays from the bundle."""
LIVE_PROVIDER_REQUIRED.set(item.get_closest_marker("provider_live") is not None)
if item.get_closest_marker("e2e") is None:
return
reason = _proxy_fail_reason()
@ -235,6 +238,7 @@ def pytest_runtest_teardown(item: pytest.Item) -> Generator[None, None, None]:
yield so fixture finalizers replay their recorded calls first. Failed tests
are left alone - their own failure already explains any unconsumed tail."""
result = yield
LIVE_PROVIDER_REQUIRED.set(False)
if not item.stash.get(_CALL_PASSED, False):
return result
reason = replay_leftover_error(

View file

@ -853,6 +853,7 @@ def _stream_steps(resp: requests.Response) -> Generator[StreamStep, None, None]:
the chunks already delivered are exactly what makes a mid-stream failure
different from a request that never streamed at all."""
try:
yield StreamChunk(b"")
for piece in cast("Iterator[bytes]", resp.iter_content(chunk_size=None)):
if piece:
yield StreamChunk(data=piece)
@ -862,6 +863,44 @@ def _stream_steps(resp: requests.Response) -> Generator[StreamStep, None, None]:
resp.close()
def primed_steps(steps: Generator[StreamStep, None, None]) -> Generator[StreamStep, None, None]:
first: Final = next(steps)
assert isinstance(first, StreamChunk) and first.data == b""
return steps
@dataclass(frozen=True, slots=True, repr=False)
class PreparedForward:
request: requests.PreparedRequest
url: str
headers: dict[str, str]
def prepare_forward(
method: str, url: str, headers: dict[str, str], body: bytes | None,
) -> PreparedForward | NetworkError:
try:
with requests.Session() as session:
request: Final = session.prepare_request(requests.Request(method, url, headers=headers, data=body))
except requests.RequestException as exc:
return NetworkError(message=str(exc))
assert request.url is not None
return PreparedForward(request, request.url, dict(request.headers))
def forward_prepared_stream(prepared: PreparedForward, timeout: float) -> StreamHead | NetworkError:
try:
with requests.Session() as session:
settings: Final = session.merge_environment_settings(prepared.url, {}, True, None, None)
resp: Final = session.send(prepared.request, timeout=timeout, allow_redirects=False, **settings)
except requests.RequestException as exc:
return NetworkError(message=str(exc))
return StreamHead(
resp.status_code, {name.lower(): value for name, value in resp.headers.items()},
primed_steps(_stream_steps(resp)),
)
def open_stream(url: URL, *, headers: BaseModel, json: BaseModel, timeout: float = 60.0) -> StreamHead | NetworkError:
"""POST a streaming request and return the moment its response head arrives,
leaving the body unread behind ``StreamHead.steps``. For a test that must keep
@ -907,5 +946,5 @@ def forward_stream(
return StreamHead(
status_code=resp.status_code,
headers={name.lower(): value for name, value in resp.headers.items()},
steps=_stream_steps(resp),
steps=primed_steps(_stream_steps(resp)),
)

View file

@ -44,7 +44,7 @@ from models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody, Usage
from passthrough_client import PassthroughClient
import os
pytestmark = pytest.mark.e2e
pytestmark = [pytest.mark.e2e, pytest.mark.provider_live]
BEDROCK_MODEL = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0"
VERTEX_MODEL = "vertex_ai/gemini-2.5-flash"

View file

@ -24,10 +24,10 @@ from models import (
AnthropicAssistantTurn,
AnthropicContentBlock,
AnthropicCustomTool,
AnthropicMessagesBody,
AnthropicToolChoice,
AnthropicToolResultBlock,
AnthropicToolResultTurn,
AnthropicMessagesBody,
ChatMessage,
JsonSchemaProperty,
LiteLLMParamsBody,
@ -165,6 +165,7 @@ class TestAnthropicMessages:
)
@pytest.mark.covers("llm.messages.anthropic.basic.stream.works")
@pytest.mark.provider_live
def test_messages_streams_completion(
self, endpoints_client: EndpointsClient, resources: ResourceManager
) -> None:

View file

@ -744,6 +744,7 @@ class TestTogetherMessages:
assert "22" in text, f"the model never saw the tool result: {response.content}"
@pytest.mark.covers("llm.messages.together_ai.basic.stream.works")
@pytest.mark.provider_live
def test_streams_text_deltas(
self, client: PassthroughClient, resources: ResourceManager, reasoning_tool_backend: str
) -> None:

299
tests/e2e/provider_cache.py Normal file
View file

@ -0,0 +1,299 @@
from __future__ import annotations
import base64
import hashlib
import hmac
import io
import threading
import time
from collections.abc import Callable, Generator, Mapping
from contextlib import closing
from dataclasses import dataclass, field
from typing import Final, Literal, Protocol
from urllib.parse import urlsplit
from e2e_http import (
NetworkError,
StreamChunk,
StreamHead,
StreamStep,
StreamTruncation,
forward_prepared_stream,
forward_stream,
prepare_forward,
primed_steps,
)
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
LIFETIME_SECONDS: Final = 86_400
MAX_REQUEST_BYTES: Final = 256 * 1024
MAX_RESPONSE_BYTES: Final = 8 * 1024 * 1024
UNRECORDED_RESPONSE_HEADERS: Final = frozenset({"set-cookie"})
JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
@dataclass(frozen=True, slots=True)
class CacheHit:
payload: bytes
valid_until: float
@dataclass(frozen=True, slots=True)
class CaptureLease:
token: str
captured_at_ms: int
expires_at_ms: int
@dataclass(frozen=True, slots=True)
class CacheBusy:
pass
@dataclass(frozen=True, slots=True)
class CacheUnavailable:
pass
type CacheLookup = CacheHit | CaptureLease | CacheBusy | CacheUnavailable
class ResponseStore(Protocol):
def lookup(self, key: str) -> CacheLookup: ...
def publish(self, key: str, lease: CaptureLease, payload: bytes) -> bool: ...
def release(self, key: str, lease: CaptureLease) -> bool: ...
def discard(self, key: str, payload: bytes) -> bool: ...
class CachedResponse(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid", strict=True)
format_version: Literal[1] = 1
request_key: str
status_code: int
headers: dict[str, str]
chunks: tuple[str, ...]
class SignedResponse(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid", strict=True)
response: str
signature: str
def exact_key(secret: bytes, method: str, url: str, headers: Mapping[str, str], body: bytes | None) -> str:
fields: Final = (
b"provider-cache-exact-v1", method.encode(), url.encode(),
*(part.encode() for pair in sorted(headers.items()) for part in pair),
b"no-body" if body is None else b"body", b"" if body is None else body,
)
encoded: Final = b"".join(len(part).to_bytes(8, "big") + part for part in fields)
return hmac.new(secret, encoded, hashlib.sha256).hexdigest()
def cacheable_endpoint(method: str, url: str, body: bytes | None) -> bool:
return (
method == "POST"
and urlsplit(url).path in {"/v1/chat/completions", "/v1/messages"}
and body is not None
and len(body) <= MAX_REQUEST_BYTES
)
def successful_response(url: str, status: int, headers: Mapping[str, str], body: bytes) -> bool:
if not 200 <= status < 300 or len(body) > MAX_RESPONSE_BYTES:
return False
streaming: Final = "text/event-stream" in headers.get("content-type", "").lower()
if streaming:
try:
text: Final = body.decode("utf-8").replace("\r\n", "\n")
if not text.endswith("\n\n"):
return False
events: Final = tuple(
"\n".join(line[5:].removeprefix(" ") for line in event.split("\n") if line.startswith("data:"))
for event in text.split("\n\n") if any(line.startswith("data:") for line in event.split("\n"))
)
values: Final = tuple(JSON_VALUE.validate_json(event) for event in events if event != "[DONE]")
except (UnicodeDecodeError, ValidationError):
return False
if not values or any(not isinstance(value, dict) or "error" in value or value.get("type") == "error" for value in values):
return False
if urlsplit(url).path == "/v1/chat/completions":
return events[-1] == "[DONE]" and "[DONE]" not in events[:-1] and complete_chat_stream(values)
return (
"[DONE]" not in events
and isinstance(values[0], dict) and values[0].get("type") == "message_start"
and isinstance(values[-1], dict) and values[-1].get("type") == "message_stop"
and any(
isinstance(value, dict) and value.get("type") == "message_delta"
and isinstance(delta := value.get("delta"), dict) and isinstance(delta.get("stop_reason"), str)
for value in values
)
)
try:
value: Final = JSON_VALUE.validate_json(body)
except ValidationError:
return False
if not isinstance(value, dict) or "error" in value:
return False
if urlsplit(url).path == "/v1/messages":
return value.get("type") == "message" and isinstance(value.get("content"), list) and isinstance(value.get("stop_reason"), str)
choices: Final = value.get("choices")
return isinstance(choices, list) and bool(choices) and all(
isinstance(choice, dict) and isinstance(choice.get("message"), dict) and isinstance(choice.get("finish_reason"), str)
for choice in choices
)
def complete_chat_stream(values: tuple[JsonValue, ...]) -> bool:
if any(not isinstance(value, dict) or not isinstance(value.get("choices"), list) for value in values):
return False
choices: Final = tuple(
choice for value in values if isinstance(value, dict)
if isinstance(items := value.get("choices"), list) for choice in items
)
if not choices or any(
not isinstance(choice, dict) or type(choice.get("index")) is not int
or not isinstance(choice.get("delta"), dict)
for choice in choices
):
return False
indices: Final = frozenset(choice["index"] for choice in choices if isinstance(choice, dict))
return all(
isinstance(tuple(choice for choice in choices if isinstance(choice, dict) and choice["index"] == index)[-1].get("finish_reason"), str)
for index in indices
)
def encode_response(secret: bytes, response: CachedResponse) -> bytes:
raw: Final = response.model_dump_json()
return SignedResponse(response=raw, signature=hmac.new(secret, raw.encode(), hashlib.sha256).hexdigest()).model_dump_json().encode()
def decode_response(secret: bytes, key: str, payload: bytes, url: str) -> CachedResponse | None:
if len(payload) > 2 * MAX_RESPONSE_BYTES:
return None
try:
signed: Final = SignedResponse.model_validate_json(payload)
if not hmac.compare_digest(signed.signature.encode(), hmac.new(secret, signed.response.encode(), hashlib.sha256).hexdigest().encode()):
return None
response: Final = CachedResponse.model_validate_json(signed.response)
chunks: Final = tuple(base64.b64decode(chunk, validate=True) for chunk in response.chunks)
except (ValidationError, ValueError):
return None
if response.request_key != key or not successful_response(url, response.status_code, response.headers, b"".join(chunks)):
return None
return response
@dataclass(slots=True)
class CacheCounters:
counts: tuple[tuple[str, int], ...] = ()
lock: threading.Lock = field(default_factory=threading.Lock)
def increment(self, name: str) -> None:
with self.lock:
current: Final = dict(self.counts)
self.counts = tuple((current | {name: current.get(name, 0) + 1}).items())
@dataclass(slots=True)
class ResponseCapture:
buffer: io.BytesIO = field(default_factory=io.BytesIO)
size: int = 0
eligible: bool = True
def observe(self, step: StreamStep) -> None:
if not self.eligible:
return
if isinstance(step, StreamTruncation) or self.size + len(step.data) + 8 > MAX_RESPONSE_BYTES:
self.eligible = False
self.buffer.close()
return
self.buffer.write(len(step.data).to_bytes(8, "big"))
self.buffer.write(step.data)
self.size += len(step.data) + 8
def chunks(self) -> tuple[bytes, ...]:
self.buffer.seek(0)
return tuple(self.buffer.read(int.from_bytes(size, "big")) for size in iter(lambda: self.buffer.read(8), b""))
def response_steps(response: CachedResponse) -> Generator[StreamStep, None, None]:
for chunk in response.chunks:
yield StreamChunk(base64.b64decode(chunk, validate=True))
@dataclass(frozen=True, slots=True)
class CacheEdge:
store: ResponseStore
secret: bytes = field(repr=False)
counters: CacheCounters = field(default_factory=CacheCounters)
wait_seconds: float = 2.0
clock: Callable[[], float] = time.monotonic
sleep: Callable[[float], None] = time.sleep
def lookup(self, key: str) -> CacheLookup:
deadline: Final = self.clock() + self.wait_seconds
while isinstance(result := self.store.lookup(key), CacheBusy) and self.clock() < deadline:
self.sleep(min(0.05, max(0, deadline - self.clock())))
return result
def forward(self, method: str, url: str, headers: dict[str, str], body: bytes | None, timeout: float) -> StreamHead | NetworkError:
if not cacheable_endpoint(method, url, body):
self.counters.increment("bypass")
self.counters.increment("upstream_attempts")
return forward_stream(method, url, headers=headers, body=body, timeout=timeout)
prepared: Final = prepare_forward(method, url, headers, body)
if isinstance(prepared, NetworkError):
self.counters.increment("rejected")
return prepared
key: Final = exact_key(self.secret, method, url, prepared.headers, body)
found: Final = self.lookup(key)
if isinstance(found, CacheHit):
response: Final = decode_response(self.secret, key, found.payload, url)
if response is not None and self.clock() < found.valid_until:
self.counters.increment("hits")
return StreamHead(response.status_code, response.headers, response_steps(response))
self.counters.increment("corrupt" if response is None else "expired")
self.store.discard(key, found.payload)
capture_slot: Final = self.lookup(key) if isinstance(found, CacheHit) else found
self.counters.increment("misses")
if isinstance(capture_slot, CacheUnavailable):
self.counters.increment("cache_errors")
self.counters.increment("upstream_attempts")
head: Final = forward_prepared_stream(prepared, timeout)
if not isinstance(capture_slot, CaptureLease):
return head
if isinstance(head, NetworkError):
self.store.release(key, capture_slot)
self.counters.increment("rejected")
return head
return StreamHead(head.status_code, head.headers, primed_steps(self.capture(key, capture_slot, url, head)))
def capture(self, key: str, lease: CaptureLease, url: str, head: StreamHead) -> Generator[StreamStep, None, None]:
capture: Final = ResponseCapture()
try:
with closing(head.steps):
yield StreamChunk(b"")
for step in head.steps:
yield step
capture.observe(step)
chunks: Final = capture.chunks() if capture.eligible else ()
headers: Final = {
name: value for name, value in head.headers.items() if name.lower() not in UNRECORDED_RESPONSE_HEADERS
}
if not capture.eligible or not successful_response(url, head.status_code, headers, b"".join(chunks)):
self.counters.increment("rejected")
return
response: Final = CachedResponse(
request_key=key, status_code=head.status_code, headers=headers,
chunks=tuple(base64.b64encode(chunk).decode("ascii") for chunk in chunks),
)
published: Final = self.store.publish(key, lease, encode_response(self.secret, response))
self.counters.increment("writes" if published else "write_failures")
finally:
self.store.release(key, lease)
capture.buffer.close()

View file

@ -0,0 +1,154 @@
from __future__ import annotations
import atexit
import functools
import json
import logging
import os
import re
import time
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import Final, Protocol, cast
from provider_cache import LIFETIME_SECONDS, CacheBusy, CacheEdge, CacheHit, CacheLookup, CacheUnavailable, CaptureLease
from pydantic import TypeAdapter, ValidationError
from redis import Redis
from redis.exceptions import RedisError
REDIS_ARRAY: Final[TypeAdapter[list[bytes]]] = TypeAdapter(list[bytes])
LOOKUP: Final = """
local clock = redis.call('TIME')
local now = clock[1] * 1000 + math.floor(clock[2] / 1000)
local row = redis.call('HMGET', KEYS[1], 'captured', 'expires', 'payload')
if row[3] then
local captured = tonumber(row[1])
local expires = tonumber(row[2])
if captured and expires and captured <= now and expires > now
and expires - captured == tonumber(ARGV[2]) then
return {'hit', row[3], tostring(expires - now)}
end
redis.call('DEL', KEYS[1])
end
if redis.call('SET', KEYS[2], ARGV[1], 'NX', 'PX', ARGV[3]) then
return {'lease', tostring(now), tostring(now + tonumber(ARGV[2]))}
end
return {'busy'}
"""
PUBLISH: Final = """
if redis.call('GET', KEYS[2]) ~= ARGV[1] then return 0 end
local clock = redis.call('TIME')
local now = clock[1] * 1000 + math.floor(clock[2] / 1000)
local captured = tonumber(ARGV[2])
local expires = tonumber(ARGV[3])
if captured > now or expires <= now or expires - captured ~= tonumber(ARGV[5]) then return 0 end
if redis.call('EXISTS', KEYS[1]) == 1 then return 0 end
redis.call('HSET', KEYS[1], 'captured', ARGV[2], 'expires', ARGV[3], 'payload', ARGV[4])
redis.call('PEXPIREAT', KEYS[1], expires)
redis.call('DEL', KEYS[2])
return 1
"""
RELEASE: Final = """
if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end
return redis.call('DEL', KEYS[1])
"""
DISCARD: Final = """
if redis.call('HGET', KEYS[1], 'payload') ~= ARGV[1] then return 0 end
return redis.call('DEL', KEYS[1])
"""
class RedisCommands(Protocol):
def eval(self, script: str, numkeys: int, *args: str | bytes | int) -> object: ...
@dataclass(frozen=True, slots=True)
class RedisResponseStore:
client: RedisCommands
namespace: str
lifetime_ms: int = LIFETIME_SECONDS * 1000
lease_ms: int = 120_000
def keys(self, key: str) -> tuple[str, str]:
prefix: Final = f"e2e-provider-cache:v1:{self.namespace}:{{{key}}}"
return prefix + ":response", prefix + ":lease"
def lookup(self, key: str) -> CacheLookup:
token: Final = uuid.uuid4().hex
started: Final = time.monotonic()
try:
result: Final = self.client.eval(LOOKUP, 2, *self.keys(key), token, self.lifetime_ms, self.lease_ms)
except (RedisError, OSError):
return CacheUnavailable()
try:
parts: Final = tuple(REDIS_ARRAY.validate_python(result, strict=True))
except ValidationError:
return CacheUnavailable()
if len(parts) == 3 and parts[0] == b"hit" and parts[2].isdigit():
return CacheHit(parts[1], started + int(parts[2]) / 1000)
if len(parts) == 3 and parts[0] == b"lease" and parts[1].isdigit() and parts[2].isdigit():
return CaptureLease(token, int(parts[1]), int(parts[2]))
if parts == (b"busy",):
return CacheBusy()
return CacheUnavailable()
def publish(self, key: str, lease: CaptureLease, payload: bytes) -> bool:
try:
result: Final = self.client.eval(
PUBLISH, 2, *self.keys(key), lease.token, lease.captured_at_ms, lease.expires_at_ms, payload, self.lifetime_ms,
)
except (RedisError, OSError):
return False
return result == 1
def release(self, key: str, lease: CaptureLease) -> bool:
try:
result: Final = self.client.eval(RELEASE, 1, self.keys(key)[1], lease.token)
except (RedisError, OSError):
return False
return result == 1
def discard(self, key: str, payload: bytes) -> bool:
try:
result: Final = self.client.eval(DISCARD, 1, self.keys(key)[0], payload)
except (RedisError, OSError):
return False
return result == 1
def redis_store(url: str, namespace: str) -> RedisResponseStore:
client: Final = Redis.from_url(url, socket_timeout=0.25, socket_connect_timeout=0.25, decode_responses=False)
return RedisResponseStore(cast(RedisCommands, client), namespace)
def write_metrics(cache: CacheEdge) -> None:
report: Final = json.dumps({"provider_cache": dict(cache.counters.counts)})
directory: Final = os.environ.get("E2E_PROVIDER_CACHE_METRICS_DIR")
if directory:
try:
root: Final = Path(directory)
root.mkdir(parents=True, exist_ok=True)
(root / f"{os.getpid()}.json").write_text(report + "\n")
except OSError:
logging.getLogger(__name__).warning("provider cache metrics artifact unavailable")
logging.getLogger(__name__).info("%s", report)
@functools.lru_cache(maxsize=1)
def configured_cache() -> CacheEdge | None:
if os.environ.get("E2E_PROVIDER_CACHE", "0") == "0":
return None
if os.environ.get("E2E_PROVIDER_CACHE") != "1":
raise ValueError("E2E_PROVIDER_CACHE must be 0 or 1")
secret: Final = os.environ.get("E2E_PROVIDER_CACHE_HMAC_KEY", "").encode()
namespace: Final = os.environ.get("E2E_PROVIDER_CACHE_NAMESPACE", "")
if len(secret) < 32 or re.fullmatch(r"[a-zA-Z0-9_-]{1,64}", namespace) is None:
raise ValueError("provider cache requires a dedicated key and namespace")
cache: Final = CacheEdge(redis_store(os.environ["E2E_PROVIDER_CACHE_REDIS_URL"], namespace), secret)
atexit.register(write_metrics, cache)
return cache

View file

@ -0,0 +1,23 @@
from __future__ import annotations
from collections.abc import Callable
from contextvars import ContextVar
from typing import Final
from models import LiteLLMParamsBody, ModelMode
LIVE_PROVIDER_REQUIRED: Final[ContextVar[bool]] = ContextVar("live_provider_required", default=False)
def route_cache_model(
params: LiteLLMParamsBody, base_for: Callable[[str], str | None], *, enabled: bool, mode: ModelMode | None = None,
) -> LiteLLMParamsBody:
if not enabled or mode == "realtime" or LIVE_PROVIDER_REQUIRED.get() or params.api_base is not None or params.mock_response is not None:
return params
provider: Final = params.model.partition("/")[0]
if provider not in {"openai", "anthropic"} or params.litellm_credential_name is not None:
return params
base: Final = base_for(provider)
if base is None:
return params
return params.model_copy(update={"api_base": f"{base}/v1" if provider == "openai" else base})

View file

@ -42,6 +42,7 @@ import base64
import difflib
import functools
import hashlib
import os
import re
import threading
from collections import deque
@ -93,6 +94,8 @@ from fixture_mode import (
parse_fixture_mode,
)
from fixture_profile import IneligibleRequest, MatchProfile, match_profile, strict_identity
from provider_cache import CacheEdge
from provider_cache_routing import LIVE_PROVIDER_REQUIRED
from pydantic import JsonValue, TypeAdapter
EDGE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType(
@ -506,7 +509,7 @@ class LiveEdge:
pass
type EdgeBackend = RecordEdge | ReplayEdge | LiveEdge
type EdgeBackend = RecordEdge | ReplayEdge | LiveEdge | CacheEdge
@dataclass(slots=True)
@ -750,12 +753,16 @@ def _handle_record(
def _handle_live(
method: str, url: str, headers: Mapping[str, str], body: bytes | None, timeout: float
method: str, url: str, headers: Mapping[str, str], body: bytes | None, timeout: float,
cache: CacheEdge | None = None,
) -> EdgeOutcome:
forwarded: Final = {
name: value for name, value in headers.items() if name.lower() not in _REQUEST_DROPPED_HEADERS
}
head: Final = forward_stream(method, url, headers=forwarded, body=body, timeout=timeout)
head: Final = (
forward_stream(method, url, headers=forwarded, body=body, timeout=timeout)
if cache is None else cache.forward(method, url, forwarded, body, timeout)
)
match head:
case NetworkError(message=message):
return _recorded_outcome(_network_error_response(message))
@ -821,6 +828,10 @@ def handle_edge_request(
else edge_request(method, split.path, split.query, body, _header_value(headers, "content-type"))
)
match backend:
case CacheEdge():
return _handle_live(
method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout, backend,
)
case LiveEdge():
return _handle_live(
method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout
@ -871,11 +882,19 @@ class _EdgeHandler(BaseHTTPRequestHandler):
or isinstance(edge_server.backend, ReplayEdge)
and edge_server.backend.source.bundle.manifest.match_profile == "stateless_v1"
)
if strict and len({name.lower() for name in self.headers.keys()}) != len(self.headers):
if strict and len({name.lower() for name in self.headers}) != len(self.headers):
self._write_reply(_text_reply(REPLAY_MISS_STATUS, "stateless_v1 eligibility error: duplicate headers"))
return
duplicate_headers: Final = len({name.lower() for name in self.headers}) != len(self.headers)
selected_backend: Final = (
LiveEdge() if isinstance(edge_server.backend, CacheEdge) and duplicate_headers else edge_server.backend
)
if isinstance(edge_server.backend, CacheEdge) and duplicate_headers:
edge_server.backend.counters.increment("duplicate_header_bypass")
if urlsplit(self.path).path.lstrip("/").partition("/")[0] in edge_server.mounts:
edge_server.backend.counters.increment("upstream_attempts")
outcome: Final = handle_edge_request(
edge_server.backend,
selected_backend,
edge_server.mounts,
self.command,
self.path,
@ -908,12 +927,12 @@ class _EdgeHandler(BaseHTTPRequestHandler):
shuts down write-side first: the proxy sees a graceful close mid-message,
which is the incomplete chunked read a provider hanging up produces, and not
the reset that could discard the chunks already in flight."""
self.send_response(stream.status_code)
for name, value in stream.headers.items():
self.send_header(name, value)
self.send_header("transfer-encoding", "chunked")
self.end_headers()
with closing(stream.steps) as steps:
self.send_response(stream.status_code)
for name, value in stream.headers.items():
self.send_header(name, value)
self.send_header("transfer-encoding", "chunked")
self.end_headers()
for step in steps:
match step:
case StreamChunk(data=data):
@ -923,7 +942,7 @@ class _EdgeHandler(BaseHTTPRequestHandler):
return
case _:
assert_never(step)
self.wfile.write(b"0\r\n\r\n")
self.wfile.write(b"0\r\n\r\n")
def log_message(self, format: str, *args: object) -> None:
"""Silence the per-request stderr line BaseHTTPRequestHandler emits."""
@ -1056,6 +1075,8 @@ def provider_edge_api_base(
case InvalidFixtureMode(value=value):
raise ValueError(f"E2E_FIXTURE_MODE={value!r} is not one of {', '.join(FIXTURE_MODES)}")
case "live":
if configured_cache_backend() is not None:
return _shared_cache_edge(bind_host, advertise_host, forward_timeout).api_base(mount)
return None
case "record" | "replay":
if mount not in EDGE_MOUNTS:
@ -1073,7 +1094,7 @@ def _observed_backend(mode_raw: str, bundle_dir: Path) -> EdgeBackend:
case InvalidFixtureMode(value=value):
raise ValueError(f"E2E_FIXTURE_MODE={value!r} is not one of {', '.join(FIXTURE_MODES)}")
case "live":
return LiveEdge()
return configured_cache_backend() or LiveEdge()
case "record":
return RecordEdge(_shared_recorder(bundle_dir, match_profile()), threading.Lock())
case "replay":
@ -1082,6 +1103,24 @@ def _observed_backend(mode_raw: str, bundle_dir: Path) -> EdgeBackend:
assert_never(mode)
def configured_cache_backend() -> CacheEdge | None:
if LIVE_PROVIDER_REQUIRED.get() or os.environ.get("E2E_PROVIDER_CACHE", "0") == "0":
return None
from provider_cache_redis import configured_cache
return configured_cache()
@functools.lru_cache(maxsize=8)
def _shared_cache_edge(bind_host: str, advertise_host: str, forward_timeout: float) -> ProviderEdge:
backend: Final = configured_cache_backend()
assert backend is not None
return start_provider_edge(
backend, mounts=EDGE_MOUNTS, bind_host=bind_host,
advertise_host=advertise_host, forward_timeout=forward_timeout,
).edge
@contextmanager
def observed_provider_edge(
observation: ProviderRequestObservation,

View file

@ -8,6 +8,7 @@ ProxyClient's key/customer methods for cleanup. Read-backs are eventually consis
from __future__ import annotations
import os
import time
import warnings
from collections.abc import Callable, Mapping
@ -26,6 +27,7 @@ from e2e_config import (
PROXY_REPLICA_URLS,
REQUEST_TIMEOUT,
SLOW_PROVIDER_TIMEOUT_SECONDS,
provider_edge_base,
settle_propagation,
)
from e2e_http import (
@ -93,6 +95,7 @@ from models import (
UserDeleteBody,
UserDeleteResponse,
)
from provider_cache_routing import route_cache_model
from pydantic import BaseModel
from transport import HttpTransport, SplitTransport, Transport, is_control_plane_path
@ -645,7 +648,10 @@ class ProxyClient:
self.transport.post(
"/model/new",
headers=self.management_headers(),
json=body,
json=body.model_copy(update={"litellm_params": route_cache_model(
body.litellm_params, provider_edge_base,
enabled=os.environ.get("E2E_PROVIDER_CACHE", "0") == "1", mode=body.model_info.mode,
)}),
response_type=ModelNewResponse,
)
).model_id

View file

@ -26,7 +26,7 @@ from models import (
)
from quota_client import QuotaClient
pytestmark = pytest.mark.e2e
pytestmark = [pytest.mark.e2e, pytest.mark.provider_live]
# Anthropic prompt caching (host has ANTHROPIC_API_KEY; Bedrock was "Operation not allowed").
ANTHROPIC_MODEL = "anthropic/claude-haiku-4-5-20251001"

View file

@ -1254,7 +1254,8 @@ class TestHandleEdgeRequestPure:
class TestApiBaseSeam:
def test_live_mode_returns_none(self, tmp_path: Path) -> None:
def test_live_mode_returns_none(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("E2E_PROVIDER_CACHE", raising=False)
for mode_raw in ("live", ""):
assert (
provider_edge_api_base(

View file

@ -7,6 +7,7 @@ import {
import {
E2E_TEAM_CRUD_ALIAS,
E2E_TEAM_ORG_ALIAS,
E2E_TEAM_ORG_ID,
INTERNAL_USER_STORAGE_PATH,
} from "../../constants";
import { Page } from "../../fixtures/pages";
@ -179,18 +180,28 @@ test.describe("Models and Endpoints for an internal user", () => {
`switching to ${ALL_MODELS_VIEW} leaves the table populated rather than blanking it`,
).toHaveCount(1, { timeout: 15_000 });
await expect(page).toHaveURL((url) =>
url.searchParams.get("filter_team") === E2E_TEAM_ORG_ID &&
url.searchParams.get("view_mode") === "all",
);
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 });
"the selected team is restored from the URL after a reload",
).toContainText(E2E_TEAM_ORG_ALIAS, { timeout: 15_000 });
await expect(
viewSelector(page),
"the view selection is not persisted across a reload either",
).toContainText(CURRENT_TEAM_VIEW, { timeout: 15_000 });
"the selected view is restored from the URL after a reload",
).toContainText(ALL_MODELS_VIEW, { timeout: 15_000 });
await expect(modelRow(page, CHAT_MODEL_A)).toHaveCount(1, { timeout: 15_000 });
await expect(page.getByTestId("pagination-range")).toHaveText("Showing 1-1 of 1");
await expect(modelRow(page, CHAT_MODEL_B)).toHaveCount(0);
await expect(modelRow(page, ungrantedModelName)).toHaveCount(0);
await chooseOption(page, teamSelector(page), PERSONAL_TEAM);
await expect(
modelRow(page, ungrantedModelName),
"the personal view still renders models after a reload rather than coming back empty",
"switching back to the personal team restores models outside the selected team",
).toHaveCount(1, { timeout: 30_000 });
});
});

View file

@ -163,6 +163,9 @@ async def test_reset_budget_keys_partial_failure():
key1, key2, key3, key4, key5, key6 = (
_attrify(k) for k in [key1, key2, key3, key4, key5, key6]
)
pre_reset_spend = {
k["token"]: k["spend"] for k in [key2, key3, key4, key5, key6]
}
prisma_client.get_data = AsyncMock(
return_value=[key1, key2, key3, key4, key5, key6]
)
@ -201,7 +204,7 @@ async def test_reset_budget_keys_partial_failure():
# And every write must carry only {spend, budget_reset_at} — never the full row.
for c in key_writes:
assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
assert c["data"]["spend"] == 0
assert c["data"]["spend"] == {"decrement": pre_reset_spend[c["where"]["token"]]}
# Verify that the failure logging hook was scheduled (due to the failure for key1)
failure_hook_calls = (
@ -252,6 +255,9 @@ async def test_reset_budget_users_partial_failure():
user1, user2, user3, user4, user5, user6 = (
_attrify(u) for u in [user1, user2, user3, user4, user5, user6]
)
pre_reset_spend = {
u["user_id"]: u["spend"] for u in [user2, user3, user4, user5, user6]
}
prisma_client.get_data = AsyncMock(
return_value=[user1, user2, user3, user4, user5, user6]
)
@ -280,7 +286,9 @@ async def test_reset_budget_users_partial_failure():
assert written_ids == ["user2", "user3", "user4", "user5", "user6"]
for c in user_writes:
assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
assert c["data"]["spend"] == 0
assert c["data"]["spend"] == {
"decrement": pre_reset_spend[c["where"]["user_id"]]
}
failure_hook_calls = (
proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args_list
@ -441,6 +449,7 @@ async def test_reset_budget_teams_partial_failure():
for t in [team1, team2]:
t.setdefault("team_id", t["id"])
team1, team2 = _attrify(team1), _attrify(team2)
pre_reset_spend = team2["spend"]
prisma_client.get_data = AsyncMock(return_value=[team1, team2])
async def fake_reset_team(team, current_time, reset_settings=None):
@ -465,7 +474,7 @@ async def test_reset_budget_teams_partial_failure():
assert len(team_writes) == 1
assert team_writes[0]["where"] == {"team_id": "team2"}
assert set(team_writes[0]["data"].keys()) == {"spend", "budget_reset_at"}
assert team_writes[0]["data"]["spend"] == 0
assert team_writes[0]["data"]["spend"] == {"decrement": pre_reset_spend}
failure_hook_calls = (
proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args_list
@ -542,6 +551,11 @@ async def test_reset_budget_continues_other_categories_on_failure():
user1, user2 = _attrify(user1), _attrify(user2)
team1, team2 = _attrify(team1), _attrify(team2)
enduser1 = _attrify(enduser1)
pre_reset_spend = {
**{k["token"]: k["spend"] for k in [key1, key2]},
**{u["user_id"]: u["spend"] for u in [user2]},
**{t["team_id"]: t["spend"] for t in [team1, team2]},
}
_wire_cascade_reads_for_test(prisma_client)
proxy_logging_obj = MagicMock()
@ -618,7 +632,9 @@ async def test_reset_budget_continues_other_categories_on_failure():
# Every batched write must carry only the two reset fields, never the full row.
for c in key_writes + user_writes + team_writes:
assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
assert c["data"]["spend"] == 0
assert c["data"]["spend"] == {
"decrement": pre_reset_spend[next(iter(c["where"].values()))]
}
# ---------------------------------------------------------------------------

View file

@ -86,6 +86,7 @@ def test_async_fallbacks(caplog):
if "Task exception was never retrieved" not in log
and "Task was destroyed but it is pending" not in log
and "get_available_deployment" not in log
and "Selected deployment for model" not in log
and "in the Langfuse queue" not in log
and "Unclosed client session" not in log
and "Unclosed connector" not in log

View file

@ -1069,6 +1069,7 @@ async def test_jwt_non_admin_team_route_access(monkeypatch):
mock_jwt_response = {
"is_proxy_admin": False,
"jwt_claims": {},
"team_id": None,
"team_object": None,
"user_id": None,

View file

@ -1,6 +1,7 @@
import json
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
@ -579,6 +580,20 @@ def test_text_only_streaming_has_index_zero():
), f"Expected index=0, got {parsed.choices[0].index}"
def test_message_delta_without_usage_returns_chunk_with_no_usage():
iterator: Final = ModelResponseIterator(None, sync_stream=True)
model_response: Final = iterator.chunk_parser(
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
}
)
assert model_response.choices[0].finish_reason == "stop"
assert model_response.usage is None
def test_streaming_thinking_deltas_count_reasoning_tokens_in_usage():
"""Anthropic streaming usage should account for emitted thinking deltas."""
chunks = [

View file

@ -1,14 +1,20 @@
import math
from datetime import datetime, timezone
from typing import Final
import pytest
import litellm
from litellm.llms.fireworks_ai.cost_calculator import cost_per_token
from litellm.types.utils import OffPeakPricing, PromptTokensDetailsWrapper, Usage
from litellm.types.utils import (
CompletionTokensDetailsWrapper,
OffPeakPricing,
PromptTokensDetailsWrapper,
Usage,
)
MODEL = "accounts/fireworks/models/glm-5p2"
INPUT_COST = 1.4e-06
# Read the cached rate from the price map so this test tracks the shipped value
# (glm-5p2 is $0.14/1M) instead of hardcoding a number that breaks when it changes.
CACHE_READ_COST = litellm.get_model_info(model=MODEL, custom_llm_provider="fireworks_ai")["cache_read_input_token_cost"]
OUTPUT_COST = 4.4e-06
@ -44,13 +50,16 @@ STANDARD_CACHE_READ_COST = 1.5e-08
def _register_off_peak_model(
off_peak_pricing: OffPeakPricing, cache_read_cost: float | None = STANDARD_CACHE_READ_COST
) -> None:
litellm.model_cost[f"fireworks_ai/{OFF_PEAK_MODEL}"] = {
"litellm_provider": "fireworks_ai",
"mode": "chat",
"input_cost_per_token": STANDARD_INPUT_COST,
"output_cost_per_token": STANDARD_OUTPUT_COST,
"off_peak_pricing": off_peak_pricing,
**({} if cache_read_cost is None else {"cache_read_input_token_cost": cache_read_cost}),
litellm.model_cost = { # test-quality-ok: the save/restore conftest returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing
**litellm.model_cost,
f"fireworks_ai/{OFF_PEAK_MODEL}": {
"litellm_provider": "fireworks_ai",
"mode": "chat",
"input_cost_per_token": STANDARD_INPUT_COST,
"output_cost_per_token": STANDARD_OUTPUT_COST,
"off_peak_pricing": off_peak_pricing,
**({} if cache_read_cost is None else {"cache_read_input_token_cost": cache_read_cost}),
},
}
@ -125,3 +134,75 @@ def test_off_peak_defaults_to_the_current_time():
assert math.isclose(prompt_cost, 1000 * 1e-08, rel_tol=1e-10)
assert math.isclose(completion_cost, 200 * 2e-08, rel_tol=1e-10)
COMPONENT_MODEL = "accounts/fireworks/models/cost-components-test"
COMPONENT_INPUT_COST = 1e-06
COMPONENT_OUTPUT_COST = 2e-06
COMPONENT_CACHE_READ_COST = 1e-07
COMPONENT_CACHE_CREATION_COST = 3e-06
COMPONENT_REASONING_COST = 4e-06
COMPONENT_AUDIO_IN_COST = 5e-06
COMPONENT_AUDIO_OUT_COST = 6e-06
def test_cache_write_reasoning_and_audio_tokens_are_billed_at_their_component_rates():
litellm.model_cost = { # test-quality-ok: the save/restore conftest returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing
**litellm.model_cost,
f"fireworks_ai/{COMPONENT_MODEL}": {
"litellm_provider": "fireworks_ai",
"mode": "chat",
"input_cost_per_token": COMPONENT_INPUT_COST,
"output_cost_per_token": COMPONENT_OUTPUT_COST,
"cache_read_input_token_cost": COMPONENT_CACHE_READ_COST,
"cache_creation_input_token_cost": COMPONENT_CACHE_CREATION_COST,
"output_cost_per_reasoning_token": COMPONENT_REASONING_COST,
"input_cost_per_audio_token": COMPONENT_AUDIO_IN_COST,
"output_cost_per_audio_token": COMPONENT_AUDIO_OUT_COST,
},
}
usage = Usage(
prompt_tokens=1000,
completion_tokens=500,
total_tokens=1500,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=300,
cache_creation_tokens=200,
audio_tokens=100,
),
completion_tokens_details=CompletionTokensDetailsWrapper(
reasoning_tokens=200,
audio_tokens=50,
),
)
prompt_cost, completion_cost = cost_per_token(model=COMPONENT_MODEL, usage=usage)
expected_prompt_cost = (
400 * COMPONENT_INPUT_COST
+ 300 * COMPONENT_CACHE_READ_COST
+ 200 * COMPONENT_CACHE_CREATION_COST
+ 100 * COMPONENT_AUDIO_IN_COST
)
expected_completion_cost = (
250 * COMPONENT_OUTPUT_COST + 200 * COMPONENT_REASONING_COST + 50 * COMPONENT_AUDIO_OUT_COST
)
assert prompt_cost == pytest.approx(expected_prompt_cost)
assert completion_cost == pytest.approx(expected_completion_cost)
def test_an_entry_without_an_input_rate_gets_no_cache_read_fallback():
litellm.model_cost = { # test-quality-ok: the save/restore conftest returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing
**litellm.model_cost, # pyright: ignore[reportUnknownMemberType] # the SDK types model_cost as dict[Unknown, Unknown]
"fireworks_ai/accounts/fireworks/models/no-input-rate-test": {
"litellm_provider": "fireworks_ai",
"mode": "chat",
"output_cost_per_token": 2e-06,
},
}
usage: Final = _usage(prompt_tokens=1000, cached_tokens=300, completion_tokens=200)
prompt_cost, completion_cost = cost_per_token(model="accounts/fireworks/models/no-input-rate-test", usage=usage)
assert prompt_cost == 0
assert completion_cost == 200 * 2e-06

View file

@ -5836,3 +5836,126 @@ def test_supported_reasoning_efforts_still_map(model):
drop_params=False,
)
assert "thinkingConfig" in result
def _generate_content_body() -> dict:
return {
"candidates": [
{
"content": {"role": "model", "parts": [{"text": "hi"}]},
"finishReason": "STOP",
"index": 0,
}
],
"usageMetadata": {
"promptTokenCount": 5,
"candidatesTokenCount": 7,
"totalTokenCount": 12,
},
}
def test_generate_content_transform_uses_reported_model_version():
"""The served modelVersion must win over the requested name so downstream
pricing sees what actually ran."""
import httpx
body = {**_generate_content_body(), "modelVersion": "gemini-x-served"}
response: Final = VertexGeminiConfig()._transform_google_generate_content_to_openai_model_response(
completion_response=body,
model_response=ModelResponse(),
model="gemini-x",
logging_obj=MagicMock(),
raw_response=httpx.Response(200, headers={}),
)
assert response.model == "gemini-x-served"
def test_generate_content_transform_falls_back_to_requested_model():
import httpx
response: Final = VertexGeminiConfig()._transform_google_generate_content_to_openai_model_response(
completion_response=_generate_content_body(),
model_response=ModelResponse(),
model="gemini-x",
logging_obj=MagicMock(),
raw_response=httpx.Response(200, headers={}),
)
assert response.model == "gemini-x"
def test_streaming_chunk_carries_model_version():
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator,
)
chunk = {**_generate_content_body(), "modelVersion": "gemini-x-served"}
iterator: Final = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock())
streaming_chunk: Final = iterator.chunk_parser(chunk)
assert streaming_chunk.model == "gemini-x-served"
def test_served_model_version_reaches_assembled_stream_through_custom_stream_wrapper():
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator,
)
served_model: Final = "gemini-3.8-flash-001"
iterator: Final = ModelResponseIterator(
streaming_response=iter(
[json.dumps({**_generate_content_body(), "modelVersion": served_model}) for _ in range(3)]
),
sync_stream=True,
logging_obj=MagicMock(),
)
wrapper: Final = CustomStreamWrapper(
completion_stream=iter(iterator),
model="gemini/gemini-3.8-flash",
custom_llm_provider="gemini",
logging_obj=MagicMock(),
)
chunks: Final = list(wrapper)
assert len(chunks) >= 3
for chunk in chunks[:-1]:
assert chunk._hidden_params["provider_response_model"] == served_model
assembled: Final = litellm.stream_chunk_builder(chunks=list(chunks), messages=[{"role": "user", "content": "hi"}])
assert assembled._hidden_params["provider_response_model"] == served_model
def test_generate_content_transform_strips_version_suffix_from_model_version():
import httpx
body: Final = {**_generate_content_body(), "modelVersion": "gemini-3.8-flash-001@default"}
response: Final = VertexGeminiConfig()._transform_google_generate_content_to_openai_model_response(
completion_response=body,
model_response=ModelResponse(),
model="gemini-3.8-flash",
logging_obj=MagicMock(),
raw_response=httpx.Response(200, headers={}),
)
assert response.model == "gemini-3.8-flash-001"
def test_prompt_blocked_chunk_keeps_served_model_version():
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator,
)
chunk: Final = {
"promptFeedback": {"blockReason": "SAFETY", "blockReasonMessage": "prompt was blocked"},
"modelVersion": "gemini-3.8-flash-001",
"responseId": "resp-1",
}
iterator: Final = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock())
streaming_chunk: Final = iterator.chunk_parser(chunk)
assert streaming_chunk.model == "gemini-3.8-flash-001"
assert streaming_chunk.choices[0].finish_reason == "content_filter"

View file

@ -51,23 +51,23 @@ class TestXAIResponsesAPITransformation:
assert result["tools"][0]["type"] == "code_interpreter"
assert "container" not in result["tools"][0], "Container field should be removed"
def test_instructions_parameter_dropped(self):
"""Test that instructions parameter is dropped for XAI"""
def test_instructions_parameter_forwarded(self):
"""xAI supports 'instructions' on /v1/responses, so it must survive param mapping"""
config = XAIResponsesAPIConfig()
params = ResponsesAPIOptionalRequestParams(instructions="You are a helpful assistant.", temperature=0.7)
result = config.map_openai_params(response_api_optional_params=params, model="grok-4-fast", drop_params=False)
assert "instructions" not in result, "Instructions should be dropped"
assert result.get("instructions") == "You are a helpful assistant."
assert result.get("temperature") == 0.7, "Other params should be preserved"
def test_supported_params_excludes_instructions(self):
"""Test that get_supported_openai_params excludes instructions"""
def test_supported_params_includes_instructions(self):
"""A system message bridged to 'instructions' must not be rejected for xAI"""
config = XAIResponsesAPIConfig()
supported = config.get_supported_openai_params("grok-4-fast")
assert "instructions" not in supported, "instructions should not be supported"
assert "instructions" in supported, "instructions should be supported"
assert "tools" in supported, "tools should be supported"
assert "temperature" in supported, "temperature should be supported"
assert "model" in supported, "model should be supported"

View file

@ -53,8 +53,8 @@ class TestXAIResponsesAPITransformation:
"container" not in result["tools"][0]
), "Container field should be removed"
def test_instructions_parameter_dropped(self):
"""Test that instructions parameter is dropped for XAI"""
def test_instructions_parameter_forwarded(self):
"""xAI supports 'instructions' on /v1/responses, so it must survive param mapping"""
config = XAIResponsesAPIConfig()
params = ResponsesAPIOptionalRequestParams(
@ -65,15 +65,15 @@ class TestXAIResponsesAPITransformation:
response_api_optional_params=params, model="grok-4-fast", drop_params=False
)
assert "instructions" not in result, "Instructions should be dropped"
assert result.get("instructions") == "You are a helpful assistant."
assert result.get("temperature") == 0.7, "Other params should be preserved"
def test_supported_params_excludes_instructions(self):
"""Test that get_supported_openai_params excludes instructions"""
def test_supported_params_includes_instructions(self):
"""A system message bridged to 'instructions' must not be rejected for xAI"""
config = XAIResponsesAPIConfig()
supported = config.get_supported_openai_params("grok-4-fast")
assert "instructions" not in supported, "instructions should not be supported"
assert "instructions" in supported, "instructions should be supported"
assert "tools" in supported, "tools should be supported"
assert "temperature" in supported, "temperature should be supported"
assert "model" in supported, "model should be supported"

View file

@ -5,7 +5,7 @@ import json
import time
from base64 import urlsafe_b64encode
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -15,6 +15,9 @@ from litellm.types.mcp import MCPAuth
if TYPE_CHECKING:
import httpx
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.types.mcp_server.mcp_server_manager import MCPServer
@ -6977,6 +6980,11 @@ async def _exchange_persistence_attempted_for_auth_type(auth_type) -> bool:
new_callable=AsyncMock,
return_value="admin-user",
),
patch( # test-quality-ok: this control tests persistence by auth mode; write-policy behavior is covered separately
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.authorize_oauth_credential_request",
new_callable=AsyncMock,
return_value="admin-user",
),
patch(
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._store_per_user_token_server_side",
new_callable=AsyncMock,
@ -7124,12 +7132,12 @@ async def test_build_oauth_protected_resource_response_obo_end_to_end():
global_mcp_server_manager.registry.clear()
def _token_request(headers):
def _token_request(headers, path="/token"):
"""A real Starlette request with case-insensitive headers (matches production)."""
from starlette.requests import Request
raw = [(k.lower().encode(), v.encode()) for k, v in headers.items()]
return Request({"type": "http", "method": "POST", "path": "/token", "headers": raw, "query_string": b""})
return Request({"type": "http", "method": "POST", "path": path, "headers": raw, "query_string": b""})
@pytest.fixture
@ -11162,14 +11170,14 @@ async def test_identity_bound_authorization_carries_nonce_and_caller_through_cal
),
)
request = Request({"type": "http", "scheme": "https", "server": ("proxy.example.com", 443),
"path": "/authorize", "query_string": b"", "headers": []})
"path": "/authorize", "query_string": b"", "headers": [(b"authorization", b"Bearer sk-alice")]})
with (
patch( # test-quality-ok: isolate authenticated request resolution from the real encrypted OAuth round trip
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_user_id_from_request",
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.authorize_oauth_credential_request",
new=AsyncMock(return_value="alice")),
patch( # test-quality-ok: isolate user access lookup while testing nonce and caller preservation
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._bridge_authorize_access_denial",
new=AsyncMock(return_value=None)),
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._user_can_reach_mcp_server",
new=AsyncMock(return_value=True)),
):
authorized = await authorize_with_server(
request, server, "client", "http://127.0.0.1:6274/callback", state="client-state",
@ -11374,3 +11382,858 @@ with TestClient(app) as client:
assert responses[path]["status"] == 200, responses[path]
assert responses[path]["body"]["issuer"] == f"http://testserver/gateway/{path}"
assert responses["example/mcp"]["body"]["token_endpoint"] == "http://testserver/gateway/example/token"
@pytest.fixture
def jwt_oauth_identity(monkeypatch: pytest.MonkeyPatch) -> tuple["JWTHandler", "RSAPrivateKey"]:
import jwt
from cryptography.hazmat.primitives.asymmetric import rsa
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy import proxy_server
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
signing_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
cache: Final = UserApiKeyCache()
cache.set_cache(
"litellm_jwt_auth_keys_https://idp.example.test/jwks",
[json.loads(jwt.algorithms.RSAAlgorithm.to_jwk(signing_key.public_key()))],
)
cache.set_cache("jwt-owner", LiteLLM_UserTable(user_id="jwt-owner", user_email="owner@example.test"))
handler: Final = JWTHandler()
handler.update_environment(
prisma_client=None,
user_api_key_cache=cache,
litellm_jwtauth=LiteLLM_JWTAuth(user_id_jwt_field="identity.user_id"),
)
monkeypatch.setenv("JWT_PUBLIC_KEY_URL", "https://idp.example.test/jwks")
monkeypatch.setenv("JWT_ISSUER", "https://idp.example.test")
monkeypatch.setenv("JWT_AUDIENCE", "litellm-proxy")
monkeypatch.setattr(proxy_server, "jwt_handler", handler)
monkeypatch.setattr(proxy_server, "general_settings", {"enable_jwt_auth": True})
monkeypatch.setattr(proxy_server, "premium_user", True)
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
return handler, signing_key
def _oauth_identity_jwt(
signing_key: "RSAPrivateKey",
*,
expires_in: int = 300,
audience: str = "litellm-proxy",
issuer: str = "https://idp.example.test",
owner: str | None = "jwt-owner",
scope: str = "",
claims: dict[str, object] | None = None,
) -> str:
import jwt
return jwt.encode(
{
"sub": "not-the-configured-user-id",
"identity": {"user_id": owner},
"email": "owner@example.test",
"iss": issuer,
"aud": audience,
"exp": int(time.time()) + expires_in,
"scope": scope,
**(claims or {}),
},
signing_key,
algorithm="RS256",
)
@pytest.mark.asyncio
@pytest.mark.parametrize("header", ["Authorization", "x-litellm-api-key"])
@pytest.mark.parametrize("policy_allowed", [False, True])
@pytest.mark.parametrize("server_allowed", [False, True])
@pytest.mark.parametrize("admin", [False, True])
@pytest.mark.parametrize("owner_state", ["active", "missing", "inactive", "database_error"])
async def test_oauth_exchange_stores_token_for_validated_jwt_user(
jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"],
header: str,
policy_allowed: bool,
server_allowed: bool,
admin: bool,
owner_state: str,
monkeypatch: pytest.MonkeyPatch,
) -> None:
import httpx
from litellm.proxy._experimental.mcp_server import discoverable_endpoints
from litellm.proxy._types import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
handler, signing_key = jwt_oauth_identity
handler.litellm_jwtauth.custom_validate = lambda claims: policy_allowed
from litellm.proxy._experimental.mcp_server import mcp_server_manager
manager: Final = MagicMock()
manager.get_allowed_mcp_servers = AsyncMock(return_value=["jwt-oauth-server"] if server_allowed else [])
manager.invalidate_user_oauth_token_cache = AsyncMock()
monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager)
bearer: Final = _oauth_identity_jwt(signing_key, scope="litellm_proxy_admin" if admin else "")
request: Final = _token_request({header: f"Bearer {bearer}"}, path="/jwt-oauth-server/token")
server: Final = MCPServer(
server_id="jwt-oauth-server",
name="jwt-oauth-server",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow="authorization_code",
authorization_url="https://upstream.example.test/authorize",
token_url="https://upstream.example.test/token",
client_id="registered-client",
)
import litellm
from litellm.caching.llm_caching_handler import LLMClientCache
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy import proxy_server
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
from litellm.types.llms.custom_http import httpxSpecialProvider
def upstream_response(outbound: httpx.Request) -> httpx.Response:
assert outbound.url == server.token_url
assert bearer not in str(outbound.headers)
assert bearer.encode() not in outbound.content
return httpx.Response(200, json={"access_token": "upstream-token", "token_type": "Bearer"})
database: Final = MagicMock()
users: Final = database.db.litellm_usertable
users.find_unique = AsyncMock(return_value=None)
users.find_first = AsyncMock(return_value=None)
users.create = AsyncMock()
if owner_state in ("missing", "database_error"):
handler.user_api_key_cache.delete_cache("jwt-owner")
if owner_state == "database_error":
users.find_unique.side_effect = RuntimeError("database unavailable")
if owner_state == "inactive":
handler.user_api_key_cache.set_cache(
"jwt-owner", LiteLLM_UserTable(user_id="jwt-owner", metadata={"scim_active": False})
)
table: Final = database.db.litellm_mcpusercredentials
table.find_unique = AsyncMock(return_value=None)
table.upsert = AsyncMock()
monkeypatch.setattr(proxy_server, "prisma_client", database)
monkeypatch.setenv("LITELLM_SALT_KEY", "oauth-jwt-test-encryption-key")
clients: Final = LLMClientCache()
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", clients)
async with httpx.AsyncClient(transport=httpx.MockTransport(upstream_response)) as transport:
upstream: Final = AsyncHTTPHandler()
await upstream.client.aclose()
upstream.client = transport
clients.set_cache("async_httpx_client" + httpxSpecialProvider.Oauth2Check, upstream)
response: Final = await discoverable_endpoints.exchange_token_with_server(
request=request,
mcp_server=server,
grant_type="authorization_code",
code="upstream-code",
redirect_uri="http://localhost/callback",
client_id="registered-client",
client_secret=None,
code_verifier=None,
)
assert response.status_code == 200
assert json.loads(response.body)["access_token"] == "upstream-token"
users.create.assert_not_awaited()
if (
not server_allowed
or not policy_allowed
or owner_state in ("inactive", "database_error")
or (owner_state == "missing" and not admin)
):
table.upsert.assert_not_awaited()
return
table.upsert.assert_awaited_once()
stored: Final = table.upsert.call_args.kwargs
assert stored["where"] == {"user_id_server_id": {"user_id": "jwt-owner", "server_id": server.server_id}}
credential: Final = stored["data"]["create"]["credential_b64"]
assert "upstream-token" not in credential
decoded: Final = decrypt_value_helper(credential, key="mcp_user_credential")
assert json.loads(decoded)["access_token"] == "upstream-token"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"rejection",
[
"expired",
"audience",
"issuer",
"signature",
"missing_user",
"unknown_user",
"disabled",
"not_premium",
"scim_inactive",
"custom_validate",
"missing_database",
],
)
@pytest.mark.parametrize("credential_write", [False, True])
async def test_oauth_jwt_identity_rejects_untrusted_or_inactive_owner(
jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"],
monkeypatch: pytest.MonkeyPatch,
rejection: str,
credential_write: bool,
) -> None:
from cryptography.hazmat.primitives.asymmetric import rsa
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import mcp_server_manager
from litellm.proxy._experimental.mcp_server.bridge_token_flow import (
_extract_user_id_from_request, authorize_oauth_credential_request,
)
allowed_servers: Final = AsyncMock(return_value=["server-a"])
monkeypatch.setattr(mcp_server_manager.global_mcp_server_manager, "get_allowed_mcp_servers", allowed_servers)
handler, signing_key = jwt_oauth_identity
key: Final = (
rsa.generate_private_key(public_exponent=65537, key_size=2048) if rejection == "signature" else signing_key
)
bearer: Final = _oauth_identity_jwt(
key,
expires_in=-60 if rejection == "expired" else 300,
audience="upstream-only" if rejection == "audience" else "litellm-proxy",
issuer="https://untrusted.example.test" if rejection == "issuer" else "https://idp.example.test",
owner=None if rejection == "missing_user" else "unknown" if rejection == "unknown_user" else "jwt-owner",
)
if rejection == "disabled":
monkeypatch.setattr(proxy_server, "general_settings", {"enable_jwt_auth": False})
if rejection == "not_premium":
monkeypatch.setattr(proxy_server, "premium_user", False)
if rejection == "missing_database":
monkeypatch.setattr(proxy_server, "prisma_client", None)
if rejection == "scim_inactive":
handler.user_api_key_cache.set_cache(
"jwt-owner", LiteLLM_UserTable(user_id="jwt-owner", metadata={"scim_active": False})
)
if rejection == "custom_validate":
handler.litellm_jwtauth.custom_validate = lambda claims: False
request: Final = _token_request({"Authorization": f"Bearer {bearer}"})
result: Final = (
await authorize_oauth_credential_request(request, "server-a")
if credential_write else await _extract_user_id_from_request(request)
)
assert result is None
allowed_servers.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize("blocked", [False, True])
async def test_oauth_jwt_cannot_override_explicit_litellm_key(
jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"],
blocked: bool,
) -> None:
from litellm.proxy._experimental.mcp_server.bridge_token_flow import _extract_user_id_from_request
from litellm.proxy._types import UserAPIKeyAuth, hash_token
handler, signing_key = jwt_oauth_identity
key: Final = "sk-explicit-key"
handler.user_api_key_cache.set_cache(hash_token(key), UserAPIKeyAuth(user_id="key-owner", blocked=blocked))
request: Final = _token_request(
{
"Authorization": f"Bearer {_oauth_identity_jwt(signing_key)}",
"x-litellm-api-key": key,
}
)
assert await _extract_user_id_from_request(request) == (None if blocked else "key-owner")
@pytest.mark.asyncio
@pytest.mark.parametrize(
"mapping", ["active", "blocked", "inactive_owner", "fallback", "pending", "reject", "custom_reject"]
)
async def test_oauth_jwt_uses_configured_virtual_key_owner(
jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"],
mapping: str,
) -> None:
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy._experimental.mcp_server.bridge_token_flow import _extract_user_id_from_request
from litellm.proxy._types import UserAPIKeyAuth, UnregisteredJWTClientBehavior, hash_token
from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key
handler, signing_key = jwt_oauth_identity
handler.litellm_jwtauth.virtual_key_claim_field = "sub"
if mapping == "custom_reject":
handler.litellm_jwtauth.custom_validate = lambda claims: False
handler.litellm_jwtauth.unregistered_jwt_client_behavior = (
UnregisteredJWTClientBehavior.AUTO_REGISTER
if mapping == "pending"
else UnregisteredJWTClientBehavior.REJECT
if mapping == "reject"
else UnregisteredJWTClientBehavior.FALLBACK_TEAM_MAPPING
)
key_hash: Final = hash_token("sk-mapped-oauth-owner")
handler.user_api_key_cache.set_cache(
jwt_key_mapping_cache_key("sub", "not-the-configured-user-id"),
"__NO_MAPPING__" if mapping in ("fallback", "pending", "reject") else key_hash,
)
handler.user_api_key_cache.set_cache(
key_hash, UserAPIKeyAuth(token=key_hash, user_id="mapped-owner", blocked=mapping == "blocked")
)
handler.user_api_key_cache.set_cache(
"mapped-owner", LiteLLM_UserTable(user_id="mapped-owner", metadata={"scim_active": mapping != "inactive_owner"})
)
request: Final = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key)}"})
expected: Final = "jwt-owner" if mapping == "fallback" else "mapped-owner" if mapping == "active" else None
assert await _extract_user_id_from_request(request) == expected
@pytest.mark.asyncio
@pytest.mark.parametrize("allowed_domain", [None, "allowed.example.test"])
async def test_oauth_jwt_respects_custom_validation_and_email_policy(
jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"],
allowed_domain: str | None,
) -> None:
from litellm.proxy._experimental.mcp_server.bridge_token_flow import _extract_user_id_from_request
handler, signing_key = jwt_oauth_identity
handler.litellm_jwtauth.custom_validate = lambda claims: True
handler.litellm_jwtauth.user_allowed_email_domain = allowed_domain
request: Final = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key)}"})
assert await _extract_user_id_from_request(request) == (None if allowed_domain else "jwt-owner")
@pytest.mark.asyncio
@pytest.mark.parametrize("route_allowed", [False, True])
async def test_oauth_jwt_identity_preserves_separate_mcp_route_authorization(
jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"],
monkeypatch: pytest.MonkeyPatch,
route_allowed: bool,
) -> None:
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server.bridge_token_flow import _extract_user_id_from_request
from litellm.proxy._types import LitellmUserRoles, RoleBasedPermissions, RoleMapping
from litellm.proxy.auth.handle_jwt import JWTAuthManager
handler, signing_key = jwt_oauth_identity
handler.litellm_jwtauth.user_id_jwt_field = "sub"
handler.litellm_jwtauth.roles_jwt_field = "aud"
handler.litellm_jwtauth.object_id_jwt_field = "identity.user_id"
handler.litellm_jwtauth.role_mappings = [
RoleMapping(role="litellm-proxy", internal_role=LitellmUserRoles.INTERNAL_USER)
]
handler.litellm_jwtauth.enforce_rbac = True
monkeypatch.setattr(
proxy_server,
"general_settings",
{
"enable_jwt_auth": True,
"role_permissions": [
RoleBasedPermissions(
role=LitellmUserRoles.INTERNAL_USER,
routes=["mcp_routes"] if route_allowed else ["/models"],
)
],
},
)
bearer: Final = _oauth_identity_jwt(signing_key)
request: Final = _token_request({"Authorization": f"Bearer {bearer}"}, path="/example/token")
assert await _extract_user_id_from_request(request) == "jwt-owner"
admission: Final = JWTAuthManager.auth_builder(
api_key=bearer,
jwt_handler=handler,
request_data={},
general_settings=proxy_server.general_settings,
route="/mcp/example",
prisma_client=proxy_server.prisma_client,
user_api_key_cache=handler.user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_server.proxy_logging_obj,
request_method="POST",
)
if route_allowed:
assert (await admission)["user_id"] == "jwt-owner"
else:
with pytest.raises(HTTPException) as denial:
await admission
assert denial.value.status_code == 403
@pytest.mark.asyncio
@pytest.mark.parametrize("identity", ["sso", "email"])
@pytest.mark.parametrize("inactive", [False, True])
@pytest.mark.parametrize("admin", [False, True])
async def test_oauth_jwt_resolves_canonical_owner_without_cached_identity(
jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"],
monkeypatch: pytest.MonkeyPatch,
identity: str,
inactive: bool,
admin: bool,
) -> None:
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server.bridge_token_flow import _extract_user_id_from_request
from litellm.proxy.auth.handle_jwt import JWTAuthManager
handler, signing_key = jwt_oauth_identity
external_id: Final = f"external-{identity}-{inactive}-{admin}"
handler.litellm_jwtauth.user_email_jwt_field = "email"
handler.litellm_jwtauth.admin_allowed_routes = ["mcp_routes"]
owner: Final = LiteLLM_UserTable(
user_id="canonical-oauth-owner",
user_email="owner@example.test",
metadata={"scim_active": not inactive},
organization_memberships=[],
)
database: Final = MagicMock()
table: Final = database.db.litellm_usertable
table.find_unique = AsyncMock(side_effect=[None, owner if identity == "sso" else None])
table.find_first = AsyncMock(return_value=owner)
table.update = AsyncMock(return_value=owner)
monkeypatch.setattr(proxy_server, "prisma_client", database)
bearer: Final = _oauth_identity_jwt(signing_key, owner=external_id, scope="litellm_proxy_admin" if admin else "")
request: Final = _token_request({"Authorization": f"Bearer {bearer}"})
stored_owner: Final = await _extract_user_id_from_request(request)
assert stored_owner == (None if inactive else external_id if admin else "canonical-oauth-owner")
assert table.find_unique.await_count == 2
if identity == "email":
table.find_first.assert_awaited_once()
if not inactive:
admission: Final = await JWTAuthManager.auth_builder(
api_key=bearer,
jwt_handler=handler,
request_data={},
general_settings=proxy_server.general_settings,
route="/mcp/example",
prisma_client=database,
user_api_key_cache=handler.user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_server.proxy_logging_obj,
)
assert stored_owner == admission["user_id"]
@pytest.mark.asyncio
async def test_oauth_jwt_identity_does_not_provision_or_synchronize_teams(
jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"],
) -> None:
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server.bridge_token_flow import _extract_user_id_from_request
handler, signing_key = jwt_oauth_identity
handler.litellm_jwtauth.enforce_team_based_model_access = True
handler.litellm_jwtauth.team_id_default = "new-team"
handler.litellm_jwtauth.team_id_upsert = True
handler.litellm_jwtauth.sync_user_role_and_teams = True
owner: Final = LiteLLM_UserTable(user_id="jwt-owner", teams=["existing-team"])
handler.user_api_key_cache.set_cache("jwt-owner", owner)
request: Final = _token_request(
{"Authorization": f"Bearer {_oauth_identity_jwt(signing_key)}"}, path="/example/token"
)
assert await _extract_user_id_from_request(request) == "jwt-owner"
assert owner.teams == ["existing-team"]
proxy_server.prisma_client.db.litellm_teamtable.find_unique.assert_not_called()
proxy_server.prisma_client.db.litellm_teamtable.upsert.assert_not_called()
proxy_server.prisma_client.db.litellm_usertable.update.assert_not_called()
@pytest.mark.asyncio
@pytest.mark.parametrize("state", ["active", "inactive", "missing_database"])
async def test_oauth_refresh_revalidates_the_same_active_user_rule(
jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"],
monkeypatch: pytest.MonkeyPatch,
state: str,
) -> None:
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server.bridge_token_flow import _reload_active_user_by_id
handler, _ = jwt_oauth_identity
handler.user_api_key_cache.set_cache(
"jwt-owner", LiteLLM_UserTable(user_id="jwt-owner", metadata={"scim_active": state != "inactive"})
)
if state == "missing_database":
monkeypatch.setattr(proxy_server, "prisma_client", None)
expected: Final = None if state == "active" else "no_active_key" if state == "inactive" else "unresolvable"
assert await _reload_active_user_by_id("jwt-owner") == expected
@pytest.mark.asyncio
@pytest.mark.parametrize("mapped", [False, True])
@pytest.mark.parametrize("state", ["allowed", "route_denied", "server_denied", "blocked", "expired", "lookup_error", "cancelled"])
async def test_oauth_credential_write_keeps_virtual_key_permissions(
jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"],
monkeypatch: pytest.MonkeyPatch,
mapped: bool,
state: str,
) -> None:
import asyncio
from litellm.proxy._experimental.mcp_server import mcp_server_manager
from litellm.proxy._experimental.mcp_server.bridge_token_flow import authorize_oauth_credential_request
from litellm.proxy._types import UserAPIKeyAuth, hash_token
from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key
handler, signing_key = jwt_oauth_identity
key: Final = "sk-oauth-permission-test"
hashed: Final = hash_token(key)
credential: Final = UserAPIKeyAuth(
token=hashed,
user_id="jwt-owner",
blocked=state == "blocked",
expires=datetime.now(timezone.utc) - timedelta(seconds=60) if state == "expired" else None,
allowed_routes=["openai_routes"] if state == "route_denied" else ["mcp_routes"],
agent_id="agent-scope",
org_id="org-scope",
end_user_id="end-user-scope",
)
handler.user_api_key_cache.set_cache(hashed, credential)
if mapped:
handler.litellm_jwtauth.virtual_key_claim_field = "sub"
handler.user_api_key_cache.set_cache(jwt_key_mapping_cache_key("sub", "not-the-configured-user-id"), hashed)
manager: Final = MagicMock()
manager.get_allowed_mcp_servers = AsyncMock(
return_value=[] if state == "server_denied" else ["server-a"],
side_effect=(asyncio.CancelledError() if state == "cancelled" else RuntimeError("permission lookup unavailable") if state == "lookup_error" else None),
)
monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager)
bearer: Final = _oauth_identity_jwt(signing_key) if mapped else key
request: Final = _token_request({"Authorization": f"Bearer {bearer}"}, path="/server-a/token")
if state == "cancelled":
with pytest.raises(asyncio.CancelledError):
await authorize_oauth_credential_request(request, "server-a")
manager.get_allowed_mcp_servers.assert_awaited_once()
return
assert await authorize_oauth_credential_request(request, "server-a") == ("jwt-owner" if state == "allowed" else None)
if state in ("allowed", "server_denied", "lookup_error"):
manager.get_allowed_mcp_servers.assert_awaited_once()
writer: Final = manager.get_allowed_mcp_servers.call_args.args[0]
assert (writer.user_id, writer.token, writer.org_id, writer.agent_id, writer.end_user_id) == (
"jwt-owner",
hashed,
"org-scope",
"agent-scope",
"end-user-scope",
)
@pytest.mark.asyncio
@pytest.mark.parametrize("server_id", ["team-a-server", "team-b-server"])
async def test_oauth_writer_preserves_claimed_team_instead_of_expanding_user_roster(
jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"],
monkeypatch: pytest.MonkeyPatch,
server_id: str,
) -> None:
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import mcp_server_manager
from litellm.proxy._experimental.mcp_server.bridge_token_flow import authorize_oauth_credential_request
from litellm.proxy._types import LiteLLM_TeamTable, Member
handler, signing_key = jwt_oauth_identity
handler.litellm_jwtauth.team_id_jwt_field = "team"
handler.litellm_jwtauth.team_id_upsert = True
handler.litellm_jwtauth.user_id_upsert = True
handler.litellm_jwtauth.sync_user_role_and_teams = True
handler.user_api_key_cache.set_cache("jwt-owner", LiteLLM_UserTable(user_id="jwt-owner", teams=["a", "b"]))
handler.user_api_key_cache.set_cache(
"team_id:a",
LiteLLM_TeamTable(team_id="a", models=[], members_with_roles=[Member(user_id="jwt-owner", role="user")]),
)
manager: Final = MagicMock()
manager.get_allowed_mcp_servers = AsyncMock(return_value=["team-a-server"])
monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager)
bearer: Final = _oauth_identity_jwt(signing_key, claims={"team": "a"})
request: Final = _token_request({"Authorization": f"Bearer {bearer}"}, path=f"/{server_id}/token")
assert await authorize_oauth_credential_request(request, server_id) == (
"jwt-owner" if server_id == "team-a-server" else None
)
manager.get_allowed_mcp_servers.assert_awaited_once()
writer: Final = manager.get_allowed_mcp_servers.call_args.args[0]
assert writer.team_id == "a"
assert not writer.mcp_admitted_user_subject
proxy_server.prisma_client.db.litellm_teamtable.create.assert_not_called()
proxy_server.prisma_client.db.litellm_usertable.create.assert_not_called()
proxy_server.prisma_client.db.litellm_usertable.update.assert_not_called()
assert handler.litellm_jwtauth.user_id_upsert and handler.litellm_jwtauth.team_id_upsert
assert handler.litellm_jwtauth.sync_user_role_and_teams
@pytest.mark.asyncio
async def test_oauth_write_denial_does_not_erase_identity_binding(
jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"], monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.proxy._experimental.mcp_server import discoverable_endpoints, mcp_server_manager
from litellm.proxy._types import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer
_, signing_key = jwt_oauth_identity
monkeypatch.setenv("LITELLM_SALT_KEY", "oauth-identity-binding-test-salt")
manager: Final = MagicMock()
manager.get_allowed_mcp_servers = AsyncMock(return_value=[])
monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager)
server: Final = MCPServer(
server_id="bound-server", name="bound-server", transport=MCPTransport.http,
auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", client_id="client",
token_url="https://upstream.example.test/token",
oauth_identity_binding=MCPOAuthIdentityBinding(
mode="enforce", issuer="https://upstream.example.test", audiences=["client"],
),
)
request: Final = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key)}"})
code: Final = discoverable_endpoints.seal_bridge_authorization_code(
"upstream-code", "another-owner", server.server_id, "bound-nonce",
)
with pytest.raises(HTTPException) as denied:
await discoverable_endpoints.exchange_token_with_server(
request=request, mcp_server=server, grant_type="authorization_code", code=code,
redirect_uri="http://localhost/callback", client_id="client", client_secret=None, code_verifier="verifier",
)
assert denied.value.status_code == 403
assert denied.value.detail == {"error": "oauth_principal_mismatch"}
manager.get_allowed_mcp_servers.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize("admin_only", [False, True])
async def test_signed_oauth_callback_honors_credential_write_policy(
jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"],
monkeypatch: pytest.MonkeyPatch,
admin_only: bool,
) -> None:
import httpx
import litellm
from litellm.caching.llm_caching_handler import LLMClientCache
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import discoverable_endpoints, mcp_server_manager
from litellm.proxy._types import MCPTransport
from litellm.types.llms.custom_http import httpxSpecialProvider
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server: Final = MCPServer(
server_id="signed-server", name="signed-server", transport=MCPTransport.http,
auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", client_id="client",
token_url="https://upstream.example.test/token",
)
monkeypatch.setattr(proxy_server, "general_settings", {
"enable_jwt_auth": True,
"admin_only_routes": [f"/v1/mcp/server/{server.server_id}/oauth-user-credential"] if admin_only else [],
})
monkeypatch.setenv("LITELLM_SALT_KEY", "signed-oauth-test-salt")
manager: Final = MagicMock()
manager.get_allowed_mcp_servers = AsyncMock(return_value=[server.server_id])
manager.invalidate_user_oauth_token_cache = AsyncMock()
monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager)
table: Final = proxy_server.prisma_client.db.litellm_mcpusercredentials
table.find_unique = AsyncMock(return_value=None)
table.upsert = AsyncMock()
clients: Final = LLMClientCache()
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", clients)
def upstream_response(outbound: httpx.Request) -> httpx.Response:
assert outbound.url == server.token_url
assert b"code=upstream-code" in outbound.content
return httpx.Response(200, json={"access_token": "upstream-token", "token_type": "Bearer"})
async with httpx.AsyncClient(transport=httpx.MockTransport(upstream_response)) as transport:
upstream: Final = AsyncHTTPHandler()
await upstream.client.aclose()
upstream.client = transport
clients.set_cache("async_httpx_client" + httpxSpecialProvider.Oauth2Check, upstream)
response: Final = await discoverable_endpoints.exchange_token_with_server(
request=_token_request({}, path="/signed-server/token"), mcp_server=server,
grant_type="authorization_code",
code=discoverable_endpoints.seal_bridge_authorization_code("upstream-code", "jwt-owner", server.server_id),
redirect_uri="http://localhost/callback", client_id="client", client_secret=None, code_verifier=None,
)
assert response.status_code == 200
assert json.loads(response.body)["access_token"] == "upstream-token"
if admin_only:
table.upsert.assert_not_awaited()
else:
table.upsert.assert_awaited_once()
assert table.upsert.call_args.kwargs["where"]["user_id_server_id"] == {
"user_id": "jwt-owner", "server_id": server.server_id,
}
@pytest.mark.asyncio
@pytest.mark.parametrize("allowed", [False, True])
@pytest.mark.parametrize("credential", [
"jwt", "key", "expired_jwt", "wrong_audience", "bad_signature", "malformed_jwt", "missing_issuer",
"foreign_explicit", "blank_explicit", "unknown_key", "blocked_key", "expired_key", "opaque_record",
"opaque_outage", "opaque_oidc", "opaque_custom", "foreign_unscoped", "foreign_configured", "encrypted", "invalid_encrypted", "envelope", "master",
])
async def test_identity_bound_authorize_preserves_presented_jwt_permissions(
jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"],
monkeypatch: pytest.MonkeyPatch,
allowed: bool,
credential: str,
) -> None:
import jwt
from datetime import datetime, timedelta, timezone
from urllib.parse import parse_qs, urlparse
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy import proxy_server
from litellm.proxy._types import JWTIssuerConfig, UserAPIKeyAuth, hash_token
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
from litellm.proxy._experimental.mcp_server import discoverable_endpoints, mcp_server_manager
from litellm.proxy._types import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer
handler, signing_key = jwt_oauth_identity
master: Final = "browser-session-test-signing-key-123456789"
monkeypatch.setattr(proxy_server, "master_key", master)
monkeypatch.setattr(proxy_server, "user_custom_auth", (lambda: None) if credential == "opaque_custom" else None)
handler.litellm_jwtauth.oidc_userinfo_enabled = credential == "opaque_oidc"
if credential == "foreign_unscoped":
monkeypatch.delenv("JWT_ISSUER")
if credential == "foreign_configured":
handler.litellm_jwtauth.issuers = [JWTIssuerConfig(
issuer="https://unrelated.example.test", jwks_url="https://idp.example.test/jwks",
audience="litellm-proxy", user_id_jwt_field="identity.user_id",
)]
proxy_server.prisma_client.get_data = AsyncMock(
return_value=None, side_effect=RuntimeError("database unavailable") if credential == "opaque_outage" else None,
)
handler.user_api_key_cache.set_cache("cookie-owner", LiteLLM_UserTable(user_id="cookie-owner"))
key: Final = "opaque-record" if credential == "opaque_record" else "sk-browser-gateway-key"
if credential in ("key", "blocked_key", "expired_key", "opaque_record"):
handler.user_api_key_cache.set_cache(hash_token(key), UserAPIKeyAuth(
token=hash_token(key), user_id="jwt-owner", blocked=credential in ("blocked_key", "opaque_record"),
expires=datetime.now(timezone.utc) - timedelta(seconds=60) if credential == "expired_key" else None,
))
monkeypatch.setenv("LITELLM_SALT_KEY", "authorize-policy-test-salt")
server: Final = MCPServer(
server_id="bound-server", name="bound-server", transport=MCPTransport.http,
auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", client_id="client",
authorization_url="https://upstream.example.test/authorize", token_url="https://upstream.example.test/token",
oauth_identity_binding=MCPOAuthIdentityBinding(
mode="enforce", issuer="https://upstream.example.test", audiences=["client"],
),
)
manager: Final = MagicMock()
# The full user roster permits the server; the presented JWT may have narrower access.
manager.get_allowed_mcp_servers = AsyncMock(
side_effect=lambda auth: [server.server_id] if allowed or auth.mcp_admitted_user_subject else [],
)
monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager)
bearer: Final = (
key if credential in ("key", "blocked_key", "expired_key", "opaque_record", "unknown_key")
else "opaque-bearer" if credential in ("opaque_outage", "opaque_oidc", "opaque_custom")
else "not.a.jwt" if credential == "malformed_jwt"
else "llm_env_invalid" if credential == "envelope"
else "v2:gcm:invalid" if credential == "invalid_encrypted"
else master if credential == "master"
else ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(
LiteLLM_UserTable(user_id="jwt-owner", user_role="internal_user"),
) if credential == "encrypted"
else jwt.encode({"iss": "https://idp.example.test"}, "wrong-signing-key-at-least-32-bytes", algorithm="HS256")
if credential == "bad_signature"
else jwt.encode({"sub": "jwt-owner"}, signing_key, algorithm="RS256") if credential == "missing_issuer"
else _oauth_identity_jwt(
signing_key,
expires_in=-60 if credential == "expired_jwt" else 300,
audience="another-service" if credential == "wrong_audience" else "litellm-proxy",
issuer="https://unrelated.example.test" if credential.startswith("foreign_") or credential == "blank_explicit" else "https://idp.example.test",
)
)
cookie: Final = jwt.encode(
{"user_id": "cookie-owner", "login_method": "sso", "exp": int(time.time()) + 300}, master, algorithm="HS256",
)
response: Final = await discoverable_endpoints.authorize_with_server(
request=_token_request({
"Authorization": f"Bearer {bearer}", "Cookie": f"token={cookie}",
**({"x-litellm-api-key": bearer} if credential == "foreign_explicit" else {}),
**({"x-litellm-api-key": ""} if credential == "blank_explicit" else {}),
}),
mcp_server=server, client_id="client", redirect_uri="http://127.0.0.1:6274/callback",
state="client-state", code_challenge="pkce-challenge", code_challenge_method="S256",
)
redirect: Final = urlparse(response.headers["location"])
query: Final = parse_qs(redirect.query)
if allowed and credential in ("jwt", "key", "foreign_unscoped", "foreign_configured"):
assert redirect.hostname == "upstream.example.test"
assert query["nonce"] and response.headers.get("set-cookie")
assert all(call.args[0].user_id == "jwt-owner" for call in manager.get_allowed_mcp_servers.await_args_list)
else:
assert redirect.hostname == "127.0.0.1"
assert query["error"] == ["access_denied"]
assert query["state"] == ["client-state"]
assert "set-cookie" not in response.headers
proxy_server.prisma_client.db.litellm_mcpusercredentials.upsert.assert_not_called()
proxy_server.prisma_client.db.litellm_usertable.create.assert_not_called()
proxy_server.prisma_client.db.litellm_teamtable.create.assert_not_called()
@pytest.mark.asyncio
@pytest.mark.parametrize("credential", ["none", "opaque", "foreign_jwt"])
@pytest.mark.parametrize("cookie_state", ["allowed", "server_denied", "expired", "missing"])
async def test_identity_bound_authorize_unrelated_bearer_uses_browser_session(
jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"],
monkeypatch: pytest.MonkeyPatch,
credential: str,
cookie_state: str,
) -> None:
import jwt
from urllib.parse import parse_qs, urlparse
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import discoverable_endpoints, mcp_server_manager
from litellm.proxy._types import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer
handler, signing_key = jwt_oauth_identity
master: Final = "browser-session-test-signing-key-123456789"
monkeypatch.setattr(proxy_server, "master_key", master)
monkeypatch.setattr(proxy_server, "user_custom_auth", None)
monkeypatch.setenv("LITELLM_SALT_KEY", "authorize-policy-test-salt")
handler.user_api_key_cache.set_cache("cookie-owner", LiteLLM_UserTable(user_id="cookie-owner"))
proxy_server.prisma_client.get_data = AsyncMock(return_value=None)
server: Final = MCPServer(
server_id="bound-server", name="bound-server", transport=MCPTransport.http,
auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", client_id="client",
authorization_url="https://upstream.example.test/authorize", token_url="https://upstream.example.test/token",
oauth_identity_binding=MCPOAuthIdentityBinding(
mode="enforce", issuer="https://upstream.example.test", audiences=["client"],
),
)
manager: Final = MagicMock()
manager.get_allowed_mcp_servers = AsyncMock(return_value=[] if cookie_state == "server_denied" else [server.server_id])
monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager)
bearer: Final = (
_oauth_identity_jwt(signing_key, issuer="https://unrelated.example.test")
if credential == "foreign_jwt" else "unrelated-upstream-bearer"
)
cookie: Final = jwt.encode(
{"user_id": "cookie-owner", "login_method": "sso", "exp": int(time.time()) + (-60 if cookie_state == "expired" else 300)},
master, algorithm="HS256",
)
response: Final = await discoverable_endpoints.authorize_with_server(
request=_token_request({
**({"Authorization": f"Bearer {bearer}"} if credential != "none" else {}),
**({"Cookie": f"token={cookie}"} if cookie_state != "missing" else {}),
}),
mcp_server=server, client_id="client", redirect_uri="http://127.0.0.1:6274/callback",
state="client-state", code_challenge="pkce-challenge", code_challenge_method="S256",
)
redirect: Final = urlparse(response.headers["location"])
query: Final = parse_qs(redirect.query)
if cookie_state == "allowed":
assert redirect.hostname == "upstream.example.test"
assert query["nonce"] and response.headers.get("set-cookie")
manager.get_allowed_mcp_servers.assert_awaited_once()
assert manager.get_allowed_mcp_servers.call_args.args[0].user_id == "cookie-owner"
elif cookie_state == "server_denied":
assert query["error"] == ["access_denied"]
assert query["state"] == ["client-state"]
else:
assert redirect.path == "/sso/key/generate"
manager.get_allowed_mcp_servers.assert_not_awaited()
proxy_server.prisma_client.db.litellm_mcpusercredentials.upsert.assert_not_called()
proxy_server.prisma_client.db.litellm_usertable.create.assert_not_called()
proxy_server.prisma_client.db.litellm_teamtable.create.assert_not_called()

View file

@ -2,7 +2,7 @@ import asyncio
import re
import time
from collections.abc import Mapping, Sequence
from typing import Optional
from typing import Final, Optional
from unittest.mock import AsyncMock, MagicMock, patch
from fastapi import HTTPException
@ -6790,6 +6790,88 @@ async def test_sync_user_role_and_teams_singular_claim_only_recognized_under_fla
assert user.teams == []
@pytest.mark.asyncio
@pytest.mark.parametrize("operation", ["identity", "authorize", "admit"])
@pytest.mark.parametrize("existing_user", [False, True])
@pytest.mark.parametrize("model_allowed", [False, True])
async def test_jwt_identity_and_authorization_keep_provisioning_in_admission(
monkeypatch: pytest.MonkeyPatch, operation: str, existing_user: bool, model_allowed: bool
) -> None:
from litellm.proxy._types import ScopeMapping
from litellm.proxy.auth.auth_checks import UserNotFoundError
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
private_key, jwk = _get_rsa_key_and_jwk("identity-mode")
cache: Final = UserApiKeyCache()
cache.set_cache("litellm_jwt_auth_keys_https://identity.example/jwks", [jwk])
user_id: Final = f"identity-mode-{operation}-{existing_user}-{model_allowed}"
user: Final = LiteLLM_UserTable(user_id=user_id, organization_memberships=[])
if existing_user:
cache.set_cache(user_id, user)
database: Final = MagicMock()
users: Final = database.db.litellm_usertable
users.find_unique = AsyncMock(return_value=None)
users.find_first = AsyncMock(return_value=None)
users.create = AsyncMock(return_value=user)
handler: Final = JWTHandler()
handler.update_environment(
prisma_client=database,
user_api_key_cache=cache,
litellm_jwtauth=LiteLLM_JWTAuth(
user_id_jwt_field="sub",
user_id_upsert=True,
enforce_scope_based_access=True,
scope_mappings=[ScopeMapping(scope="allowed", models=["allowed-model"])],
),
)
monkeypatch.setenv("JWT_PUBLIC_KEY_URL", "https://identity.example/jwks")
monkeypatch.setenv("JWT_ISSUER", "https://identity.example")
monkeypatch.setenv("JWT_AUDIENCE", "gateway")
token: Final = _encode_rsa_jwt(
private_key, "https://identity.example", "gateway", "identity-mode", {"sub": user_id, "scope": "allowed"}
)
common: Final = {
"api_key": token,
"jwt_handler": handler,
"prisma_client": database,
"user_api_key_cache": cache,
"parent_otel_span": None,
"proxy_logging_obj": MagicMock(),
}
if operation == "identity":
if not existing_user:
with pytest.raises(UserNotFoundError):
await JWTAuthManager.resolve_identity(**common)
else:
identity: Final = await JWTAuthManager.resolve_identity(**common)
assert identity.user_id == user_id
assert identity.user_object is not None and identity.user_object.user_id == user_id
users.create.assert_not_awaited()
return
authorize: Final = JWTAuthManager.auth_builder if operation == "admit" else JWTAuthManager.authorize_jwt
pending: Final = authorize(
**common,
request_data={"model": "allowed-model" if model_allowed else "forbidden-model"},
general_settings={},
route="/mcp/example",
)
if not model_allowed:
with pytest.raises(HTTPException) as denial:
await pending
assert denial.value.status_code == 403
users.create.assert_not_awaited()
return
if operation == "authorize" and not existing_user:
with pytest.raises(UserNotFoundError):
await pending
else:
result: Final = await pending
assert result["user_id"] == user_id
assert result["user_object"] is not None
assert result["user_object"].user_id == user_id
assert users.create.await_count == (0 if operation == "authorize" or existing_user else 1)
def _entra_agent_registry() -> AgentRegistry:
registry = AgentRegistry()
registry.register_agent(
@ -6916,7 +6998,8 @@ def _entra_signed_app_token(monkeypatch, azp: str, scope: str) -> tuple[JWTHandl
@pytest.mark.asyncio
@pytest.mark.parametrize("is_admin_token", [False, True], ids=["standard_jwt", "proxy_admin_jwt"])
async def test_auth_builder_propagates_agent_id_from_jwt_claim(monkeypatch, is_admin_token: bool):
@pytest.mark.parametrize("identity_only", [False, True])
async def test_auth_builder_propagates_agent_id_from_jwt_claim(monkeypatch, is_admin_token: bool, identity_only: bool):
"""auth_builder carries the resolved agent id into JWTAuthBuilderResult on both the admin and standard paths."""
jwt_handler, token = _entra_signed_app_token(
monkeypatch,
@ -6925,6 +7008,14 @@ async def test_auth_builder_propagates_agent_id_from_jwt_claim(monkeypatch, is_a
)
jwt_handler.bind_agent_lookup(_entra_agent_registry())
if identity_only:
identity = await JWTAuthManager.resolve_identity(
api_key=token, jwt_handler=jwt_handler, prisma_client=None,
user_api_key_cache=None, parent_otel_span=None, proxy_logging_obj=None,
)
assert identity.agent_id == "canonical-agent-id"
return
result = await JWTAuthManager.auth_builder(
api_key=token,
jwt_handler=jwt_handler,
@ -6942,7 +7033,8 @@ async def test_auth_builder_propagates_agent_id_from_jwt_claim(monkeypatch, is_a
@pytest.mark.asyncio
async def test_auth_builder_denies_jwt_naming_unregistered_agent_before_admin_check(monkeypatch):
@pytest.mark.parametrize("identity_only", [False, True])
async def test_auth_builder_denies_jwt_naming_unregistered_agent_before_admin_check(monkeypatch, identity_only: bool):
"""An unknown agent claim is rejected even when the token would otherwise be a proxy admin."""
jwt_handler, token = _entra_signed_app_token(
monkeypatch,
@ -6951,6 +7043,14 @@ async def test_auth_builder_denies_jwt_naming_unregistered_agent_before_admin_ch
)
jwt_handler.bind_agent_lookup(_entra_agent_registry())
if identity_only:
with pytest.raises(HTTPException) as denial:
await JWTAuthManager.resolve_identity(
api_key=token, jwt_handler=jwt_handler, prisma_client=None,
user_api_key_cache=None, parent_otel_span=None, proxy_logging_obj=None,
)
assert denial.value.status_code == 403
return
with pytest.raises(HTTPException) as exc_info:
await JWTAuthManager.auth_builder(
api_key=token,
@ -6965,3 +7065,36 @@ async def test_auth_builder_denies_jwt_naming_unregistered_agent_before_admin_ch
)
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
@pytest.mark.parametrize("admission", [False, True])
async def test_admin_jwt_team_header_only_provisions_during_admission(monkeypatch, admission: bool):
from litellm.proxy.management_endpoints import team_endpoints
handler, token = _entra_signed_app_token(
monkeypatch, azp="canonical-agent-id", scope=LiteLLM_JWTAuth().admin_jwt_scope,
)
handler.bind_agent_lookup(_entra_agent_registry())
handler.litellm_jwtauth.team_id_upsert = True
handler.litellm_jwtauth.admin_allowed_routes = ["openai_routes"]
database = MagicMock()
database.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
create_team = AsyncMock(return_value=LiteLLM_TeamTable(team_id="new-team").model_dump())
monkeypatch.setattr(team_endpoints, "new_team", create_team)
resolve = JWTAuthManager.auth_builder if admission else JWTAuthManager.authorize_jwt
result = await resolve(
api_key=token, jwt_handler=handler, request_data={}, general_settings={},
route="/chat/completions", prisma_client=database,
user_api_key_cache=handler.user_api_key_cache, parent_otel_span=None,
proxy_logging_obj=MagicMock(), request_headers={"x-litellm-team-id": "new-team"},
)
assert result["is_proxy_admin"] is True
if admission:
create_team.assert_awaited_once()
assert result["team_id"] == "new-team"
else:
create_team.assert_not_awaited()
assert result["team_id"] is None

View file

@ -209,6 +209,8 @@ async def test_get_aggregated_daily_spend_update_transactions_same_key():
"prompt_caching_savings_spend": 0,
"gateway_injected_caching_savings_spend": 0,
"autorouter_savings_spend": 0,
"total_response_time_ms": 0,
"timed_requests": 0,
}
updates = [{test_key: test_transaction1}, {test_key: test_transaction2}]
@ -261,6 +263,8 @@ async def test_flush_and_get_aggregated_daily_spend_update_transactions(
"prompt_caching_savings_spend": 0,
"gateway_injected_caching_savings_spend": 0,
"autorouter_savings_spend": 0,
"total_response_time_ms": 0,
"timed_requests": 0,
}
# Add updates to queue
@ -550,7 +554,7 @@ async def test_every_optional_daily_metric_aggregates(daily_spend_update_queue):
numeric_fields = [
name for name, annotation in BaseDailySpendTransaction.__annotations__.items() if _numeric(annotation)
]
assert "autorouter_savings_spend" in numeric_fields
assert {"autorouter_savings_spend", "total_response_time_ms", "timed_requests"} <= set(numeric_fields)
increments = {field: index + 1 for index, field in enumerate(numeric_fields)}
await daily_spend_update_queue.add_update({test_key: dict(increments)})
@ -579,8 +583,12 @@ async def test_optional_metric_missing_from_an_older_payload_still_aggregates(
}
await daily_spend_update_queue.add_update({test_key: dict(base)})
await daily_spend_update_queue.add_update({test_key: {**base, "autorouter_savings_spend": 0.25}})
await daily_spend_update_queue.add_update(
{test_key: {**base, "autorouter_savings_spend": 0.25, "total_response_time_ms": 900, "timed_requests": 1}}
)
await daily_spend_update_queue.aggregate_queue_updates()
updates = await daily_spend_update_queue.flush_all_updates_from_in_memory_queue()
assert updates[0][test_key]["autorouter_savings_spend"] == pytest.approx(0.25)
assert updates[0][test_key]["total_response_time_ms"] == 900
assert updates[0][test_key]["timed_requests"] == 1

View file

@ -85,10 +85,10 @@ def test_one_statement_carries_every_row_in_the_batch():
assert sql.count("INSERT INTO") == 1
assert len(re.findall(r"ON CONFLICT", sql)) == 1
# 23 bound columns per row plus the inlined updated_at, so the row count is what
# 25 bound columns per row plus the inlined updated_at, so the row count is what
# separates one multi-row statement from a hundred single-row ones.
assert len(params) == 100 * 23
assert "$2300::text" in sql
assert len(params) == 100 * 25
assert "$2500::text" in sql
assert sql.count("(NOW() AT TIME ZONE 'UTC')") == 100 + 1
@ -104,7 +104,16 @@ def test_conflict_target_is_the_full_unique_constraint():
@pytest.mark.parametrize(
"column",
["prompt_tokens", "completion_tokens", "spend", "api_requests", "successful_requests", "failed_requests"],
[
"prompt_tokens",
"completion_tokens",
"spend",
"api_requests",
"successful_requests",
"failed_requests",
"total_response_time_ms",
"timed_requests",
],
)
def test_counters_increment_rather_than_overwrite(column):
"""An overwrite would silently discard every earlier flush's spend for that row."""

View file

@ -2864,6 +2864,76 @@ async def test_daily_transaction_internal_call_keeps_spend_but_not_request_count
assert user_sent["successful_requests"] == 1
def _response_time_payload(request_duration_ms: object, metadata: dict | None = None) -> dict:
return {
"request_id": "req-timed-1",
"user": "test-user",
"startTime": "2026-09-15T00:00:00",
"api_key": "test-key",
"model": "gpt-5.5",
"custom_llm_provider": "openai",
"model_group": "gpt-5.5",
"call_type": "acompletion",
"prompt_tokens": 10,
"completion_tokens": 5,
"spend": 0.01,
"request_duration_ms": request_duration_ms,
"metadata": json.dumps(metadata or {}),
}
@pytest.mark.asyncio
@pytest.mark.parametrize("request_duration_ms", [1234, 0])
async def test_daily_transaction_rolls_up_response_time_for_successful_requests(request_duration_ms: int):
"""A successful user-sent request contributes its request_duration_ms to the daily
response-time sum and counts as one timed request, including a 0 ms duration."""
writer = DBSpendUpdateWriter()
mock_prisma = MagicMock()
mock_prisma.get_request_status = MagicMock(return_value="success")
transaction = await writer._common_add_spend_log_transaction_to_daily_transaction(
payload=_response_time_payload(request_duration_ms),
prisma_client=mock_prisma,
type="user",
)
assert transaction is not None
assert transaction["total_response_time_ms"] == request_duration_ms
assert transaction["timed_requests"] == 1
@pytest.mark.asyncio
@pytest.mark.parametrize(
("request_status", "request_duration_ms", "metadata"),
[
("failure", 1234, {}),
("success", None, {}),
("success", -5, {}),
("success", "1234", {}),
("success", 1234, {"internal_call_origin": "shadow_eval_judge"}),
],
ids=["failed", "missing", "negative", "non_int", "internal_call"],
)
async def test_daily_transaction_excludes_untimed_requests_from_response_time(
request_status: str, request_duration_ms: object, metadata: dict
):
"""Failed, internal, and missing/invalid-duration requests never enter the response-time
average: both the duration sum and the timed_requests denominator stay at zero."""
writer = DBSpendUpdateWriter()
mock_prisma = MagicMock()
mock_prisma.get_request_status = MagicMock(return_value=request_status)
transaction = await writer._common_add_spend_log_transaction_to_daily_transaction(
payload=_response_time_payload(request_duration_ms, metadata),
prisma_client=mock_prisma,
type="user",
)
assert transaction is not None
assert transaction["total_response_time_ms"] == 0
assert transaction["timed_requests"] == 0
def _deadlock_error():
from prisma.errors import RawQueryError

View file

@ -1,11 +1,14 @@
"""Unit tests for the LLM-as-a-Judge guardrail hook."""
import json
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
import litellm
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge import (
LLMAsAJudgeGuardrail,
_build_judge_prompt,
@ -13,7 +16,8 @@ from litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge import (
_parse_judge_verdict,
initialize_guardrail,
)
from litellm.types.guardrails import GuardrailEventHooks, Mode
from litellm.types.utils import LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN
# ---------------------------------------------------------------------------
# Helpers
@ -136,17 +140,314 @@ def test_initialize_guardrail_invalid_on_failure():
initialize_guardrail(lp, g)
@pytest.mark.parametrize(
("mode", "runs_pre_call", "runs_post_call"),
[
("pre_call", True, False),
(["pre_call", "post_call"], True, True),
(Mode(tags={"judge": ["pre_call"]}, default="post_call"), True, False),
(None, False, True),
],
ids=["scalar", "list", "tagged", "missing"],
)
def test_initialize_guardrail_preserves_every_mode_shape(
mode: str | list[str] | Mode | None,
runs_pre_call: bool,
runs_post_call: bool,
):
lp: Final = _make_litellm_params(mode=mode)
instance: Final = initialize_guardrail(lp, _make_guardrail_dict())
request_data: Final[dict[str, object]] = {"metadata": {"guardrails": ["g"], "tags": ["judge"]}}
premium: Final = patch("litellm.proxy.proxy_server.premium_user", True) # test-quality-ok: no seam for Mode tags
try:
with premium:
assert instance.should_run_guardrail(request_data, GuardrailEventHooks.pre_call) is runs_pre_call
assert instance.should_run_guardrail(request_data, GuardrailEventHooks.post_call) is runs_post_call
finally:
litellm.logging_callback_manager.remove_callback_from_all_lists(instance)
def test_initialize_guardrail_rejects_unknown_mode():
lp: Final = _make_litellm_params(mode="sometimes")
with pytest.raises(ValueError, match="sometimes"):
initialize_guardrail(lp, _make_guardrail_dict())
# ---------------------------------------------------------------------------
# apply_guardrail — enforcement paths
# ---------------------------------------------------------------------------
def _judge_router(overall_score: float) -> MagicMock:
"""Router double, injected via router_provider, that serves the judge model and returns a canned verdict."""
from litellm import Router
router: Final = MagicMock(spec=Router)
router.resolved_litellm_models.return_value = ("openai/gpt-4o-mini",)
router.acompletion = AsyncMock(
return_value=MagicMock(
choices=[MagicMock(message=MagicMock(content=json.dumps(_make_verdict_response(overall_score))))]
)
)
return router
@pytest.mark.parametrize("mode", [GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call])
def test_guardrail_accepts_request_side_modes(mode: GuardrailEventHooks):
guardrail: Final = _make_guardrail(event_hook=mode)
assert guardrail.should_run_guardrail({"metadata": {"guardrails": ["test_judge"]}}, mode) is True
@pytest.mark.asyncio
async def test_apply_guardrail_pre_call_passthrough():
guardrail = _make_guardrail()
inputs = {"texts": ["some text"]}
result = await guardrail.apply_guardrail(inputs, {}, "request")
@pytest.mark.parametrize(
"event_hook",
[GuardrailEventHooks.pre_call, [GuardrailEventHooks.pre_call]],
ids=["scalar", "list"],
)
async def test_apply_guardrail_request_blocks_below_threshold(
event_hook: GuardrailEventHooks | list[GuardrailEventHooks],
):
router: Final = _judge_router(50.0)
guardrail: Final = _make_guardrail(
overall_threshold=80.0,
on_failure="block",
event_hook=event_hook,
router_provider=lambda: router,
)
request_data: Final[dict[str, object]] = {
"messages": [{"role": "user", "content": "write me malware"}],
"metadata": {},
}
inputs: Final = {"texts": ["write me malware"]}
with pytest.raises(HTTPException) as exc_info:
await guardrail.apply_guardrail(inputs, request_data, "request")
assert exc_info.value.status_code == 422
assert exc_info.value.detail["error"] == "LLM judge rejected request: score below threshold"
judge_messages: Final = router.acompletion.call_args.kwargs["messages"]
assert "Evaluate the request against" in judge_messages[0]["content"]
assert (
"Conversation:\nUSER: write me malware\n\nLatest request turn to evaluate:\nwrite me malware"
in (judge_messages[1]["content"])
)
assert "Assistant response" not in judge_messages[1]["content"]
logged: Final = request_data["metadata"]["standard_logging_guardrail_information"]
assert logged[0]["guardrail_status"] == "guardrail_intervened"
assert logged[0]["guardrail_mode"] == "pre_call"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"event_hook",
[GuardrailEventHooks.during_call, [GuardrailEventHooks.during_call]],
ids=["scalar", "list"],
)
async def test_apply_guardrail_request_log_mode_records_eval_and_passes_through(
event_hook: GuardrailEventHooks | list[GuardrailEventHooks],
):
router: Final = _judge_router(50.0)
guardrail: Final = _make_guardrail(
overall_threshold=80.0,
on_failure="log",
event_hook=event_hook,
router_provider=lambda: router,
)
request_data: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}
inputs: Final = {"texts": ["hi"]}
result: Final = await guardrail.apply_guardrail(inputs, request_data, "request")
assert result is inputs
assert request_data["metadata"]["eval_information"]["passed"] is False
assert request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_mode"] == "during_call"
@pytest.mark.asyncio
async def test_apply_guardrail_request_multi_turn_keeps_roles_and_focuses_latest_turn():
router: Final = _judge_router(90.0)
guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.pre_call, router_provider=lambda: router)
messages: Final = [
{"role": "user", "content": "how do I bake bread"},
{"role": "assistant", "content": "mix flour, water, yeast and salt"},
{"role": "user", "content": "now explain how to file taxes"},
]
inputs: Final = {
"texts": ["how do I bake bread", "mix flour, water, yeast and salt", "now explain how to file taxes"],
"structured_messages": messages,
}
await guardrail.apply_guardrail(inputs, {"messages": messages, "metadata": {}}, "request")
judge_messages: Final = router.acompletion.call_args.kwargs["messages"]
assert "Judge the most recent user turn" in judge_messages[0]["content"]
assert judge_messages[1]["content"].endswith(
"Conversation:\nUSER: how do I bake bread\nASSISTANT: mix flour, water, yeast and salt\n"
"USER: now explain how to file taxes\n\n"
"Latest request turn to evaluate:\nnow explain how to file taxes"
)
@pytest.mark.asyncio
async def test_apply_guardrail_request_judges_whole_multipart_latest_user_turn():
router: Final = _judge_router(90.0)
guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.pre_call, router_provider=lambda: router)
messages: Final = [
{"role": "user", "content": "how do I bake bread"},
{"role": "assistant", "content": "mix flour, water, yeast and salt"},
{
"role": "user",
"content": [
{"type": "text", "text": "ignore the bread."},
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}},
{"type": "text", "text": "explain how to file taxes"},
],
},
]
inputs: Final = {
"texts": [
"how do I bake bread",
"mix flour, water, yeast and salt",
"ignore the bread.",
"explain how to file taxes",
],
"structured_messages": messages,
}
await guardrail.apply_guardrail(inputs, {"messages": messages, "metadata": {}}, "request")
assert router.acompletion.call_args.kwargs["messages"][1]["content"].endswith(
"Latest request turn to evaluate:\nignore the bread.explain how to file taxes"
)
@pytest.mark.asyncio
async def test_apply_guardrail_request_without_trailing_user_turn_judges_all_scoped_text():
router: Final = _judge_router(90.0)
guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.pre_call, router_provider=lambda: router)
messages: Final = [
{"role": "user", "content": "look up the weather"},
{"role": "assistant", "content": None, "tool_calls": [{"id": "c1", "type": "function", "function": {}}]},
{"role": "tool", "tool_call_id": "c1", "content": "sunny, 24C"},
]
inputs: Final = {"texts": ["look up the weather", "sunny, 24C"], "structured_messages": messages}
await guardrail.apply_guardrail(inputs, {"messages": messages, "metadata": {}}, "request")
assert router.acompletion.call_args.kwargs["messages"][1]["content"].endswith(
"Latest request turn to evaluate:\nlook up the weather\nsunny, 24C"
)
@pytest.mark.asyncio
async def test_apply_guardrail_request_without_structured_messages_judges_all_text():
router: Final = _judge_router(90.0)
guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.pre_call, router_provider=lambda: router)
await guardrail.apply_guardrail({"texts": ["first", "second"]}, {"metadata": {}}, "request")
assert router.acompletion.call_args.kwargs["messages"][1]["content"].endswith(
"Latest request turn to evaluate:\nfirst\nsecond"
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("modes", "input_type"),
[
([GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call], "request"),
([GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only], "request"),
([GuardrailEventHooks.post_call, GuardrailEventHooks.logging_only], "response"),
],
)
async def test_apply_guardrail_with_ambiguous_modes_logs_configured_mode(
modes: list[GuardrailEventHooks], input_type: str
):
router: Final = _judge_router(90.0)
guardrail: Final = _make_guardrail(event_hook=modes, router_provider=lambda: router)
request_data: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}
await guardrail.apply_guardrail({"texts": ["hi"]}, request_data, input_type)
assert request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_mode"] == [
mode.value for mode in modes
]
@pytest.mark.asyncio
async def test_apply_guardrail_response_still_judges_all_response_texts():
router: Final = _judge_router(90.0)
guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.post_call, router_provider=lambda: router)
await guardrail.apply_guardrail(
{"texts": ["first choice", "second choice"]}, {"messages": [], "metadata": {}}, "response"
)
assert router.acompletion.call_args.kwargs["messages"][1]["content"].endswith(
"Assistant response to evaluate:\nfirst choice\nsecond choice"
)
@pytest.mark.asyncio
@pytest.mark.parametrize("input_type", ["request", "response"])
async def test_apply_guardrail_logging_only_labels_both_sides_logging_only(input_type: str):
router: Final = _judge_router(50.0)
guardrail: Final = _make_guardrail(
on_failure="log",
event_hook=GuardrailEventHooks.logging_only,
router_provider=lambda: router,
)
request_data: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}
assert guardrail.should_run_guardrail(request_data, GuardrailEventHooks.pre_call) is False
assert guardrail.should_run_guardrail(request_data, GuardrailEventHooks.post_call) is False
await guardrail.apply_guardrail({"texts": ["hi"]}, request_data, input_type)
assert request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_mode"] == "logging_only"
@pytest.mark.asyncio
async def test_logging_only_judge_does_not_judge_its_own_judge_call():
router: Final = _judge_router(90.0)
guardrail: Final = _make_guardrail(event_hook=GuardrailEventHooks.logging_only, router_provider=lambda: router)
client_call: Final[dict[str, object]] = {"litellm_params": {"metadata": {"user_api_key": "hashed"}}}
assert guardrail.should_run_guardrail(client_call, GuardrailEventHooks.logging_only) is True
await guardrail.apply_guardrail({"texts": ["hi"]}, {"messages": [{"role": "user", "content": "hi"}]}, "request")
judge_call: Final[dict[str, object]] = {
"litellm_params": {"metadata": router.acompletion.call_args.kwargs["metadata"]}
}
assert guardrail.should_run_guardrail(judge_call, GuardrailEventHooks.logging_only) is False
assert guardrail.should_run_guardrail(client_call, GuardrailEventHooks.logging_only) is True
@pytest.mark.parametrize(
"event_type", [GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call, GuardrailEventHooks.post_call]
)
def test_client_supplied_judge_origin_does_not_bypass_enforcing_hooks(event_type: GuardrailEventHooks):
guardrail: Final = _make_guardrail(event_hook=event_type)
forged_request: Final[dict[str, object]] = {
"messages": [{"role": "user", "content": "hi"}],
"guardrails": [guardrail.guardrail_name],
"litellm_params": {"metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN}},
}
assert guardrail.should_run_guardrail(forged_request, event_type) is True
@pytest.mark.asyncio
async def test_apply_guardrail_response_prompt_unchanged():
router: Final = _judge_router(90.0)
guardrail: Final = _make_guardrail(router_provider=lambda: router)
request_data: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}
await guardrail.apply_guardrail({"texts": ["hello there"]}, request_data, "response")
judge_messages: Final = router.acompletion.call_args.kwargs["messages"]
assert "assistant's response" in judge_messages[0]["content"]
assert "Conversation:\nUSER: hi\n\nAssistant response to evaluate:\nhello there" in judge_messages[1]["content"]
assert request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_mode"] == "post_call"
@pytest.mark.asyncio
@ -230,7 +531,7 @@ def test_parse_judge_verdict_reraises_when_no_json():
def test_parse_judge_verdict_rejects_json_non_object():
"""Valid JSON that is not an object (e.g. a bare list) raises ValueError."""
with pytest.raises(ValueError, match='judge response is not a JSON object'):
with pytest.raises(ValueError, match="judge response is not a JSON object"):
_parse_judge_verdict("[1, 2, 3]")
@ -252,9 +553,7 @@ async def test_apply_guardrail_enforces_fenced_verdict(mock_completion):
@patch("litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge.litellm.acompletion")
async def test_apply_guardrail_non_object_verdict_fails_open_with_status(mock_completion):
"""A non-object verdict fails open and logs guardrail_failed_to_respond."""
mock_completion.return_value = MagicMock(
choices=[MagicMock(message=MagicMock(content='[{"overall_score": 50}]'))]
)
mock_completion.return_value = MagicMock(choices=[MagicMock(message=MagicMock(content='[{"overall_score": 50}]'))])
guardrail = _make_guardrail(overall_threshold=80.0, on_failure="block", router_provider=lambda: None)
inputs = {"texts": ["response"]}
request_data: dict = {"messages": [], "metadata": {}}
@ -314,7 +613,12 @@ def _real_router(model_list, **router_kwargs):
"model_list, router_kwargs, judge_model",
[
(
[{"model_name": "my-judge-alias", "litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"}}],
[
{
"model_name": "my-judge-alias",
"litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"},
}
],
{},
"my-judge-alias",
),
@ -324,12 +628,22 @@ def _real_router(model_list, **router_kwargs):
"anthropic/claude-sonnet-4-6",
),
(
[{"model_name": "backing-group", "litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"}}],
[
{
"model_name": "backing-group",
"litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"},
}
],
{"model_group_alias": {"my-judge-alias": "backing-group"}},
"my-judge-alias",
),
(
[{"model_name": "backing-group", "litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"}}],
[
{
"model_name": "backing-group",
"litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"},
}
],
{"model_group_alias": {"my-judge-alias": {"model": "backing-group", "hidden": True}}},
"my-judge-alias",
),
@ -412,7 +726,12 @@ async def test_judge_resolves_router_lazily_per_call(mock_sdk_completion):
mock_sdk_completion.assert_awaited_once()
holder["router"] = _real_router(
[{"model_name": "my-judge-alias", "litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"}}]
[
{
"model_name": "my-judge-alias",
"litellm_params": {"model": "anthropic/claude-sonnet-4-6", "api_key": "sk-ant-test"},
}
]
)
await guardrail.apply_guardrail({"texts": ["r"]}, {"messages": [], "metadata": {}}, "response")
holder["router"].acompletion.assert_awaited_once()

View file

@ -157,6 +157,8 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown():
"prompt_caching_savings_spend": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"autorouter_savings_spend": 0.0,
"total_response_time_ms": 0,
"timed_requests": 0,
"failed_requests": 0,
}
mock_rows = [
@ -647,6 +649,8 @@ def test_update_breakdown_metrics_includes_user_email():
prompt_caching_savings_spend=0,
gateway_injected_caching_savings_spend=0,
autorouter_savings_spend=0,
total_response_time_ms=0,
timed_requests=0,
total_tokens=2,
api_requests=1,
successful_requests=1,
@ -722,6 +726,8 @@ async def test_tag_daily_activity_metadata_totals_not_zero():
mock_record_1.prompt_caching_savings_spend = 0.0
mock_record_1.gateway_injected_caching_savings_spend = 0.0
mock_record_1.autorouter_savings_spend = 0.0
mock_record_1.total_response_time_ms = 18_000
mock_record_1.timed_requests = 9
mock_record_1.api_requests = 10
mock_record_1.successful_requests = 9
mock_record_1.failed_requests = 1
@ -746,6 +752,8 @@ async def test_tag_daily_activity_metadata_totals_not_zero():
mock_record_2.prompt_caching_savings_spend = 0.0
mock_record_2.gateway_injected_caching_savings_spend = 0.0
mock_record_2.autorouter_savings_spend = 0.0
mock_record_2.total_response_time_ms = 2_500
mock_record_2.timed_requests = 5
mock_record_2.api_requests = 5
mock_record_2.successful_requests = 5
mock_record_2.failed_requests = 0
@ -778,6 +786,8 @@ async def test_tag_daily_activity_metadata_totals_not_zero():
assert result.metadata.total_successful_requests == 14 # 9 + 5
assert result.metadata.total_failed_requests == 1
assert result.metadata.total_tokens == 1100 # (500+200) + (300+100)
assert result.metadata.total_response_time_ms == 20_500
assert result.metadata.total_timed_requests == 14
# Verify breakdown still works
assert len(result.results) == 1
@ -786,6 +796,10 @@ async def test_tag_daily_activity_metadata_totals_not_zero():
assert "staging" in daily.breakdown.entities
assert daily.breakdown.entities["production"].metrics.spend == 25.0
assert daily.breakdown.entities["staging"].metrics.spend == 5.0
assert daily.breakdown.models["gpt-4"].metrics.total_response_time_ms == 18_000
assert daily.breakdown.models["gpt-4"].metrics.timed_requests == 9
assert daily.breakdown.models["gpt-3.5-turbo"].metrics.total_response_time_ms == 2_500
assert daily.breakdown.models["gpt-3.5-turbo"].metrics.timed_requests == 5
@pytest.mark.asyncio
@ -810,6 +824,8 @@ async def test_aggregated_activity_preserves_metadata_for_deleted_keys():
"prompt_caching_savings_spend": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"autorouter_savings_spend": 0.0,
"total_response_time_ms": 0,
"timed_requests": 0,
"failed_requests": 0,
}
mock_rows = [
@ -900,6 +916,8 @@ def _daily_user_spend_record(*, user_id, api_key, spend, model="gpt-4", model_gr
prompt_caching_savings_spend=0.0,
gateway_injected_caching_savings_spend=0.0,
autorouter_savings_spend=0.0,
total_response_time_ms=0,
timed_requests=0,
api_requests=1,
successful_requests=1,
failed_requests=0,
@ -1333,6 +1351,8 @@ async def test_get_daily_activity_aggregated_empty_result_set():
"prompt_caching_savings_spend": None,
"gateway_injected_caching_savings_spend": None,
"autorouter_savings_spend": None,
"total_response_time_ms": None,
"timed_requests": None,
"api_requests": None,
"successful_requests": None,
"failed_requests": None,
@ -1378,6 +1398,8 @@ def _no_spend_record():
prompt_caching_savings_spend=None,
gateway_injected_caching_savings_spend=None,
autorouter_savings_spend=None,
total_response_time_ms=None,
timed_requests=None,
api_requests=None,
successful_requests=None,
failed_requests=None,
@ -1465,6 +1487,55 @@ class TestEverySavingsDriverSurvivesTheReadPath:
)
class TestResponseTimeSurvivesTheReadPath:
"""The dashboard averages total_response_time_ms over timed_requests, so both halves
of the pair must be summed by the rollup query, accumulated across rows, carried by
a single-row conversion, and coalesced when a NULL aggregate comes back."""
_FIELDS = ("total_response_time_ms", "timed_requests")
def test_both_halves_are_summed_by_the_rollup_query(self):
sql, _ = _build_aggregated_sql_query(
table_name="litellm_dailyuserspend",
entity_id_field="user_id",
entity_id="user-1",
start_date="2026-09-01",
end_date="2026-09-30",
model=None,
api_key=None,
timezone_offset_minutes=None,
)
for field in self._FIELDS:
assert f"SUM({field})" in sql, f"{field} is never summed, so the average reads as zero"
def test_accumulating_rows_keeps_sum_and_count_paired(self):
first = _no_spend_record()
first.total_response_time_ms = 1500
first.timed_requests = 2
second = _no_spend_record()
second.total_response_time_ms = 500
second.timed_requests = 1
metrics = update_metrics(update_metrics(SpendMetrics(), first), second)
assert metrics.total_response_time_ms == 2000
assert metrics.timed_requests == 3
def test_single_row_conversion_carries_both_halves(self):
record = _no_spend_record()
record.total_response_time_ms = 1234
record.timed_requests = 4
metrics = _record_to_spend_metrics(record)
assert metrics.total_response_time_ms == 1234
assert metrics.timed_requests == 4
def test_null_aggregates_read_as_zero(self):
metrics = _record_to_spend_metrics(_no_spend_record())
assert metrics.total_response_time_ms == 0
assert metrics.timed_requests == 0
accumulated = update_metrics(SpendMetrics(), _no_spend_record())
assert accumulated.total_response_time_ms == 0
assert accumulated.timed_requests == 0
@pytest.fixture
def ptu_cost_attribution_enabled(monkeypatch):
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
@ -1488,6 +1559,8 @@ def _spend_record(api_key, *, model="gpt-4o-mini-ptu", spend=0.0, ptu_flat_cost=
prompt_caching_savings_spend=0,
gateway_injected_caching_savings_spend=0,
autorouter_savings_spend=0,
total_response_time_ms=0,
timed_requests=0,
total_tokens=0,
api_requests=0,
successful_requests=0,
@ -1554,6 +1627,8 @@ def _grouping_row(
prompt_caching_savings_spend=0.0,
gateway_injected_caching_savings_spend=0.0,
autorouter_savings_spend=0.0,
total_response_time_ms=0,
timed_requests=0,
api_requests=0,
successful_requests=0,
failed_requests=0,
@ -1714,6 +1789,8 @@ def test_update_breakdown_metrics_covers_mcp_endpoint_and_entity(ptu_cost_attrib
prompt_caching_savings_spend=0,
gateway_injected_caching_savings_spend=0,
autorouter_savings_spend=0,
total_response_time_ms=0,
timed_requests=0,
total_tokens=0,
api_requests=0,
successful_requests=0,
@ -2118,6 +2195,8 @@ async def test_get_daily_activity_aggregated_with_entity_breakdown():
"prompt_caching_savings_spend": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"autorouter_savings_spend": 0.0,
"total_response_time_ms": 0,
"timed_requests": 0,
"failed_requests": 0,
"prompt_tokens": 0,
"completion_tokens": 0,

View file

@ -4,6 +4,7 @@ import pytest
import litellm
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
from litellm.llms.anthropic.cost_calculation import cost_per_token as anthropic_cost_per_token
from litellm.proxy.spend_tracking.savings import (
_baseline_usage,
_resolve_model,
@ -17,6 +18,34 @@ from litellm.types.utils import Usage
pytestmark = pytest.mark.usefixtures("local_model_cost_map")
@pytest.mark.parametrize("modifier", [{"speed": "fast"}, {"inference_geo": "us"}])
@pytest.mark.parametrize("continuing", [False, True])
def test_baseline_preserves_anthropic_pricing_fields(modifier: dict[str, str], continuing: bool) -> None:
usage: Final = _usage(1000, 0, 1000, 100).model_copy(update=modifier)
expected: Final = (_usage(1000, 1000, 0, 100) if continuing else usage).model_copy(update=modifier)
normalized: Final = _baseline_usage(usage, continuing)
cache_fields: Final = {"prompt_tokens_details", "cache_read_input_tokens", "cache_creation_input_tokens"}
assert normalized.model_dump(exclude=cache_fields) == usage.model_dump(exclude=cache_fields)
assert usage.prompt_tokens_details.cached_tokens == 0
selected_cost: Final = 0.013
assert compute_autorouter_savings(
"claude-opus-5", "claude-sonnet-5", "anthropic", usage, conversation_continuing=continuing,
cost_breakdown={"input_cost": 0.01, "output_cost": 0.003},
) == pytest.approx(sum(anthropic_cost_per_token("claude-opus-5", expected)) - selected_cost)
def test_anthropic_baseline_keeps_negotiated_prices_with_provider_multiplier() -> None:
info: Final = {
**litellm.get_model_info("claude-opus-5", "anthropic"),
"input_cost_per_token": 1e-6, "output_cost_per_token": 2e-6, "cache_read_input_token_cost": 3e-7,
}
usage: Final = _usage(1000, 1000, 0, 100).model_copy(update={"speed": "fast"})
assert compute_autorouter_savings(
"claude-opus-5", "claude-sonnet-5", "anthropic", usage, baseline_info=info,
cost_breakdown={"input_cost": 0.01, "output_cost": 0.003},
) == pytest.approx(0.0015 * 2 - 0.013)
def _anthropic_costs(model: str) -> tuple[float, float]:
info = litellm.get_model_info(model=model, custom_llm_provider="anthropic")
input_cost = info["input_cost_per_token"] or 0.0

View file

@ -689,6 +689,36 @@ async def test_during_call_hook_records_latency_metric(proxy_logging, make_user_
assert recorded["status"] == "success"
class _RecordingApplyGuardrail(CustomGuardrail):
def __init__(self, guardrail_name: str, applied: list[str]) -> None:
super().__init__(
guardrail_name=guardrail_name,
event_hook=GuardrailEventHooks.during_call,
default_on=True,
)
self._applied = applied
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
await asyncio.sleep(0)
self._applied.append(self.guardrail_name or "")
return inputs
@pytest.mark.asyncio
async def test_during_call_hook_runs_every_unified_guardrail(proxy_logging, make_user_api_key_auth, monkeypatch):
applied: list[str] = []
guardrails = [_RecordingApplyGuardrail(f"judge-{i}", applied) for i in range(3)]
monkeypatch.setattr(litellm, "callbacks", guardrails)
await proxy_logging.during_call_hook(
data={"model": "m", "messages": [{"role": "user", "content": "hi"}], "metadata": {}},
user_api_key_dict=make_user_api_key_auth(),
call_type="completion",
)
assert sorted(applied) == ["judge-0", "judge-1", "judge-2"]
@pytest.mark.asyncio
async def test_post_call_success_hook_records_latency_metric(proxy_logging, make_user_api_key_auth, monkeypatch):
cb = _moderation_guardrail()

View file

@ -5,17 +5,20 @@ completion_start_time = end_time."""
import json
from datetime import datetime
from typing import Optional
from typing import Final, Optional
from unittest.mock import Mock, patch
import httpx
import pytest
from pydantic_core import PydanticSerializationError
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.responses.streaming_iterator import (
ResponsesAPIStreamingIterator,
SyncResponsesAPIStreamingIterator,
_estimate_usage_from_text,
)
from litellm.types.llms.openai import (
ResponseAPIUsage,
@ -31,16 +34,23 @@ def _sse_event(payload: dict) -> bytes:
def _mock_config() -> Mock:
mock_config = Mock(spec=BaseResponsesAPIConfig)
mock_responses_api_response = Mock(spec=ResponsesAPIResponse)
mock_responses_api_response.id = "resp_ttft"
mock_responses_api_response = ResponsesAPIResponse(
id="resp_ttft",
created_at=0,
status="completed",
model="gpt-4o-mini",
object="response",
output=[],
usage=ResponseAPIUsage(input_tokens=1, output_tokens=1, total_tokens=2),
)
def _transform(model, parsed_chunk, logging_obj):
evt_type = parsed_chunk.get("type")
if evt_type == "response.completed":
completed = Mock(spec=ResponseCompletedEvent)
completed.type = ResponsesAPIStreamEvents.RESPONSE_COMPLETED
completed.response = mock_responses_api_response
return completed
return ResponseCompletedEvent(
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
response=mock_responses_api_response,
)
stub = Mock()
stub.type = evt_type
return stub
@ -54,6 +64,8 @@ def _make_iterator(
sse_events: list[bytes],
logging_obj: LiteLLMLoggingObj,
trailing_error: Optional[Exception] = None,
config: Mock | None = None,
request_data: dict | None = None,
) -> ResponsesAPIStreamingIterator:
async def aiter_bytes():
for evt in sse_events:
@ -68,10 +80,11 @@ def _make_iterator(
return ResponsesAPIStreamingIterator(
response=mock_response,
model="gpt-4o-mini",
responses_api_provider_config=_mock_config(),
responses_api_provider_config=config or _mock_config(),
logging_obj=logging_obj,
litellm_metadata={},
custom_llm_provider="openai",
request_data=request_data,
)
@ -329,6 +342,88 @@ def test_run_post_success_hooks_does_not_report_generation_time_as_overhead():
assert "litellm_overhead_time_ms" not in iterator.completed_response._hidden_params
def _mock_config_with_completed_response(response: ResponsesAPIResponse) -> Mock:
mock_config = Mock(spec=BaseResponsesAPIConfig)
def _transform(model, parsed_chunk, logging_obj):
evt_type = parsed_chunk.get("type")
if evt_type == "response.completed":
return ResponseCompletedEvent(
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
response=response,
)
stub = Mock()
stub.type = evt_type
if "delta" in parsed_chunk:
stub.delta = parsed_chunk.get("delta")
if "item" in parsed_chunk:
stub.item = parsed_chunk.get("item")
return stub
mock_config.transform_streaming_response.side_effect = _transform
return mock_config
def _responses_api_response_without_usage() -> ResponsesAPIResponse:
return ResponsesAPIResponse(
id="resp_no_usage",
created_at=int(datetime(2025, 1, 1).timestamp()),
status="completed",
model="gpt-4o-mini",
object="response",
output=[],
usage=None,
)
@pytest.mark.asyncio
async def test_completed_event_without_usage_gets_text_estimate():
"""A response.completed event carrying usage: null still bills: the
iterator estimates usage from the request input and generated text."""
response = _responses_api_response_without_usage()
iterator = _make_iterator(
sse_events=[
_sse_event({"type": "response.output_text.delta", "delta": "hello world"}),
_sse_event({"type": "response.completed", "response": {}}),
],
logging_obj=_logging_obj_stub(),
config=_mock_config_with_completed_response(response),
request_data={"input": "count these input tokens please"},
)
async for _ in iterator:
pass
usage = iterator.completed_response.response.usage
assert usage is not None
assert usage.input_tokens > 0
assert usage.output_tokens > 0
assert usage.total_tokens == usage.input_tokens + usage.output_tokens
@pytest.mark.asyncio
async def test_completed_event_with_usage_is_left_untouched():
"""Provider-reported usage on response.completed wins over the estimate."""
response = _responses_api_response_with_usage()
iterator = _make_iterator(
sse_events=[
_sse_event({"type": "response.output_text.delta", "delta": "hello world"}),
_sse_event({"type": "response.completed", "response": {}}),
],
logging_obj=_logging_obj_stub(),
config=_mock_config_with_completed_response(response),
request_data={"input": "count these input tokens please"},
)
async for _ in iterator:
pass
usage = iterator.completed_response.response.usage
assert usage.input_tokens == 20
assert usage.output_tokens == 60
assert usage.total_tokens == 80
def _responses_api_response_with_usage() -> ResponsesAPIResponse:
return ResponsesAPIResponse(
id="resp_lit6427",
@ -628,3 +723,222 @@ async def test_streaming_logging_copy_keeps_client_usage_when_response_fails_val
assert isinstance(client_usage, ResponseAPIUsage)
assert client_usage.input_tokens == 29
assert client_usage.cost == pytest.approx(0.0001)
@pytest.mark.asyncio
async def test_completed_event_without_usage_counts_tool_call_arguments():
"""A function-call-only stream still bills output tokens: streamed
function_call_arguments deltas feed the text estimate."""
response = _responses_api_response_without_usage()
iterator = _make_iterator(
sse_events=[
_sse_event(
{
"type": "response.output_item.added",
"item": {"type": "function_call", "name": "get_weather", "call_id": "call_1"},
}
),
_sse_event(
{
"type": "response.function_call_arguments.delta",
"delta": '{"location": "San Francisco", "unit": "celsius"}',
}
),
_sse_event({"type": "response.completed", "response": {}}),
],
logging_obj=_logging_obj_stub(),
config=_mock_config_with_completed_response(response),
request_data={"input": "what is the weather in san francisco"},
)
async for _ in iterator:
pass
usage = iterator.completed_response.response.usage
assert usage is not None
assert usage.output_tokens > 0
assert usage.total_tokens == usage.input_tokens + usage.output_tokens
@pytest.mark.asyncio
async def test_completed_event_without_usage_counts_multimodal_input_as_messages():
"""Multimodal request input is counted as chat messages, not as a JSON blob:
a huge base64 image must not inflate the estimated input tokens."""
image_input: Final = [
{
"role": "user",
"content": [
{"type": "input_text", "text": "what is in this image"},
{
"type": "input_image",
"image_url": "data:image/png;base64," + "A" * 4000,
},
],
}
]
json_count: Final = litellm.token_counter(model="gpt-4o-mini", text=json.dumps(image_input))
response = _responses_api_response_without_usage()
iterator = _make_iterator(
sse_events=[
_sse_event({"type": "response.output_text.delta", "delta": "it is a cat"}),
_sse_event({"type": "response.completed", "response": {}}),
],
logging_obj=_logging_obj_stub(),
config=_mock_config_with_completed_response(response),
request_data={"input": image_input},
)
async for _ in iterator:
pass
usage = iterator.completed_response.response.usage
assert usage is not None
assert usage.input_tokens < json_count / 2
@pytest.mark.asyncio
async def test_completed_event_survives_a_failing_usage_estimate():
"""A malformed request input that makes the message transformer raise must not
break a stream that previously completed: the estimate is best-effort and
falls back to usage None."""
malformed_input: Final = [{"type": "message", "role": "user", "content": 42}]
with pytest.raises(ValueError, match="Invalid content type"):
_estimate_usage_from_text("gpt-4o-mini", malformed_input, {"input": malformed_input}, "hello world")
response = _responses_api_response_without_usage()
iterator = _make_iterator(
sse_events=[
_sse_event({"type": "response.output_text.delta", "delta": "hello world"}),
_sse_event({"type": "response.completed", "response": {}}),
],
logging_obj=_logging_obj_stub(),
config=_mock_config_with_completed_response(response),
request_data={"input": malformed_input},
)
yielded: list = []
async for chunk in iterator:
yielded.append(chunk)
assert yielded
assert iterator.completed_response.response.usage is None
@pytest.mark.asyncio
@pytest.mark.parametrize(
"tool_delta_event_type",
["response.custom_tool_call_input.delta", "response.mcp_call_arguments.delta"],
)
async def test_completed_event_without_usage_counts_tool_input_deltas(tool_delta_event_type):
"""Custom-tool and MCP argument deltas feed the streamed usage fallback the
same way function_call_arguments deltas do."""
response = _responses_api_response_without_usage()
iterator = _make_iterator(
sse_events=[
_sse_event({"type": tool_delta_event_type, "delta": '{"query": "weather in sf"}'}),
_sse_event({"type": "response.completed", "response": {}}),
],
logging_obj=_logging_obj_stub(),
config=_mock_config_with_completed_response(response),
request_data={"input": "what is the weather in san francisco"},
)
async for _ in iterator:
pass
usage = iterator.completed_response.response.usage
assert usage is not None
assert usage.output_tokens > 0
assert usage.total_tokens == usage.input_tokens + usage.output_tokens
@pytest.mark.asyncio
async def test_completed_event_with_a_dict_response_is_typed_and_billed():
"""transform_streaming_response can model_construct a terminal event whose
response stays a plain dict; the iterator must type it so the estimated
usage reaches the cost stamping path."""
dict_response: Final = {
"id": "resp_dict",
"model": "gpt-4o-mini",
"object": "response",
"output": [],
"usage": None,
}
def _transform(model, parsed_chunk, logging_obj):
if parsed_chunk.get("type") == "response.completed":
return ResponseCompletedEvent.model_construct(type="response.completed", response=dict_response)
stub: Final = Mock()
stub.type = parsed_chunk.get("type")
if "delta" in parsed_chunk:
stub.delta = parsed_chunk.get("delta")
return stub
config: Final = Mock(spec=BaseResponsesAPIConfig)
config.transform_streaming_response.side_effect = _transform
logging_obj: Final = _logging_obj_stub()
logging_obj._response_cost_calculator.return_value = 0.000704
iterator: Final = _make_iterator(
sse_events=[
_sse_event({"type": "response.output_text.delta", "delta": "hello world"}),
_sse_event({"type": "response.completed", "response": {}}),
],
logging_obj=logging_obj,
config=config,
request_data={"input": "count these input tokens please"},
)
yielded: Final = [chunk async for chunk in iterator]
terminal_event: Final = iterator.completed_response
assert yielded[-1] is terminal_event
completed_response: Final = terminal_event.response
assert isinstance(completed_response, ResponsesAPIResponse)
usage: Final = completed_response.usage
assert usage is not None
assert usage.input_tokens > 0
assert usage.output_tokens > 0
assert usage.cost == pytest.approx(0.000704)
logging_obj._response_cost_calculator.assert_any_call(result=completed_response)
def test_billed_terminal_response_keeps_a_response_that_already_has_usage():
from litellm.responses.streaming_iterator import _billed_terminal_response
response: Final = _responses_api_response_with_usage()
assert _billed_terminal_response(response, None) is response
def test_billed_terminal_response_copies_when_estimating_and_leaves_the_original_untouched():
from litellm.responses.streaming_iterator import _billed_terminal_response
response: Final = _responses_api_response_without_usage()
estimated: Final = ResponseAPIUsage(input_tokens=3, output_tokens=4, total_tokens=7)
billed: Final = _billed_terminal_response(response, lambda: estimated)
assert billed is not response
assert billed.usage is estimated
assert response.usage is None
def test_persist_completed_response_to_cache_survives_an_unserializable_response(monkeypatch):
bad_response: Final = ResponsesAPIResponse.model_construct(id="r", output=[object()], usage=None)
with pytest.raises(PydanticSerializationError):
bad_response.model_dump_json()
logging_obj: Final = _logging_obj_stub()
caching_handler: Final = Mock()
caching_handler.request_kwargs = {"stream": True}
logging_obj._llm_caching_handler = caching_handler
iterator: Final = _make_iterator(sse_events=[], logging_obj=logging_obj)
iterator.completed_response = ResponseCompletedEvent.model_construct(
type="response.completed", response=bad_response
)
cache: Final = Mock()
monkeypatch.setattr(litellm, "cache", cache)
iterator._persist_completed_response_to_cache(is_async=False)
cache.add_cache.assert_not_called()

View file

@ -49,6 +49,20 @@ CI = [".github/workflows/test-litellm-ui-unit.yml"]
@pytest.mark.parametrize(
"category,changed,expected",
[
("provider-harness", ["tests/e2e/provider_cache.py"], "run"),
("provider-harness", ["tests/e2e/conftest.py"], "run"),
("provider-harness", ["tests/e2e/e2e_http.py"], "run"),
("provider-harness", ["tests/code_coverage_tests/test_provider_cache.py"], "run"),
("provider-harness", ["tests/code_coverage_tests/test_provider_replay_harness.py"], "run"),
("provider-harness", [".circleci/config.yml"], "run"),
("provider-harness", [".circleci/scripts/classify_changes.sh"], "run"),
("provider-harness", ["pyproject.toml"], "run"),
("provider-harness", ["uv.lock"], "run"),
("provider-harness", ["tests/e2e/PROVIDER_CACHE.md"], "skip"),
("provider-harness", ["tests/e2e/ui/test_example.py"], "skip"),
("provider-harness", ["tests/e2e/quota_management/test_quota.py"], "skip"),
("provider-harness", ["litellm/main.py"], "skip"),
("provider-harness", ["ui/litellm-dashboard/src/App.tsx"], "skip"),
# docs-only: skip everything
("backend", DOCS, "skip"),
("client", DOCS, "skip"),

View file

@ -1,11 +1,73 @@
import asyncio
import time
from collections.abc import Callable, Mapping
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.router import Router
from litellm.router import _silent_experiment_kwargs_snapshot
from litellm.router import _silent_experiment_targets
class _RecordingLogger(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.success_kwargs: list[dict[str, object]] = []
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.success_kwargs.append(kwargs)
def shadow_successes(self) -> list[dict[str, object]]:
return [
call
for call in self.success_kwargs
if call.get("litellm_params", {}).get("metadata", {}).get("is_silent_experiment") is True
]
@pytest.fixture
def recording_logger():
original_callbacks: Final = litellm.callbacks
logger: Final = _RecordingLogger()
litellm.callbacks = [logger]
try:
yield logger
finally:
litellm.callbacks = original_callbacks
async def _wait_for_shadow_successes(logger: _RecordingLogger, expected: int, timeout: float = 5.0) -> None:
deadline: Final = time.monotonic() + timeout
while len(logger.shadow_successes()) < expected and time.monotonic() < deadline:
await asyncio.sleep(0.05)
def _wait_for_shadow_successes_sync(logger: _RecordingLogger, expected: int, timeout: float = 5.0) -> None:
deadline: Final = time.monotonic() + timeout
while len(logger.shadow_successes()) < expected and time.monotonic() < deadline:
time.sleep(0.05)
def _streaming_model_list(silent_model: object) -> list[dict[str, object]]:
return [
{
"model_name": "primary-model",
"litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "fake-key", "silent_model": silent_model},
},
{
"model_name": "shadow-a",
"litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "fake-key", "silent_model": "shadow-b"},
},
{
"model_name": "shadow-b",
"litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "fake-key"},
},
]
class _NonCopyableSpan:
@ -65,8 +127,7 @@ def test_get_silent_experiment_kwargs():
assert result["metadata"]["is_silent_experiment"] is True
assert result["metadata"]["foo"] == "bar"
assert "litellm_call_id" not in result
# stream must be forced to False so callbacks fire in background
assert result["stream"] is False
assert result["stream"] is True
# proxy_server_request must be preserved for spend log metadata
assert "proxy_server_request" in result
# CRITICAL: metadata must be a DIFFERENT dict object than the original,
@ -86,6 +147,247 @@ def test_get_silent_experiment_kwargs():
assert result["metadata"]["user_api_key_auth"] is mock_auth
def test_get_silent_experiment_kwargs_without_stream_stays_non_streaming():
router = Router(model_list=[{"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "k"}}])
result = router._get_silent_experiment_kwargs(metadata={"foo": "bar"}, stream=False)
assert result["stream"] is False
assert "stream" not in router._get_silent_experiment_kwargs(metadata={"foo": "bar"})
@pytest.mark.parametrize(
"silent_model, expected",
[
("shadow-a", ("shadow-a",)),
(["shadow-a", "shadow-b"], ("shadow-a", "shadow-b")),
([], ()),
(None, ()),
(42, ()),
(["shadow-a", 42], ()),
],
)
def test_silent_experiment_targets(silent_model, expected):
assert _silent_experiment_targets(silent_model) == expected
@pytest.mark.asyncio
async def test_streaming_shadow_is_streamed_and_drained_async(recording_logger):
router = Router(model_list=_streaming_model_list("shadow-a"))
response = await router.acompletion(
model="primary-model",
messages=[{"role": "user", "content": "hi"}],
stream=True,
stream_options={"include_usage": True},
mock_response="pong",
metadata={"foo": "bar"},
)
chunks = [chunk async for chunk in response]
assert chunks
await _wait_for_shadow_successes(recording_logger, expected=1)
shadow_successes = recording_logger.shadow_successes()
assert len(shadow_successes) == 1
shadow = shadow_successes[0]
assert shadow["stream"] is True
assert shadow["stream_options"] == {"include_usage": True}
assert shadow["litellm_params"]["metadata"]["model_group"] == "shadow-a"
assert shadow["async_complete_streaming_response"] is not None
def test_streaming_shadow_is_streamed_and_drained_sync(recording_logger):
router = Router(model_list=_streaming_model_list("shadow-a"))
response = router.completion(
model="primary-model",
messages=[{"role": "user", "content": "hi"}],
stream=True,
mock_response="pong",
metadata={"foo": "bar"},
)
chunks = list(response)
assert chunks
_wait_for_shadow_successes_sync(recording_logger, expected=1)
shadow_successes = recording_logger.shadow_successes()
assert len(shadow_successes) == 1
assert shadow_successes[0]["stream"] is True
assert shadow_successes[0]["litellm_params"]["metadata"]["model_group"] == "shadow-a"
assert shadow_successes[0]["async_complete_streaming_response"] is not None
@pytest.mark.asyncio
async def test_multiple_shadow_targets_fan_out_async(recording_logger):
router = Router(model_list=_streaming_model_list(["shadow-a", "shadow-b"]))
metadata = {"foo": "bar"}
response = await router.acompletion(
model="primary-model",
messages=[{"role": "user", "content": "hi"}],
stream=True,
mock_response="pong",
metadata=metadata,
)
assert [chunk async for chunk in response]
await _wait_for_shadow_successes(recording_logger, expected=2)
shadow_successes = recording_logger.shadow_successes()
model_groups = sorted(call["litellm_params"]["metadata"]["model_group"] for call in shadow_successes)
assert model_groups == ["shadow-a", "shadow-b"]
shadow_metadatas = [call["litellm_params"]["metadata"] for call in shadow_successes]
assert shadow_metadatas[0] is not shadow_metadatas[1]
assert all(call["stream"] is True for call in shadow_successes)
assert "is_silent_experiment" not in metadata
assert metadata.get("model_group") != "shadow-a"
primary_successes = [call for call in recording_logger.success_kwargs if call not in shadow_successes]
assert len(primary_successes) == 1
assert primary_successes[0]["litellm_params"]["metadata"]["model_group"] == "primary-model"
def test_multiple_shadow_targets_fan_out_sync(recording_logger):
router = Router(model_list=_streaming_model_list(["shadow-a", "shadow-b"]))
response = router.completion(
model="primary-model",
messages=[{"role": "user", "content": "hi"}],
mock_response="pong",
metadata={"foo": "bar"},
)
assert response.choices[0].message.content == "pong"
_wait_for_shadow_successes_sync(recording_logger, expected=2)
shadow_successes = recording_logger.shadow_successes()
model_groups = sorted(call["litellm_params"]["metadata"]["model_group"] for call in shadow_successes)
assert model_groups == ["shadow-a", "shadow-b"]
assert all(call["stream"] is False for call in shadow_successes)
def _tagged_primary_model_list() -> list[dict[str, object]]:
return [
{
"model_name": "primary-model",
"litellm_params": {
"model": "openai/gpt-5.4-mini",
"api_key": "fake-key",
"silent_model": "shadow-b",
"tags": ["primary-only"],
},
},
{
"model_name": "shadow-b",
"litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "fake-key"},
},
]
def test_silent_experiment_kwargs_snapshot_is_isolated_from_later_primary_mutations():
metadata = {"foo": "bar"}
kwargs: dict[str, object] = {"metadata": metadata, "stream": True}
snapshot = _silent_experiment_kwargs_snapshot(kwargs)
kwargs["messages"] = [{"role": "user", "content": "added by the primary"}]
metadata["tags"] = ["primary-only"]
assert dict(snapshot) == {"metadata": {"foo": "bar"}, "stream": True}
assert dict(_silent_experiment_kwargs_snapshot({"stream": False, "metadata": None})) == {
"stream": False,
"metadata": None,
}
def test_sync_shadow_gets_kwargs_snapshot_taken_before_primary_mutates_them(recording_logger):
deferred: list[Callable[[], None]] = []
class _DeferredThread:
def __init__(self, target, args, kwargs, daemon) -> None:
deferred.append(lambda: target(*args, **kwargs))
def start(self) -> None:
return None
router = Router(model_list=_tagged_primary_model_list())
with patch( # test-quality-ok: Router has no thread factory to inject; deferring start is the only deterministic way to expose the race
"litellm.router.threading", SimpleNamespace(Thread=_DeferredThread)
):
response = router.completion(
model="primary-model",
messages=[{"role": "user", "content": "hi"}],
mock_response="pong",
metadata={"foo": "bar"},
)
assert response.choices[0].message.content == "pong"
assert len(deferred) == 1
deferred[0]()
_wait_for_shadow_successes_sync(recording_logger, expected=1)
shadow_successes = recording_logger.shadow_successes()
assert len(shadow_successes) == 1
shadow_metadata = shadow_successes[0]["litellm_params"]["metadata"]
assert shadow_metadata["model_group"] == "shadow-b"
assert "primary-only" not in shadow_metadata.get("tags", [])
def test_sync_shadow_workers_do_not_share_metadata_with_each_other(recording_logger):
workers: list[tuple[Mapping[str, object], Callable[[], None]]] = []
class _DeferredThread:
def __init__(self, target, args, kwargs, daemon) -> None:
workers.append((kwargs, lambda: target(*args, **kwargs)))
def start(self) -> None:
return None
router = Router(model_list=_streaming_model_list(["shadow-a", "shadow-b"]))
with patch( # test-quality-ok: Router has no thread factory to inject; deferring start is the only deterministic way to expose the race
"litellm.router.threading", SimpleNamespace(Thread=_DeferredThread)
):
router.completion(
model="primary-model",
messages=[{"role": "user", "content": "hi"}],
mock_response="pong",
metadata={"foo": "bar"},
)
assert len(workers) == 2
(first_kwargs, run_first), (_, run_second) = workers
first_kwargs["metadata"].pop("foo")
run_second()
run_first()
_wait_for_shadow_successes_sync(recording_logger, expected=2)
metadata_by_group = {
call["litellm_params"]["metadata"]["model_group"]: call["litellm_params"]["metadata"]
for call in recording_logger.shadow_successes()
}
assert metadata_by_group["shadow-b"]["foo"] == "bar"
assert "foo" not in metadata_by_group["shadow-a"]
@pytest.mark.asyncio
async def test_async_shadow_does_not_inherit_primary_deployment_tags(recording_logger):
router = Router(model_list=_tagged_primary_model_list())
response = await router.acompletion(
model="primary-model",
messages=[{"role": "user", "content": "hi"}],
mock_response="pong",
metadata={"foo": "bar"},
)
assert response.choices[0].message.content == "pong"
await _wait_for_shadow_successes(recording_logger, expected=1)
shadow_successes = recording_logger.shadow_successes()
assert len(shadow_successes) == 1
assert "primary-only" not in shadow_successes[0]["litellm_params"]["metadata"].get("tags", [])
@pytest.mark.asyncio
async def test_shadow_of_a_shadow_is_not_launched(recording_logger):
router = Router(model_list=_streaming_model_list(["shadow-a"]))
response = await router.acompletion(
model="primary-model",
messages=[{"role": "user", "content": "hi"}],
mock_response="pong",
)
assert response.choices[0].message.content == "pong"
await _wait_for_shadow_successes(recording_logger, expected=2, timeout=1.0)
model_groups = [call["litellm_params"]["metadata"]["model_group"] for call in recording_logger.shadow_successes()]
assert model_groups == ["shadow-a"]
def test_silent_experiment_completion_direct():
"""
Test _silent_experiment_completion directly (for router code coverage).
@ -127,6 +429,25 @@ async def test_silent_experiment_acompletion_direct():
)
@pytest.mark.asyncio
async def test_run_silent_experiment_drains_stream_so_callbacks_fire(recording_logger):
router = Router(model_list=_streaming_model_list(None))
silent_kwargs: Final = {
"stream": True,
"stream_options": {"include_usage": True},
"mock_response": "pong",
"metadata": {"is_silent_experiment": True, "model_group": "shadow-b"},
}
await router._run_silent_experiment("shadow-b", [{"role": "user", "content": "hi"}], silent_kwargs)
await _wait_for_shadow_successes(recording_logger, expected=1)
shadow_successes = recording_logger.shadow_successes()
assert len(shadow_successes) == 1
assert shadow_successes[0]["stream"] is True
assert shadow_successes[0]["async_complete_streaming_response"] is not None
assert silent_kwargs["stream"] is True
@pytest.mark.asyncio
async def test_router_silent_experiment_acompletion():
"""

View file

@ -2,14 +2,30 @@
Test automatic routing to xAI Responses API when tools are present
"""
import json
from collections.abc import Mapping
from typing import Final
from unittest.mock import MagicMock, patch
import httpx
import pytest
import litellm
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.main import responses_api_bridge_check
class _RecordingResponsesHandler:
"""MockTransport handler that serves a canned /responses reply and keeps the body xAI would have received"""
def __init__(self, reply: Mapping[str, object]) -> None:
self.reply: Final = reply
self.request_body: Mapping[str, object] | None = None
def __call__(self, request: httpx.Request) -> httpx.Response:
self.request_body = json.loads(request.content)
return httpx.Response(200, json=dict(self.reply), request=request)
class TestXAIResponsesAutoRouting:
"""Test that xAI requests with tools automatically route to Responses API"""
@ -254,6 +270,44 @@ class TestXAIResponsesAutoRouting:
# Note: This test may need adjustment based on actual mock_response behavior
# The key is that the responses_api_bridge_check logic routes correctly
def test_system_message_survives_web_search_bridge(self):
"""A system message becomes 'instructions' on the bridged /responses call, and xAI accepts it"""
handler: Final = _RecordingResponsesHandler(
reply={
"id": "resp_test",
"object": "response",
"created_at": 0,
"status": "completed",
"model": "grok-4.6",
"output": [
{
"type": "message",
"id": "msg_test",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "1.0.0", "annotations": []}],
}
],
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
}
)
response: Final = litellm.completion(
model="xai/grok-4.6",
messages=[
{"role": "system", "content": "Answer briefly."},
{"role": "user", "content": "newest litellm version?"},
],
web_search_options={"search_context_size": "medium"},
api_key="fake-key",
client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))),
)
assert response.choices[0].message.content == "1.0.0"
assert handler.request_body is not None
assert handler.request_body["instructions"] == "Answer briefly."
assert handler.request_body["tools"] == [{"type": "web_search"}]
if __name__ == "__main__":
pytest.main([__file__, "-v"])

View file

@ -87,8 +87,8 @@ const LLMJudgeFields: React.FC<LLMJudgeFieldsProps> = ({ availableModels, contro
return (
<FieldGroup>
<div className="rounded-md border border-success/20 bg-success/10 px-3.5 py-2.5 text-[13px] text-success">
After each LLM response, the <strong>Judge Model</strong> scores it 0100 against your criteria. If the weighted
average falls below the threshold, the response is blocked (or logged).
The <strong>Judge Model</strong> scores the user request (pre_call, during_call) or the LLM response (post_call)
0100 against your criteria. If the weighted average falls below the threshold, it is blocked (or logged).
</div>
<GuardrailField

View file

@ -104,8 +104,8 @@ export function AutoRoutersPanel({
<DialogHeader>
<DialogTitle>Add Auto Router</DialogTitle>
<DialogDescription>
Routes each request to a model by classifying its complexity. Called like any other model, so clients keep
using a single model name.
Choose a classifier to route each request to a model. Called like any other model, so clients keep using a
single model name.
</DialogDescription>
</DialogHeader>
<AddAutoRouterTab

View file

@ -57,6 +57,8 @@ const dedupe = (models: string[]): string[] => Array.from(new Set(models));
const COMPLEXITY_TYPE_LABELS: Record<string, string> = {
llm: "LLM Classifier",
capability: "Capability",
llm_v2: "Fuse v2",
heuristic_first: "Heuristic first",
hybrid: "Hybrid",
custom: "Custom classifier",

View file

@ -43,6 +43,8 @@ describe("sumMetadata", () => {
total_cache_read_input_tokens: 1,
total_cache_creation_input_tokens: 1,
total_flat_cost: 1,
total_response_time_ms: 1,
total_timed_requests: 1,
};
const merged = sumMetadata(page, page);

View file

@ -30,6 +30,8 @@ const SUMMABLE_METADATA_KEYS = [
"total_cache_read_input_tokens",
"total_cache_creation_input_tokens",
"total_flat_cost",
"total_response_time_ms",
"total_timed_requests",
] as const;
interface DailyActivityResponse {
@ -78,6 +80,8 @@ const EMPTY_DATA: DailyActivityResponse = {
total_failed_requests: 0,
total_cache_read_input_tokens: 0,
total_cache_creation_input_tokens: 0,
total_response_time_ms: 0,
total_timed_requests: 0,
total_pages: 1,
has_more: false,
page: 1,

View file

@ -14,6 +14,8 @@ export interface SpendMetrics {
prompt_caching_savings_spend?: number;
gateway_injected_caching_savings_spend?: number;
autorouter_savings_spend?: number;
total_response_time_ms?: number;
timed_requests?: number;
}
export type DailyData = {
@ -81,6 +83,8 @@ export interface ModelActivityData {
prompt_tokens: number;
completion_tokens: number;
total_spend: number;
total_response_time_ms?: number;
total_timed_requests?: number;
top_api_keys: TopApiKeyData[];
top_models: TopModelData[];
daily_data: {
@ -95,6 +99,7 @@ export interface ModelActivityData {
failed_requests: number;
cache_read_input_tokens: number;
cache_creation_input_tokens: number;
avg_response_time_ms?: number | null;
};
}[];
}

View file

@ -1,5 +1,36 @@
import { describe, expect, it } from "vitest";
import { valueFormatter, valueFormatterSpend } from "./value_formatters";
import { averageResponseTimeMs, formatResponseTime, valueFormatter, valueFormatterSpend } from "./value_formatters";
describe("averageResponseTimeMs", () => {
it("divides the summed duration by the number of timed requests", () => {
expect(averageResponseTimeMs(6000, 4)).toBe(1500);
expect(averageResponseTimeMs(0, 3)).toBe(0);
});
it("returns null instead of dividing by zero when nothing was timed", () => {
expect(averageResponseTimeMs(0, 0)).toBeNull();
expect(averageResponseTimeMs(1200, 0)).toBeNull();
});
});
describe("formatResponseTime", () => {
it("shows sub-second durations in whole milliseconds", () => {
expect(formatResponseTime(0)).toBe("0ms");
expect(formatResponseTime(412.6)).toBe("413ms");
expect(formatResponseTime(999)).toBe("999ms");
});
it("shows durations of a second or more in seconds with two decimals", () => {
expect(formatResponseTime(1000)).toBe("1.00s");
expect(formatResponseTime(1500)).toBe("1.50s");
expect(formatResponseTime(12345)).toBe("12.35s");
});
it("shows a dash when there is no average to display", () => {
expect(formatResponseTime(null)).toBe("-");
expect(formatResponseTime(undefined)).toBe("-");
});
});
describe("valueFormatter", () => {
it("should format numbers >= 1,000,000 as millions with 2 decimal places", () => {

View file

@ -11,6 +11,17 @@ export function valueFormatter(number: number) {
return number.toString();
}
export function averageResponseTimeMs(totalResponseTimeMs: number, timedRequests: number): number | null {
if (timedRequests <= 0) return null;
return totalResponseTimeMs / timedRequests;
}
export function formatResponseTime(ms: number | null | undefined) {
if (ms == null) return "-";
if (ms < 1000) return `${Math.round(ms)}ms`;
return `${(ms / 1000).toFixed(2)}s`;
}
export function valueFormatterSpend(number: number) {
if (number === 0) return "$0";
if (number >= 1_000_000_000) {

View file

@ -1,7 +1,8 @@
import { fireEvent, render, screen } from "@testing-library/react";
import React from "react";
import { beforeAll, describe, expect, it, vi } from "vitest";
import { ActivityMetrics, formatKeyLabel, processActivityData } from "./activity_metrics";
import { ActivityMetrics, formatKeyLabel, processActivityData, ResponseTimeTooltip } from "./activity_metrics";
import type { ChartTooltipProps } from "@/components/shared/charts";
import { Team } from "./key_team_helpers/key_list";
import { DailyData, KeyMetricWithMetadata, ModelActivityData } from "./UsagePage/types";
@ -1424,6 +1425,144 @@ describe("processActivityData", () => {
expect(result).toEqual({});
});
it("sums response time per model and derives a per-day average over timed requests", () => {
const dayWithModel = (date: string, metrics: Partial<typeof EMPTY_SPEND_METRICS> & Record<string, number>) =>
createMockDailyData(date, EMPTY_SPEND_METRICS, {
...EMPTY_BREAKDOWN,
models: {
"gpt-5.5": { metrics: { ...EMPTY_SPEND_METRICS, ...metrics }, metadata: {}, api_key_breakdown: {} },
},
});
const fourTimedRequests = {
api_requests: 4,
successful_requests: 4,
total_response_time_ms: 6000,
timed_requests: 4,
};
const oneTimedOneFailed = {
api_requests: 2,
successful_requests: 1,
failed_requests: 1,
total_response_time_ms: 500,
timed_requests: 1,
};
const onlyFailures = { api_requests: 1, successful_requests: 0, failed_requests: 1 };
const activity: { results: DailyData[] } = {
results: [
dayWithModel("2025-01-02", fourTimedRequests),
dayWithModel("2025-01-01", oneTimedOneFailed),
dayWithModel("2025-01-03", onlyFailures),
],
};
const result = processActivityData(activity, "models");
expect(result["gpt-5.5"].total_response_time_ms).toBe(6500);
expect(result["gpt-5.5"].total_timed_requests).toBe(5);
expect(result["gpt-5.5"].daily_data.map((day) => day.metrics.avg_response_time_ms)).toEqual([500, 1500, null]);
});
it("treats rollups written before response time existed as zero timed requests", () => {
const activity: { results: DailyData[] } = {
results: [
createMockDailyData("2025-01-01", EMPTY_SPEND_METRICS, {
...EMPTY_BREAKDOWN,
models: {
"gpt-5.5": {
metrics: { ...EMPTY_SPEND_METRICS, api_requests: 3, successful_requests: 3 },
metadata: {},
api_key_breakdown: {},
},
},
}),
],
};
const result = processActivityData(activity, "models");
expect(result["gpt-5.5"].total_response_time_ms).toBe(0);
expect(result["gpt-5.5"].total_timed_requests).toBe(0);
expect(result["gpt-5.5"].daily_data[0].metrics.avg_response_time_ms).toBeNull();
});
});
describe("ActivityMetrics response time", () => {
const timedModel = createMockModelActivityData("GPT-5.5", {
total_response_time_ms: 6000,
total_timed_requests: 4,
daily_data: [
{
date: "2025-01-01",
metrics: {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15,
api_requests: 3,
spend: 1,
successful_requests: 3,
failed_requests: 0,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
avg_response_time_ms: 2000,
},
},
{
date: "2025-01-02",
metrics: {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15,
api_requests: 1,
spend: 1,
successful_requests: 1,
failed_requests: 0,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
avg_response_time_ms: 1000,
},
},
],
});
it("shows the model's average response time in the summary card and the collapsed header", () => {
render(<ActivityMetrics modelMetrics={{ "gpt-5.5": timedModel }} />);
expect(screen.getByText("Avg Response Time")).toBeInTheDocument();
expect(screen.getByRole("heading", { name: "1.50s" })).toBeInTheDocument();
expect(screen.getByText("over 4 timed successful requests")).toBeInTheDocument();
expect(screen.getByText("1.50s avg response")).toBeInTheDocument();
});
it("renders the per-day response time chart with duration-formatted axis ticks", () => {
render(<ActivityMetrics modelMetrics={{ "gpt-5.5": timedModel }} />);
expect(screen.getByText("Avg Response Time per day")).toBeInTheDocument();
expect(screen.getByText("Avg Response Time Ms")).toBeInTheDocument();
expect(screen.getAllByText(/^\d+(\.\d+)?(ms|s)$/).length).toBeGreaterThan(1);
});
it("labels the chart tooltip with the readable series name and a formatted duration", () => {
const payload = [
{ dataKey: "metrics.avg_response_time_ms", value: 1500, color: "#f59e0b", payload: timedModel.daily_data[0] },
] as NonNullable<ChartTooltipProps["payload"]>;
render(<ResponseTimeTooltip active={true} payload={payload} label="2025-01-01" />);
expect(screen.getByText("Avg Response Time Ms")).toBeInTheDocument();
expect(screen.getByText("1.50s")).toBeInTheDocument();
expect(screen.queryByText("metrics.avg_response_time_ms")).not.toBeInTheDocument();
});
it("shows a dash and no response time chart when the model has no timed requests", () => {
render(<ActivityMetrics modelMetrics={{ "gpt-5.5": createMockModelActivityData("GPT-5.5") }} />);
expect(screen.getByText("Avg Response Time")).toBeInTheDocument();
expect(screen.getByRole("heading", { name: "-" })).toBeInTheDocument();
expect(screen.getByText("over 0 timed successful requests")).toBeInTheDocument();
expect(screen.queryByText(/avg response$/)).not.toBeInTheDocument();
expect(screen.queryByText("Avg Response Time per day")).not.toBeInTheDocument();
expect(screen.queryByText("Avg Response Time Ms")).not.toBeInTheDocument();
});
});
describe("formatKeyLabel", () => {

View file

@ -1,4 +1,13 @@
import { AreaChart, BarChart, CustomLegend, CustomTooltip } from "@/components/shared/charts";
import {
AreaChart,
BarChart,
type ChartTooltipProps,
CustomLegend,
CustomTooltip,
formatCategoryName,
LineChart,
ValueTooltip,
} from "@/components/shared/charts";
import { formatNumberWithCommas } from "@/utils/dataUtils";
import { resolveTeamAliasFromTeamID } from "@/utils/teamUtils";
import { Card, CardContent } from "@/components/ui/card";
@ -9,13 +18,25 @@ import { Team } from "./key_team_helpers/key_list";
import KeyModelUsageView from "./UsagePage/components/KeyModelUsageView";
import { keyActivityLabel } from "./UsagePage/keyActivityLabel";
import { DailyData, KeyMetricWithMetadata, ModelActivityData, TopApiKeyData, TopModelData } from "./UsagePage/types";
import { valueFormatter } from "./UsagePage/utils/value_formatters";
import { averageResponseTimeMs, formatResponseTime, valueFormatter } from "./UsagePage/utils/value_formatters";
interface ActivityMetricsProps {
modelMetrics: Record<string, ModelActivityData>;
hidePromptCachingMetrics?: boolean;
}
const modelAverageResponseTimeMs = (metrics: ModelActivityData): number | null =>
averageResponseTimeMs(metrics.total_response_time_ms ?? 0, metrics.total_timed_requests ?? 0);
export const ResponseTimeTooltip = ({ active, payload, label }: ChartTooltipProps) => (
<ValueTooltip
active={active}
payload={payload?.map((item) => ({ ...item, name: formatCategoryName(String(item.dataKey ?? "")) }))}
label={label}
valueFormatter={formatResponseTime}
/>
);
const ModelSection = ({
modelName,
metrics,
@ -28,7 +49,7 @@ const ModelSection = ({
return (
<div className="space-y-2">
{/* Summary Cards */}
<div className="grid grid-cols-4 gap-4">
<div className="grid grid-cols-5 gap-4">
<Card>
<CardContent>
<p className="text-sm text-muted-foreground">Total Requests</p>
@ -62,6 +83,17 @@ const ModelSection = ({
</p>
</CardContent>
</Card>
<Card>
<CardContent>
<p className="text-sm text-muted-foreground">Avg Response Time</p>
<h3 className="text-lg font-medium text-foreground">
{formatResponseTime(modelAverageResponseTimeMs(metrics))}
</h3>
<p className="text-sm text-muted-foreground">
over {(metrics.total_timed_requests ?? 0).toLocaleString()} timed successful requests
</p>
</CardContent>
</Card>
</div>
{metrics.top_api_keys && metrics.top_api_keys.length > 0 && (
@ -154,6 +186,28 @@ const ModelSection = ({
</CardContent>
</Card>
{(metrics.total_timed_requests ?? 0) > 0 && (
<Card>
<CardContent>
<div className="flex justify-between items-center">
<h3 className="text-lg font-medium text-foreground">Avg Response Time per day</h3>
<CustomLegend categories={["metrics.avg_response_time_ms"]} colors={["amber"]} />
</div>
<LineChart
className="mt-4"
data={metrics.daily_data}
index="date"
categories={["metrics.avg_response_time_ms"]}
colors={["amber"]}
valueFormatter={formatResponseTime}
customTooltip={ResponseTimeTooltip}
connectNulls={true}
showLegend={false}
/>
</CardContent>
</Card>
)}
<Card>
<CardContent>
<div className="flex justify-between items-center">
@ -416,6 +470,9 @@ export const ActivityMetrics: React.FC<ActivityMetricsProps> = ({ modelMetrics,
<div className="flex space-x-4 text-sm text-muted-foreground">
<span>${formatNumberWithCommas(modelMetrics[modelName].total_spend, 2)}</span>
<span>{modelMetrics[modelName].total_requests.toLocaleString()} requests</span>
{modelAverageResponseTimeMs(modelMetrics[modelName]) != null && (
<span>{formatResponseTime(modelAverageResponseTimeMs(modelMetrics[modelName]))} avg response</span>
)}
</div>
</div>
}
@ -471,11 +528,15 @@ export const processActivityData = (
total_spend: 0,
total_cache_read_input_tokens: 0,
total_cache_creation_input_tokens: 0,
total_response_time_ms: 0,
total_timed_requests: 0,
top_api_keys: [],
top_models: [],
daily_data: [],
};
}
const dayResponseTimeMs = modelData.metrics.total_response_time_ms || 0;
const dayTimedRequests = modelData.metrics.timed_requests || 0;
// Update totals
modelMetrics[model].total_requests += modelData.metrics.api_requests;
modelMetrics[model].prompt_tokens += modelData.metrics.prompt_tokens;
@ -486,6 +547,9 @@ export const processActivityData = (
modelMetrics[model].total_failed_requests += modelData.metrics.failed_requests;
modelMetrics[model].total_cache_read_input_tokens += modelData.metrics.cache_read_input_tokens || 0;
modelMetrics[model].total_cache_creation_input_tokens += modelData.metrics.cache_creation_input_tokens || 0;
modelMetrics[model].total_response_time_ms =
(modelMetrics[model].total_response_time_ms ?? 0) + dayResponseTimeMs;
modelMetrics[model].total_timed_requests = (modelMetrics[model].total_timed_requests ?? 0) + dayTimedRequests;
// Add daily data
modelMetrics[model].daily_data.push({
@ -500,6 +564,7 @@ export const processActivityData = (
failed_requests: modelData.metrics.failed_requests,
cache_read_input_tokens: modelData.metrics.cache_read_input_tokens || 0,
cache_creation_input_tokens: modelData.metrics.cache_creation_input_tokens || 0,
avg_response_time_ms: averageResponseTimeMs(dayResponseTimeMs, dayTimedRequests),
},
});
});

View file

@ -0,0 +1,75 @@
import React, { useState } from "react";
import { describe, expect, it, vi } from "vitest";
import { fireEvent, renderWithProviders, screen } from "../../../tests/test-utils";
import AutoRouterClassifierTabs from "./AutoRouterClassifierTabs";
import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
const initial: ComplexityRouterConfigValue = {
classifier_type: "llm",
tiers: { SIMPLE: ["efficient"], MEDIUM: [], COMPLEX: [], REASONING: ["capable"] },
};
function Form({ initialValue = initial }: { initialValue?: ComplexityRouterConfigValue }) {
const [value, setValue] = useState(initialValue);
return (
<AutoRouterClassifierTabs value={value} onChange={setValue}>
<output aria-label="Classifier type">{value.classifier_type}</output>
</AutoRouterClassifierTabs>
);
}
describe("AutoRouterClassifierTabs", () => {
it.each(["heuristic", "heuristic_v2", "llm", "heuristic_first", "hybrid"] as const)(
"groups %s under Complexity without resetting its configuration",
(classifier_type) => {
const onChange = vi.fn();
renderWithProviders(
<AutoRouterClassifierTabs value={{ ...initial, classifier_type }} onChange={onChange}>
Existing classifier settings
</AutoRouterClassifierTabs>,
);
expect(screen.getByRole("tab", { name: "Complexity" })).toHaveAttribute("aria-selected", "true");
expect(screen.getByRole("tabpanel", { name: "Complexity" })).toHaveTextContent("Existing classifier settings");
fireEvent.click(screen.getByRole("tab", { name: "Complexity" }));
expect(onChange).not.toHaveBeenCalled();
},
);
it.each([
["capability", "Capability"],
["llm_v2", "Fuse v2"],
] as const)("opens saved %s settings and switches back to local Complexity", (classifier_type, label) => {
renderWithProviders(<Form initialValue={{ ...initial, classifier_type }} />);
expect(screen.getByRole("tab", { name: label })).toHaveAttribute("aria-selected", "true");
fireEvent.click(screen.getByRole("tab", { name: "Complexity" }));
expect(screen.getByRole("tab", { name: "Complexity" })).toHaveAttribute("aria-selected", "true");
expect(screen.getByRole("status", { name: "Classifier type" })).toHaveTextContent("heuristic");
});
it("keeps custom tiers editable under Complexity and explains why forecast tabs are disabled", () => {
const onChange = vi.fn();
renderWithProviders(
<AutoRouterClassifierTabs
value={{
...initial,
custom_tier_set: {
tiers: [{ id: "review", name: "REVIEW", definition: "Code reviews", models: ["capable"] }],
fallback_tier_id: "review",
},
}}
onChange={onChange}
>
Custom tiers
</AutoRouterClassifierTabs>,
);
expect(screen.getByRole("tabpanel", { name: "Complexity" })).toHaveTextContent("Custom tiers");
for (const name of ["Capability", "Fuse v2"]) {
const tab = screen.getByRole("tab", { name });
expect(tab).toHaveAttribute("aria-disabled", "true");
expect(tab).toHaveAccessibleDescription("Restore standard tiers to use Capability or Fuse v2.");
fireEvent.click(tab);
}
expect(onChange).not.toHaveBeenCalled();
expect(screen.getByText("Restore standard tiers to use Capability or Fuse v2.")).toBeVisible();
});
});

View file

@ -0,0 +1,58 @@
import React, { useId } from "react";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { effectiveClassifierType, type ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
import { transitionClassifierType } from "./classifier_type_transition";
import { isForecastClassifier } from "./forecast_classifier_config";
interface AutoRouterClassifierTabsProps {
value: ComplexityRouterConfigValue;
onChange: (value: ComplexityRouterConfigValue) => void;
children: React.ReactNode;
}
const AutoRouterClassifierTabs: React.FC<AutoRouterClassifierTabsProps> = ({ value, onChange, children }) => {
const restrictionId = useId();
const classifierType = effectiveClassifierType(value);
const selected = isForecastClassifier(classifierType) ? classifierType : "complexity";
const hasCustomTiers = Boolean(value.custom_tier_set);
const handleChange = (tab: unknown) => {
if (tab === selected) return;
if (tab === "complexity") {
onChange(transitionClassifierType(value, isForecastClassifier(classifierType) ? "heuristic" : classifierType));
} else if (!hasCustomTiers && (tab === "capability" || tab === "llm_v2")) {
onChange(transitionClassifierType(value, tab));
}
};
return (
<Tabs value={selected} onValueChange={handleChange}>
<p className="text-sm font-medium">Classifier type</p>
<TabsList aria-label="Classifier type" className="w-full">
<TabsTrigger value="complexity">Complexity</TabsTrigger>
<TabsTrigger
value="capability"
disabled={hasCustomTiers}
aria-describedby={hasCustomTiers ? restrictionId : undefined}
>
Capability
</TabsTrigger>
<TabsTrigger
value="llm_v2"
disabled={hasCustomTiers}
aria-describedby={hasCustomTiers ? restrictionId : undefined}
>
Fuse v2
</TabsTrigger>
</TabsList>
{hasCustomTiers && (
<p id={restrictionId} className="text-sm text-muted-foreground">
Restore standard tiers to use Capability or Fuse v2.
</p>
)}
<TabsContent value={selected}>{children}</TabsContent>
</Tabs>
);
};
export default AutoRouterClassifierTabs;

View file

@ -1,3 +1,4 @@
import { transitionClassifierType } from "./classifier_type_transition";
import { Info } from "lucide-react";
import { SimpleTooltip } from "@/components/ui/tooltip";
import { MultiSelect } from "@/components/shared/MultiSelect";
@ -17,7 +18,6 @@ import ClassifierReasoningEffortSelect from "./ClassifierReasoningEffortSelect";
import ClassifierCircuitBreakerConfig from "./ClassifierCircuitBreakerConfig";
import ClassifierVisionConfig from "./ClassifierVisionConfig";
import type { ReasoningEffort } from "./complexity_router_tiers";
import { nonReasoningTierFields } from "./nonReasoningTierFields";
import { useComplexityScorerDefaults } from "@/app/(dashboard)/hooks/autoRouter/useComplexityScorerDefaults";
import {
ClassificationFrequency,
@ -33,12 +33,10 @@ import {
DEFAULT_CLASSIFIER_FALLBACK,
DEFAULT_CLASSIFIER_TIMEOUT_MS,
DEFAULT_CLASSIFICATION_RUBRIC,
NEW_CLASSIFIER_CLASSIFICATION_RUBRIC,
ClassificationRubric,
effectiveTierLabel,
heuristicScoringRole,
usesLlmClassifier,
DEFAULT_HEURISTIC_FIRST_MAX_TIER,
DEFAULT_HYBRID_BOUNDARY_MARGIN,
HEURISTIC_FIRST_MAX_TIER_KEYS,
effectiveClassifierType,
@ -263,35 +261,7 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
const explicitlySupportedClassifierEfforts = effortOptionsByModel[classifierModel];
const handleClassifierTypeChange = (classifierType: ClassifierType) => {
const nextValue: ComplexityRouterConfigValue = {
...value,
classifier_type: classifierType,
classifier_llm_config: usesLlmClassifier(classifierType)
? value.classifier_llm_config ?? {
model: "",
timeout_ms: DEFAULT_CLASSIFIER_TIMEOUT_MS,
classification_rubric: NEW_CLASSIFIER_CLASSIFICATION_RUBRIC,
}
: undefined,
classifier_context_window_size: usesLlmClassifier(classifierType)
? value.classifier_context_window_size ?? DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE
: undefined,
classifier_context_budget_chars: usesLlmClassifier(classifierType)
? value.classifier_context_budget_chars ?? DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS
: undefined,
classifier_context_include_assistant_turns: usesLlmClassifier(classifierType)
? value.classifier_context_include_assistant_turns
: undefined,
classifier_fallback: usesLlmClassifier(classifierType) ? value.classifier_fallback : undefined,
heuristic_first_max_tier:
classifierType === "heuristic_first"
? value.heuristic_first_max_tier ?? DEFAULT_HEURISTIC_FIRST_MAX_TIER
: undefined,
hybrid_boundary_margin:
classifierType === "hybrid" ? value.hybrid_boundary_margin ?? DEFAULT_HYBRID_BOUNDARY_MARGIN : undefined,
...nonReasoningTierFields(classifierType, value),
};
onChange(nextValue);
onChange(transitionClassifierType(value, classifierType));
};
const handleHeuristicFirstMaxTierChange = (tier: string) => {
@ -433,27 +403,6 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
});
};
if (classifierType === "capability") {
return (
<p className="text-sm text-muted-foreground">
This router uses capability forecasting. Configure its classifier, threshold, and calibration through YAML or
the API. Saving preserves those settings
</p>
);
}
if (classifierType === "llm_v2") {
return (
<div className="rounded-md border p-4 text-sm">
<strong>LLM V2 classifier (experimental)</strong>
<p className="mt-2 text-muted-foreground">
Combines task demands and model capability in one forecast. Its solver profiles and quality allowance are
configured through the API. Saving this router preserves those settings
</p>
</div>
);
}
return (
<>
<ClassifierTypeRadios value={value} classifierType={classifierType} onTypeChange={handleClassifierTypeChange} />

View file

@ -1,8 +1,11 @@
import RoutingOptions from "./RoutingOptions";
import PlanModeOverrideControls from "./PlanModeOverrideControls";
import ForecastClassifierConfig, { ForecastSolverModels } from "./ForecastClassifierConfig";
import { isForecastClassifier, type CapabilitySettings, type FuseSettings } from "./forecast_classifier_config";
import { SimpleTooltip } from "@/components/ui/tooltip";
import { MultiSelect } from "@/components/shared/MultiSelect";
import { SearchSelect } from "@/components/shared/SearchSelect";
import DefaultModelField from "./DefaultModelField";
import { ChevronRight, Info, Plus, Trash2, X } from "lucide-react";
import { Switch } from "@/components/ui/switch";
import { AffinityControls } from "./AffinityControls";
import NonReasoningTierToggle from "./NonReasoningTierToggle";
@ -204,11 +207,6 @@ const rowOrigin = (row: TierRow, editing: boolean): string => {
return isBuiltInTierName(row.name) ? "built-in" : "custom";
};
const defaultModelPlaceholderFor = (derivedDefaultModel: string | undefined, isCustomSet: boolean): string => {
if (derivedDefaultModel) return `Derived from tiers: ${derivedDefaultModel}`;
return isCustomSet ? "Add a model to your fallback tier" : "Add a model to the Simple or Medium tier";
};
const builtInTierInfo = (rowId: string): { label: string; description: string; examples: string } | undefined => {
const builtIn = ALL_BUILT_IN_TIERS.find((tier) => tier === rowId);
return builtIn ? TIER_DESCRIPTIONS[builtIn] : undefined;
@ -376,6 +374,8 @@ export interface ComplexityRouterConfigValue {
/** An explicit pin. Unset means the default tracks the tiers - see resolveComplexityDefaultModel. */
default_model?: string;
classifier_type: ClassifierType;
capability_classifier_config?: CapabilitySettings;
llm_v2_config?: FuseSettings;
classifier_llm_config?: ClassifierLLMConfig;
classifier_context_window_size?: number;
classifier_context_budget_chars?: number;
@ -535,44 +535,6 @@ export const DEFAULT_HYBRID_BOUNDARY_MARGIN = 0.03;
*/
export const HEURISTIC_FIRST_MAX_TIER_KEYS = TIER_ORDER.slice(0, -1);
const PlanModeOverrideControls: React.FC<{
value: ComplexityRouterConfigValue;
onChange: (value: ComplexityRouterConfigValue) => void;
planModeTierOptions: { value: string; label: string }[];
}> = ({ value, onChange, planModeTierOptions }) => (
<>
<div className="flex items-center gap-2 mb-2">
<Switch
checked={value.plan_mode_min_tier !== undefined}
disabled={planModeTierOptions.length === 0}
onCheckedChange={(enabled) =>
onChange({
...value,
plan_mode_min_tier: enabled ? planModeTierOptions.at(-1)?.value : undefined,
})
}
aria-label="Route plan-mode requests to a minimum tier"
/>
<strong className="font-semibold">Route plan-mode requests to a minimum tier</strong>
</div>
<span className="block text-xs mb-3 text-muted-foreground">
Requests from coding agents in plan mode (Claude Code, GitHub Copilot) route to at least this tier. The classifier
still wins when it picks higher, and the override only lasts while plan mode is active.
{planModeTierOptions.length === 0 && " Add models to a tier to enable this."}
</span>
{value.plan_mode_min_tier !== undefined && (
<div style={{ maxWidth: 320 }}>
<TierRowSelect
label="Plan-mode minimum tier"
options={planModeTierOptions}
value={value.plan_mode_min_tier ?? null}
onValueChange={(tier) => onChange({ ...value, plan_mode_min_tier: tier })}
/>
</div>
)}
</>
);
const ComplexityRouterConfig: React.FC<ComplexityRouterConfigProps> = ({
modelInfo,
value,
@ -596,6 +558,7 @@ const ComplexityRouterConfig: React.FC<ComplexityRouterConfigProps> = ({
onAutoRouterCompressionChange,
showValidationErrors = false,
}) => {
const forecast = isForecastClassifier(value.classifier_type);
const customTierSet = value.custom_tier_set;
const tierRows = activeTierRows(value);
const tierRowsError = customTierSet ? getCustomTierRowsError(customTierSet) : null;
@ -605,8 +568,6 @@ const ComplexityRouterConfig: React.FC<ComplexityRouterConfigProps> = ({
value: row.id,
label: tierRowLabel(row, value.tier_labels),
}));
const derivedDefaultModel = resolveComplexityDefaultModel(value);
const defaultModelPlaceholder = defaultModelPlaceholderFor(derivedDefaultModel, Boolean(customTierSet));
const defaultModel = resolveComplexityDefaultModel(value, value.default_model);
const dispatch = (action: TierSetAction) => {
@ -641,298 +602,319 @@ const ComplexityRouterConfig: React.FC<ComplexityRouterConfigProps> = ({
tier_model_params: setTierModelParam(value.tier_model_params, tier, model, change),
});
// Clearing the select drops the key entirely rather than storing "", so an emptied pin reads as
// "track the tiers" everywhere downstream instead of as a blank model name.
const handleDefaultModelChange = (model: string | null | undefined) => {
onChange({ ...value, default_model: model || undefined });
};
const handleTierLabelChange = (tier: keyof ComplexityTiers, label: string) => {
onChange({
...value,
tier_labels: { ...value.tier_labels, [tier]: label },
});
};
const handleTierLabelChange = (tier: keyof ComplexityTiers, label: string) =>
onChange({ ...value, tier_labels: { ...value.tier_labels, [tier]: label } });
return (
<div className="w-full max-w-none">
<div className="inline-flex items-center gap-2 mb-4">
<h4 className="m-0 text-xl font-semibold text-foreground">Complexity Tier Configuration</h4>
<SimpleTooltip content="Map each complexity tier to one or more models. Simple queries use cheaper/faster models, complex queries use more capable models.">
<Info className="size-4 text-muted-foreground" />
</SimpleTooltip>
<h4 className="m-0 text-xl font-semibold text-foreground">
{forecast ? "Solver models" : "Complexity Tier Configuration"}
</h4>
{!forecast && (
<SimpleTooltip content="Map each complexity tier to one or more models. Simple queries use cheaper/faster models, complex queries use more capable models.">
<Info className="size-4 text-muted-foreground" />
</SimpleTooltip>
)}
</div>
<TierConfigIntro value={value} />
<Card>
<CardContent>
{!customTierSet && (
<NonReasoningTierToggle value={value} onChange={onChange} available={value.classifier_type === "llm"} />
)}
{tierRows.map((row, index) => {
const tierInfo = builtInTierInfo(row.id);
const label = tierRowLabel(row, value.tier_labels);
const tierMissing = showValidationErrors && row.models.length === 0;
const needsDefinition = Boolean(customTierSet) && !row.definition.trim() && !isBuiltInTierName(row.name);
const definitionMissing = showValidationErrors && needsDefinition;
const showsDisplayName = !customTierSet && !editingTiers;
return (
<div key={row.id}>
{index > 0 && <Separator className="my-4" />}
<div className="mb-4">
<TierRowHeader
row={row}
index={index}
rowCount={tierRows.length}
label={label}
description={tierInfo?.description}
editing={editingTiers}
isCustomSet={Boolean(customTierSet)}
onRemove={() => removeTierRow(row.id)}
/>
{tierInfo && !customTierSet && (
<span className="block mb-2 text-xs text-muted-foreground">Examples: {tierInfo.examples}</span>
)}
{editingTiers && (
<TierRowEditFields
row={row}
index={index}
definitionMissing={definitionMissing}
onPatch={(patch) => updateTierRow(row.id, patch)}
/>
)}
{showsDisplayName && tierInfo && (
<InputGroup className="mb-2">
<InputGroupInput
value={value.tier_labels?.[row.id as keyof ComplexityTiers] ?? ""}
onChange={(event) => handleTierLabelChange(row.id as keyof ComplexityTiers, event.target.value)}
placeholder={`Display name (default: ${tierInfo.label})`}
aria-label={`Display name for the ${tierInfo.label} tier`}
/>
{value.tier_labels?.[row.id as keyof ComplexityTiers] && (
<InputGroupAddon align="inline-end">
<InputGroupButton
size="icon-xs"
aria-label={`Clear display name for the ${tierInfo.label} tier`}
onClick={() => handleTierLabelChange(row.id as keyof ComplexityTiers, "")}
>
<X />
</InputGroupButton>
</InputGroupAddon>
)}
</InputGroup>
)}
<MultiSelect
options={modelOptions}
value={row.models}
onValueChange={(models: string[]) => setRowModels(row, models)}
placeholder={`Select model(s) for ${label.toLowerCase()} queries`}
emptyText="No models found"
className={tierMissing ? "w-full border-destructive" : "w-full"}
/>
<TierModelEffortRows
tierLabel={label}
models={row.models}
effortOptionsByModel={tierEffortOptionsByModel}
paramsByModel={row.params}
fastModeByModel={fastModeByModel}
onEffortChange={(model, effort) =>
handleTierModelParamChange(row.id, model, ["reasoning_effort", effort])
}
onFastModeChange={(model, enabled) =>
handleTierModelParamChange(row.id, model, ["speed", enabled ? "fast" : undefined])
}
/>
{row.models.length > 1 && (
<span className="text-xs text-muted-foreground">
Multiple models selected: the router randomly picks among them per request (or Thompson-samples
within the pool when adaptive routing is on).
</span>
)}
{tierMissing && <span className="text-xs text-destructive">The {label} tier is required</span>}
</div>
</div>
);
})}
<TierSetToolbar
editing={editingTiers}
isCustomSet={Boolean(customTierSet)}
rowCount={tierRows.length}
rowsError={tierRowsError}
keywordRulesError={keywordRulesError}
onEditingChange={onEditingTiersChange}
onAdd={addCustomTier}
onRestore={exitToBuiltInTiers}
{forecast ? (
<>
<ForecastSolverModels
value={value}
onChange={onChange}
modelOptions={modelOptions}
effortOptionsByModel={tierEffortOptionsByModel}
fastModeByModel={fastModeByModel}
/>
<ForecastClassifierConfig
value={value}
onChange={onChange}
modelOptions={modelOptions}
effortOptionsByModel={classifierEffortOptionsByModel}
/>
</>
) : (
<>
<TierConfigIntro value={value} />
{customTierSet && (
<FallbackTierField
rows={tierRows}
fallbackTierId={customTierSet.fallback_tier_id}
onValueChange={(fallbackTierId) => onChange(setFallbackTier(value, fallbackTierId))}
/>
)}
<Card>
<CardContent>
{!customTierSet && (
<NonReasoningTierToggle value={value} onChange={onChange} available={value.classifier_type === "llm"} />
)}
<Separator className="my-4" />
{tierRows.map((row, index) => {
const tierInfo = builtInTierInfo(row.id);
const label = tierRowLabel(row, value.tier_labels);
const tierMissing = showValidationErrors && row.models.length === 0;
const needsDefinition =
Boolean(customTierSet) && !row.definition.trim() && !isBuiltInTierName(row.name);
const definitionMissing = showValidationErrors && needsDefinition;
const showsDisplayName = !customTierSet && !editingTiers;
return (
<div key={row.id}>
{index > 0 && <Separator className="my-4" />}
<div className="mb-4">
<TierRowHeader
row={row}
index={index}
rowCount={tierRows.length}
label={label}
description={tierInfo?.description}
editing={editingTiers}
isCustomSet={Boolean(customTierSet)}
onRemove={() => removeTierRow(row.id)}
/>
{tierInfo && !customTierSet && (
<span className="block mb-2 text-xs text-muted-foreground">Examples: {tierInfo.examples}</span>
)}
{editingTiers && (
<TierRowEditFields
row={row}
index={index}
definitionMissing={definitionMissing}
onPatch={(patch) => updateTierRow(row.id, patch)}
/>
)}
{showsDisplayName && tierInfo && (
<InputGroup className="mb-2">
<InputGroupInput
value={value.tier_labels?.[row.id as keyof ComplexityTiers] ?? ""}
onChange={(event) =>
handleTierLabelChange(row.id as keyof ComplexityTiers, event.target.value)
}
placeholder={`Display name (default: ${tierInfo.label})`}
aria-label={`Display name for the ${tierInfo.label} tier`}
/>
{value.tier_labels?.[row.id as keyof ComplexityTiers] && (
<InputGroupAddon align="inline-end">
<InputGroupButton
size="icon-xs"
aria-label={`Clear display name for the ${tierInfo.label} tier`}
onClick={() => handleTierLabelChange(row.id as keyof ComplexityTiers, "")}
>
<X />
</InputGroupButton>
</InputGroupAddon>
)}
</InputGroup>
)}
<MultiSelect
options={modelOptions}
value={row.models}
onValueChange={(models: string[]) => setRowModels(row, models)}
placeholder={`Select model(s) for ${label.toLowerCase()} queries`}
emptyText="No models found"
className={tierMissing ? "w-full border-destructive" : "w-full"}
/>
<TierModelEffortRows
tierLabel={label}
models={row.models}
effortOptionsByModel={tierEffortOptionsByModel}
paramsByModel={row.params}
fastModeByModel={fastModeByModel}
onEffortChange={(model, effort) =>
handleTierModelParamChange(row.id, model, ["reasoning_effort", effort])
}
onFastModeChange={(model, enabled) =>
handleTierModelParamChange(row.id, model, ["speed", enabled ? "fast" : undefined])
}
/>
{row.models.length > 1 && (
<span className="text-xs text-muted-foreground">
Multiple models selected: the router randomly picks among them per request (or
Thompson-samples within the pool when adaptive routing is on).
</span>
)}
{tierMissing && <span className="text-xs text-destructive">The {label} tier is required</span>}
</div>
</div>
);
})}
<div className="mb-2">
<div className="flex items-center gap-2 mb-2">
<strong className="text-base font-semibold">Default Model</strong>
<SimpleTooltip content="Leave empty to follow the tiers. A model chosen here is pinned: it stays the default however the tiers change.">
<Info className="size-4 text-muted-foreground" />
</SimpleTooltip>
</div>
<SearchSelect
options={modelOptions}
value={value.default_model ?? ""}
onValueChange={handleDefaultModelChange}
placeholder={defaultModelPlaceholder}
emptyText="No models found"
aria-label="Default model"
/>
<span className="block mt-1 text-xs text-muted-foreground">
Used when the tier the request lands in has no model, and when the classifier fails with &quot;Route to
the default model&quot; selected.
</span>
</div>
</CardContent>
</Card>
<TierSetToolbar
editing={editingTiers}
isCustomSet={Boolean(customTierSet)}
rowCount={tierRows.length}
rowsError={tierRowsError}
keywordRulesError={keywordRulesError}
onEditingChange={onEditingTiersChange}
onAdd={addCustomTier}
onRestore={exitToBuiltInTiers}
/>
{customTierSet && (
<FallbackTierField
rows={tierRows}
fallbackTierId={customTierSet.fallback_tier_id}
onValueChange={(fallbackTierId) => onChange(setFallbackTier(value, fallbackTierId))}
/>
)}
</CardContent>
</Card>
</>
)}
{!forecast && <DefaultModelField value={value} onChange={onChange} modelOptions={modelOptions} />}
<Separator className="my-6" />
<div className="rounded-lg border border-border bg-muted">
{[
{
key: "classifier",
label: <strong className="text-foreground font-semibold">Advanced: Classification Method</strong>,
children: (
<ClassificationMethodConfig
value={value}
onChange={onChange}
modelOptions={modelOptions}
effortOptionsByModel={classifierEffortOptionsByModel}
customTechnicalKeywords={customTechnicalKeywords}
onCustomTechnicalKeywordsChange={onCustomTechnicalKeywordsChange}
showValidationErrors={showValidationErrors}
defaultModel={defaultModel}
/>
),
},
{
key: "adaptive",
label: <strong className="text-foreground font-semibold">Advanced: Adaptive Routing</strong>,
children: (
<Restricted by={restrictedBy(value, "adaptive")}>
<AdaptiveRoutingConfig value={value} onChange={onChange} />
</Restricted>
),
},
{
key: "affinity",
label: <strong className="text-foreground font-semibold">Advanced: Affinity</strong>,
children: <AffinityControls value={value} onChange={onChange} />,
},
{
key: "modality",
label: <strong className="text-foreground font-semibold">Advanced: Modality Routing</strong>,
children: <ModalityRoutingControls value={value} onChange={onChange} />,
},
{
key: "plan-mode",
label: <strong className="text-foreground font-semibold">Advanced: Plan-Mode Override</strong>,
children: (
<PlanModeOverrideControls value={value} onChange={onChange} planModeTierOptions={planModeTierOptions} />
),
},
{
key: "context-window",
label: <strong className="text-foreground font-semibold">Advanced: Context Window Escalation</strong>,
children: <ContextWindowEscalationConfig value={value} onChange={onChange} />,
},
{
key: "stall-escalation",
label: <strong className="text-foreground font-semibold">Advanced: Stalled Task Escalation</strong>,
children: (
<Restricted by={restrictedBy(value, "stallEscalation")}>
<StallEscalationConfig value={value} onChange={onChange} />
</Restricted>
),
},
{
key: "response",
label: <strong className="text-foreground font-semibold">Advanced: Response Format</strong>,
children: <ResponseFormatControls value={value} onChange={onChange} />,
},
...(onEscalationKeywordsChange
? [
{
key: "escalation",
label: <strong className="text-foreground font-semibold">Advanced: Escalation Keywords</strong>,
children: (
<Restricted by={restrictedBy(value, "escalation")}>
<EscalationKeywords keywords={escalationKeywords} onChange={onEscalationKeywordsChange} />
</Restricted>
),
},
]
: []),
...(onAutoRouterCompressionChange
? [
{
key: "compression",
label: <strong className="text-foreground font-semibold">Advanced: Compression</strong>,
children: (
<CompressionControls value={autoRouterCompression} onChange={onAutoRouterCompressionChange} />
),
},
]
: []),
...(onKeywordTierRulesChange || onSemanticMatchingEnabledChange
? [
{
key: "keyword-semantic",
label: <strong className="text-foreground font-semibold">Advanced: Keyword/Semantic Matching</strong>,
children: (
<>
{onKeywordTierRulesChange && (
<KeywordTierRules
rules={keywordTierRules}
onChange={onKeywordTierRulesChange}
tierLabels={value.tier_labels}
tierNames={customTierSet && tierRows.map(activeTierName).filter(Boolean)}
/>
)}
{onKeywordTierRulesChange && onSemanticMatchingEnabledChange && <Separator className="my-4" />}
{onSemanticMatchingEnabledChange && (
<SemanticKeywordMatching
enabled={semanticMatchingEnabled}
onEnabledChange={onSemanticMatchingEnabledChange}
embeddingModel={embeddingModel}
onEmbeddingModelChange={onEmbeddingModelChange}
matchThreshold={matchThreshold}
onMatchThresholdChange={onMatchThresholdChange}
modelInfo={modelInfo}
showValidationErrors={showValidationErrors}
/>
)}
</>
),
},
]
: []),
].map(({ key, label, children }) => (
<Collapsible key={key} className="border-b border-border last:border-b-0">
<CollapsibleTrigger className="group flex w-full items-center gap-2 px-4 py-3 text-left">
<ChevronRight className="size-4 shrink-0 text-muted-foreground transition-transform group-data-panel-open:rotate-90" />
{label}
</CollapsibleTrigger>
<CollapsibleContent className="px-4 pb-4">{children}</CollapsibleContent>
</Collapsible>
))}
</div>
<RoutingOptions forecast={forecast}>
{forecast && (
<>
<DefaultModelField value={value} onChange={onChange} modelOptions={modelOptions} />
<ForecastSolverModels
additionalPoolsOnly
value={value}
onChange={onChange}
modelOptions={modelOptions}
effortOptionsByModel={tierEffortOptionsByModel}
fastModeByModel={fastModeByModel}
/>
</>
)}
<div className="rounded-lg border border-border bg-muted">
{[
...(!forecast
? [
{
key: "classifier",
label: <strong className="text-foreground font-semibold">Advanced: Classification Method</strong>,
children: (
<ClassificationMethodConfig
value={value}
onChange={onChange}
modelOptions={modelOptions}
effortOptionsByModel={classifierEffortOptionsByModel}
customTechnicalKeywords={customTechnicalKeywords}
onCustomTechnicalKeywordsChange={onCustomTechnicalKeywordsChange}
showValidationErrors={showValidationErrors}
defaultModel={defaultModel}
/>
),
},
]
: []),
{
key: "adaptive",
label: <strong className="text-foreground font-semibold">Advanced: Adaptive Routing</strong>,
children: (
<Restricted by={restrictedBy(value, "adaptive")}>
<AdaptiveRoutingConfig value={value} onChange={onChange} />
</Restricted>
),
},
{
key: "affinity",
label: <strong className="text-foreground font-semibold">Advanced: Affinity</strong>,
children: <AffinityControls value={value} onChange={onChange} />,
},
{
key: "modality",
label: <strong className="text-foreground font-semibold">Advanced: Modality Routing</strong>,
children: <ModalityRoutingControls value={value} onChange={onChange} />,
},
{
key: "plan-mode",
label: <strong className="text-foreground font-semibold">Advanced: Plan-Mode Override</strong>,
children: (
<PlanModeOverrideControls value={value} onChange={onChange} planModeTierOptions={planModeTierOptions} />
),
},
{
key: "context-window",
label: <strong className="text-foreground font-semibold">Advanced: Context Window Escalation</strong>,
children: <ContextWindowEscalationConfig value={value} onChange={onChange} />,
},
{
key: "stall-escalation",
label: <strong className="text-foreground font-semibold">Advanced: Stalled Task Escalation</strong>,
children: (
<Restricted by={restrictedBy(value, "stallEscalation")}>
<StallEscalationConfig value={value} onChange={onChange} />
</Restricted>
),
},
{
key: "response",
label: <strong className="text-foreground font-semibold">Advanced: Response Format</strong>,
children: <ResponseFormatControls value={value} onChange={onChange} />,
},
...(onEscalationKeywordsChange
? [
{
key: "escalation",
label: <strong className="text-foreground font-semibold">Advanced: Escalation Keywords</strong>,
children: (
<Restricted by={restrictedBy(value, "escalation")}>
<EscalationKeywords keywords={escalationKeywords} onChange={onEscalationKeywordsChange} />
</Restricted>
),
},
]
: []),
...(onAutoRouterCompressionChange
? [
{
key: "compression",
label: <strong className="text-foreground font-semibold">Advanced: Compression</strong>,
children: (
<CompressionControls value={autoRouterCompression} onChange={onAutoRouterCompressionChange} />
),
},
]
: []),
...(onKeywordTierRulesChange || onSemanticMatchingEnabledChange
? [
{
key: "keyword-semantic",
label: (
<strong className="text-foreground font-semibold">Advanced: Keyword/Semantic Matching</strong>
),
children: (
<>
{onKeywordTierRulesChange && (
<KeywordTierRules
rules={keywordTierRules}
onChange={onKeywordTierRulesChange}
tierLabels={value.tier_labels}
tierNames={
customTierSet || isForecastClassifier(value.classifier_type)
? tierRows.map(activeTierName).filter(Boolean)
: undefined
}
/>
)}
{onKeywordTierRulesChange && onSemanticMatchingEnabledChange && <Separator className="my-4" />}
{onSemanticMatchingEnabledChange && (
<SemanticKeywordMatching
enabled={semanticMatchingEnabled}
onEnabledChange={onSemanticMatchingEnabledChange}
embeddingModel={embeddingModel}
onEmbeddingModelChange={onEmbeddingModelChange}
matchThreshold={matchThreshold}
onMatchThresholdChange={onMatchThresholdChange}
modelInfo={modelInfo}
showValidationErrors={showValidationErrors}
/>
)}
</>
),
},
]
: []),
]
.filter(({ key }) => !forecast || !["adaptive", "context-window", "escalation"].includes(key))
.map(({ key, label, children }) => (
<Collapsible key={key} className="border-b border-border last:border-b-0">
<CollapsibleTrigger className="group flex w-full items-center gap-2 px-4 py-3 text-left">
<ChevronRight className="size-4 shrink-0 text-muted-foreground transition-transform group-data-panel-open:rotate-90" />
{label}
</CollapsibleTrigger>
<CollapsibleContent className="px-4 pb-4">{children}</CollapsibleContent>
</Collapsible>
))}
</div>
</RoutingOptions>
</div>
);
};

View file

@ -1,11 +1,12 @@
import userEvent from "@testing-library/user-event";
import React from "react";
import { describe, expect, it, vi } from "vitest";
import { renderWithProviders, screen } from "../../../tests/test-utils";
import { renderWithProviders, screen, within } from "../../../tests/test-utils";
import {
buildUpdatedComplexityRouterConfig,
hydrateComplexityRouterConfig,
} from "../edit_auto_router/edit_auto_router_modal";
import type { KeywordTierRule } from "./KeywordTierRules";
import type { ModelGroup } from "../llm_calls/fetch_models";
import ComplexityRouterConfig, { type ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
@ -50,8 +51,8 @@ it.each([false, true])("edits and round-trips independent model settings with cu
const view = renderWithProviders(editor(initial));
const fast = () => screen.getByRole("switch", { name: `Fast mode for primary in the ${label} tier` });
expect(screen.getAllByRole("switch", { name: /^Fast mode for/ })).toHaveLength(3);
expect(screen.queryByRole("switch", { name: /^Fast mode for (blocked|missing)/ })).not.toBeInTheDocument();
expect(screen.getAllByRole("switch", { name: /^Fast mode for/ })).toHaveLength(4);
expect(screen.queryByRole("switch", { name: /^Fast mode for missing/ })).not.toBeInTheDocument();
expect(screen.queryByRole("combobox", { name: /^Reasoning effort for secondary/ })).not.toBeInTheDocument();
expect(screen.getByRole("switch", { name: `Fast mode for secondary in the ${label} tier` })).toBeChecked();
expect(fast()).not.toBeChecked();
@ -125,19 +126,170 @@ it.each([false, true])("edits and round-trips independent model settings with cu
);
});
describe("Fast mode metadata", () => {
it("offers nothing before model capabilities load and leaves stored speed untouched", () => {
const value: ComplexityRouterConfigValue = {
tiers: { SIMPLE: ["primary"], MEDIUM: [], COMPLEX: [], REASONING: [] },
classifier_type: "heuristic",
tier_model_params: { SIMPLE: { primary: { speed: "fast" } } },
};
const onChange = vi.fn();
renderWithProviders(<ComplexityRouterConfig modelInfo={[]} value={value} onChange={onChange} />);
expect(screen.queryByRole("switch", { name: /^Fast mode for/ })).not.toBeInTheDocument();
expect(onChange).not.toHaveBeenCalled();
expect(buildUpdatedComplexityRouterConfig({}, value).tier_model_configs).toEqual({
SIMPLE: [{ model_name: "primary", litellm_params: { speed: "fast" } }],
});
it.each(["capability", "llm_v2"] as const)("preserves Fast mode controls for %s solvers", async (classifierType) => {
const user = userEvent.setup();
const initial: ComplexityRouterConfigValue = {
classifier_type: classifierType,
classifier_llm_config: { model: "primary", timeout_ms: 3000 },
tiers: { SIMPLE: ["primary"], MEDIUM: [], COMPLEX: [], REASONING: ["blocked"] },
capability_classifier_config: { efficient_tier: "SIMPLE", capable_tier: "REASONING", base_threshold: 0.7 },
llm_v2_config: {
efficient_profile: "Small solver",
capable_profile: "Large solver",
harness: "One attempt",
max_quality_gap: 0.05,
},
tier_model_params: { SIMPLE: { primary: { reasoning_effort: "high", max_tokens: 1024 } } },
};
const onChange = vi.fn<(value: ComplexityRouterConfigValue) => void>();
const editor = (value: ComplexityRouterConfigValue) => (
<ComplexityRouterConfig modelInfo={modelInfo} value={value} onChange={onChange} />
);
const view = renderWithProviders(editor(initial));
const fast = () => screen.getByRole("switch", { name: "Fast mode for primary in the Efficient solver tier" });
expect(screen.getAllByRole("switch", { name: /^Fast mode for/ })).toHaveLength(1);
expect(fast()).not.toBeChecked();
await user.click(fast());
const enabled = onChange.mock.lastCall![0];
expect(enabled.tier_model_params?.SIMPLE.primary).toEqual({
reasoning_effort: "high",
max_tokens: 1024,
speed: "fast",
});
const saved = buildUpdatedComplexityRouterConfig({}, enabled);
expect(saved.tier_model_configs).toEqual({
SIMPLE: [{ model_name: "primary", litellm_params: { reasoning_effort: "high", max_tokens: 1024, speed: "fast" } }],
});
view.rerender(editor(hydrateComplexityRouterConfig(saved, undefined)));
expect(fast()).toBeChecked();
await user.click(fast());
expect(onChange.mock.lastCall![0].tier_model_params?.SIMPLE.primary).toEqual({
reasoning_effort: "high",
max_tokens: 1024,
});
});
describe("Fast mode metadata", () => {
it.each(["heuristic", "capability", "llm_v2"] as const)(
"can clear stored Fast mode without current capability metadata for %s",
async (classifier_type) => {
const user = userEvent.setup();
const value: ComplexityRouterConfigValue = {
tiers: { SIMPLE: ["primary"], MEDIUM: [], COMPLEX: [], REASONING: ["secondary"] },
classifier_type,
tier_model_params: { SIMPLE: { primary: { speed: "fast", max_tokens: 512 } } },
};
const onChange = vi.fn<(value: ComplexityRouterConfigValue) => void>();
const editor = (current: ComplexityRouterConfigValue, info: ModelGroup[]) => (
<ComplexityRouterConfig modelInfo={info} value={current} onChange={onChange} />
);
const view = renderWithProviders(editor(value, []));
const fast = () => screen.getByRole("switch", { name: /^Fast mode for primary/ });
expect(fast()).toBeChecked();
expect(onChange).not.toHaveBeenCalled();
view.rerender(editor(value, [{ model_group: "primary", supports_fast_mode: false }]));
expect(fast()).toBeChecked();
await user.click(fast());
const cleared = onChange.mock.lastCall![0];
expect(cleared.tier_model_params?.SIMPLE.primary).toEqual({ max_tokens: 512 });
const saved = buildUpdatedComplexityRouterConfig({}, cleared);
expect(saved.tier_model_configs).toEqual({
SIMPLE: [{ model_name: "primary", litellm_params: { max_tokens: 512 } }],
});
view.rerender(editor(hydrateComplexityRouterConfig(saved, undefined), []));
expect(screen.queryByRole("switch", { name: /^Fast mode for primary/ })).not.toBeInTheDocument();
view.rerender(editor(cleared, modelInfo));
expect(fast()).not.toBeChecked();
},
);
});
it.each(["MEDIUM", "REASONING"])("clears a legacy Capability pool while reconciling plan floor %s", async (floor) => {
const user = userEvent.setup();
const stored = {
classifier_type: "capability" as const,
plan_mode_min_tier: floor,
tiers: { SIMPLE: ["primary"], MEDIUM: ["secondary"], COMPLEX: [], REASONING: ["blocked"] },
tier_model_configs: { MEDIUM: [{ model_name: "secondary", litellm_params: { speed: "fast" } }] },
};
const value = hydrateComplexityRouterConfig(stored, undefined);
const onChange = vi.fn<(value: ComplexityRouterConfigValue) => void>();
renderWithProviders(<ComplexityRouterConfig modelInfo={modelInfo} value={value} onChange={onChange} />);
await user.click(screen.getByRole("button", { name: "Advanced routing options" }));
expect(screen.getByRole("switch", { name: "Fast mode for secondary in the Medium routing pool tier" })).toBeChecked();
expect(onChange).not.toHaveBeenCalled();
await user.click(screen.getByRole("combobox", { name: "Select medium routing pool models" }));
await user.click(await screen.findByRole("option", { name: "secondary" }));
await user.keyboard("{Escape}");
const cleared = onChange.mock.lastCall![0];
expect(cleared.tiers.MEDIUM).toEqual([]);
expect(cleared.plan_mode_min_tier).toBe(floor === "MEDIUM" ? undefined : floor);
expect(cleared.tier_model_params).toBeUndefined();
expect(buildUpdatedComplexityRouterConfig(stored, cleared).tiers).toEqual({
SIMPLE: ["primary"],
REASONING: ["blocked"],
});
});
it.each(["capability", "llm_v2"] as const)(
"shows and clears a persisted default model in %s",
async (classifier_type) => {
const user = userEvent.setup();
const stored = {
classifier_type,
default_model: "legacy-default",
tiers: { SIMPLE: ["primary"], REASONING: ["secondary"] },
};
const onChange = vi.fn<(value: ComplexityRouterConfigValue) => void>();
const editor = (value: ComplexityRouterConfigValue) => (
<ComplexityRouterConfig modelInfo={modelInfo} value={value} onChange={onChange} />
);
const view = renderWithProviders(editor(hydrateComplexityRouterConfig(stored, undefined)));
await user.click(screen.getByRole("button", { name: "Advanced routing options" }));
const select = () => screen.getByRole("combobox", { name: "Default model" });
expect(select()).toHaveValue("legacy-default");
expect(onChange).not.toHaveBeenCalled();
await user.click(select());
await user.click(await screen.findByRole("option", { name: "blocked" }));
const changed = onChange.mock.lastCall![0];
expect(buildUpdatedComplexityRouterConfig(stored, changed).default_model).toBe("blocked");
view.rerender(editor(changed));
await user.click(
within(screen.getByRole("group", { name: "Default model configuration" })).getByRole("button", { name: "Clear" }),
);
const cleared = onChange.mock.lastCall![0];
expect(cleared.default_model).toBeUndefined();
const saved = buildUpdatedComplexityRouterConfig(stored, cleared);
expect(saved).not.toHaveProperty("default_model");
view.rerender(editor(hydrateComplexityRouterConfig(saved, undefined)));
expect(select()).toHaveValue("");
expect(select()).toHaveAttribute("placeholder", expect.stringContaining("primary"));
},
);
it.each(["capability", "llm_v2"] as const)("offers only populated keyword targets for %s", async (classifier_type) => {
const user = userEvent.setup();
const value: ComplexityRouterConfigValue = {
classifier_type,
tiers: { SIMPLE: ["primary"], MEDIUM: [], COMPLEX: [], REASONING: ["secondary"] },
};
const onRulesChange = vi.fn<(rules: KeywordTierRule[]) => void>();
const editor = (rules: KeywordTierRule[]) => (
<ComplexityRouterConfig
value={value}
onChange={vi.fn()}
modelInfo={modelInfo}
keywordTierRules={rules}
onKeywordTierRulesChange={onRulesChange}
/>
);
const view = renderWithProviders(editor([]));
await user.click(screen.getByRole("button", { name: "Advanced routing options" }));
await user.click(screen.getByText("Advanced: Keyword/Semantic Matching"));
await user.click(screen.getByRole("button", { name: "Add keyword rule" }));
const rules = onRulesChange.mock.lastCall![0];
expect(rules[0].tier).toBe("SIMPLE");
view.rerender(editor(rules));
await user.click(screen.getByRole("combobox", { name: "Route keyword rule 1 to tier" }));
expect((await screen.findAllByRole("option")).map((option) => option.textContent)).toEqual(["Simple", "Reasoning"]);
});

View file

@ -0,0 +1,56 @@
import React from "react";
import { Info } from "lucide-react";
import { SearchSelect } from "@/components/shared/SearchSelect";
import { SimpleTooltip } from "@/components/ui/tooltip";
import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
import { isForecastClassifier } from "./forecast_classifier_config";
import { resolveComplexityDefaultModel } from "./tier_rows";
interface DefaultModelFieldProps {
value: ComplexityRouterConfigValue;
onChange: (value: ComplexityRouterConfigValue) => void;
modelOptions: { value: string; label: string }[];
}
const defaultModelPlaceholderFor = (derivedDefaultModel: string | undefined, isCustomSet: boolean): string => {
if (derivedDefaultModel) return `Derived from tiers: ${derivedDefaultModel}`;
return isCustomSet ? "Add a model to your fallback tier" : "Add a model to the Simple or Medium tier";
};
const DefaultModelField = ({ value, onChange, modelOptions }: DefaultModelFieldProps) => {
const defaultModelPlaceholder = defaultModelPlaceholderFor(
resolveComplexityDefaultModel(value),
Boolean(value.custom_tier_set),
);
// Clearing the select drops the key entirely rather than storing "", so an emptied pin reads as
// "track the tiers" everywhere downstream instead of as a blank model name.
const handleDefaultModelChange = (model: string | null | undefined) => {
onChange({ ...value, default_model: model || undefined });
};
return (
<div className="mt-4 mb-2" role="group" aria-label="Default model configuration">
<div className="flex items-center gap-2 mb-2">
<strong className="text-base font-semibold">Default Model</strong>
<SimpleTooltip content="Leave empty to follow the tiers. A model chosen here is pinned: it stays the default however the tiers change.">
<Info className="size-4 text-muted-foreground" />
</SimpleTooltip>
</div>
<SearchSelect
options={modelOptions}
value={value.default_model ?? ""}
onValueChange={handleDefaultModelChange}
placeholder={defaultModelPlaceholder}
emptyText="No models found"
aria-label="Default model"
/>
<span className="block mt-1 text-xs text-muted-foreground">
{isForecastClassifier(value.classifier_type)
? "Used when routing cannot find a suitable model. Classifier failures route to the capable solver."
: 'Used when the tier the request lands in has no model, and when the classifier fails with "Route to the default model" selected.'}
</span>
</div>
);
};
export default DefaultModelField;

View file

@ -0,0 +1,224 @@
import React, { useState } from "react";
import { describe, expect, it, vi } from "vitest";
import userEvent from "@testing-library/user-event";
import { fireEvent, renderWithProviders, screen } from "../../../tests/test-utils";
import ClassificationMethodConfig from "./ClassificationMethodConfig";
import AutoRouterClassifierTabs from "./AutoRouterClassifierTabs";
import ForecastClassifierConfig from "./ForecastClassifierConfig";
import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
import { getForecastConfigError, isForecastClassifier } from "./forecast_classifier_config";
import { buildUpdatedComplexityRouterConfig } from "../edit_auto_router/edit_auto_router_modal";
vi.mock("@/components/networking", async (importOriginal) => ({
...(await importOriginal<typeof import("@/components/networking")>()),
getComplexityScorerDefaults: vi.fn(async () => ({
tier_boundaries: {},
token_thresholds: {},
dimension_weights: {},
})),
}));
const initial: ComplexityRouterConfigValue = {
classifier_type: "capability",
classifier_llm_config: { model: "judge", timeout_ms: 20000 },
tiers: { SIMPLE: ["efficient"], MEDIUM: [], COMPLEX: [], REASONING: ["capable"] },
capability_classifier_config: { efficient_tier: "SIMPLE", capable_tier: "REASONING", base_threshold: 0.7 },
};
const fuseInitial: ComplexityRouterConfigValue = {
...initial,
classifier_type: "llm_v2",
capability_classifier_config: undefined,
adaptive: false,
llm_v2_config: {
efficient_profile: "Small solver",
capable_profile: "Larger solver",
harness: "One attempt",
max_quality_gap: 0.05,
},
};
const options = ["judge", "efficient", "capable"].map((model) => ({ value: model, label: model }));
function Form({ initialValue = initial }: { initialValue?: ComplexityRouterConfigValue }) {
const [value, setValue] = useState(initialValue);
const [saved, setSaved] = useState("");
return (
<>
<AutoRouterClassifierTabs value={value} onChange={setValue}>
{isForecastClassifier(value.classifier_type) ? (
<ForecastClassifierConfig
value={value}
onChange={setValue}
modelOptions={options}
effortOptionsByModel={{}}
/>
) : (
<ClassificationMethodConfig
value={value}
onChange={setValue}
modelOptions={options}
effortOptionsByModel={{}}
/>
)}
</AutoRouterClassifierTabs>
<button
disabled={Boolean(getForecastConfigError(value))}
onClick={() => setSaved(JSON.stringify(buildUpdatedComplexityRouterConfig({}, value)))}
>
Save configuration
</button>
<output aria-label="Saved configuration">{saved}</output>
</>
);
}
describe("forecast classifier form", () => {
it("switches a populated standard router to Capability without saving hidden pools or their overrides", () => {
renderWithProviders(
<Form
initialValue={{
classifier_type: "llm",
classifier_llm_config: { model: "judge", timeout_ms: 20000, classification_rubric: "agentic" },
adaptive: true,
plan_mode_min_tier: "MEDIUM",
tiers: {
SIMPLE: ["efficient", "second-efficient"],
MEDIUM: ["leftover-medium"],
COMPLEX: ["leftover-complex"],
REASONING: ["capable"],
},
tier_model_params: {
SIMPLE: { efficient: { reasoning_effort: "low", speed: "fast", max_tokens: 1024 } },
MEDIUM: { "leftover-medium": { speed: "fast" } },
COMPLEX: { "leftover-complex": { max_tokens: 4096 } },
REASONING: { capable: { reasoning_effort: "high" } },
},
}}
/>,
);
fireEvent.click(screen.getByRole("tab", { name: "Capability" }));
fireEvent.change(screen.getByLabelText("Solve probability threshold"), { target: { value: "0.7" } });
expect(screen.getByRole("button", { name: "Save configuration" })).toBeEnabled();
fireEvent.click(screen.getByRole("button", { name: "Save configuration" }));
const output = screen.getByRole("status", { name: "Saved configuration" });
expect(output).toHaveTextContent('"classifier_type":"capability"');
expect(output).toHaveTextContent('"SIMPLE":["efficient","second-efficient"]');
expect(output).toHaveTextContent('"REASONING":["capable"]');
expect(output).toHaveTextContent('"reasoning_effort":"low","speed":"fast","max_tokens":1024');
expect(output).toHaveTextContent('"reasoning_effort":"high"');
expect(output).toHaveTextContent('"adaptive":false');
expect(output).not.toHaveTextContent("leftover-medium");
expect(output).not.toHaveTextContent("leftover-complex");
expect(output).not.toHaveTextContent('"plan_mode_min_tier"');
});
it.each(["capability", "llm_v2"] as const)(
"carries non-default solver assignments when switching away from %s",
(source) => {
const pair = { efficient_tier: "MEDIUM", capable_tier: "COMPLEX" };
const previous: ComplexityRouterConfigValue = {
...(source === "capability" ? initial : fuseInitial),
tiers: { SIMPLE: [], MEDIUM: ["efficient"], COMPLEX: ["capable"], REASONING: [] },
capability_classifier_config:
source === "capability" ? { ...initial.capability_classifier_config!, ...pair } : undefined,
llm_v2_config: source === "llm_v2" ? { ...fuseInitial.llm_v2_config!, ...pair } : undefined,
plan_mode_min_tier: "COMPLEX",
tier_model_params: { MEDIUM: { efficient: { max_tokens: 128 } }, COMPLEX: { capable: { speed: "fast" } } },
};
renderWithProviders(<Form initialValue={previous} />);
fireEvent.click(screen.getByRole("tab", { name: source === "capability" ? "Fuse v2" : "Capability" }));
if (source === "capability") {
fireEvent.change(screen.getByLabelText("Efficient solver profile"), { target: { value: "Small solver" } });
fireEvent.change(screen.getByLabelText("Capable solver profile"), { target: { value: "Large solver" } });
fireEvent.change(screen.getByLabelText("Harness and budget"), { target: { value: "One attempt" } });
fireEvent.change(screen.getByLabelText("Maximum quality gap"), { target: { value: "0.05" } });
} else {
fireEvent.change(screen.getByLabelText("Solve probability threshold"), { target: { value: "0.7" } });
}
expect(screen.getByRole("button", { name: "Save configuration" })).toBeEnabled();
fireEvent.click(screen.getByRole("button", { name: "Save configuration" }));
const output = screen.getByRole("status", { name: "Saved configuration" });
expect(output).toHaveTextContent('"efficient_tier":"MEDIUM","capable_tier":"COMPLEX"');
expect(output).toHaveTextContent('"tiers":{"MEDIUM":["efficient"],"COMPLEX":["capable"]}');
expect(output).toHaveTextContent('"plan_mode_min_tier":"COMPLEX"');
expect(output).toHaveTextContent('"max_tokens":128');
expect(output).toHaveTextContent('"speed":"fast"');
},
);
it("keeps decimal and negative numbers when entered one character at a time", async () => {
const user = userEvent.setup();
renderWithProviders(<Form />);
const threshold = screen.getByLabelText("Solve probability threshold");
await user.clear(threshold);
await user.type(threshold, "0.65");
expect(threshold).toHaveValue(0.65);
await user.click(screen.getByRole("button", { name: "Classifier options" }));
await user.click(screen.getByRole("switch", { name: "Use fitted calibration" }));
await user.type(screen.getByLabelText("Efficient intercept"), "-0.3");
expect(screen.getByLabelText("Efficient intercept")).toHaveValue(-0.3);
});
it.each([
["capability", "LLM Classifier"],
["capability", "Heuristic first"],
["capability", "Hybrid"],
["llm_v2", "LLM Classifier"],
["llm_v2", "Heuristic first"],
["llm_v2", "Hybrid"],
] as const)("restores the current rubric when switching %s through Complexity to %s", async (source, target) => {
const user = userEvent.setup();
renderWithProviders(<Form initialValue={source === "capability" ? initial : fuseInitial} />);
fireEvent.click(screen.getByRole("tab", { name: "Complexity" }));
fireEvent.click(screen.getByRole("radio", { name: new RegExp(`^${target}`) }));
await user.click(screen.getByRole("combobox", { name: "Classifier Model" }));
await user.click(screen.getByRole("option", { name: "judge", exact: true }));
fireEvent.click(screen.getByRole("button", { name: "Save configuration" }));
const output = screen.getByRole("status", { name: "Saved configuration" });
expect(output).toHaveTextContent('"classification_rubric":"agentic"');
expect(output).toHaveTextContent('"model":"judge"');
expect(output).toHaveTextContent('"timeout_ms":3000');
expect(output).not.toHaveTextContent('"capability_classifier_config"');
expect(output).not.toHaveTextContent('"llm_v2_config"');
});
it("saves capability threshold edits together with fitted calibration", () => {
renderWithProviders(<Form />);
fireEvent.change(screen.getByLabelText("Solve probability threshold"), { target: { value: "0.6" } });
fireEvent.click(screen.getByRole("button", { name: "Classifier options" }));
fireEvent.click(screen.getByRole("switch", { name: "Use fitted calibration" }));
expect(screen.getByRole("button", { name: "Save configuration" })).toBeDisabled();
fireEvent.change(screen.getByLabelText("Calibration version"), { target: { value: "eval-a" } });
fireEvent.change(screen.getByLabelText("Efficient slope"), { target: { value: "1.2" } });
fireEvent.change(screen.getByLabelText("Efficient intercept"), { target: { value: "-0.3" } });
fireEvent.click(screen.getByRole("button", { name: "Save configuration" }));
const output = screen.getByRole("status", { name: "Saved configuration" });
expect(output).toHaveTextContent('"base_threshold":0.6');
expect(output).toHaveTextContent('"calibration":{"version":"eval-a","slope":1.2,"intercept":-0.3}');
fireEvent.change(screen.getByLabelText("Solve probability threshold"), { target: { value: "" } });
expect(screen.getByRole("button", { name: "Save configuration" })).toBeDisabled();
});
it("switches to Fuse, requires solver context, and saves the filled fields", () => {
renderWithProviders(<Form />);
fireEvent.click(screen.getByRole("tab", { name: "Fuse v2" }));
expect(screen.queryByLabelText("Solve probability threshold")).not.toBeInTheDocument();
expect(screen.getByRole("button", { name: "Save configuration" })).toBeDisabled();
fireEvent.change(screen.getByLabelText("Efficient solver profile"), {
target: { value: "Short reasoning budget" },
});
fireEvent.change(screen.getByLabelText("Capable solver profile"), { target: { value: "Larger reasoning budget" } });
fireEvent.change(screen.getByLabelText("Harness and budget"), {
target: { value: "Shell and test runner, one attempt" },
});
fireEvent.change(screen.getByLabelText("Maximum quality gap"), { target: { value: "0.05" } });
fireEvent.click(screen.getByRole("button", { name: "Save configuration" }));
const output = screen.getByRole("status", { name: "Saved configuration" });
expect(output).toHaveTextContent('"classifier_type":"llm_v2"');
expect(output).toHaveTextContent('"efficient_profile":"Short reasoning budget"');
expect(output).toHaveTextContent('"capable_profile":"Larger reasoning budget"');
expect(output).toHaveTextContent('"harness":"Shell and test runner, one attempt"');
expect(output).toHaveTextContent('"max_quality_gap":0.05');
expect(output).toHaveTextContent('"adaptive":false');
expect(output).not.toHaveTextContent('"capability_classifier_config"');
});
});

View file

@ -0,0 +1,433 @@
import React from "react";
import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
import { ChevronRight } from "lucide-react";
import { Input } from "@/components/ui/input";
import { Label } from "@/components/ui/label";
import { Textarea } from "@/components/ui/textarea";
import { Switch } from "@/components/ui/switch";
import { SearchSelect } from "@/components/shared/SearchSelect";
import { MultiSelect } from "@/components/shared/MultiSelect";
import {
type ComplexityRouterConfigValue,
type ClassificationFrequency,
classificationFrequency,
withClassificationFrequency,
DEFAULT_CLASSIFIER_TIMEOUT_MS,
} from "./ComplexityRouterConfig";
import {
forecastTierNames,
forecastModels,
getForecastConfigError,
newCapabilitySettings,
newFuseSettings,
type CapabilitySettings,
type FuseSettings,
} from "./forecast_classifier_config";
import ClassifierReasoningEffortSelect from "./ClassifierReasoningEffortSelect";
import ClassifierCircuitBreakerConfig from "./ClassifierCircuitBreakerConfig";
import ClassifierVisionConfig from "./ClassifierVisionConfig";
import TierModelEffortRows from "./TierModelEffortRows";
import { activeTierRows } from "./tier_rows";
import { setTierModels } from "./tier_set_actions";
import { tierRowLabel, setTierModelParam, setTierModelReasoningEffort } from "./complexity_router_tiers";
interface Props {
value: ComplexityRouterConfigValue;
onChange: (value: ComplexityRouterConfigValue) => void;
modelOptions: { value: string; label: string }[];
effortOptionsByModel: Record<string, string[] | null | undefined>;
}
const NumberField = ({
label,
value,
onChange,
min,
max,
step = "any",
help,
}: {
label: string;
value: number;
onChange: (value: number) => void;
min?: number;
max?: number;
step?: number | "any";
help?: string;
}) => {
const id = React.useId();
return (
<div className="space-y-1">
<Label htmlFor={id}>{label}</Label>
<Input
id={id}
type="number"
min={min}
max={max}
step={step}
value={Number.isFinite(value) ? value : ""}
onChange={(event) => onChange(event.target.value === "" ? Number.NaN : Number(event.target.value))}
/>
{help && <p className="text-xs text-muted-foreground">{help}</p>}
</div>
);
};
export const ForecastSolverModels = ({
value,
onChange,
modelOptions,
effortOptionsByModel,
fastModeByModel,
additionalPoolsOnly = false,
}: Props & { fastModeByModel: Record<string, boolean>; additionalPoolsOnly?: boolean }) => {
const id = React.useId();
const names = forecastTierNames(value);
const additionalRows =
value.classifier_type === "capability"
? activeTierRows(value)
.filter((row) => !names.includes(row.id) && row.models.length > 0)
.map((row) => ({ tier: row.id, label: `${tierRowLabel(row, value.tier_labels)} routing pool` }))
: [];
const rows = additionalPoolsOnly
? additionalRows
: names.map((tier, index) => ({ tier, label: index === 0 ? "Efficient solver" : "Capable solver" }));
if (rows.length === 0) return null;
return (
<div className="rounded-lg border p-4 space-y-4">
{rows.map(({ tier, label }) => {
const models = forecastModels(value.tiers, tier);
const setModels = (next: string[]) => onChange(setTierModels(value, tier, next));
return (
<div key={tier} className="space-y-2">
<Label htmlFor={`${id}-${tier}`} className="block text-sm font-semibold">
{label}
</Label>
{value.classifier_type === "llm_v2" ? (
<SearchSelect
options={modelOptions}
inputId={`${id}-${tier}`}
value={models[0] ?? ""}
aria-label={label}
placeholder={`Select ${label.toLowerCase()}`}
onValueChange={(model) => setModels(model ? [model] : [])}
/>
) : (
<MultiSelect
options={modelOptions}
id={`${id}-${tier}`}
value={models}
onValueChange={setModels}
placeholder={`Select ${label.toLowerCase()} models`}
/>
)}
<TierModelEffortRows
tierLabel={label}
models={models}
effortOptionsByModel={Object.fromEntries(
Object.entries(effortOptionsByModel).map(([model, efforts]) => [model, efforts ?? []]),
)}
paramsByModel={value.tier_model_params?.[tier] ?? {}}
fastModeByModel={fastModeByModel}
onFastModeChange={(model, enabled) =>
onChange({
...value,
tier_model_params: setTierModelParam(value.tier_model_params, tier, model, [
"speed",
enabled ? "fast" : undefined,
]),
})
}
onEffortChange={(model, effort) =>
onChange({
...value,
tier_model_params: setTierModelReasoningEffort(value.tier_model_params, tier, model, effort),
})
}
/>
</div>
);
})}
{!additionalPoolsOnly && (
<p className="text-sm text-muted-foreground">
Invalid forecasts and classifier failures route to the capable solver
</p>
)}
</div>
);
};
const CalibrationFields = ({
label,
value,
onChange,
bounded = false,
}: {
label: string;
bounded?: boolean;
value: { slope: number; intercept: number };
onChange: (value: { slope: number; intercept: number }) => void;
}) => (
<div className="grid gap-3 sm:grid-cols-2">
<NumberField
label={`${label} slope`}
value={value.slope}
min={0}
max={bounded ? 20 : undefined}
onChange={(slope) => onChange({ ...value, slope })}
/>
<NumberField
label={`${label} intercept`}
value={value.intercept}
min={bounded ? -20 : undefined}
max={bounded ? 20 : undefined}
onChange={(intercept) => onChange({ ...value, intercept })}
/>
</div>
);
const emptyCoefficients = () => ({ slope: Number.NaN, intercept: Number.NaN });
const ForecastClassifierConfig = ({ value, onChange, modelOptions, effortOptionsByModel }: Props) => {
const id = React.useId();
const isCapability = value.classifier_type === "capability";
const capability = value.capability_classifier_config ?? newCapabilitySettings();
const fuse = value.llm_v2_config ?? newFuseSettings();
const config = isCapability ? capability : fuse;
const llm = value.classifier_llm_config ?? { model: "", timeout_ms: DEFAULT_CLASSIFIER_TIMEOUT_MS };
const updateCapability = (next: CapabilitySettings) => onChange({ ...value, capability_classifier_config: next });
const updateFuse = (next: FuseSettings) => onChange({ ...value, llm_v2_config: next });
const updateTransport = (patch: { max_output_tokens?: number; response_format?: "json_schema" | "json_object" }) =>
isCapability ? updateCapability({ ...capability, ...patch }) : updateFuse({ ...fuse, ...patch });
const setCalibrationVersion = (version: string) => {
if (isCapability && capability.calibration)
updateCapability({ ...capability, calibration: { ...capability.calibration, version } });
if (!isCapability && fuse.calibration) updateFuse({ ...fuse, calibration: { ...fuse.calibration, version } });
};
const error = getForecastConfigError(value);
return (
<div className="mt-4 space-y-4">
<p className="text-sm text-muted-foreground">
{isCapability
? "Forecasts whether the efficient solver can complete the task using the bundled capability card"
: "Forecasts success for both solvers and selects efficient when the estimated quality gap is within your allowance"}
</p>
<div className="space-y-1">
<Label htmlFor={`${id}-judge`}>Judge model</Label>
<SearchSelect
inputId={`${id}-judge`}
aria-label="Judge model"
options={modelOptions}
value={llm.model}
placeholder="Select the judge model"
onValueChange={(model) => {
if (model === llm.model) return;
onChange({ ...value, classifier_llm_config: { ...llm, model: model ?? "", reasoning_effort: undefined } });
}}
/>
</div>
{isCapability ? (
<>
<NumberField
label="Solve probability threshold"
value={capability.base_threshold}
min={0}
max={1}
help="Minimum estimated chance of whole-task success required to use the efficient solver"
onChange={(base_threshold) => updateCapability({ ...capability, base_threshold })}
/>
</>
) : (
<>
{(["efficient_profile", "capable_profile", "harness"] as const).map((field) => {
const label = {
efficient_profile: "Efficient solver profile",
capable_profile: "Capable solver profile",
harness: "Harness and budget",
}[field];
return (
<div key={field} className="space-y-1">
<Label htmlFor={`${id}-${field}`}>{label}</Label>
<Textarea
id={`${id}-${field}`}
value={fuse[field]}
maxLength={4000}
placeholder={
field === "harness"
? "Tools, execution environment, verification, and budget available to each solver"
: "Describe this solver's strengths, limitations, and settings"
}
onChange={(event) => updateFuse({ ...fuse, [field]: event.target.value })}
/>
</div>
);
})}
<NumberField
label="Maximum quality gap"
value={fuse.max_quality_gap}
min={0}
max={1}
help="Allowed difference between capable and efficient success probabilities, from 0 to 1. This is an estimate, not a measured quality guarantee"
onChange={(max_quality_gap) => updateFuse({ ...fuse, max_quality_gap })}
/>
</>
)}
<Collapsible className="rounded-lg border">
<CollapsibleTrigger className="group flex w-full items-center gap-2 px-4 py-3 text-left font-medium">
<ChevronRight className="size-4 transition-transform group-data-panel-open:rotate-90" />
Classifier options
</CollapsibleTrigger>
<CollapsibleContent className="space-y-4 px-4 pb-4">
<ClassifierReasoningEffortSelect
model={llm.model}
value={llm.reasoning_effort}
explicitlySupported={effortOptionsByModel[llm.model]}
onChange={(reasoning_effort) => onChange({ ...value, classifier_llm_config: { ...llm, reasoning_effort } })}
/>
<NumberField
label="Timeout (ms)"
min={1}
step={1}
value={llm.timeout_ms}
help="Allow enough time for the judge to produce its forecast"
onChange={(timeout_ms) => onChange({ ...value, classifier_llm_config: { ...llm, timeout_ms } })}
/>
<ClassifierCircuitBreakerConfig
value={llm}
onChange={(classifier_llm_config) => onChange({ ...value, classifier_llm_config })}
/>
<ClassifierVisionConfig
value={llm}
onChange={(classifier_llm_config) => onChange({ ...value, classifier_llm_config })}
/>
<div className="space-y-1">
<Label htmlFor={`${id}-frequency`}>How often to classify</Label>
<SearchSelect
inputId={`${id}-frequency`}
aria-label="How often to classify"
value={classificationFrequency(value)}
allowClear={false}
options={[
{ value: "every_request", label: "Every request" },
{ value: "user_turn", label: "Every new user message" },
{ value: "session", label: "Once per session" },
]}
onValueChange={(frequency) => {
if (frequency) onChange(withClassificationFrequency(value, frequency as ClassificationFrequency));
}}
/>
</div>
{isCapability && (
<NumberField
label="Capability boundary step"
value={capability.threshold_step ?? 0}
min={0}
max={0.5}
help="Added once for uncertain or unmatched tasks and twice for unsupported tasks; the final threshold cannot exceed 1"
onChange={(threshold_step) => updateCapability({ ...capability, threshold_step })}
/>
)}
<NumberField
label="Classifier output token limit"
min={1}
step={1}
value={config.max_output_tokens ?? (isCapability ? 4096 : 1024)}
onChange={(max_output_tokens) => updateTransport({ max_output_tokens })}
/>
<div className="space-y-1">
<Label htmlFor={`${id}-format`}>Forecast response format</Label>
<SearchSelect
inputId={`${id}-format`}
aria-label="Forecast response format"
value={config.response_format ?? "json_schema"}
allowClear={false}
options={[
{ value: "json_schema", label: "Strict JSON schema" },
{ value: "json_object", label: "JSON object (for judges without strict schema support)" },
]}
onValueChange={(response_format) => {
if (response_format === "json_schema" || response_format === "json_object")
updateTransport({ response_format });
}}
/>
</div>
<div className="space-y-3 rounded-md border p-3">
<Label>
<Switch
checked={Boolean(config.calibration)}
onCheckedChange={(enabled) =>
isCapability
? updateCapability({
...capability,
calibration: enabled ? { version: "", ...emptyCoefficients() } : undefined,
})
: updateFuse({
...fuse,
calibration: enabled
? {
version: "",
prompt_version: "llm-v2-1",
efficient: emptyCoefficients(),
capable: emptyCoefficients(),
}
: undefined,
})
}
/>
Use fitted calibration
</Label>
<p className="text-xs text-muted-foreground">
Optional coefficients fitted for your judge, solvers, and harness. Leave off to use raw forecasts
</p>
{config.calibration && (
<div className="space-y-1">
<Label htmlFor={`${id}-version`}>Calibration version</Label>
<Input
id={`${id}-version`}
value={config.calibration.version}
maxLength={isCapability ? 128 : 512}
onChange={(event) => setCalibrationVersion(event.target.value)}
/>
</div>
)}
{isCapability && capability.calibration && (
<CalibrationFields
label="Efficient"
bounded
value={capability.calibration}
onChange={(next) =>
updateCapability({
...capability,
calibration: { version: capability.calibration?.version ?? "", ...next },
})
}
/>
)}
{!isCapability &&
fuse.calibration &&
(["efficient", "capable"] as const).map((role) => (
<CalibrationFields
key={role}
label={role === "efficient" ? "Efficient" : "Capable"}
value={fuse.calibration![role]}
onChange={(next) => {
if (fuse.calibration) updateFuse({ ...fuse, calibration: { ...fuse.calibration, [role]: next } });
}}
/>
))}
</div>
</CollapsibleContent>
</Collapsible>
<p className="text-xs text-muted-foreground">
The classifier uses its bundled prompt and always falls back to the capable solver
</p>
{error && (
<p role="alert" className="text-sm text-destructive">
{error}
</p>
)}
</div>
);
};
export default ForecastClassifierConfig;

View file

@ -0,0 +1,44 @@
import React from "react";
import { Switch } from "@/components/ui/switch";
import TierRowSelect from "./TierRowSelect";
import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
const PlanModeOverrideControls: React.FC<{
value: ComplexityRouterConfigValue;
onChange: (value: ComplexityRouterConfigValue) => void;
planModeTierOptions: { value: string; label: string }[];
}> = ({ value, onChange, planModeTierOptions }) => (
<>
<div className="flex items-center gap-2 mb-2">
<Switch
checked={value.plan_mode_min_tier !== undefined}
disabled={planModeTierOptions.length === 0}
onCheckedChange={(enabled) =>
onChange({
...value,
plan_mode_min_tier: enabled ? planModeTierOptions.at(-1)?.value : undefined,
})
}
aria-label="Route plan-mode requests to a minimum tier"
/>
<strong className="font-semibold">Route plan-mode requests to a minimum tier</strong>
</div>
<span className="block text-xs mb-3 text-muted-foreground">
Requests from coding agents in plan mode (Claude Code, GitHub Copilot) route to at least this tier. The classifier
still wins when it picks higher, and the override only lasts while plan mode is active.
{planModeTierOptions.length === 0 && " Add models to a tier to enable this."}
</span>
{value.plan_mode_min_tier !== undefined && (
<div style={{ maxWidth: 320 }}>
<TierRowSelect
label="Plan-mode minimum tier"
options={planModeTierOptions}
value={value.plan_mode_min_tier ?? null}
onValueChange={(tier) => onChange({ ...value, plan_mode_min_tier: tier })}
/>
</div>
)}
</>
);
export default PlanModeOverrideControls;

View file

@ -0,0 +1,23 @@
import React from "react";
import { ChevronRight } from "lucide-react";
import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
interface RoutingOptionsProps {
forecast: boolean;
children: React.ReactNode;
}
const RoutingOptions = ({ forecast, children }: RoutingOptionsProps) =>
forecast ? (
<Collapsible className="rounded-lg border">
<CollapsibleTrigger className="group flex w-full items-center gap-2 px-4 py-3 text-left font-medium">
<ChevronRight className="size-4 transition-transform group-data-panel-open:rotate-90" />
Advanced routing options
</CollapsibleTrigger>
<CollapsibleContent className="space-y-4 px-4 pb-4">{children}</CollapsibleContent>
</Collapsible>
) : (
<>{children}</>
);
export default RoutingOptions;

View file

@ -23,6 +23,12 @@ interface TierModelEffortRowsProps {
onFastModeChange: (model: string, enabled: boolean) => void;
}
const canEditFastMode = (
model: string,
fastModeByModel: TierModelEffortRowsProps["fastModeByModel"],
paramsByModel: TierModelEffortRowsProps["paramsByModel"],
): boolean => fastModeByModel?.[model] === true || paramsByModel?.[model]?.speed === "fast";
export interface TierEffortRow {
model: string;
effort: ReasoningEffort | undefined;
@ -49,7 +55,7 @@ export const tierEffortRows = ({
const listed = effort !== undefined && !supported.includes(effort) ? [...supported, effort] : supported;
return { model, effort, options: Array.from(new Set(listed)) };
})
.filter(({ model, options }) => options.length > 0 || fastModeByModel?.[model] === true);
.filter(({ model, options }) => options.length > 0 || canEditFastMode(model, fastModeByModel, paramsByModel));
const TierModelEffortRows: React.FC<TierModelEffortRowsProps> = (props) => {
const { tierLabel, paramsByModel, onEffortChange, fastModeByModel, onFastModeChange } = props;
@ -101,7 +107,7 @@ const TierModelEffortRows: React.FC<TierModelEffortRowsProps> = (props) => {
</SelectContent>
</Select>
)}
{fastModeByModel?.[model] === true && (
{canEditFastMode(model, fastModeByModel, paramsByModel) && (
<SimpleTooltip content="Fast mode has higher pricing and requires an eligible provider account. Off removes this tier's speed override and inherits the request or provider default">
<label
className="flex items-center gap-2 text-xs"

View file

@ -2,9 +2,17 @@ import React from "react";
import { CUSTOM_TIER_RESTRICTIONS, CustomTierSet, TierRestriction } from "./tier_rows";
export const restrictedBy = (
value: { custom_tier_set?: CustomTierSet },
value: { custom_tier_set?: CustomTierSet; classifier_type?: string },
key: keyof typeof CUSTOM_TIER_RESTRICTIONS,
): TierRestriction | undefined => (value.custom_tier_set ? CUSTOM_TIER_RESTRICTIONS[key] : undefined);
): TierRestriction | undefined => {
if (value.custom_tier_set) return CUSTOM_TIER_RESTRICTIONS[key];
if (value.classifier_type === "llm_v2" && key === "adaptive")
return {
omit: ["adaptive", "adaptive_weights", "adaptive_eligible", "tier_distance_penalty"],
reason: "Fuse v2 uses its quality-gap decision directly; adaptive routing is unavailable",
};
return undefined;
};
export const Restricted: React.FC<{ by: TierRestriction | undefined; children: React.ReactNode }> = ({
by,

View file

@ -167,6 +167,124 @@ describe("AddAutoRouterTab", () => {
mockFetchAllModelDeployments.mockResolvedValue([]);
});
it.each(["Capability", "Fuse v2"])(
"creates %s from its dedicated tab without complexity templates",
async (label) => {
const user = userEvent.setup();
mockFetchAvailableModels.mockResolvedValue([
{ model_group: "efficient", mode: "chat" },
{ model_group: "capable", mode: "chat" },
{ model_group: "judge", mode: "chat" },
]);
renderWithProviders(<Harness />);
await user.type(screen.getByLabelText("Auto Router Name"), "forecast-router");
await user.click(screen.getByRole("tab", { name: label, exact: true }));
expect(screen.getByLabelText("Auto Router Name")).toHaveValue("forecast-router");
expect(screen.queryByTestId("template-selector")).not.toBeInTheDocument();
expect(screen.queryByTestId("configure-automatically-button")).not.toBeInTheDocument();
expect(screen.queryByTestId("detailed-configuration-toggle")).not.toBeInTheDocument();
expect(screen.queryByText("Complexity Tier Configuration")).not.toBeInTheDocument();
expect(screen.queryByText("Advanced: Classification Method")).not.toBeInTheDocument();
expect(screen.queryByText("Advanced: Adaptive Routing")).not.toBeInTheDocument();
const capability = label === "Capability";
for (const [role, model] of [
["Efficient", "efficient"],
["Capable", "capable"],
]) {
await user.click(
screen.getByRole("combobox", {
name: capability ? `Select ${role.toLowerCase()} solver models` : `${role} solver`,
}),
);
await user.click(await screen.findByRole("option", { name: model, exact: true }));
if (capability) await user.keyboard("{Escape}");
}
await user.click(screen.getByRole("combobox", { name: "Judge model" }));
await user.click(await screen.findByRole("option", { name: "judge", exact: true }));
expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeDisabled();
if (capability) {
fireEvent.change(screen.getByLabelText("Solve probability threshold"), { target: { value: "0.7" } });
} else {
fireEvent.change(screen.getByLabelText("Efficient solver profile"), { target: { value: "Small solver" } });
fireEvent.change(screen.getByLabelText("Capable solver profile"), { target: { value: "Large solver" } });
fireEvent.change(screen.getByLabelText("Harness and budget"), { target: { value: "One attempt" } });
fireEvent.change(screen.getByLabelText("Maximum quality gap"), { target: { value: "0.05" } });
}
expect(screen.queryByRole("alert")).not.toBeInTheDocument();
expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeEnabled();
await user.click(screen.getByRole("button", { name: "Advanced routing options" }));
for (const label of ["Adaptive Routing", "Context Window Escalation", "Escalation Keywords"]) {
expect(screen.queryByText(`Advanced: ${label}`)).not.toBeInTheDocument();
}
expect(screen.getByText("Advanced: Stalled Task Escalation")).toBeInTheDocument();
expect(screen.getByText("Advanced: Response Format")).toBeInTheDocument();
expect(screen.queryByText("Advanced: Classification Method")).not.toBeInTheDocument();
expect(screen.getByText("Advanced: Affinity")).toBeInTheDocument();
await user.click(screen.getByRole("button", { name: "Add Auto Router" }));
await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalledTimes(1));
const expected = {
classifier_type: capability ? "capability" : "llm_v2",
adaptive: false,
enable_context_window_escalation: false,
escalation_keywords: [],
tiers: { SIMPLE: ["efficient"], REASONING: ["capable"] },
classifier_llm_config: { model: "judge" },
};
expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls[0][0].complexity_router_config).toMatchObject(expected);
},
);
it("restores the automatic/template/detail flow on the Complexity tab", async () => {
const user = userEvent.setup();
mockFetchAvailableModels.mockResolvedValue(ALL_FAMILY_MODELS);
renderWithProviders(<Harness />);
await screen.findByTestId("configure-automatically-button");
await user.click(screen.getByRole("tab", { name: "Capability", exact: true }));
expect(screen.queryByTestId("configure-automatically-button")).not.toBeInTheDocument();
await user.click(screen.getByRole("tab", { name: "Complexity", exact: true }));
expect(screen.getByTestId("configure-automatically-button")).toBeInTheDocument();
expect(screen.getByTestId("template-selector")).toBeInTheDocument();
expandDetailedConfiguration();
expect(screen.getByText("Complexity Tier Configuration")).toBeInTheDocument();
for (const label of ["Adaptive Routing", "Context Window Escalation", "Escalation Keywords"]) {
expect(screen.getByText(`Advanced: ${label}`)).toBeInTheDocument();
}
await user.click(screen.getByText("Advanced: Classification Method"));
expect(screen.queryByRole("radio", { name: /^Capability/ })).not.toBeInTheDocument();
expect(screen.queryByRole("radio", { name: /^Fuse v2/ })).not.toBeInTheDocument();
expect(screen.getByRole("radio", { name: /^Heuristic \(default/ })).toBeChecked();
});
it.each(["Capability", "Fuse v2"])(
"retries failed model loading on %s without losing entered settings",
async (label) => {
const user = userEvent.setup();
mockFetchAvailableModels.mockRejectedValueOnce(new Error("Model list unavailable")).mockResolvedValue([
{ model_group: "efficient", mode: "chat" },
{ model_group: "capable", mode: "chat" },
{ model_group: "judge", mode: "chat" },
]);
renderWithProviders(<Harness />);
await user.click(screen.getByRole("tab", { name: label, exact: true }));
await user.type(screen.getByLabelText("Auto Router Name"), "forecast-retry");
const capability = label === "Capability";
const policyField = capability ? "Solve probability threshold" : "Efficient solver profile";
fireEvent.change(screen.getByLabelText(policyField), {
target: { value: capability ? "0.7" : "Small solver" },
});
expect(await screen.findByText("Could not load available models.")).toBeVisible();
expect(screen.queryByTestId("template-selector")).not.toBeInTheDocument();
await user.click(screen.getByRole("button", { name: "Retry", exact: true }));
await waitFor(() => expect(screen.queryByText("Could not load available models.")).not.toBeInTheDocument());
expect(screen.getByRole("tab", { name: label, exact: true })).toHaveAttribute("aria-selected", "true");
expect(screen.getByLabelText("Auto Router Name")).toHaveValue("forecast-retry");
expect(screen.getByLabelText(policyField)).toHaveValue(capability ? 0.7 : "Small solver");
await user.click(screen.getByRole("combobox", { name: "Judge model" }));
expect(await screen.findByRole("option", { name: "judge", exact: true })).toBeVisible();
expect(mockFetchAvailableModels).toHaveBeenCalledTimes(2);
},
);
// Detailed Configuration starts collapsed so the modal opens onto just Name + Template; a caller
// opts into the full tier/classifier form rather than always seeing it up front.
it("keeps Detailed Configuration collapsed until a caller opens it", () => {

View file

@ -1,3 +1,5 @@
import AutoRouterClassifierTabs from "./AutoRouterClassifierTabs";
import { getForecastConfigError, isForecastClassifier } from "../add_model/forecast_classifier_config";
import React, { useEffect, useState } from "react";
import { useQuery } from "@tanstack/react-query";
import { useWatch } from "react-hook-form";
@ -138,7 +140,9 @@ export const getSubmitBlockedReason = (
(config.custom_tier_set
? getCustomTierRowsError(config.custom_tier_set)
: getTierLabelsError(config.tier_labels)) ??
getMissingTiersError(activeTierRows(config)) ??
(isForecastClassifier(config.classifier_type)
? getForecastConfigError(config)
: getMissingTiersError(activeTierRows(config))) ??
getPlanModeTierError(config.plan_mode_min_tier, activeTierRows(config)) ??
getKeywordTierRulesError(keywordTierRules, activeTierRows(config)) ??
getClassifierModelError(config) ??
@ -401,6 +405,8 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
classificationMode: complexityRouterConfig.classification_mode,
tierLabels: complexityRouterConfig.tier_labels,
classifierType: complexityRouterConfig.classifier_type,
capabilityClassifierConfig: complexityRouterConfig.capability_classifier_config,
llmV2Config: complexityRouterConfig.llm_v2_config,
classifierLlmConfig: complexityRouterConfig.classifier_llm_config,
classifierContextWindowSize: complexityRouterConfig.classifier_context_window_size,
classifierContextBudgetChars: complexityRouterConfig.classifier_context_budget_chars,
@ -543,81 +549,138 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
setIsTestModalVisible(true);
};
const configurationForm = (
<ComplexityRouterConfig
editingTiers={editingTiers}
onEditingTiersChange={setEditingTiers}
modelInfo={modelInfo}
value={complexityRouterConfig}
onChange={setComplexityRouterConfig}
customTechnicalKeywords={customTechnicalKeywords}
onCustomTechnicalKeywordsChange={setCustomTechnicalKeywords}
keywordTierRules={keywordTierRules}
onKeywordTierRulesChange={setKeywordTierRules}
keywordRulesError={getKeywordTierRulesError(keywordTierRules, activeTierRows(complexityRouterConfig))}
semanticMatchingEnabled={semanticMatchingEnabled}
onSemanticMatchingEnabledChange={setSemanticMatchingEnabled}
embeddingModel={embeddingModel}
onEmbeddingModelChange={setEmbeddingModel}
matchThreshold={matchThreshold}
onMatchThresholdChange={setMatchThreshold}
escalationKeywords={escalationKeywords}
onEscalationKeywordsChange={setEscalationKeywords}
autoRouterCompression={autoRouterCompression}
onAutoRouterCompressionChange={isMemberManaged ? undefined : setAutoRouterCompression}
showValidationErrors={showValidationErrors}
/>
);
const forecast = isForecastClassifier(complexityRouterConfig.classifier_type);
return (
<TooltipProvider>
<Card>
<CardContent>
<form onSubmit={form.handleSubmit(() => handleAutoRouterSubmit())} noValidate>
<FieldGroup>
<div>
<FormField
control={form.control}
name="auto_router_name"
label={labelWithHint("Auto Router Name", "Unique name for this auto router configuration")}
>
{({ ref, ...field }) => (
<Input {...field} ref={ref} placeholder="e.g., smart_router, auto_router_1" />
)}
</FormField>
{!automaticSetupLoading && automaticRouterConfig && (
<div className="mt-5 flex flex-wrap items-center justify-between gap-3 rounded-lg border border-border bg-muted px-4 py-3">
<div>
<p className="text-sm font-medium text-foreground">Not sure where to start?</p>
<p className="text-sm text-muted-foreground">Let us pick models for each complexity tier.</p>
</div>
<Button type="button" data-testid="configure-automatically-button" onClick={handleAutomaticSetup}>
Configure automatically
</Button>
</div>
)}
<div className="mt-5">
<label className="block text-sm font-medium text-foreground mb-2">Template</label>
<Select
items={templateItems}
value={selectedPreset ?? null}
onValueChange={(presetKey: string | null) => handlePresetChange(presetKey ?? undefined)}
>
<SelectTrigger data-testid="template-selector" className="w-full">
<SelectValue placeholder="Choose a template or select Custom to define your own" />
</SelectTrigger>
<SelectContent>
{sortedPresetOptions.map(({ preset, availability: presetState }) => {
const disabledHint = presetDisabledHint(presetState);
const hintClass = isPresetHintAlarming(presetState)
? "text-destructive"
: "text-muted-foreground";
const matchedHint =
presetState.kind === "available" && presetState.viaDeployments
? "Matches your deployments"
: null;
return (
<SelectItem
key={preset.key}
value={preset.key}
label={preset.label}
disabled={disabledHint !== null}
title={disabledHint ?? preset.description}
<div className="mb-6">
<FormField
control={form.control}
name="auto_router_name"
label={labelWithHint("Auto Router Name", "Unique name for this auto router configuration")}
>
{({ ref, ...field }) => <Input {...field} ref={ref} placeholder="e.g., smart_router, auto_router_1" />}
</FormField>
</div>
<AutoRouterClassifierTabs
value={complexityRouterConfig}
onChange={(config) => {
setSelectedPreset(undefined);
setComplexityRouterConfig(config);
}}
>
<FieldGroup>
<div>
{!forecast && (
<>
{!automaticSetupLoading && automaticRouterConfig && (
<div className="mt-5 flex flex-wrap items-center justify-between gap-3 rounded-lg border border-border bg-muted px-4 py-3">
<div>
<p className="text-sm font-medium text-foreground">Not sure where to start?</p>
<p className="text-sm text-muted-foreground">
Let us pick models for each complexity tier.
</p>
</div>
<Button
type="button"
data-testid="configure-automatically-button"
onClick={handleAutomaticSetup}
>
<div>
<div className="font-medium">{preset.label}</div>
<div className="text-xs text-muted-foreground">{preset.description}</div>
{disabledHint && <div className={`text-xs mt-1 ${hintClass}`}>{disabledHint}</div>}
{matchedHint && <div className="text-xs mt-1 text-success">{matchedHint}</div>}
</div>
</SelectItem>
);
})}
<SelectItem value="custom" label="Custom Configuration">
<div>
<div className="font-medium">Custom Configuration</div>
<div className="text-xs text-muted-foreground">Define your auto router from scratch</div>
Configure automatically
</Button>
</div>
</SelectItem>
</SelectContent>
</Select>
)}
<div className="mt-5">
<label className="block text-sm font-medium text-foreground mb-2">Template</label>
<Select
items={templateItems}
value={selectedPreset ?? null}
onValueChange={(presetKey: string | null) => handlePresetChange(presetKey ?? undefined)}
>
<SelectTrigger data-testid="template-selector" className="w-full">
<SelectValue placeholder="Choose a template or select Custom to define your own" />
</SelectTrigger>
<SelectContent>
{sortedPresetOptions.map(({ preset, availability: presetState }) => {
const disabledHint = presetDisabledHint(presetState);
const hintClass = isPresetHintAlarming(presetState)
? "text-destructive"
: "text-muted-foreground";
const matchedHint =
presetState.kind === "available" && presetState.viaDeployments
? "Matches your deployments"
: null;
return (
<SelectItem
key={preset.key}
value={preset.key}
label={preset.label}
disabled={disabledHint !== null}
title={disabledHint ?? preset.description}
>
<div>
<div className="font-medium">{preset.label}</div>
<div className="text-xs text-muted-foreground">{preset.description}</div>
{disabledHint && <div className={`text-xs mt-1 ${hintClass}`}>{disabledHint}</div>}
{matchedHint && <div className="text-xs mt-1 text-success">{matchedHint}</div>}
</div>
</SelectItem>
);
})}
<SelectItem value="custom" label="Custom Configuration">
<div>
<div className="font-medium">Custom Configuration</div>
<div className="text-xs text-muted-foreground">
Define your auto router from scratch
</div>
</div>
</SelectItem>
</SelectContent>
</Select>
{presetsPending && (
<div className="text-xs mt-1 text-muted-foreground">Loading templates...</div>
)}
{presetsUnavailable && (
<div className="text-xs mt-1 text-destructive">
Could not load templates, so only Custom Configuration is shown.{" "}
<button type="button" className="underline" onClick={() => void refetchPresets()}>
Retry
</button>
</div>
)}
</div>
</>
)}
{modelsUnverifiable && (
<div className="text-xs mt-1 text-destructive">
Could not load available models.{" "}
@ -626,163 +689,129 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
</button>
</div>
)}
{presetsPending && <div className="text-xs mt-1 text-muted-foreground">Loading templates...</div>}
{presetsUnavailable && (
<div className="text-xs mt-1 text-destructive">
Could not load templates, so only Custom Configuration is shown.{" "}
<button type="button" className="underline" onClick={() => void refetchPresets()}>
Retry
</button>
</div>
)}
</div>
</div>
{requiresTeamScope && (
<FormField
control={form.control}
name="team_id"
label={labelWithHint(
"Select Team",
"Select the team this auto router belongs to. Only keys for this team will be able to call it.",
)}
>
{({ id, value, onChange }) => (
<TeamDropdown
id={id}
value={value}
onChange={onChange}
filterTeam={(team) => canCreateAutoRouterForTeam(actor, team)}
/>
)}
</FormField>
)}
<div className="border border-border rounded-lg">
<button
type="button"
onClick={() => setDetailsExpanded((expanded) => !expanded)}
className="w-full flex flex-col gap-1 px-4 py-3 text-left hover:bg-muted"
data-testid="detailed-configuration-toggle"
>
<span className="flex items-center gap-2 font-medium text-foreground">
{detailsExpanded ? (
<ChevronDown className="size-3 text-muted-foreground" />
) : (
<ChevronRight className="size-3 text-muted-foreground" />
{requiresTeamScope && (
<FormField
control={form.control}
name="team_id"
label={labelWithHint(
"Select Team",
"Select the team this auto router belongs to. Only keys for this team will be able to call it.",
)}
Detailed Configuration
</span>
{!detailsExpanded && (
<span className="text-xs text-muted-foreground line-clamp-2">
{tierConfigSummary(complexityRouterConfig)}
</span>
)}
</button>
{detailsExpanded && (
<div className="px-4 pb-4">
<ComplexityRouterConfig
editingTiers={editingTiers}
onEditingTiersChange={setEditingTiers}
modelInfo={modelInfo}
value={complexityRouterConfig}
onChange={setComplexityRouterConfig}
customTechnicalKeywords={customTechnicalKeywords}
onCustomTechnicalKeywordsChange={setCustomTechnicalKeywords}
keywordTierRules={keywordTierRules}
onKeywordTierRulesChange={setKeywordTierRules}
keywordRulesError={getKeywordTierRulesError(
keywordTierRules,
activeTierRows(complexityRouterConfig),
>
{({ id, value, onChange }) => (
<TeamDropdown
id={id}
value={value}
onChange={onChange}
filterTeam={(team) => canCreateAutoRouterForTeam(actor, team)}
/>
)}
</FormField>
)}
{forecast ? (
configurationForm
) : (
<div className="border border-border rounded-lg">
<button
type="button"
onClick={() => setDetailsExpanded((expanded) => !expanded)}
className="w-full flex flex-col gap-1 px-4 py-3 text-left hover:bg-muted"
data-testid="detailed-configuration-toggle"
>
<span className="flex items-center gap-2 font-medium text-foreground">
{detailsExpanded ? (
<ChevronDown className="size-3 text-muted-foreground" />
) : (
<ChevronRight className="size-3 text-muted-foreground" />
)}
Detailed Configuration
</span>
{!detailsExpanded && (
<span className="text-xs text-muted-foreground line-clamp-2">
{tierConfigSummary(complexityRouterConfig)}
</span>
)}
semanticMatchingEnabled={semanticMatchingEnabled}
onSemanticMatchingEnabledChange={setSemanticMatchingEnabled}
embeddingModel={embeddingModel}
onEmbeddingModelChange={setEmbeddingModel}
matchThreshold={matchThreshold}
onMatchThresholdChange={setMatchThreshold}
escalationKeywords={escalationKeywords}
onEscalationKeywordsChange={setEscalationKeywords}
autoRouterCompression={autoRouterCompression}
onAutoRouterCompressionChange={isMemberManaged ? undefined : setAutoRouterCompression}
showValidationErrors={showValidationErrors}
/>
</button>
{detailsExpanded && <div className="px-4 pb-4">{configurationForm}</div>}
</div>
)}
</div>
{isAdmin && (
<FormField
control={form.control}
name="model_access_group"
label={labelWithHint(
"Model Access Group",
"Use model access groups to control who can access this auto router",
)}
>
{({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => (
<AccessGroupTagsCombobox
id={id}
value={value}
onChange={onChange}
options={modelAccessGroups}
ariaInvalid={ariaInvalid}
ariaDescribedBy={ariaDescribedBy}
{isAdmin && (
<FormField
control={form.control}
name="model_access_group"
label={labelWithHint(
"Model Access Group",
"Use model access groups to control who can access this auto router",
)}
>
{({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => (
<AccessGroupTagsCombobox
id={id}
value={value}
onChange={onChange}
options={modelAccessGroups}
ariaInvalid={ariaInvalid}
ariaDescribedBy={ariaDescribedBy}
/>
)}
</FormField>
)}
<div className="flex justify-between items-center">
<Tooltip>
<TooltipTrigger
render={
<a
href="https://github.com/BerriAI/litellm/issues"
className="text-sm text-primary underline-offset-4 hover:underline"
>
Need Help?
</a>
}
/>
)}
</FormField>
)}
<div className="flex justify-between items-center">
<Tooltip>
<TooltipTrigger
render={
<a
href="https://github.com/BerriAI/litellm/issues"
className="text-sm text-primary underline-offset-4 hover:underline"
<TooltipContent>Get help on our github</TooltipContent>
</Tooltip>
<div className="flex gap-2">
<BlockedReasonTooltip reason={submitBlockedReason}>
<Button
type="button"
variant="outline"
data-testid="auto-router-test-routing-btn"
disabled={submitBlockedReason !== null || isSubmitting}
onClick={() => setIsRoutingTestVisible(true)}
>
Need Help?
</a>
}
/>
<TooltipContent>Get help on our github</TooltipContent>
</Tooltip>
<div className="flex gap-2">
<BlockedReasonTooltip reason={submitBlockedReason}>
Test Routing
</Button>
</BlockedReasonTooltip>
<Button
type="button"
variant="outline"
data-testid="auto-router-test-routing-btn"
disabled={submitBlockedReason !== null || isSubmitting}
onClick={() => setIsRoutingTestVisible(true)}
data-testid="auto-router-test-connect-btn"
onClick={handleTestConnection}
disabled={isTestingConnection}
>
Test Routing
{isTestingConnection && <UiLoadingSpinner className="size-4" />}
Test Connection
</Button>
</BlockedReasonTooltip>
<Button
type="button"
variant="outline"
data-testid="auto-router-test-connect-btn"
onClick={handleTestConnection}
disabled={isTestingConnection}
>
{isTestingConnection && <UiLoadingSpinner className="size-4" />}
Test Connection
</Button>
<BlockedReasonTooltip reason={submitBlockedReason}>
<Button
type="button"
disabled={submitBlockedReason !== null || isSubmitting}
onClick={() => {
void handleAutoRouterSubmit();
}}
>
Add Auto Router
</Button>
</BlockedReasonTooltip>
<BlockedReasonTooltip reason={submitBlockedReason}>
<Button
type="button"
disabled={submitBlockedReason !== null || isSubmitting}
onClick={() => {
void handleAutoRouterSubmit();
}}
>
Add Auto Router
</Button>
</BlockedReasonTooltip>
</div>
</div>
</div>
</FieldGroup>
</FieldGroup>
</AutoRouterClassifierTabs>
</form>
</CardContent>
</Card>

View file

@ -48,6 +48,37 @@ const baseParams: BuildComplexityRouterConfigParams = {
};
describe("buildComplexityRouterConfig", () => {
it.each(["capability", "llm_v2", "heuristic"] as const)(
"disables the removed overrides only for forecast creates: %s",
(classifierType) => {
const forecast = classifierType !== "heuristic";
const params = {
...baseParams,
classifierType,
adaptive: true,
enableContextWindowEscalation: true,
contextWindowEscalationBuffer: 0.9,
};
const config = buildComplexityRouterConfig(params);
expect(config.adaptive).toBe(!forecast);
expect(config.enable_context_window_escalation).toBe(!forecast);
expect(config.escalation_keywords).toEqual(forecast ? [] : ["LITELLM ESCALATE"]);
for (const key of [
"adaptive_weights",
"adaptive_eligible",
"tier_distance_penalty",
"context_window_escalation_buffer",
]) {
expect(Object.hasOwn(config, key)).toBe(!forecast);
}
if (forecast) {
const untouched = buildComplexityRouterConfig({ ...baseParams, classifierType });
expect(untouched.enable_context_window_escalation).toBe(false);
expect(untouched.escalation_keywords).toEqual([]);
}
},
);
it("carries Fast and reasoning overrides independently into a new router payload", () => {
const params = { speed: "fast", reasoning_effort: "high", max_tokens: 1024 };
const config = buildComplexityRouterConfig({

View file

@ -1,3 +1,9 @@
import {
isForecastClassifier,
withoutForecastPromptOverrides,
type CapabilitySettings,
type FuseSettings,
} from "./forecast_classifier_config";
import type { ModelGroup } from "../llm_calls/fetch_models";
import { KeywordTierRule } from "./KeywordTierRules";
import {
@ -126,6 +132,48 @@ const scorerKnobPayload = ({
};
};
export interface StoredComplexityRouterConfig {
tiers?: Partial<Record<keyof ComplexityTiers, unknown>>;
enable_non_reasoning_tier?: boolean;
tier_model_configs?: unknown;
default_model?: string | null;
plan_mode_min_tier?: unknown;
classification_prompt?: unknown;
classification_examples?: unknown;
heuristic_first_max_tier?: unknown;
hybrid_boundary_margin?: unknown;
tier_labels?: unknown;
classifier_type?: ClassifierType;
capability_classifier_config?: unknown;
llm_v2_config?: unknown;
classifier_llm_config?: ClassifierLLMConfig;
classifier_context_window_size?: unknown;
classifier_context_budget_chars?: unknown;
classifier_context_include_assistant_turns?: unknown;
classifier_fallback?: unknown;
classification_mode?: unknown;
tier_boundaries?: unknown;
token_thresholds?: unknown;
dimension_weights?: unknown;
custom_dimensions?: unknown;
reasoning_override_min_score?: unknown;
session_affinity?: unknown;
session_affinity_ttl_seconds?: unknown;
modality_routing?: unknown;
modality_pin_override?: unknown;
deployment_affinity?: unknown;
adaptive?: boolean;
adaptive_weights?: AdaptiveRouterWeights;
tier_distance_penalty?: number;
adaptive_eligible?: AdaptiveEligible;
return_raw_model_name?: boolean;
enable_context_window_escalation?: unknown;
context_window_escalation_buffer?: unknown;
stall_escalation_enabled?: unknown;
stall_escalation_window?: unknown;
stall_escalation_repeat_threshold?: unknown;
}
export interface BuildComplexityRouterConfigParams {
tiers: ComplexityTiers;
enableNonReasoningTier?: boolean;
@ -134,6 +182,8 @@ export interface BuildComplexityRouterConfigParams {
planModeMinTier: string | undefined;
tierLabels: ComplexityTierLabels | undefined;
classifierType: ClassifierType;
capabilityClassifierConfig?: CapabilitySettings;
llmV2Config?: FuseSettings;
classifierLlmConfig: ClassifierLLMConfigWire | undefined;
classifierContextWindowSize: number | undefined;
classifierContextBudgetChars: number | undefined;
@ -198,6 +248,8 @@ export interface ComplexityRouterConfigPayload {
plan_mode_min_tier?: string;
tier_labels?: ComplexityTierLabels;
classifier_type: ClassifierType;
capability_classifier_config?: CapabilitySettings;
llm_v2_config?: FuseSettings;
classifier_llm_config?: ClassifierLLMConfig;
classifier_context_window_size?: number;
classifier_context_budget_chars?: number;
@ -468,31 +520,34 @@ const classifierWireFields = (
| "classifierContextBudgetChars"
| "classifierContextIncludeAssistantTurns"
>,
): Partial<ComplexityRouterConfigPayload> => ({
...(usesLlmClassifier(effectiveType) &&
classifierLlmConfig && {
classifier_llm_config:
effectiveType === "capability" ? classifierLlmConfig : normalizeClassifierLlmConfig(classifierLlmConfig),
}),
...(usesLlmClassifier(effectiveType) &&
classifierFallback !== undefined && { classifier_fallback: classifierFallback }),
...(effectiveType === "heuristic_first" &&
heuristicFirstMaxTier?.trim() && { heuristic_first_max_tier: heuristicFirstMaxTier }),
...(effectiveType === "hybrid" &&
hybridBoundaryMargin !== undefined && { hybrid_boundary_margin: hybridBoundaryMargin }),
...(usesLlmClassifier(effectiveType) &&
classifierContextWindowSize !== undefined && {
classifier_context_window_size: classifierContextWindowSize,
}),
...(usesLlmClassifier(effectiveType) &&
classifierContextBudgetChars !== undefined && {
classifier_context_budget_chars: classifierContextBudgetChars,
}),
...(usesLlmClassifier(effectiveType) &&
classifierContextIncludeAssistantTurns !== undefined && {
classifier_context_include_assistant_turns: classifierContextIncludeAssistantTurns,
}),
});
): Partial<ComplexityRouterConfigPayload> => {
const supportsFallback = usesLlmClassifier(effectiveType) && !isForecastClassifier(effectiveType);
return {
...(usesLlmClassifier(effectiveType) &&
classifierLlmConfig && {
classifier_llm_config: isForecastClassifier(effectiveType)
? withoutForecastPromptOverrides(classifierLlmConfig)
: normalizeClassifierLlmConfig(classifierLlmConfig),
}),
...(supportsFallback && classifierFallback !== undefined && { classifier_fallback: classifierFallback }),
...(effectiveType === "heuristic_first" &&
heuristicFirstMaxTier?.trim() && { heuristic_first_max_tier: heuristicFirstMaxTier }),
...(effectiveType === "hybrid" &&
hybridBoundaryMargin !== undefined && { hybrid_boundary_margin: hybridBoundaryMargin }),
...(usesLlmClassifier(effectiveType) &&
classifierContextWindowSize !== undefined && {
classifier_context_window_size: classifierContextWindowSize,
}),
...(usesLlmClassifier(effectiveType) &&
classifierContextBudgetChars !== undefined && {
classifier_context_budget_chars: classifierContextBudgetChars,
}),
...(usesLlmClassifier(effectiveType) &&
classifierContextIncludeAssistantTurns !== undefined && {
classifier_context_include_assistant_turns: classifierContextIncludeAssistantTurns,
}),
};
};
export const buildComplexityRouterConfig = ({
tiers,
@ -502,6 +557,8 @@ export const buildComplexityRouterConfig = ({
planModeMinTier,
tierLabels,
classifierType,
capabilityClassifierConfig,
llmV2Config,
classifierLlmConfig,
classifierContextWindowSize,
classifierContextBudgetChars,
@ -571,9 +628,11 @@ export const buildComplexityRouterConfig = ({
// An edited tier set forces the LLM classifier, so llm-only inputs must survive a classifier_type
// the form never rewrote. The UI gates the same controls on this, not on the raw value.
const effectiveType: ClassifierType = customTierSet ? "llm" : classifierType;
const forecast = isForecastClassifier(effectiveType);
const supportsOpeningPrompt = !customTierSet && !forecast && usesLlmClassifier(effectiveType);
const payload: ComplexityRouterConfigPayload = {
tiers,
tiers: forecast ? Object.fromEntries(Object.entries(tiers).filter(([, models]) => models.length > 0)) : tiers,
// The backend rejects the flag beside a custom tier set.
...(!customTierSet && enableNonReasoningTier && { enable_non_reasoning_tier: true }),
...(serializedTierModelConfigs && { tier_model_configs: serializedTierModelConfigs }),
@ -582,10 +641,13 @@ export const buildComplexityRouterConfig = ({
...(cleanedTierLabels && { tier_labels: cleanedTierLabels }),
classifier_type: classifierType,
...classifierWireFields(effectiveType, classifierInputs),
...(effectiveType === "capability" &&
capabilityClassifierConfig && { capability_classifier_config: capabilityClassifierConfig }),
...(effectiveType === "llm_v2" && { llm_v2_config: llmV2Config }),
...(forecast && { adaptive: false }),
// A built-in router's opening instructions. Suppressed beside a legacy whole-prompt override,
// which the backend rejects as a second override of the same prompt.
...(!customTierSet &&
usesLlmClassifier(effectiveType) &&
...(supportsOpeningPrompt &&
!classifierLlmConfig?.system_prompt?.trim() && {
...(classificationPrompt?.trim() && { classification_prompt: classificationPrompt.trim() }),
...(classificationExamples?.trim() && { classification_examples: classificationExamples.trim() }),
@ -597,7 +659,7 @@ export const buildComplexityRouterConfig = ({
modality_pin_override: modalityPinOverride ?? false,
...(customTechnicalKeywords.length > 0 && { custom_technical_keywords: customTechnicalKeywords }),
...(cleanedKeywordTierRules.length > 0 && { keyword_tier_rules: cleanedKeywordTierRules }),
escalation_keywords: cleanedEscalationKeywords,
escalation_keywords: forecast ? [] : cleanedEscalationKeywords,
// Only written when on: the backend rejects it alongside session_affinity, user_turn mode and
// a custom tier set, so an off router must not carry the key into any of those saves.
...(stallEscalationEnabled && {
@ -612,19 +674,22 @@ export const buildComplexityRouterConfig = ({
embedding_model: embeddingModel,
match_threshold: matchThreshold,
}),
...(adaptive && {
adaptive: true,
adaptive_weights: adaptiveWeights,
...(adaptiveEligible === "all" && { tier_distance_penalty: tierDistancePenalty }),
adaptive_eligible: adaptiveEligible,
}),
...(adaptive &&
!forecast && {
adaptive: true,
adaptive_weights: adaptiveWeights,
...(adaptiveEligible === "all" && { tier_distance_penalty: tierDistancePenalty }),
adaptive_eligible: adaptiveEligible,
}),
...(returnRawModelName && { return_raw_model_name: true }),
...(enableContextWindowEscalation !== undefined && {
enable_context_window_escalation: enableContextWindowEscalation,
}),
...(contextWindowEscalationBuffer !== undefined && {
context_window_escalation_buffer: contextWindowEscalationBuffer,
// Omission enables the backend default, so hidden forecast controls need an explicit opt-out.
...((forecast || enableContextWindowEscalation !== undefined) && {
enable_context_window_escalation: forecast ? false : enableContextWindowEscalation,
}),
...(!forecast &&
contextWindowEscalationBuffer !== undefined && {
context_window_escalation_buffer: contextWindowEscalationBuffer,
}),
...(sessionAffinityTtlSeconds !== undefined && {
session_affinity_ttl_seconds: sessionAffinityTtlSeconds,
}),

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