Merge remote-tracking branch 'origin/litellm_internal_staging' into shadow-eval-pre-adoption

This commit is contained in:
Abhimanyu Kapur 2026-08-08 13:17:18 -07:00
commit d3020685cf
40 changed files with 1825 additions and 364 deletions

View file

@ -9,7 +9,7 @@ Don't assume that the existing code is correct or the right way of doing things
- easy to maintain/change
- modern
In that order of importance
In descending order of importance
When adding new features, add meaningful tests. Don't add tests that don't check anything substantial and is there just to make the code coverage pass. Yes, code coverage is important, but I'd rather have no signal whether the code is working than tests that don't fail when code is broken. The goal is to have tests that would fail before the feature was added/if the code was mutated in a way that breaks the feature and succeed only when the feature is fully working. I should run mutation testing and see > 90% kill rate
@ -41,7 +41,7 @@ Python max line length is 120, not 88
When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing
`make pre-commit` saves its complete output to a log file in .git (overwriting previous pre-commit logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice
`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice
If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in

View file

@ -8,7 +8,7 @@
lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \
lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \
install-dev install-proxy-dev install-test-deps install-hooks \
install-helm-unittest check-circular-imports check-import-safety pre-commit \
install-helm-unittest check-circular-imports check-import-safety check pre-commit \
lint-install lint-fetch-base bootstrap
# Default target
@ -22,7 +22,8 @@ help:
@echo " make install-test-deps - Install the full local test environment"
@echo " make install-helm-unittest - Install helm unittest plugin"
@echo " make install-hooks - Install git hooks (Conventional Commits + Branches)"
@echo " make pre-commit - Run CI-equivalent lint on staged files (run before committing)"
@echo " make check - Run CI-equivalent lint on staged files, or on the diff vs the base branch when nothing is staged"
@echo " make pre-commit - Legacy alias for make check"
@echo " make format - Apply ruff format code formatting"
@echo " make format-check - Check ruff format code formatting (matches CI)"
@echo " make lint - Run all linting (Ruff, basedpyright, format check, circular imports, import safety)"
@ -236,13 +237,20 @@ lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline
# Faster linting for local development (only checks changed code)
lint-dev: lint-format-changed check-circular-imports check-import-safety
# Run the gating CI checks against your staged files right before committing. Mirrors
# Run the gating CI checks against your changes. Scopes to staged files when anything
# is staged (warning about changed files left unstaged); with nothing staged it falls
# back to the working tree's diff against the merge base with the base branch, so a
# fresh merge commit or an unstaged working tree still gets checked. Mirrors
# test-linting.yml (Python), test-litellm-ui-build.yml's frontend-lint (dashboard), and
# check-ui-api-types.yml (API-type drift), skipping any whose files you didn't stage.
# check-ui-api-types.yml (API-type drift), skipping any whose files aren't in scope.
# Not auto-installed as a git hook so it never slows an unrelated human commit.
pre-commit: bootstrap
check: bootstrap
./scripts/pre_commit_lint.sh
pre-commit:
@echo "make pre-commit is a legacy alias; use make check" >&2
@$(MAKE) check
# Testing targets
test: install-test-deps
$(UV_RUN) pytest tests/

View file

@ -1322,6 +1322,7 @@ X_LITELLM_DISABLE_CALLBACKS: Final = "x-litellm-disable-callbacks"
LITELLM_METADATA_FIELD: Final = "litellm_metadata"
OLD_LITELLM_METADATA_FIELD: Final = "metadata"
RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name"
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated"
LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = (

View file

@ -33,6 +33,7 @@ from litellm.integrations.otel.model.payloads import (
is_mcp_list_tools,
is_mcp_tool_call,
)
from litellm.integrations.otel.model.semconv import Error
from litellm.integrations.otel.model.spans import SpanRole, span_role_for_service
from litellm.integrations.otel.model.utils import to_ns
from litellm.integrations.otel.plumbing.context import (
@ -634,18 +635,23 @@ class OpenTelemetryV2(CustomLogger):
"""Stamp the v2 error.* attributes on the FastAPI-owned SERVER span for a
failure that dies before any LLM-call span exists (malformed body, auth /
validation rejection). Called from the proxy's global exception handler via
``_close_dangling_otel_server_span``. The instrumentor still owns the span's
status and lifecycle, so this only decorates it never sets status, never
ends it and emits no exception event, matching v1's SERVER-span behavior
and avoiding a duplicate of the event ``async_post_call_failure_hook`` or
the ``auth`` phase span already records."""
``_close_dangling_otel_server_span``, which swallows the exception into a
``JSONResponse`` so the instrumentor never sees it and leaves the span
``UNSET``; the status is set here instead (v1 did the same from the handler)
so a failed request reads as failed and not merely as a span carrying an
error message. The instrumentor still owns the span's lifecycle, so this
never ends it. The exception event is recorded only when nothing stamped
this span already ``async_post_call_failure_hook`` and the ``auth`` phase
span record their own, and a second event would duplicate it while the
attributes are always restamped so ``error.code`` stays pinned to the real
response status."""
if span is None or not is_recordable_span(span):
return
already_stamped: Final = Error.TYPE in (getattr(span, "attributes", None) or ())
stamp_error(
span,
_span_error_from_exception(exception, status_code=status_code),
record_event=False,
set_status=False,
record_event=not already_stamped,
)
async def async_post_call_failure_hook(

View file

@ -42,7 +42,12 @@ from litellm.types.integrations.websearch_interception import (
WebSearchInterceptionConfig,
)
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import AgenticLoopParams, CallTypes, LlmProviders
from litellm.types.utils import (
AgenticLoopParams,
CallTypes,
LlmProviders,
StandardLoggingUserAPIKeyMetadata,
)
from litellm.utils import ProviderConfigManager
if TYPE_CHECKING:
@ -1318,6 +1323,7 @@ class WebSearchInterceptionLogger(CustomLogger):
search_tool: Final = self._select_search_tool_from_router(llm_router=llm_router)
search_provider: str | None = None
search_litellm_params: dict[str, Any] = {}
search_tool_name: Final = self._selected_search_tool_name(search_tool=search_tool)
if search_tool is not None:
await self._authorize_search_tool(search_tool=search_tool, kwargs=kwargs)
search_litellm_params = dict(search_tool.get("litellm_params", {}) or {})
@ -1334,12 +1340,30 @@ class WebSearchInterceptionLogger(CustomLogger):
verbose_logger.debug(
"WebSearchInterception: Executing search for '%s' using provider '%s'", query, search_provider
)
user_api_key_auth: Final = self._get_user_api_key_auth_from_kwargs(kwargs)
search_metadata: Final = (
None
if user_api_key_auth is None
else self._build_search_request_metadata(
user_api_key_auth=user_api_key_auth,
search_tool_name=search_tool_name,
)
)
search_kwargs: Final = {
key: value
for key, value in search_litellm_params.items()
if key != "search_provider" and value is not None
}
result: Final = await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs)
result: Final = (
await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs)
if search_metadata is None
else await litellm.asearch(
query=query,
search_provider=search_provider,
litellm_metadata=search_metadata,
**search_kwargs,
)
)
# Format using transformation function
search_result_text: Final = WebSearchTransformation.format_search_response(result)
@ -1396,6 +1420,35 @@ class WebSearchInterceptionLogger(CustomLogger):
team_object=team_object,
)
@staticmethod
def _build_search_request_metadata(
user_api_key_auth: "UserAPIKeyAuth",
search_tool_name: str | None,
) -> Mapping[str, object]:
"""
Spend-tracking metadata for the intercepted search, so its provider cost is logged
and billed against the key/user/team that made the originating LLM request instead
of being dropped by the proxy's spend hook for lack of an owner.
"""
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
user_api_key_metadata: Final[StandardLoggingUserAPIKeyMetadata] = (
LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_auth)
)
return { # mutable-ok: litellm's metadata channel is a plain dict its logging path reads and enriches
**user_api_key_metadata,
"model_group": search_tool_name,
"user_api_key": user_api_key_auth.api_key,
"user_api_key_auth": user_api_key_auth,
}
@staticmethod
def _selected_search_tool_name(search_tool: Mapping[str, object] | None) -> str | None:
if search_tool is None:
return None
search_tool_name: Final = search_tool.get("search_tool_name")
return search_tool_name if isinstance(search_tool_name, str) and search_tool_name else None
@staticmethod
def _get_user_api_key_auth_from_kwargs(kwargs: Mapping[str, object] | None) -> "UserAPIKeyAuth | None":
if not kwargs:

View file

@ -4557,10 +4557,10 @@ class DefaultInternalUserParams(LiteLLMPydanticObjectBase):
user_role: (
Literal[
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
]
| None
) = Field(

View file

@ -1044,6 +1044,22 @@ def _ensure_parent_otel_span_on_request_state(request: Request) -> None:
request.state.parent_otel_span = parent_otel_span
async def _read_request_body_deferring_parse_failure(
request: Request,
) -> tuple[dict, ProxyException | None]:
"""Parse the body, returning a parse failure instead of raising it.
A body that fails to parse is still a request from a known caller, so auth
must run (resolving identity onto the request's trace) before the 400 goes
out; the caller re-raises the returned exception once identity is seeded.
"""
try:
parsed_body: Final = await _read_request_body(request=request)
except ProxyException as parse_exception:
return {}, parse_exception # mutable-ok: request_data is a plain dict across the whole auth path
return populate_request_with_path_params(request_data=parsed_body, request=request), None
async def _user_api_key_auth_builder(
request: Request,
api_key: str,
@ -2516,6 +2532,72 @@ def _resolve_request_principal(request: Request, valid_token: UserAPIKeyAuth) ->
)
async def _authorize_authenticated_request(
user_api_key_auth_obj: UserAPIKeyAuth,
request: Request,
request_data: dict,
route: str,
api_key: str,
) -> UserAPIKeyAuth | None:
"""Authorize an already-authenticated request: disabled-route check, the single
``common_checks`` gate (which also reserves budget), and end-user fallback
resolution. Returns the auth object the exception handler recovered when a check
failed but the request may proceed anyway, else ``None``.
"""
## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ##
RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request)
# Single authorization point. Builder paths MUST NOT call common_checks.
# Route through the same exception handler the builder uses so
# authorization failures (ProxyException, or plain Exception from
# admin-only-route / model-access / budget checks) surface as
# ProxyException consistently with pre-refactor behavior.
try:
await _run_centralized_common_checks(
user_api_key_auth_obj=user_api_key_auth_obj,
request=request,
request_data=request_data,
route=route,
)
except Exception as e:
return await UserAPIKeyAuthExceptionHandler._handle_authentication_error(
e=e,
request=request,
request_data=request_data,
route=route,
parent_otel_span=user_api_key_auth_obj.parent_otel_span,
api_key=api_key,
resolved_identity=user_api_key_auth_obj,
)
# Defense-in-depth: ``_user_api_key_auth_builder`` has multiple early-return
# paths (no master key, /user/auth route, JWT short-circuits) that bypass
# the end-user resolution block. If those paths produced an auth obj
# without an ``end_user_id`` set, fall back to extracting from the request
# body so spend logs are still attributed correctly. Validation honours
# ``litellm.validate_end_user_id_in_db``.
if user_api_key_auth_obj.end_user_id is None:
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
raw_end_user_id: Final = get_end_user_id_from_request_body(request_data, _safe_get_request_headers(request))
if raw_end_user_id is not None:
resolved_end_user_id: Final = await resolve_and_validate_end_user_id(
raw_end_user_id=raw_end_user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_auth_obj.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
route=route,
)
if resolved_end_user_id is not None:
user_api_key_auth_obj.end_user_id = resolved_end_user_id
return None
@tracer.wrap()
async def user_api_key_auth(
request: Request,
@ -2536,8 +2618,7 @@ async def user_api_key_auth(
# close, and the trace never reaches the backend.
_ensure_parent_otel_span_on_request_state(request)
request_data = await _read_request_body(request=request)
request_data = populate_request_with_path_params(request_data=request_data, request=request)
request_data, body_parse_exception = await _read_request_body_deferring_parse_failure(request=request)
route: Final[str] = get_request_route(request=request)
## CHECK IF ROUTE IS ALLOWED
@ -2545,69 +2626,41 @@ async def user_api_key_auth(
# triggers (key/user/team object reads) nest under it instead of flattening
# onto the server span. No-op when OTel V2 isn't active.
with phase_span(f"auth {route}"):
user_api_key_auth_obj: Final = await _user_api_key_auth_builder(
request=request,
api_key=api_key,
azure_api_key_header=azure_api_key_header,
anthropic_api_key_header=anthropic_api_key_header,
google_ai_studio_api_key_header=google_ai_studio_api_key_header,
azure_apim_header=azure_apim_header,
request_data=request_data,
custom_litellm_key_header=custom_litellm_key_header,
)
try:
user_api_key_auth_obj: Final = await _user_api_key_auth_builder(
request=request,
api_key=api_key,
azure_api_key_header=azure_api_key_header,
anthropic_api_key_header=anthropic_api_key_header,
google_ai_studio_api_key_header=google_ai_studio_api_key_header,
azure_apim_header=azure_apim_header,
request_data=request_data,
custom_litellm_key_header=custom_litellm_key_header,
)
except Exception:
# The body was read first, so a caller who sent both a malformed body and
# a rejected key used to get the 400; the response is unchanged, and the
# auth failure is still recorded on the trace by the handler that ran.
if body_parse_exception is not None:
raise body_parse_exception
raise
user_api_key_auth_obj.budget_reservation = None
## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ##
RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request)
# Single authorization point. Builder paths MUST NOT call common_checks.
# Route through the same exception handler the builder uses so
# authorization failures (ProxyException, or plain Exception from
# admin-only-route / model-access / budget checks) surface as
# ProxyException consistently with pre-refactor behavior.
try:
await _run_centralized_common_checks(
# A body that never parsed is authenticated (so the trace carries identity
# and this ``auth`` span) but not authorized: there is no model to check it
# against, and budget reservation would increment live spend counters that
# only the endpoint's post-call path releases; the endpoint never runs, since
# the parse failure is raised below.
if body_parse_exception is None:
recovered_auth_obj: Final = await _authorize_authenticated_request(
user_api_key_auth_obj=user_api_key_auth_obj,
request=request,
request_data=request_data,
route=route,
)
except Exception as e:
return await UserAPIKeyAuthExceptionHandler._handle_authentication_error(
e=e,
request=request,
request_data=request_data,
route=route,
parent_otel_span=user_api_key_auth_obj.parent_otel_span,
api_key=api_key,
resolved_identity=user_api_key_auth_obj,
)
# Defense-in-depth: ``_user_api_key_auth_builder`` has multiple early-return
# paths (no master key, /user/auth route, JWT short-circuits) that bypass
# the end-user resolution block. If those paths produced an auth obj
# without an ``end_user_id`` set, fall back to extracting from the request
# body so spend logs are still attributed correctly. Validation honours
# ``litellm.validate_end_user_id_in_db``.
if user_api_key_auth_obj.end_user_id is None:
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
raw_end_user_id: Final = get_end_user_id_from_request_body(request_data, _safe_get_request_headers(request))
if raw_end_user_id is not None:
resolved_end_user_id: Final = await resolve_and_validate_end_user_id(
raw_end_user_id=raw_end_user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_auth_obj.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
route=route,
)
if resolved_end_user_id is not None:
user_api_key_auth_obj.end_user_id = resolved_end_user_id
if recovered_auth_obj is not None:
return recovered_auth_obj
# Identity is now resolved. Seed it AFTER the auth span closes so the Baggage
# persists on the request task (detaching the span's context token inside the
@ -2619,6 +2672,9 @@ async def user_api_key_auth(
)
user_api_key_auth_obj.request_route = normalize_request_route(route)
if body_parse_exception is not None:
raise body_parse_exception
# Resolve caller identity once, here at the seam, into a single per-request
# Principal projected off the key object the builder already fetched (no
# second lookup). Downstream consumers read identity off this instead of

View file

@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any, Final, Optional
import litellm
from litellm import get_secret
from litellm._logging import verbose_proxy_logger
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import (
get_metadata_variable_name_from_kwargs,
@ -425,6 +425,7 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset(
"_guardrail_pipelines",
"_pipeline_managed_guardrails",
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
"disable_global_guardrails",
"disable_global_guardrail",
"opted_out_global_guardrails",

View file

@ -33,6 +33,53 @@ if TYPE_CHECKING:
CACHE_TTL_5M_SECONDS: Final = 300
CACHE_TTL_1H_SECONDS: Final = 3600
AUTOROUTER_BENCHMARKS_SQL: Final = """
WITH windowed AS (
SELECT * FROM "LiteLLM_AutoRouterSession"
WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp
),
tier_maps AS (
SELECT router_name, router_type, jsonb_object_agg(tier, tier_turns) AS tier_turns
FROM (
SELECT router_name, router_type, kv.key AS tier, SUM((kv.value)::int)::int AS tier_turns
FROM windowed, LATERAL jsonb_each_text(tier_turns) AS kv
GROUP BY router_name, router_type, kv.key
) per_tier
GROUP BY router_name, router_type
)
SELECT
agg.*,
COALESCE(tier_maps.tier_turns, '{}'::jsonb) AS tier_turns
FROM (
SELECT
router_name,
router_type,
COUNT(*)::int AS sessions,
COALESCE(SUM(turns), 0)::int AS turns,
COALESCE(SUM(unordered_turns), 0)::int AS unordered_turns,
COALESCE(SUM(covered_turns), 0)::int AS covered_turns,
COALESCE(SUM(cache_hits), 0)::int AS cache_hits,
COALESCE(SUM(same_model_turns), 0)::int AS same_model_turns,
COALESCE(SUM(same_model_hits), 0)::int AS same_model_hits,
COALESCE(SUM(first_visit_turns), 0)::int AS first_visit_turns,
COALESCE(SUM(first_visit_hits), 0)::int AS first_visit_hits,
COALESCE(SUM(return_turns), 0)::int AS return_turns,
COALESCE(SUM(return_hits), 0)::int AS return_hits,
COALESCE(SUM(return_expired_misses), 0)::int AS return_expired_misses,
COALESCE(SUM(return_within_ttl_misses), 0)::int AS return_within_ttl_misses,
COALESCE(SUM(ttl_5m_turns), 0)::int AS ttl_5m_turns,
COALESCE(SUM(ttl_1h_turns), 0)::int AS ttl_1h_turns,
COALESCE(SUM(total_tokens), 0)::bigint AS total_tokens,
COALESCE(SUM(spend), 0)::float8 AS spend,
COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend,
COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds
FROM windowed
GROUP BY router_name, router_type
) agg
LEFT JOIN tier_maps USING (router_name, router_type)
ORDER BY agg.spend DESC
"""
@dataclass(frozen=True, slots=True)
class AutoRouterTurnTransaction:

View file

@ -18,6 +18,7 @@ from litellm.constants import (
INTERNAL_CALL_ORIGIN_METADATA_KEY,
LITELLM_PROXY_MASTER_KEY_ALIAS,
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
)
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
@ -226,6 +227,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = (
"applied_policies",
"policy_sources",
"routing_decision",
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
"standard_logging_object",
"proxy_server_request",

View file

@ -27,6 +27,7 @@ from litellm.proxy.auth.auth_checks import (
can_key_call_resolved_model,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.db.autorouter_session_rollup import AUTOROUTER_BENCHMARKS_SQL
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.repositories.team_repository import TeamRepository
from litellm.router_strategy.complexity_router import ComplexityRouter
@ -292,53 +293,6 @@ class _SessionAggRow(BaseModel):
_SESSION_AGG_ROWS: Final = TypeAdapter(list[_SessionAggRow])
_BENCHMARKS_SQL: Final = """
WITH windowed AS (
SELECT * FROM "LiteLLM_AutoRouterSession"
WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp
),
tier_maps AS (
SELECT router_name, router_type, jsonb_object_agg(tier, tier_turns) AS tier_turns
FROM (
SELECT router_name, router_type, kv.key AS tier, SUM((kv.value)::int)::int AS tier_turns
FROM windowed, LATERAL jsonb_each_text(tier_turns) AS kv
GROUP BY router_name, router_type, kv.key
) per_tier
GROUP BY router_name, router_type
)
SELECT
agg.*,
COALESCE(tier_maps.tier_turns, '{}'::jsonb) AS tier_turns
FROM (
SELECT
router_name,
router_type,
COUNT(*)::int AS sessions,
COALESCE(SUM(turns), 0)::int AS turns,
COALESCE(SUM(unordered_turns), 0)::int AS unordered_turns,
COALESCE(SUM(covered_turns), 0)::int AS covered_turns,
COALESCE(SUM(cache_hits), 0)::int AS cache_hits,
COALESCE(SUM(same_model_turns), 0)::int AS same_model_turns,
COALESCE(SUM(same_model_hits), 0)::int AS same_model_hits,
COALESCE(SUM(first_visit_turns), 0)::int AS first_visit_turns,
COALESCE(SUM(first_visit_hits), 0)::int AS first_visit_hits,
COALESCE(SUM(return_turns), 0)::int AS return_turns,
COALESCE(SUM(return_hits), 0)::int AS return_hits,
COALESCE(SUM(return_expired_misses), 0)::int AS return_expired_misses,
COALESCE(SUM(return_within_ttl_misses), 0)::int AS return_within_ttl_misses,
COALESCE(SUM(ttl_5m_turns), 0)::int AS ttl_5m_turns,
COALESCE(SUM(ttl_1h_turns), 0)::int AS ttl_1h_turns,
COALESCE(SUM(total_tokens), 0)::bigint AS total_tokens,
COALESCE(SUM(spend), 0)::float8 AS spend,
COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend,
COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds
FROM windowed
GROUP BY router_name, router_type
) agg
LEFT JOIN tier_maps USING (router_name, router_type)
ORDER BY agg.spend DESC
"""
def _parse_benchmark_day(value: str) -> datetime:
try:
@ -462,7 +416,7 @@ async def get_auto_router_benchmarks(
raise HTTPException(status_code=400, detail="end_date must not be earlier than start_date")
raw_rows: Final = await prisma_client.db.query_raw(
_BENCHMARKS_SQL,
AUTOROUTER_BENCHMARKS_SQL,
start_day.isoformat(),
(end_day + timedelta(days=1)).isoformat(),
)

View file

@ -46,6 +46,7 @@ from litellm.constants import (
DEFAULT_HEALTH_CHECK_INTERVAL,
DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER,
DEFAULT_MAX_LRU_CACHE_SIZE,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.asyncify import run_async_function
@ -135,6 +136,7 @@ from litellm.router_utils.handle_error import (
from litellm.router_utils.health_state_cache import DeploymentHealthCache
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
DeploymentAffinityCheck,
warn_on_unknown_model_group_affinity_flags,
)
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
build_io_token_rate_limit_headers,
@ -603,6 +605,10 @@ class Router:
# ``litellm.proxy.auth.auth_checks._is_model_cost_zero``.
self._zero_cost_cache: dict[str, bool] = {}
self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds
self.model_group_affinity_config = model_group_affinity_config
warn_on_unknown_model_group_affinity_flags(model_group_affinity_config)
if model_list is not None:
# set_model_list will build indices automatically
self.set_model_list(model_list)
@ -744,7 +750,6 @@ class Router:
litellm.failure_callback = [self.deployment_callback_on_failure]
self.routing_strategy_args = routing_strategy_args
self.provider_budget_config = provider_budget_config
self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds
self.router_budget_logger: RouterBudgetLimiting | None = None
if RouterBudgetLimiting.should_init_router_budget_limiter(
model_list=model_list, provider_budget_config=self.provider_budget_config
@ -766,7 +771,6 @@ class Router:
)
self.model_group_retry_policy: dict[str, RetryPolicy] | None = model_group_retry_policy
self.model_group_affinity_config: dict[str, list[str]] | None = model_group_affinity_config
self.allowed_fails_policy: AllowedFailsPolicy | None = None
if allowed_fails_policy is not None:
@ -789,21 +793,8 @@ class Router:
# If model_group_affinity_config is set but no global affinity checks were
# enabled, we still need the DeploymentAffinityCheck callback (with global
# flags all False) so per-group config can activate affinity per model group.
if self.model_group_affinity_config and not any(
isinstance(cb, DeploymentAffinityCheck) for cb in (self.optional_callbacks or [])
):
if self.optional_callbacks is None:
self.optional_callbacks = []
affinity_callback: Final = DeploymentAffinityCheck(
cache=self.cache,
ttl_seconds=self.deployment_affinity_ttl_seconds,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
enable_session_id_affinity=False,
model_group_affinity_config=self.model_group_affinity_config,
)
self.optional_callbacks.append(affinity_callback)
litellm.logging_callback_manager.add_litellm_callback(affinity_callback)
if self.model_group_affinity_config:
self._ensure_deployment_affinity_callback()
if self.alerting_config is not None:
self._initialize_alerting()
@ -1662,6 +1653,28 @@ class Router:
_move_before_deployment_affinity(self.optional_callbacks, ec_callback)
_move_before_deployment_affinity(litellm.callbacks, ec_callback)
def _ensure_deployment_affinity_callback(self) -> None:
"""Register the DeploymentAffinityCheck callback (global flags all False) if absent.
Needed when nothing enabled a global affinity flag but affinity can still
activate per request: per-group `model_group_affinity_config` entries, or the
session-affinity marker a complexity router stamps at pre-routing time.
"""
if any(isinstance(cb, DeploymentAffinityCheck) for cb in (self.optional_callbacks or [])):
return
if self.optional_callbacks is None:
self.optional_callbacks = []
affinity_callback: Final = DeploymentAffinityCheck(
cache=self.cache,
ttl_seconds=self.deployment_affinity_ttl_seconds,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
enable_session_id_affinity=False,
model_group_affinity_config=self.model_group_affinity_config,
)
self.optional_callbacks.append(affinity_callback)
litellm.logging_callback_manager.add_litellm_callback(affinity_callback)
def add_optional_pre_call_checks(self, optional_pre_call_checks: OptionalPreCallChecks | None):
if optional_pre_call_checks is None:
return
@ -7683,6 +7696,8 @@ class Router:
strategy=complexity_router,
strategy_label="Complexity-router",
)
if complexity_router._uses_deployment_pin:
self._ensure_deployment_affinity_callback()
def _is_adaptive_router_deployment(self, litellm_params: LiteLLM_Params) -> bool:
"""True when this deployment opts in via the `auto_router/adaptive_router` model prefix."""
@ -11190,6 +11205,9 @@ class Router:
router_strategy: Final = self._select_pre_routing_strategy(model=model, request_kwargs=request_kwargs)
if router_strategy is None:
self._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None)
self._stamp_or_clear_metadata_key(
request_kwargs=request_kwargs, key=SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, value=None
)
return None
pre_routing_hook_response: Final = await router_strategy.async_pre_routing_hook(
@ -11203,6 +11221,11 @@ class Router:
request_kwargs=request_kwargs,
routing_decision=(pre_routing_hook_response.routing_decision if pre_routing_hook_response else None),
)
self._stamp_or_clear_metadata_key(
request_kwargs=request_kwargs,
key=SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
value=(pre_routing_hook_response.session_affinity_ttl_seconds if pre_routing_hook_response else None),
)
# `model` (the alias, e.g. "smart-router") is never the deployment actually
# called - apply the alias's own litellm_params (besides `model` itself,
@ -11234,21 +11257,40 @@ class Router:
to the deployment that actually served the request. Every attempt therefore
writes or clears, never just writes.
"""
if routing_decision is None:
Router._stamp_or_clear_metadata_key(
request_kwargs=request_kwargs,
key="routing_decision",
value=(
None
if routing_decision is None
else Router._redact_prompt_text_if_needed(
request_kwargs=request_kwargs, routing_decision=routing_decision
)
),
)
@staticmethod
def _stamp_or_clear_metadata_key(request_kwargs: dict, key: str, value: object | None) -> None:
"""Write a proxy-internal metadata key for THIS routing attempt, or clear it.
Fallbacks and retries re-enter the pre-routing hook with the same
`request_kwargs`, so every attempt must write or clear, never just write;
a value left behind by an earlier attempt would be attributed to this one.
`get_or_create_metadata_bucket` is the single owner of "which dict holds
proxy-internal metadata": it picks `litellm_metadata` when present (so the
value never lands in the `metadata` dict that routes like /v1/messages
forward to the provider) and replaces a non-dict value rather than silently
skipping the write. Clearing pops from BOTH buckets so a request whose
bucket resolution changed between attempts cannot resurrect a stale value.
"""
if value is None:
for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")):
if isinstance(bucket, dict):
bucket.pop("routing_decision", None)
bucket.pop(key, None)
return
# `get_or_create_metadata_bucket` is the single owner of "which dict holds
# proxy-internal metadata": it picks `litellm_metadata` when present (so the
# decision never lands in the `metadata` dict that routes like /v1/messages
# forward to the provider) and replaces a non-dict value rather than silently
# skipping the write.
_, metadata_bucket = get_or_create_metadata_bucket(request_kwargs)
metadata_bucket["routing_decision"] = Router._redact_prompt_text_if_needed(
request_kwargs=request_kwargs, routing_decision=routing_decision
)
metadata_bucket[key] = value
@staticmethod
def _redact_prompt_text_if_needed(

View file

@ -1622,6 +1622,28 @@ class ComplexityRouter(CustomLogger):
caller_scope: Final = self._get_user_api_key_hash_from_request_kwargs(request_kwargs) or "unscoped"
return f"complexity_router_session_affinity:v1:{self.model_name}:{caller_scope}:{session_id}"
@property
def _uses_tier_pin(self) -> bool:
return bool(self.config.session_affinity and not self.config.plugins)
@property
def _uses_deployment_pin(self) -> bool:
"""session_affinity implies the deployment pin: a session frozen onto one model
group but load-balanced across its deployments would still go cache-cold, which
is the exact failure both flags exist to prevent."""
return bool((self.config.deployment_affinity or self.config.session_affinity) and not self.config.plugins)
def _with_session_deployment_affinity(
self, response: PreRoutingHookResponse | None
) -> PreRoutingHookResponse | None:
if response is None or not self._uses_deployment_pin:
return response
return response.model_copy(
update={ # mutable-ok: model_copy types update as a plain dict
"session_affinity_ttl_seconds": self.config.session_affinity_ttl_seconds
}
)
async def async_pre_routing_hook(
self,
model: str,
@ -1656,7 +1678,7 @@ class ComplexityRouter(CustomLogger):
resolved_messages: Final = self._resolve_messages(messages, request_kwargs)
conversation_continuing: Final = _conversation_is_continuing(resolved_messages)
use_session_affinity: Final = self.config.session_affinity and not self.config.plugins
use_session_affinity: Final = self._uses_tier_pin
session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs) if use_session_affinity else None
cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None
@ -1695,17 +1717,19 @@ class ComplexityRouter(CustomLogger):
"ComplexityRouter: routing decision cause=%s, routed_model=%s", cause, routed_model
)
has_original_messages: Final = messages is not None and len(messages) > 0
return PreRoutingHookResponse(
model=routed_model,
messages=messages if has_original_messages else None,
routing_decision=self._build_routing_decision(
routed_model=routed_model,
cause=cause,
tier=self._tier_for_model(routed_model),
escalation_keyword=pin_escalation_keyword,
escalated=escalated,
conversation_continuing=conversation_continuing,
),
return self._with_session_deployment_affinity(
PreRoutingHookResponse(
model=routed_model,
messages=messages if has_original_messages else None,
routing_decision=self._build_routing_decision(
routed_model=routed_model,
cause=cause,
tier=self._tier_for_model(routed_model),
escalation_keyword=pin_escalation_keyword,
escalated=escalated,
conversation_continuing=conversation_continuing,
),
)
)
response: Final = await self._classify_and_route(
@ -1723,7 +1747,7 @@ class ComplexityRouter(CustomLogger):
value=response.model,
ttl=self.config.session_affinity_ttl_seconds,
)
return response
return self._with_session_deployment_affinity(response)
async def _classify_and_route(
self,

View file

@ -508,13 +508,39 @@ class ComplexityRouterConfig(BaseModel):
"session's first turn and reuse it for every later turn, skipping re-classification. "
"Off by default so every turn is classified on its own merits and routed to the cheapest "
"adequate tier. Set True to keep a multi-turn session on one model, which preserves "
"provider prompt caches and avoids cross-model conversation-history errors."
"provider prompt caches and avoids cross-model conversation-history errors. Always "
"implies the deployment pin regardless of deployment_affinity: the session sticks to "
"one deployment of the pinned model, since freezing the model while re-shuffling its "
"deployments would still go cache-cold."
),
)
deployment_affinity: bool = Field(
default=True,
description=(
"When True and a session_id is resolvable on the request, pin the deployment chosen "
"inside each routed model group and reuse it whenever the session returns to that "
"group, without pinning which group the session routes to. Independent of "
"session_affinity, which pins the model group instead (and always carries this "
"deployment pin with it): with session_affinity off, "
"every turn is still classified on its own merits while a session that escalates to a "
"stronger tier and comes back still lands on the deployment it used before, which is "
"what keeps a provider prompt cache warm. Pins are held per model group, so switching "
"tiers does not disturb the pin left behind in the previous group. On by default "
"because re-shuffling a conversation across deployments of the same model discards "
"that cache for no benefit; set False to keep every turn load-balanced across the "
"group, which is what a deployment set with tight per-deployment rate limits wants. "
"Inert when no session_id is resolvable, since there is nothing to key a pin on, and "
"suppressed when plugins are configured, for the same reason session_affinity is."
),
)
session_affinity_ttl_seconds: int = Field(
default=3600,
gt=0,
description="TTL for the session affinity pin; refreshed on every cache hit",
description=(
"TTL for the session affinity pin; refreshed on every cache hit. Bounds both the "
"session_affinity model pin and the deployment_affinity deployment pin, so it measures "
"idle time for the session's routing decisions rather than total session length"
),
)
plugins: list[RoutingPlugin] | None = Field(

View file

@ -13,12 +13,15 @@ where routing to a consistent deployment is still beneficial.
"""
import hashlib
import json
from collections.abc import Mapping, Sequence
from typing import Any, Final, cast
from typing_extensions import TypedDict
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY
from litellm.integrations.custom_logger import CustomLogger, Span
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import AllMessageValues
@ -29,6 +32,47 @@ class DeploymentAffinityCacheValue(TypedDict):
model_id: str
VALID_MODEL_GROUP_AFFINITY_FLAGS: Final = frozenset(
{
"deployment_affinity",
"responses_api_deployment_check",
"session_affinity",
"encrypted_content_affinity",
}
)
def warn_on_unknown_model_group_affinity_flags(model_group_affinity_config: Mapping[str, Sequence[str]] | None) -> None:
"""`model_group_affinity_config` is one Router-level config consumed by two callbacks:
DeploymentAffinityCheck acts on three of the flags and EncryptedContentAffinityCheck
on the fourth, so typo detection lives here at the schema, not inside either consumer.
"""
if model_group_affinity_config is None:
return
for group, flags in model_group_affinity_config.items():
unknown = set(flags) - VALID_MODEL_GROUP_AFFINITY_FLAGS
if unknown:
verbose_router_logger.warning(
"model_group_affinity_config: unknown flag(s) %s for model group '%s'; will be ignored. Valid flags: %s",
unknown,
group,
VALID_MODEL_GROUP_AFFINITY_FLAGS,
)
_CLAIM_PIN_SCRIPT: Final = """
local current = redis.call('GET', KEYS[1])
if current == false then
redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2])
return ARGV[1]
end
if current == ARGV[1] then
redis.call('EXPIRE', KEYS[1], ARGV[2])
end
return current
"""
class DeploymentAffinityCheck(CustomLogger):
"""
Router deployment affinity callback.
@ -38,14 +82,6 @@ class DeploymentAffinityCheck(CustomLogger):
"""
CACHE_KEY_PREFIX = "deployment_affinity:v1"
VALID_FLAGS = frozenset(
{
"deployment_affinity",
"responses_api_deployment_check",
"session_affinity",
"encrypted_content_affinity",
}
)
def __init__(
self,
@ -63,15 +99,6 @@ class DeploymentAffinityCheck(CustomLogger):
self.enable_responses_api_affinity = enable_responses_api_affinity
self.enable_session_id_affinity = enable_session_id_affinity
self.model_group_affinity_config: dict[str, list[str]] = model_group_affinity_config or {}
for group, flags in self.model_group_affinity_config.items():
unknown = set(flags) - self.VALID_FLAGS
if unknown:
verbose_router_logger.warning(
"DeploymentAffinityCheck: unknown flag(s) %s for model group '%s'; will be ignored. Valid flags: %s",
unknown,
group,
self.VALID_FLAGS,
)
def _get_effective_flags(self, model_group: str) -> tuple[bool, bool, bool]:
"""
@ -218,8 +245,13 @@ class DeploymentAffinityCheck(CustomLogger):
return f"{cls.CACHE_KEY_PREFIX}:{model_group}:{hashed_user_key}"
@classmethod
def get_session_affinity_cache_key(cls, model_group: str, session_id: str) -> str:
return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{session_id}"
def get_session_affinity_cache_key(cls, model_group: str, session_id: str, user_key: str | None) -> str:
"""Session pins are scoped by the caller's hashed API key so two callers reusing
the same client-supplied session_id cannot read or steer each other's pin.
`"unscoped"` covers direct Router usage with no authenticated caller, matching
the complexity router's own session pin key."""
hashed_user_key: Final = cls._hash_user_key(user_key) if user_key is not None else "unscoped"
return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{hashed_user_key}:{session_id}"
@staticmethod
def _get_user_key_from_metadata_dict(metadata: dict) -> str | None:
@ -278,6 +310,97 @@ class DeploymentAffinityCheck(CustomLogger):
return session_id
return None
@staticmethod
def _get_marker_session_affinity_ttl(request_kwargs: dict) -> int | None:
"""TTL from the session-affinity marker the Router stamps at pre-routing time
when an auto-router routed this request with session_affinity enabled.
Marker presence enables session pinning for this request only; anything that
is not a positive int is treated as absent."""
for metadata in DeploymentAffinityCheck._iter_metadata_dicts(request_kwargs):
ttl = metadata.get(SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY)
if isinstance(ttl, int) and not isinstance(ttl, bool) and ttl > 0:
return ttl
return None
@staticmethod
def _pinned_model_id(stored: object) -> str | None:
"""Deployment id held by a stored pin, for both the dict shape this writes and the
bare string older writers left behind. None when the value is neither."""
if isinstance(stored, dict):
model_id: Final = stored.get("model_id")
return str(model_id) if model_id is not None else None
if isinstance(stored, str):
return stored
return None
def _set_local_pin(self, cache_key: str, value: object, ttl_seconds: int) -> None:
"""The one owner of authoritative local pin writes: a plain set keeps a live
key's original expiry (`allow_ttl_override`), so the entry is replaced to make
the TTL real. Every local pin write goes through here so the redis-winner sync
and the pod-local claim can never disagree about expiry again."""
self.cache.in_memory_cache.delete_cache(cache_key)
self.cache.in_memory_cache.set_cache(cache_key, value, ttl=ttl_seconds)
async def _claim_pin(self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int) -> str | None:
"""First-writer-wins pin write: store `pin_value` only when the key is absent and
return the deployment id the key holds afterwards, so a caller learns whether it won
by comparing against its own id, and None when the stored value is one no reader can
interpret. Concurrent claimers converge on the
first write instead of the last. Re-claiming with the stored value refreshes its
TTL, the same keepalive the complexity router's model pin documents: an active
session must not lose its pin mid-conversation just because it outlives the
original write, so `session_affinity_ttl_seconds` bounds idle time, not total
session length. On Redis one Lua script does the get-or-set-or-refresh
atomically (same registration seam the rate limiters use) and the in-memory
tier is synchronized to the winner; without Redis, and whenever Redis is
unreachable, the pod-local check-and-set below stands in and is atomic because it
runs synchronously on the event loop. Degrading to a pod-local claim rather than
propagating the fault is what keeps same-pod stickiness through a Redis blip: the
caller only logs this result, so an escaping error would leave the session with no
pin at all and reshuffle every turn for the outage, which is worse than losing
cross-pod agreement. The redis tier is
resolved per call because the proxy attaches it after Router construction
(`Router._update_redis_cache`); the compiled script is cached per event loop
underneath the registration seam.
"""
redis_cache: Final = self.cache.redis_cache
if redis_cache is not None:
try:
claim_script: Final = redis_cache.async_register_script(_CLAIM_PIN_SCRIPT)
raw: Final = await claim_script(keys=(cache_key,), args=(json.dumps(pin_value), int(ttl_seconds)))
decoded: Final = raw.decode("utf-8") if isinstance(raw, bytes) else raw
if not isinstance(decoded, str):
return pin_value["model_id"]
try:
winner: object = json.loads(decoded)
except json.JSONDecodeError:
winner = decoded
self._set_local_pin(cache_key=cache_key, value=winner, ttl_seconds=ttl_seconds)
return self._pinned_model_id(winner)
except Exception as e: # noqa: BLE001 # any Redis/Lua failure degrades to the pod-local claim, never unpins
verbose_router_logger.debug(
"DeploymentAffinityCheck: redis pin claim failed, falling back to pod-local claim. error=%s", e
)
return self._claim_pin_in_memory(cache_key=cache_key, pin_value=pin_value, ttl_seconds=ttl_seconds)
def _claim_pin_in_memory(
self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int
) -> str | None:
"""Pod-local half of the claim, used when no Redis tier is attached and as the
fallback when the Redis claim fails. Mirrors the Lua script exactly, including
the keepalive: re-claiming with the stored value slides the idle window through
`_set_local_pin`. Both branches stay synchronous, hence atomic on the event
loop."""
existing: Final = self.cache.in_memory_cache.get_cache(cache_key)
if existing is not None:
existing_model_id: Final = self._pinned_model_id(existing)
if existing_model_id == pin_value["model_id"]:
self._set_local_pin(cache_key=cache_key, value=pin_value, ttl_seconds=ttl_seconds)
return existing_model_id
self._set_local_pin(cache_key=cache_key, value=pin_value, ttl_seconds=ttl_seconds)
return pin_value["model_id"]
@staticmethod
def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None:
for deployment in healthy_deployments:
@ -334,12 +457,21 @@ class DeploymentAffinityCheck(CustomLogger):
if stable_model_map_key is None:
return typed_healthy_deployments
session_affinity_active: Final = (
enable_session_id or self._get_marker_session_affinity_ttl(request_kwargs=request_kwargs) is not None
)
user_key: Final = (
self._get_user_key_from_request_kwargs(request_kwargs=request_kwargs)
if (session_affinity_active or enable_user_key)
else None
)
# 2) Session-id -> deployment affinity
if enable_session_id:
if session_affinity_active:
session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs=request_kwargs)
if session_id is not None:
session_cache_key: Final = self.get_session_affinity_cache_key(
model_group=stable_model_map_key, session_id=session_id
model_group=stable_model_map_key, session_id=session_id, user_key=user_key
)
session_cache_result: Final = await self.cache.async_get_cache(key=session_cache_key)
@ -371,7 +503,6 @@ class DeploymentAffinityCheck(CustomLogger):
if not enable_user_key:
return typed_healthy_deployments
user_key: Final = self._get_user_key_from_request_kwargs(request_kwargs=request_kwargs)
if user_key is None:
return typed_healthy_deployments
@ -438,18 +569,22 @@ class DeploymentAffinityCheck(CustomLogger):
enable_session_id,
) = self._get_effective_flags(deployment_model_name)
if not enable_user_key and not enable_session_id:
marker_session_ttl: Final = self._get_marker_session_affinity_ttl(request_kwargs=kwargs)
session_affinity_active: Final = enable_session_id or marker_session_ttl is not None
if not enable_user_key and not session_affinity_active:
return None
user_key = None
if enable_user_key:
user_key = self._get_user_key_from_request_kwargs(request_kwargs=kwargs)
user_key: Final = (
self._get_user_key_from_request_kwargs(request_kwargs=kwargs)
if (enable_user_key or session_affinity_active)
else None
)
session_id: Final = (
self._get_session_id_from_request_kwargs(request_kwargs=kwargs) if session_affinity_active else None
)
session_id = None
if enable_session_id:
session_id = self._get_session_id_from_request_kwargs(request_kwargs=kwargs)
if user_key is None and session_id is None:
if not ((enable_user_key and user_key is not None) or session_id is not None):
return None
model_info = kwargs.get("model_info")
@ -473,22 +608,31 @@ class DeploymentAffinityCheck(CustomLogger):
verbose_router_logger.warning("DeploymentAffinityCheck: model_id missing; skipping affinity cache update.")
return None
if user_key is not None:
pin_value: Final = DeploymentAffinityCacheValue(model_id=str(model_id))
if enable_user_key and user_key is not None:
try:
cache_key: Final = self.get_affinity_cache_key(model_group=deployment_model_name, user_key=user_key)
await self.cache.async_set_cache(
cache_key,
DeploymentAffinityCacheValue(model_id=str(model_id)),
ttl=self.ttl_seconds,
)
verbose_router_logger.debug(
"DeploymentAffinityCheck: set affinity mapping model_map_key=%s deployment=%s ttl=%s user_key=%s",
deployment_model_name,
model_id,
self.ttl_seconds,
self._shorten_for_logs(user_key),
claimed_user_pin: Final = await self._claim_pin(
cache_key=cache_key,
pin_value=pin_value,
ttl_seconds=self.ttl_seconds,
)
if claimed_user_pin == pin_value["model_id"]:
verbose_router_logger.debug(
"DeploymentAffinityCheck: set affinity mapping model_map_key=%s deployment=%s ttl=%s user_key=%s",
deployment_model_name,
model_id,
self.ttl_seconds,
self._shorten_for_logs(user_key),
)
else:
verbose_router_logger.debug(
"DeploymentAffinityCheck: affinity pin already claimed model_map_key=%s existing=%s ours=%s",
deployment_model_name,
claimed_user_pin,
model_id,
)
except Exception as e:
# Non-blocking: affinity is a best-effort optimization.
verbose_router_logger.debug(
@ -500,21 +644,31 @@ class DeploymentAffinityCheck(CustomLogger):
# Also persist Session-ID affinity if enabled and session-id is provided
if session_id is not None:
try:
session_affinity_ttl: Final = marker_session_ttl if marker_session_ttl is not None else self.ttl_seconds
session_cache_key: Final = self.get_session_affinity_cache_key(
model_group=deployment_model_name, session_id=session_id
model_group=deployment_model_name, session_id=session_id, user_key=user_key
)
await self.cache.async_set_cache(
session_cache_key,
DeploymentAffinityCacheValue(model_id=str(model_id)),
ttl=self.ttl_seconds,
)
verbose_router_logger.debug(
"DeploymentAffinityCheck: set session affinity mapping model_map_key=%s deployment=%s ttl=%s session_id=%s",
deployment_model_name,
model_id,
self.ttl_seconds,
session_id,
claimed_session_pin: Final = await self._claim_pin(
cache_key=session_cache_key,
pin_value=pin_value,
ttl_seconds=session_affinity_ttl,
)
if claimed_session_pin == pin_value["model_id"]:
verbose_router_logger.debug(
"DeploymentAffinityCheck: set session affinity mapping model_map_key=%s deployment=%s ttl=%s session_id=%s",
deployment_model_name,
model_id,
session_affinity_ttl,
session_id,
)
else:
verbose_router_logger.debug(
"DeploymentAffinityCheck: session pin already claimed model_map_key=%s existing=%s ours=%s session_id=%s",
deployment_model_name,
claimed_session_pin,
model_id,
session_id,
)
except Exception as e:
verbose_router_logger.debug(
"DeploymentAffinityCheck: failed to set session affinity cache. model_map_key=%s error=%s",

View file

@ -816,6 +816,7 @@ class PreRoutingHookResponse(BaseModel):
model: str
messages: list[dict[str, Any]] | None
routing_decision: StandardLoggingRoutingDecision | None = None
session_affinity_ttl_seconds: int | None = None
_PreRoutingStrategyT_co = TypeVar("_PreRoutingStrategyT_co", covariant=True)

View file

@ -35,7 +35,7 @@ These hooks enforce Conventional Commits and Conventional Branches.
Bypass with --no-verify when you need to (e.g. for emergency hotfixes).
The CI-equivalent lint is deliberately not installed as an auto-firing hook
(it can take minutes); run it on demand with 'make pre-commit' before committing.
(it can take minutes); run it on demand with 'make check' before committing.
To uninstall: git config --unset core.hooksPath
EOF

View file

@ -1,18 +1,25 @@
#!/usr/bin/env bash
#
# pre_commit_lint.sh — shift CI lint left. Run it (via `make pre-commit`) right
# before `git commit`; it inspects your staged files and runs only the matching
# gating CI checks, so a clean run means a green CI lint:
# - litellm/ Python staged -> `make lint` (test-linting.yml's lint job)
# - tests/e2e Python staged -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step)
# + raw HTTP client ban (test-code-quality.yml's check_e2e_no_raw_requests)
# - dashboard staged -> prettier + eslint + lint budgets (test-litellm-ui-build.yml's frontend-lint)
# - proxy/types staged -> regenerate dashboard API types and fail on drift (check-ui-api-types.yml)
# pre_commit_lint.sh — shift CI lint left. Run it (via `make check`, formerly
# `make pre-commit`) before `git commit`, or after committing (e.g. a merge
# commit) to predict CI for the branch. It picks the files in scope and runs
# only the matching gating CI checks, so a clean run means a green CI lint:
# - anything staged -> scope is the staged files; changed-but-unstaged files
# whose checks were skipped are called out
# - nothing staged -> scope is the working tree's diff against the merge base
# with origin/litellm_internal_staging, untracked files included
# The per-area checks:
# - litellm/ Python -> `make lint` (test-linting.yml's lint job)
# - tests/e2e Python -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step)
# + raw HTTP client ban (test-code-quality.yml's check_e2e_no_raw_requests)
# - dashboard -> prettier + eslint + lint budgets (test-litellm-ui-build.yml's frontend-lint)
# - proxy/types -> regenerate dashboard API types and fail on drift (check-ui-api-types.yml)
#
# Each block is skipped when no matching files are staged, so unrelated commits stay
# fast. This is intentionally not auto-installed as a git hook (see scripts/install_git_hooks.sh):
# the dashboard and basedpyright passes can take minutes, so it's run on demand rather
# than firing on every human commit. It is hook-compatible if you want that anyway:
# Each block is skipped when no matching files are in scope, so unrelated commits
# stay fast. This is intentionally not auto-installed as a git hook (see
# scripts/install_git_hooks.sh): the dashboard and basedpyright passes can take
# minutes, so it's run on demand rather than firing on every human commit. It is
# hook-compatible if you want that anyway:
# `ln -s ../../scripts/pre_commit_lint.sh .git/hooks/pre-commit`.
set -eu
@ -20,56 +27,105 @@ set -eu
if [ -z "${PRE_COMMIT_LINT_INNER:-}" ]; then
log_file=$(git rev-parse --path-format=absolute --git-path pre_commit_lint.log)
if : > "$log_file" 2>/dev/null; then
echo "pre-commit: logging full output to $log_file"
echo "check: logging full output to $log_file"
PRE_COMMIT_LINT_INNER=1 "$0" "$@" 2>&1 | tee "$log_file"
pipe_status=("${PIPESTATUS[@]}")
if [ "${pipe_status[1]}" -eq 0 ]; then
echo "pre-commit: full log: $log_file"
echo "check: full log: $log_file"
else
echo "pre-commit: WARNING - writing $log_file failed; the log may be incomplete" >&2
echo "check: WARNING - writing $log_file failed; the log may be incomplete" >&2
fi
exit "${pipe_status[0]}"
fi
echo "pre-commit: WARNING - cannot write $log_file; output will not be saved" >&2
echo "check: WARNING - cannot write $log_file; output will not be saved" >&2
PRE_COMMIT_LINT_INNER=1 exec "$0" "$@"
fi
repo_root=$(git rev-parse --show-toplevel)
cd "$repo_root"
staged=$(git diff --cached --name-only --diff-filter=ACMR)
staged_match() { printf '%s\n' "$staged" | grep -E "$1" || true; }
staged=$(git diff --cached --name-only --diff-filter=ACMRD)
unstaged=$(git diff --name-only)
untracked=$(git ls-files --others --exclude-standard)
if [ -n "$staged" ]; then
scope=$staged
else
git fetch --quiet origin litellm_internal_staging 2>/dev/null || true
merge_base=$(git merge-base origin/litellm_internal_staging HEAD 2>/dev/null) || {
echo "check: cannot resolve the merge base with origin/litellm_internal_staging." >&2
echo " Fix: git fetch origin litellm_internal_staging" >&2
exit 1
}
scope=$(printf '%s\n' "$(git diff --name-only --diff-filter=ACMRD "$merge_base")" "$untracked" | sed '/^$/d' | sort -u)
if [ -z "$scope" ]; then
echo "check: nothing to check (no staged files, no working-tree changes, no branch changes vs origin/litellm_internal_staging)"
exit 0
fi
echo "check: nothing staged; scoping to the working tree's diff against the merge base with origin/litellm_internal_staging:"
printf '%s\n' "$scope" | sed 's/^/ /'
fi
scope_match() { printf '%s\n' "$scope" | grep -E "$1" || true; }
existing_files() {
while IFS= read -r f; do
if [ -f "$f" ]; then printf '%s\n' "$f"; fi
done
}
litellm_py_pattern='^litellm/.*\.py$'
e2e_py_pattern='^tests/e2e/.*\.py$'
spec_pattern='^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$'
ui_prettier_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|scss|md|mdx|yml|yaml|html)$'
ui_eslint_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$'
# CI's lint job (test-linting.yml) only inspects litellm/, so a tests-only or
# scripts-only commit can't turn it red; scope the trigger there to skip the slow
# make lint when it couldn't catch anything.
litellm_py_files=$(staged_match '^litellm/.*\.py$')
e2e_py_files=$(staged_match '^tests/e2e/.*\.py$')
litellm_py_files=$(scope_match "$litellm_py_pattern")
e2e_py_files=$(scope_match "$e2e_py_pattern")
# ruff format (and CI's format step) skip enterprise; the rest of make lint covers it.
fmt_files=$(printf '%s\n' "$litellm_py_files" | grep -v '^litellm/enterprise/' || true)
fmt_files=$(printf '%s\n' "$litellm_py_files" | grep -v '^litellm/enterprise/' | existing_files)
# check-ui-api-types.yml triggers on any file under litellm/proxy or litellm/types
# (Prisma schema and configs included, not just Python) plus the generator and its
# lockfiles, so match that whole trigger set rather than a Python subset.
spec_files=$(staged_match '^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$')
spec_files=$(scope_match "$spec_pattern")
# CI's frontend-lint runs prettier over a wider extension set than eslint; keep that
# split so this flags exactly what the job would.
ui_prettier_files=$(staged_match '^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|scss|md|mdx|yml|yaml|html)$')
ui_eslint_files=$(staged_match '^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$')
ui_prettier_changed=$(scope_match "$ui_prettier_pattern")
ui_eslint_changed=$(scope_match "$ui_eslint_pattern")
ui_prettier_files=$(printf '%s\n' "$ui_prettier_changed" | existing_files)
ui_eslint_files=$(printf '%s\n' "$ui_eslint_changed" | existing_files)
# CI lints the committed tree, so this script predicts CI for what you have STAGED
# (every trigger above reads `git diff --cached`). The tools it runs, though, read
# the working tree, so unstaged edits to tracked files and untracked files fold
# into the result and a green/red here won't match a commit of just the staged
# changes. There's no safe way to lint the index in place, so surface the gap
# instead of hiding it: stage everything you intend to commit before trusting a
# pass. This only warns; it never blocks or touches your changes.
unstaged=$(git diff --name-only)
untracked=$(git ls-files --others --exclude-standard)
if [ -n "$unstaged" ] || [ -n "$untracked" ]; then
echo "pre-commit: NOTE - unstaged/untracked changes are included in these checks but" >&2
echo " won't be in a commit of only your staged changes, so this result may differ from" >&2
echo " CI. Stage everything you intend to commit (git add) for an accurate prediction:" >&2
printf '%s\n' "$unstaged" "$untracked" | sed '/^$/d' | sed 's/^/ /' >&2
# CI lints the committed tree, so with staged files this script predicts CI for
# what you have STAGED (every trigger above reads `git diff --cached`). The tools
# it runs, though, read the working tree, so unstaged edits to tracked files and
# untracked files fold into the result and a green/red here won't match a commit
# of just the staged changes. There's no safe way to lint the index in place, so
# surface the gap instead of hiding it: stage everything you intend to commit
# before trusting a pass. This only warns; it never blocks or touches your changes.
if [ -n "$staged" ]; then
not_staged=$(printf '%s\n' "$unstaged" "$untracked" | sed '/^$/d' | sort -u)
if [ -n "$not_staged" ]; then
echo "check: NOTE - unstaged/untracked changes are included in these checks but" >&2
echo " won't be in a commit of only your staged changes, so this result may differ from" >&2
echo " CI. Stage everything you intend to commit (git add) for an accurate prediction:" >&2
printf '%s\n' "$not_staged" | sed 's/^/ /' >&2
fi
warn_skipped() {
local check_name=$1 pattern=$2 triggered=$3
[ -n "$triggered" ] && return 0
local missed
missed=$(printf '%s\n' "$not_staged" | grep -E "$pattern" || true)
[ -z "$missed" ] && return 0
echo "check: SKIPPED $check_name because these changed files are not staged:" >&2
printf '%s\n' "$missed" | sed 's/^/ /' >&2
}
warn_skipped "Python lint (make lint)" "$litellm_py_pattern" "$litellm_py_files"
warn_skipped "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_pattern" "$e2e_py_files"
warn_skipped "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_pattern" "$ui_prettier_changed"
warn_skipped "dashboard API-type sync (npm run gen:api)" "$spec_pattern" "$spec_files"
fi
lint_dashboard() {
@ -114,15 +170,15 @@ bootstrap_hint() {
python_checks() {
local rc=0
echo "pre-commit: linting Python (make lint)"
make lint || { echo "✗ Python lint failed. Fix the reds above, then re-run make pre-commit." >&2; rc=1; }
echo "check: linting Python (make lint)"
make lint || { echo "✗ Python lint failed. Fix the reds above, then re-run make check." >&2; rc=1; }
# `make lint` format-checks files in origin/base...HEAD, which at pre-commit time
# predates the staged change, so format-check the staged litellm files directly to
# predates the staged change, so format-check the scoped litellm files directly to
# cover a brand-new commit before it lands.
if [ -n "$fmt_files" ]; then
echo "pre-commit: ruff format --check (staged litellm files)"
echo "check: ruff format --check (scoped litellm files)"
printf '%s\n' "$fmt_files" | xargs uv run --no-sync ruff format --check --exclude '/enterprise/' \
|| { echo "✗ Unformatted staged files. Fix with: make format, then re-stage." >&2; rc=1; }
|| { echo "✗ Unformatted files in scope. Fix with: make format, then re-stage." >&2; rc=1; }
fi
return $rc
}
@ -146,18 +202,18 @@ if [ -n "$litellm_py_files" ]; then
fi
if [ -n "$e2e_py_files" ] && [ -z "$litellm_py_files" ]; then
echo "pre-commit: type-checking tests/e2e (make lint-e2e-basedpyright)"
make lint-e2e-basedpyright || { echo "✗ tests/e2e basedpyright failed. Fix the errors above, then re-run make pre-commit." >&2; status=1; }
echo "check: type-checking tests/e2e (make lint-e2e-basedpyright)"
make lint-e2e-basedpyright || { echo "✗ tests/e2e basedpyright failed. Fix the errors above, then re-run make check." >&2; status=1; }
fi
if [ -n "$e2e_py_files" ]; then
echo "pre-commit: checking tests/e2e raw HTTP client ban (check_e2e_no_raw_requests)"
echo "check: checking tests/e2e raw HTTP client ban (check_e2e_no_raw_requests)"
uv run --no-sync python tests/code_coverage_tests/check_e2e_no_raw_requests.py \
|| { echo "✗ Raw HTTP client import in tests/e2e. Route the call through tests/e2e/e2e_http.py, then re-run make pre-commit." >&2; status=1; }
|| { echo "✗ Raw HTTP client import in tests/e2e. Route the call through tests/e2e/e2e_http.py, then re-run make check." >&2; status=1; }
fi
dashboard_checks() {
echo "pre-commit: linting dashboard (prettier + eslint + lint budgets)"
echo "check: linting dashboard (prettier + eslint + lint budgets)"
if [ ! -d ui/litellm-dashboard/node_modules ]; then
echo "✗ ui/litellm-dashboard/node_modules is missing; dashboard lint cannot run." >&2
bootstrap_hint
@ -166,7 +222,7 @@ dashboard_checks() {
lint_dashboard || { echo "✗ Dashboard lint failed. See above; format with: (cd ui/litellm-dashboard && npm run format)." >&2; return 1; }
}
if [ -n "$ui_prettier_files" ] || [ -n "$ui_eslint_files" ]; then
if [ -n "$ui_prettier_changed" ] || [ -n "$ui_eslint_changed" ]; then
dash_log=$(mktemp)
set -m
dashboard_checks > "$dash_log" 2>&1 &
@ -176,7 +232,7 @@ fi
genapi_checks() {
local status=0
echo "pre-commit: checking dashboard API types are in sync (npm run gen:api)"
echo "check: checking dashboard API types are in sync (npm run gen:api)"
# gen-api-types.mjs imports litellm.proxy.proxy_server, which needs the proxy deps
# and an up-to-date Prisma client; check-ui-api-types.yml installs those and runs
# prisma generate before gen:api, so mirror that here or a stale client can mask
@ -194,7 +250,7 @@ genapi_checks() {
status=1
elif ( cd ui/litellm-dashboard && LITELLM_PYTHON="uv run --no-sync python" npm run gen:api ); then
if ! git diff --quiet -- ui/litellm-dashboard/src/lib/http/schema.d.ts; then
echo "✗ Dashboard API types are stale; regenerated src/lib/http/schema.d.ts. Stage it and commit; re-run make pre-commit only if other checks failed too." >&2
echo "✗ Dashboard API types are stale; regenerated src/lib/http/schema.d.ts. Stage it and commit; re-run make check only if other checks failed too." >&2
status=1
fi
else

View file

@ -40,7 +40,7 @@ class MockA2AClient:
name="mock-agent", url="http://mock-agent.local"
)
async def send_message(self, request):
async def send_message(self, request, *, context=None):
from a2a.compat.v0_3.conversions import pb2_v10
for text in ("hel", "hello"):

View file

@ -138,7 +138,7 @@ Mark live tests with `@pytest.mark.e2e` (on the class or the module). Pure cover
Before you push
1. Run `make lint-e2e-basedpyright` (or `make pre-commit` with your changes staged); the harness is fully typed and the gate allows zero basedpyright errors, enforced in CI on any PR touching `tests/e2e/**/*.py`
1. Run `make lint-e2e-basedpyright` (or `make check` with your changes staged); the harness is fully typed and the gate allows zero basedpyright errors, enforced in CI on any PR touching `tests/e2e/**/*.py`
2. Add the models your test needs to the config your local proxy loads

View file

@ -622,8 +622,12 @@ async def test_service_logger_keys_success():
logger success hook is called with the correct event metadata and no exception is logged.
"""
keys = [
{"id": "key1", "spend": 10.0, "budget_duration": 60, "token": "key1"},
{"id": "key2", "spend": 15.0, "budget_duration": 60, "token": "key2"},
_attrify(
{"id": "key1", "spend": 10.0, "budget_duration": 60, "token": "key1"}
),
_attrify(
{"id": "key2", "spend": 15.0, "budget_duration": 60, "token": "key2"}
),
]
prisma_client = MagicMock()
prisma_client.get_data = AsyncMock(return_value=keys)
@ -740,8 +744,12 @@ async def test_service_logger_users_success():
the correct metadata and no exception is logged.
"""
users = [
{"id": "user1", "spend": 20.0, "budget_duration": 120, "user_id": "user1"},
{"id": "user2", "spend": 25.0, "budget_duration": 120, "user_id": "user2"},
_attrify(
{"id": "user1", "spend": 20.0, "budget_duration": 120, "user_id": "user1"}
),
_attrify(
{"id": "user2", "spend": 25.0, "budget_duration": 120, "user_id": "user2"}
),
]
prisma_client = MagicMock()
prisma_client.get_data = AsyncMock(return_value=users)
@ -853,8 +861,12 @@ async def test_service_logger_teams_success():
the proper metadata and nothing is logged as an exception.
"""
teams = [
{"id": "team1", "spend": 30.0, "budget_duration": 180, "team_id": "team1"},
{"id": "team2", "spend": 35.0, "budget_duration": 180, "team_id": "team2"},
_attrify(
{"id": "team1", "spend": 30.0, "budget_duration": 180, "team_id": "team1"}
),
_attrify(
{"id": "team2", "spend": 35.0, "budget_duration": 180, "team_id": "team2"}
),
]
prisma_client = MagicMock()
prisma_client.get_data = AsyncMock(return_value=teams)

View file

@ -338,7 +338,7 @@ class BaseResponsesAPITest(ABC):
)
assert result is not None
assert result.id == response.id
assert result.output == response.output
assert result.output_text == response.output_text
else:
raise ValueError("response is not a ResponsesAPIResponse")
else:
@ -352,7 +352,7 @@ class BaseResponsesAPITest(ABC):
)
assert result is not None
assert result.id == response.id
assert result.output == response.output
assert result.output_text == response.output_text
else:
raise ValueError("response is not a ResponsesAPIResponse")

View file

@ -12,8 +12,10 @@ from typing import Final
import pytest
from litellm.proxy.db.autorouter_session_rollup import UPSERT_AUTOROUTER_SESSION_SQL
from litellm.proxy.management_endpoints.auto_router_endpoints import _BENCHMARKS_SQL
from litellm.proxy.db.autorouter_session_rollup import (
AUTOROUTER_BENCHMARKS_SQL,
UPSERT_AUTOROUTER_SESSION_SQL,
)
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -164,7 +166,7 @@ async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db):
await _turn(db, key, "A", T0 - timedelta(days=40), session_id=out_of_window, router=router)
rows = await db.query_raw(
_BENCHMARKS_SQL,
AUTOROUTER_BENCHMARKS_SQL,
(T0 - timedelta(days=1)).isoformat(),
(T0 + timedelta(days=1)).isoformat(),
)
@ -186,7 +188,7 @@ async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db
await _turn(db, key, "A", T0 + timedelta(seconds=10), session_id=f"s-{uuid.uuid4()}", router=router, router_type="quality")
rows = await db.query_raw(
_BENCHMARKS_SQL,
AUTOROUTER_BENCHMARKS_SQL,
(T0 - timedelta(days=1)).isoformat(),
(T0 + timedelta(days=1)).isoformat(),
)
@ -248,7 +250,7 @@ async def test_the_benchmarks_aggregate_sums_tier_turns_across_sessions(db):
await _turn(db, key, "C", T0 + timedelta(seconds=30), session_id=f"s-{uuid.uuid4()}", router=router, tier=None)
rows = await db.query_raw(
_BENCHMARKS_SQL,
AUTOROUTER_BENCHMARKS_SQL,
(T0 - timedelta(days=1)).isoformat(),
(T0 + timedelta(days=1)).isoformat(),
)
@ -275,7 +277,7 @@ async def test_tier_maps_stay_separate_per_router_type_on_a_reconfigured_alias(d
)
rows = await db.query_raw(
_BENCHMARKS_SQL,
AUTOROUTER_BENCHMARKS_SQL,
(T0 - timedelta(days=1)).isoformat(),
(T0 + timedelta(days=1)).isoformat(),
)
@ -289,7 +291,7 @@ async def test_a_window_with_no_tiered_turns_aggregates_to_an_empty_map(db):
await _turn(db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, tier=None)
rows = await db.query_raw(
_BENCHMARKS_SQL,
AUTOROUTER_BENCHMARKS_SQL,
(T0 - timedelta(days=1)).isoformat(),
(T0 + timedelta(days=1)).isoformat(),
)

View file

@ -1195,8 +1195,13 @@ def test_async_post_call_failure_hook_skips_a_transport_that_already_answered():
def test_record_error_attributes_on_span_decorates_without_ending():
"""PATH A: a failure that dies before any LLM-call span (malformed body,
validation) is stamped onto the instrumentor-owned SERVER span. The method must
not end the span or emit a duplicate exception event, and must pin error.code
to the real response status (not the exception's own code)."""
not end the span, and must pin error.code to the real response status (not the
exception's own code).
LIT-4780: the instrumentor never sees the exception (the proxy handler turns it
into a JSONResponse), so nothing else marks the span as failed; the status and
the exception event have to come from here or the trace shows the error message
on an otherwise successful-looking request."""
logger, exporter = _logger()
server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME)
logger.record_error_attributes_on_span(server, _proxy_exc("Invalid JSON body", 400), 422)
@ -1206,7 +1211,31 @@ def test_record_error_attributes_on_span_decorates_without_ending():
assert span.attributes["error.type"] == "ProxyException"
assert span.attributes["error.message"] == "Invalid JSON body"
assert span.attributes["litellm.provider.error.code"] == "422"
assert all(e.name != "exception" for e in span.events)
assert span.status.status_code is StatusCode.ERROR
assert [e.name for e in span.events] == ["exception"]
def test_record_error_attributes_on_span_does_not_duplicate_an_already_stamped_error():
"""A failure that already went through ``async_post_call_failure_hook`` reaches
the exception handler too; the second stamp must keep one exception event while
still repinning error.code to the real response status."""
from litellm.proxy._types import UserAPIKeyAuth
logger, exporter = _logger()
server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME)
set_request_root_span(server)
exc = _proxy_exc("Authentication Error, invalid key", 401)
asyncio.run(
logger.async_post_call_failure_hook(
request_data={}, original_exception=exc, user_api_key_dict=UserAPIKeyAuth()
)
)
logger.record_error_attributes_on_span(server, exc, 400)
server.end()
(span,) = exporter.get_finished_spans()
assert [e.name for e in span.events] == ["exception"]
assert span.attributes["litellm.provider.error.code"] == "400"
assert span.status.status_code is StatusCode.ERROR
def test_record_error_attributes_on_span_ignores_below_400_and_missing_span():

View file

@ -221,14 +221,97 @@ async def test_execute_search_passes_selected_search_tool_litellm_params(monkeyp
kwargs={"litellm_params": {"metadata": {"user_api_key_auth": user_api_key_auth}}},
)
mock_asearch.assert_awaited_once_with(
query="what is litellm",
search_provider="tavily",
api_key="fake-ui-key",
api_base="https://api.tavily.com",
timeout=10.0,
max_retries=2,
forwarded_kwargs = mock_asearch.await_args.kwargs
assert forwarded_kwargs["query"] == "what is litellm"
assert forwarded_kwargs["search_provider"] == "tavily"
assert forwarded_kwargs["api_key"] == "fake-ui-key"
assert forwarded_kwargs["api_base"] == "https://api.tavily.com"
assert forwarded_kwargs["timeout"] == 10.0
assert forwarded_kwargs["max_retries"] == 2
@pytest.mark.asyncio
async def test_execute_search_attributes_spend_to_the_calling_key(monkeypatch):
"""An intercepted search is billed and logged against the key that made the LLM request.
Without the forwarded attribution metadata the proxy's spend hook skips the search
entirely, so its provider cost never reaches SpendLogs or any budget.
"""
import litellm
from litellm.proxy import proxy_server
from litellm.proxy.hooks.proxy_track_cost_callback import _should_track_cost_callback
logger = WebSearchInterceptionLogger(
enabled_providers=["bedrock"],
search_tool_name="perplexity-sonar-pro",
)
router = MagicMock()
router.search_tools = [
{
"search_tool_name": "perplexity-sonar-pro",
"litellm_params": {"search_provider": "perplexity", "api_key": "fake-key"},
}
]
mock_asearch = AsyncMock(return_value=SearchResponse(object="search", results=[]))
user_api_key_auth = UserAPIKeyAuth(
api_key="hashed-sk-1234",
key_alias="alice-key",
user_id="user-alice",
org_id="org-1",
)
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(litellm, "asearch", mock_asearch)
await logger._execute_search(
"what is litellm",
kwargs={"litellm_params": {"metadata": {"user_api_key_auth": user_api_key_auth}}},
)
forwarded_metadata = mock_asearch.await_args.kwargs["litellm_metadata"]
assert forwarded_metadata["user_api_key"] == "hashed-sk-1234"
assert forwarded_metadata["user_api_key_hash"] == "hashed-sk-1234"
assert forwarded_metadata["user_api_key_alias"] == "alice-key"
assert forwarded_metadata["user_api_key_user_id"] == "user-alice"
assert forwarded_metadata["user_api_key_org_id"] == "org-1"
assert forwarded_metadata["model_group"] == "perplexity-sonar-pro"
assert (
_should_track_cost_callback(
user_api_key=forwarded_metadata["user_api_key"],
user_id=forwarded_metadata["user_api_key_user_id"],
team_id=forwarded_metadata["user_api_key_team_id"],
end_user_id=None,
call_type="asearch",
)
is True
)
@pytest.mark.asyncio
async def test_execute_search_without_proxy_auth_context_stays_sdk_only(monkeypatch):
"""SDK callers have no key to attribute the search to, so no proxy metadata is invented."""
import litellm
from litellm.proxy import proxy_server
logger = WebSearchInterceptionLogger(
enabled_providers=["bedrock"],
search_tool_name="perplexity-sonar-pro",
)
router = MagicMock()
router.search_tools = [
{
"search_tool_name": "perplexity-sonar-pro",
"litellm_params": {"search_provider": "perplexity", "api_key": "fake-key"},
}
]
mock_asearch = AsyncMock(return_value=SearchResponse(object="search", results=[]))
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(litellm, "asearch", mock_asearch)
await logger._execute_search("what is litellm", kwargs={"litellm_params": {}})
assert "litellm_metadata" not in mock_asearch.await_args.kwargs
@pytest.mark.asyncio

View file

@ -4786,6 +4786,117 @@ async def test_user_api_key_auth_does_not_overwrite_end_user_id_set_by_builder()
setattr(_proxy_server_mod, k, v)
@pytest.mark.asyncio
async def test_user_api_key_auth_authenticates_before_raising_malformed_body_error():
"""Regression (LIT-4780): a body that fails to parse must still be authenticated
first, so the rejected request's trace carries the caller's key / team / user
identity instead of an anonymous root span. The parse error is re-raised
unchanged once identity is seeded."""
from fastapi import Request
from starlette.datastructures import URL
import litellm.proxy.proxy_server as _proxy_server_mod
builder_token = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="team-1")
request = Request(
scope={
"type": "http",
"headers": [(b"content-type", b"application/json")],
"method": "POST",
}
)
request._url = URL(url="/chat/completions")
request._body = b'{}{"model": "gpt-4o"}'
attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None)
originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}
try:
for k, v in attrs.items():
setattr(_proxy_server_mod, k, v)
with (
patch(
"litellm.proxy.auth.user_api_key_auth._user_api_key_auth_builder",
new_callable=AsyncMock,
return_value=builder_token,
) as mock_builder,
patch(
"litellm.proxy.auth.user_api_key_auth._run_centralized_common_checks",
new_callable=AsyncMock,
) as mock_common_checks,
patch(
"litellm.proxy.auth.user_api_key_auth.RouteChecks.should_call_route",
),
patch(
"litellm.proxy.auth.user_api_key_auth.seed_request_identity",
) as mock_seed,
):
with pytest.raises(ProxyException) as exc_info:
await user_api_key_auth(request=request, api_key="Bearer sk-test")
assert "Invalid JSON payload" in str(exc_info.value.message)
assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
mock_builder.assert_awaited_once()
assert mock_seed.call_args.args[0] is builder_token
# authorization must not run for a request that is about to be rejected:
# ``common_checks`` reserves budget against live spend counters that only the
# endpoint's post-call path releases, and the endpoint never runs here
mock_common_checks.assert_not_awaited()
finally:
for k, v in originals.items():
setattr(_proxy_server_mod, k, v)
@pytest.mark.asyncio
async def test_user_api_key_auth_malformed_body_with_rejected_key_still_returns_the_parse_error():
"""The body is read before the key is authenticated, so a caller who sends both a
malformed body and a key that fails auth gets the 400. Authenticating the request
first (LIT-4780) must not turn that into the auth status code."""
from fastapi import Request
from starlette.datastructures import URL
import litellm.proxy.proxy_server as _proxy_server_mod
request = Request(
scope={
"type": "http",
"headers": [(b"content-type", b"application/json")],
"method": "POST",
}
)
request._url = URL(url="/chat/completions")
request._body = b'{}{"model": "gpt-4o"}'
attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None)
originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}
try:
for k, v in attrs.items():
setattr(_proxy_server_mod, k, v)
with (
patch(
"litellm.proxy.auth.user_api_key_auth._user_api_key_auth_builder",
new_callable=AsyncMock,
side_effect=ProxyException(
message="Authentication Error, invalid key",
type="auth_error",
param="None",
code=status.HTTP_401_UNAUTHORIZED,
),
),
patch(
"litellm.proxy.auth.user_api_key_auth.RouteChecks.should_call_route",
),
):
with pytest.raises(ProxyException) as exc_info:
await user_api_key_auth(request=request, api_key="Bearer sk-bad")
assert "Invalid JSON payload" in str(exc_info.value.message)
assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
finally:
for k, v in originals.items():
setattr(_proxy_server_mod, k, v)
def _proxy_attrs_for_db_lookup():
"""Minimal proxy_server attributes for driving the real
``_user_api_key_auth_builder`` down to the DB key lookup."""

View file

@ -676,6 +676,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields():
"applied_policies": ["spoofed-policy"],
"policy_sources": {"spoofed-policy": "request"},
"routing_decision": {"cause": "forged", "routed_model": "spoofed"},
"_session_deployment_affinity_ttl": 999999,
"internal_call_origin": "autorouter_classifier",
"_guardrail_pipelines": [{"name": "spoofed"}],
"_pipeline_managed_guardrails": ["evaded"],
@ -719,6 +720,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields():
"applied_policies",
"policy_sources",
"routing_decision",
"_session_deployment_affinity_ttl",
"internal_call_origin",
"_guardrail_pipelines",
"_pipeline_managed_guardrails",

View file

@ -3414,6 +3414,103 @@ class TestSessionAffinity:
def _request_kwargs(session_id: str) -> Dict:
return {"metadata": {"session_id": session_id}}
@pytest.mark.asyncio
async def test_hook_response_carries_session_affinity_ttl_on_classify_and_pin_paths(
self, mock_router_instance, session_affinity_config
):
"""The hook response's session_affinity_ttl_seconds is what the Router stamps as
the deployment-affinity marker, so both the classify path (turn 1) and the
session-pin path (turn 2) must carry the configured TTL."""
mock_router_instance.cache = DualCache()
router = ComplexityRouter(
model_name="test-router",
litellm_router_instance=mock_router_instance,
complexity_router_config={**session_affinity_config, "session_affinity_ttl_seconds": 321},
)
request_kwargs = self._request_kwargs("marker-session")
first = await router.async_pre_routing_hook(
model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE
)
second = await router.async_pre_routing_hook(
model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE
)
assert first.session_affinity_ttl_seconds == 321
assert second.session_affinity_ttl_seconds == 321
@pytest.mark.parametrize(
"session_affinity,deployment_affinity,plugins,tier_pinned,deployment_pinned",
[
(False, False, False, False, False),
(False, True, False, False, True),
(True, False, False, True, True),
(True, True, False, True, True),
(False, True, True, False, False),
(True, True, True, False, False),
],
)
@pytest.mark.asyncio
async def test_tier_pin_and_deployment_pin_are_independently_gated(
self,
mock_router_instance,
basic_config,
session_affinity,
deployment_affinity,
plugins,
tier_pinned,
deployment_pinned,
):
"""deployment_affinity pins the deployment inside each routed group without pinning which
group the session routes to, so with session_affinity off the tier must still reclassify
on every turn while the marker the Router stamps is still emitted. Turn 1 classifies
REASONING and turn 2 SIMPLE, so a reclassified turn 2 moves model while a tier-pinned one
does not. plugins suppress both pins, since a stale pin would bypass the plugin pipeline."""
mock_router_instance.cache = DualCache()
router = ComplexityRouter(
model_name="test-router",
litellm_router_instance=mock_router_instance,
complexity_router_config={
**basic_config,
"session_affinity": session_affinity,
"deployment_affinity": deployment_affinity,
**({"plugins": [_DummyPlugin()]} if plugins else {}),
},
)
request_kwargs = self._request_kwargs("matrix-session")
first = await router.async_pre_routing_hook(
model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE
)
second = await router.async_pre_routing_hook(
model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE
)
assert first.model == "o1-preview"
assert second.model == ("o1-preview" if tier_pinned else "gpt-4o-mini")
assert (first.session_affinity_ttl_seconds is not None) is deployment_pinned
assert (second.session_affinity_ttl_seconds is not None) is deployment_pinned
@pytest.mark.asyncio
async def test_hook_response_has_no_session_affinity_ttl_when_disabled_or_plugins(
self, mock_router_instance, basic_config, session_affinity_config
):
mock_router_instance.cache = DualCache()
disabled_router = ComplexityRouter(
model_name="test-router",
litellm_router_instance=mock_router_instance,
complexity_router_config={**basic_config, "deployment_affinity": False},
)
plugin_router = ComplexityRouter(
model_name="test-router-plugins",
litellm_router_instance=mock_router_instance,
complexity_router_config={**session_affinity_config, "plugins": [_DummyPlugin()]},
)
disabled = await disabled_router.async_pre_routing_hook(
model="test-model", request_kwargs=self._request_kwargs("s-off"), messages=self.SIMPLE_MESSAGE
)
with_plugins = await plugin_router.async_pre_routing_hook(
model="test-model", request_kwargs=self._request_kwargs("s-plugins"), messages=self.SIMPLE_MESSAGE
)
assert disabled.session_affinity_ttl_seconds is None
assert with_plugins.session_affinity_ttl_seconds is None
@pytest.mark.asyncio
async def test_disabled_by_default_reclassifies_every_turn(self, mock_router_instance, basic_config):
"""Regression: session_affinity defaults to False, so a shared session_id must NOT

View file

@ -465,8 +465,7 @@ async def test_async_pre_call_hook_uses_model_map_key_scope():
Deployment affinity caching uses (user_api_key_hash, model_map_key) -> model_id.
"""
cache = AsyncMock()
cache.async_set_cache = AsyncMock()
cache = DualCache()
callback = DeploymentAffinityCheck(
cache=cache,
@ -489,11 +488,7 @@ async def test_async_pre_call_hook_uses_model_map_key_scope():
model_group="claude-sonnet-4-5@20250929",
user_key="user-key-abc",
)
cache.async_set_cache.assert_called_once_with(
expected_cache_key,
{"model_id": "model-id-123"},
ttl=123,
)
assert await cache.async_get_cache(key=expected_cache_key) == {"model_id": "model-id-123"}
@pytest.mark.asyncio

View file

@ -1,6 +1,6 @@
import os
import sys
from unittest.mock import AsyncMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -10,6 +10,7 @@ import json
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
DeploymentAffinityCheck,
)
@ -163,7 +164,7 @@ async def test_async_session_id_affinity_priority_over_user_key():
await callback.cache.async_set_cache(
DeploymentAffinityCheck.get_session_affinity_cache_key(
"model_group", "session1"
"model_group", "session1", user_key="user1"
),
{"model_id": "deployment-2"},
)
@ -180,3 +181,439 @@ async def test_async_session_id_affinity_priority_over_user_key():
assert len(filtered) == 1
assert filtered[0]["model_info"]["id"] == "deployment-2"
MOCK_RESPONSES_API_RESPONSE = {
"id": "resp_mock-resp-456",
"object": "response",
"created_at": 1741476542,
"status": "completed",
"model": "azure/computer-use-preview",
"output": [],
"usage": {
"input_tokens": 5,
"output_tokens": 10,
"total_tokens": 15,
"output_tokens_details": {"reasoning_tokens": 0},
},
}
def _smart_router(session_affinity=True, ttl_seconds=777, deployment_affinity=True):
return litellm.Router(
model_list=[
{
"model_name": "smart-router",
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_default_model": "target-group",
"complexity_router_config": {
"session_affinity": session_affinity,
"deployment_affinity": deployment_affinity,
"session_affinity_ttl_seconds": ttl_seconds,
"tiers": {
"SIMPLE": "target-group",
"MEDIUM": "target-group",
"COMPLEX": "target-group",
"REASONING": "target-group",
},
},
},
},
{
"model_name": "target-group",
"litellm_params": {
"model": "azure/computer-use-preview-1",
"api_key": "mock-api-key-1",
"api_version": "mock-api-version",
"api_base": "https://mock-endpoint-1.openai.azure.com",
},
"model_info": {"id": "deployment-1", "base_model": "computer-use-preview"},
},
{
"model_name": "target-group",
"litellm_params": {
"model": "azure/computer-use-preview-2",
"api_key": "mock-api-key-2",
"api_version": "mock-api-version-2",
"api_base": "https://mock-endpoint-2.openai.azure.com",
},
"model_info": {"id": "deployment-2", "base_model": "computer-use-preview"},
},
],
)
def _session_pin_key(session_id, user_key):
return DeploymentAffinityCheck.get_session_affinity_cache_key(
model_group="target-group", session_id=session_id, user_key=user_key
)
def _cleanup_router_callbacks(router):
for callback in router.optional_callbacks or []:
litellm.logging_callback_manager.remove_callback_from_all_lists(callback)
async def _one_turn(router, model, session_id, key_hash):
"""One request with the shuffle forced to deployment-1, so any other landing
deployment can only come from a pin read."""
with (
patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new_callable=AsyncMock,
) as mock_post,
patch(
"litellm.router_strategy.simple_shuffle.random.choice",
side_effect=lambda seq: seq[0],
),
):
mock_post.return_value = MockResponse(MOCK_RESPONSES_API_RESPONSE, 200)
response = await router.aresponses(
model=model,
input=f"turn for {session_id} {key_hash}",
litellm_metadata={"session_id": session_id, "user_api_key_hash": key_hash},
)
return response._hidden_params["model_id"]
@pytest.mark.asyncio
async def test_auto_router_session_affinity_writes_scoped_pin_and_follows_it():
"""Turn 1 persists a key-scoped deployment pin; a pin seeded to the deployment
the shuffle would never pick is then followed, proving the read path."""
router = _smart_router()
try:
served = await _one_turn(router, "smart-router", "write-session", "key-1")
assert await router.cache.async_get_cache(key=_session_pin_key("write-session", "key-1")) == {
"model_id": served
}
assert await router.cache.async_get_cache(key=_session_pin_key("write-session", None)) is None
await router.cache.async_set_cache(
key=_session_pin_key("read-session", "key-1"), value={"model_id": "deployment-2"}
)
assert await _one_turn(router, "smart-router", "read-session", "key-1") == "deployment-2"
finally:
_cleanup_router_callbacks(router)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model,key_hash",
[
("target-group", "key-1"),
("smart-router", "key-2"),
],
ids=["direct-group-call", "different-api-key"],
)
async def test_seeded_session_pin_is_invisible_outside_its_scope(model, key_hash):
"""The pin binds (auto-routed request, api key, session): a direct call to the
group and a different key reusing the session id must both ignore it."""
router = _smart_router()
try:
await router.cache.async_set_cache(
key=_session_pin_key("scoped-session", "key-1"), value={"model_id": "deployment-2"}
)
assert await _one_turn(router, model, "scoped-session", key_hash) == "deployment-1"
finally:
_cleanup_router_callbacks(router)
@pytest.mark.asyncio
async def test_marker_write_uses_marker_ttl_and_writes_only_the_session_pin():
"""The write hook honors the marker's TTL over the callback default and writes
no user-key entry when only session affinity is active."""
import time as time_module
cache = DualCache()
callback = DeploymentAffinityCheck(
cache=cache,
ttl_seconds=3600,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
)
await callback.async_pre_call_deployment_hook(
kwargs={
"model_info": {"id": "deployment-1"},
"metadata": {
"deployment_model_name": "target-group",
"session_id": "ttl-session",
"user_api_key_hash": "key-1",
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: 777,
},
},
call_type=None,
)
session_key = _session_pin_key("ttl-session", "key-1")
assert cache.in_memory_cache.cache_dict == {session_key: {"model_id": "deployment-1"}}
assert cache.in_memory_cache.ttl_dict[session_key] == pytest.approx(time_module.time() + 777, abs=5)
@pytest.mark.asyncio
@pytest.mark.parametrize("bad_marker", ["777", True, -5, 0, None])
async def test_malformed_marker_values_do_not_enable_session_affinity(bad_marker):
cache = DualCache()
callback = DeploymentAffinityCheck(
cache=cache,
ttl_seconds=3600,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
)
await callback.async_pre_call_deployment_hook(
kwargs={
"model_info": {"id": "deployment-1"},
"metadata": {
"deployment_model_name": "target-group",
"session_id": "bad-marker-session",
"user_api_key_hash": "key-1",
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: bad_marker,
},
},
call_type=None,
)
assert cache.in_memory_cache.cache_dict == {}
@pytest.mark.asyncio
@pytest.mark.parametrize("enable_user_key", [False, True], ids=["session-pin", "user-key-pin"])
async def test_concurrent_first_requests_never_flip_a_claimed_pin(enable_user_key):
"""Two overlapping first requests select different deployments before either
write lands. Pins are first-writer-wins claims, so the second write must leave
the stored pin unchanged instead of flipping it."""
cache = DualCache()
callback = DeploymentAffinityCheck(
cache=cache,
ttl_seconds=3600,
enable_user_key_affinity=enable_user_key,
enable_responses_api_affinity=False,
)
def racing_kwargs(deployment_id):
metadata = {"deployment_model_name": "target-group", "user_api_key_hash": "key-1"}
if not enable_user_key:
metadata["session_id"] = "racing-session"
metadata[SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY] = 777
return {"model_info": {"id": deployment_id}, "metadata": metadata}
await callback.async_pre_call_deployment_hook(kwargs=racing_kwargs("deployment-1"), call_type=None)
await callback.async_pre_call_deployment_hook(kwargs=racing_kwargs("deployment-2"), call_type=None)
pinned_key = (
DeploymentAffinityCheck.get_affinity_cache_key(model_group="target-group", user_key="key-1")
if enable_user_key
else _session_pin_key("racing-session", "key-1")
)
assert await cache.async_get_cache(key=pinned_key) == {"model_id": "deployment-1"}
@pytest.mark.asyncio
@pytest.mark.parametrize(
"stored_pin",
[{"model_id": "deployment-1"}, "deployment-1"],
ids=["dict-pin", "legacy-string-pin"],
)
async def test_in_memory_reclaim_slides_idle_window_only_for_the_stored_deployment(stored_pin):
"""The pod-local claim mirrors the Lua keepalive: the winning deployment's
re-claim extends the pin's expiry, a losing deployment's claim touches neither
the value nor the expiry, so no-Redis setups keep stickiness across an active
session and ttl bounds idle time there too. Sameness is judged on the pinned
model id, so a legacy string pin written by the Redis branch slides the same."""
import time as time_module
cache = DualCache()
callback = DeploymentAffinityCheck(
cache=cache,
ttl_seconds=3600,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
)
pin_key = _session_pin_key("slide-session", "key-1")
cache.in_memory_cache.set_cache(pin_key, stored_pin, ttl=10)
first_expiry = cache.in_memory_cache.ttl_dict[pin_key]
reclaimed = await callback._claim_pin(cache_key=pin_key, pin_value={"model_id": "deployment-1"}, ttl_seconds=777)
assert reclaimed == "deployment-1"
assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5)
assert cache.in_memory_cache.ttl_dict[pin_key] > first_expiry
lost = await callback._claim_pin(cache_key=pin_key, pin_value={"model_id": "deployment-2"}, ttl_seconds=10)
assert lost == "deployment-1"
assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5)
@pytest.mark.asyncio
async def test_claim_pin_uses_redis_attached_after_construction():
"""The proxy attaches Redis via Router._update_redis_cache after the Router (and
this callback) are built. The claim must resolve the redis tier per call, or pins
silently stay pod-local and cross-pod first-writer-wins is lost."""
cache = DualCache()
callback = DeploymentAffinityCheck(
cache=cache,
ttl_seconds=3600,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
)
captured = {}
async def fake_runner(keys, args, client=None):
captured["keys"] = keys
captured["args"] = args
return b'{"model_id": "other-pod-winner"}'
late_redis = MagicMock()
late_redis.async_register_script = MagicMock(return_value=fake_runner)
cache.redis_cache = late_redis
import time as time_module
pin_key = _session_pin_key("late-redis-session", "key-1")
cache.in_memory_cache.set_cache(pin_key, {"model_id": "other-pod-winner"}, ttl=10)
claimed = await callback._claim_pin(
cache_key=pin_key,
pin_value={"model_id": "our-deployment"},
ttl_seconds=777,
)
assert claimed == "other-pod-winner"
assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5)
assert captured["keys"] == (pin_key,)
assert captured["args"] == ('{"model_id": "our-deployment"}', 777)
assert cache.in_memory_cache.get_cache(_session_pin_key("late-redis-session", "key-1")) == {
"model_id": "other-pod-winner"
}
@pytest.mark.asyncio
async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down():
"""A Redis outage must cost cross-pod agreement, never same-pod stickiness. The write
hook only logs this result, so an escaping error would leave the session unpinned and
reshuffle every turn for the whole outage. DualCache's write path, which this claim
replaced, wrote the in-memory tier before ever touching Redis."""
cache = DualCache()
callback = DeploymentAffinityCheck(
cache=cache,
ttl_seconds=3600,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
)
async def exploding_runner(keys, args, client=None):
raise ConnectionError("redis is down")
down_redis = MagicMock()
down_redis.async_register_script = MagicMock(return_value=exploding_runner)
cache.redis_cache = down_redis
key = _session_pin_key("outage-session", "key-1")
claimed = await callback._claim_pin(cache_key=key, pin_value={"model_id": "our-deployment"}, ttl_seconds=777)
assert claimed == "our-deployment"
assert cache.in_memory_cache.get_cache(key) == {"model_id": "our-deployment"}
second = await callback._claim_pin(cache_key=key, pin_value={"model_id": "another-deployment"}, ttl_seconds=777)
assert second == "our-deployment"
@pytest.mark.asyncio
async def test_marker_session_affinity_read_and_write_agree_for_wildcard_groups():
"""Wildcard deployments keep the literal pattern as model_name on both the read
path and the write path, so the marker-gated pin round-trips through one key."""
callback = DeploymentAffinityCheck(
cache=DualCache(),
ttl_seconds=3600,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
)
request_kwargs = {
"model_info": {"id": "wild-deployment-2"},
"metadata": {
"deployment_model_name": "openai/*",
"session_id": "wild-session",
"user_api_key_hash": "key-1",
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: 777,
},
}
await callback.async_pre_call_deployment_hook(kwargs=request_kwargs, call_type=None)
filtered = await callback.async_filter_deployments(
model="openai/gpt-4o",
healthy_deployments=[
{
"model_name": "openai/*",
"litellm_params": {"model": "openai/gpt-4o"},
"model_info": {"id": f"wild-deployment-{i}"},
}
for i in (1, 2)
],
messages=[],
request_kwargs=request_kwargs,
)
assert [d["model_info"]["id"] for d in filtered] == ["wild-deployment-2"]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model,session_affinity,deployment_affinity,expect_marker",
[
("smart-router", False, True, True),
("smart-router", True, False, True),
("smart-router", False, False, False),
("target-group", False, True, False),
],
ids=[
"deployment-affinity-stamps",
"session-affinity-implies-deployment-pin",
"both-off-no-stamp",
"non-auto-routed-clears",
],
)
async def test_pre_routing_hook_stamps_or_clears_the_marker_per_attempt(
model, session_affinity, deployment_affinity, expect_marker
):
"""Every routing attempt writes or clears the marker, so a fallback from an
auto-routed group to a plain group cannot carry a stale marker. session_affinity
implies the deployment pin: a session frozen onto one group must not re-shuffle
across that group's deployments."""
router = _smart_router(session_affinity=session_affinity, deployment_affinity=deployment_affinity)
try:
request_kwargs = {
"metadata": {SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: 111},
"litellm_metadata": {SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: 111},
}
await router.async_pre_routing_hook(
model=model,
request_kwargs=request_kwargs,
messages=[{"role": "user", "content": "Hello"}],
)
if expect_marker:
assert request_kwargs["litellm_metadata"][SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY] == 777
else:
assert SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY not in request_kwargs["metadata"]
assert SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY not in request_kwargs["litellm_metadata"]
finally:
_cleanup_router_callbacks(router)
def test_complexity_router_with_deployment_affinity_registers_affinity_callback():
enabled = _smart_router()
session_only = _smart_router(session_affinity=True, deployment_affinity=False)
disabled = _smart_router(session_affinity=False, deployment_affinity=False)
try:
assert [
(cb.enable_user_key_affinity, cb.enable_responses_api_affinity, cb.enable_session_id_affinity)
for cb in enabled.optional_callbacks or []
if isinstance(cb, DeploymentAffinityCheck)
] == [(False, False, False)]
assert any(isinstance(cb, DeploymentAffinityCheck) for cb in session_only.optional_callbacks or [])
assert not any(isinstance(cb, DeploymentAffinityCheck) for cb in disabled.optional_callbacks or [])
finally:
_cleanup_router_callbacks(enabled)
_cleanup_router_callbacks(session_only)
_cleanup_router_callbacks(disabled)

View file

@ -129,6 +129,126 @@ def _run(repo: Path, bin_dir: Path, extra_env: dict[str, str]) -> subprocess.Com
)
def _commit_all(repo: Path, message: str) -> None:
subprocess.run(["git", "add", "."], cwd=repo, check=True)
subprocess.run(
["git", "-c", "user.email=t@t", "-c", "user.name=t", "commit", "-qm", message],
cwd=repo,
check=True,
)
def _set_base_ref(repo: Path) -> None:
subprocess.run(
["git", "update-ref", "refs/remotes/origin/litellm_internal_staging", "HEAD"],
cwd=repo,
check=True,
)
def test_nothing_staged_scopes_to_working_tree_diff_and_runs_checks(tmp_path: Path) -> None:
repo, bin_dir = _sandbox(tmp_path)
_commit_all(repo, "base")
_set_base_ref(repo)
(repo / "litellm" / "foo.py").write_text("x = 2\n")
proc = _run(repo, bin_dir, {})
assert proc.returncode == 0, proc.stdout + proc.stderr
assert "nothing staged; scoping to the working tree's diff" in proc.stdout
assert "litellm/foo.py" in proc.stdout
assert "linting Python" in proc.stdout
def test_nothing_staged_checks_committed_branch_changes(tmp_path: Path) -> None:
repo, bin_dir = _sandbox(tmp_path)
_commit_all(repo, "base")
_set_base_ref(repo)
(repo / "litellm" / "foo.py").write_text("x = 2\n")
_commit_all(repo, "branch change")
proc = _run(repo, bin_dir, {})
assert proc.returncode == 0, proc.stdout + proc.stderr
assert "nothing staged; scoping to the working tree's diff" in proc.stdout
assert "linting Python" in proc.stdout
def test_nothing_staged_includes_untracked_files_in_scope(tmp_path: Path) -> None:
repo, bin_dir = _sandbox(tmp_path)
_commit_all(repo, "base")
_set_base_ref(repo)
(repo / "litellm" / "brand_new.py").write_text("z = 3\n")
proc = _run(repo, bin_dir, {})
assert proc.returncode == 0, proc.stdout + proc.stderr
assert "litellm/brand_new.py" in proc.stdout
assert "linting Python" in proc.stdout
def test_nothing_staged_deletion_only_branch_triggers_checks(tmp_path: Path) -> None:
repo, bin_dir = _sandbox(tmp_path)
_commit_all(repo, "base")
_set_base_ref(repo)
(repo / "litellm" / "foo.py").unlink()
_commit_all(repo, "delete module")
proc = _run(repo, bin_dir, {})
assert proc.returncode == 0, proc.stdout + proc.stderr
assert "nothing to check" not in proc.stdout
assert "litellm/foo.py" in proc.stdout
assert "linting Python" in proc.stdout
assert "ruff format --check" not in proc.stdout
def test_staged_deletion_triggers_checks_without_feeding_missing_files_to_tools(tmp_path: Path) -> None:
repo, bin_dir = _sandbox(tmp_path)
_commit_all(repo, "base")
subprocess.run(["git", "rm", "-q", "litellm/foo.py"], cwd=repo, check=True)
proc = _run(repo, bin_dir, {})
assert proc.returncode == 0, proc.stdout + proc.stderr
assert "nothing staged" not in proc.stdout
assert "linting Python" in proc.stdout
assert "ruff format --check" not in proc.stdout
def test_deleted_dashboard_file_still_triggers_dashboard_lint(tmp_path: Path) -> None:
repo, bin_dir = _sandbox(tmp_path)
_commit_all(repo, "base")
_set_base_ref(repo)
(repo / "ui" / "litellm-dashboard" / "src" / "app.ts").unlink()
_commit_all(repo, "delete dashboard file")
proc = _run(repo, bin_dir, {})
assert proc.returncode == 0, proc.stdout + proc.stderr
assert "linting dashboard" in proc.stdout
def test_nothing_staged_and_no_changes_is_an_explicit_no_op(tmp_path: Path) -> None:
repo, bin_dir = _sandbox(tmp_path)
_commit_all(repo, "base")
_set_base_ref(repo)
proc = _run(repo, bin_dir, {})
assert proc.returncode == 0, proc.stdout + proc.stderr
assert "nothing to check" in proc.stdout
assert "linting Python" not in proc.stdout
def test_nothing_staged_without_a_base_ref_fails_with_a_fetch_hint(tmp_path: Path) -> None:
repo, bin_dir = _sandbox(tmp_path)
_commit_all(repo, "base")
proc = _run(repo, bin_dir, {})
assert proc.returncode == 1
assert "cannot resolve the merge base" in proc.stdout
assert "git fetch origin litellm_internal_staging" in proc.stdout
def test_partial_staging_warns_which_checks_were_skipped(tmp_path: Path) -> None:
repo, bin_dir = _sandbox(tmp_path)
_commit_all(repo, "base")
(repo / "notes.md").write_text("hi\n")
subprocess.run(["git", "add", "notes.md"], cwd=repo, check=True)
(repo / "litellm" / "foo.py").write_text("x = 4\n")
proc = _run(repo, bin_dir, {})
assert proc.returncode == 0, proc.stdout + proc.stderr
assert "SKIPPED Python lint (make lint)" in proc.stdout
assert "litellm/foo.py" in proc.stdout
assert "linting Python" not in proc.stdout
def test_python_dashboard_and_gen_api_blocks_run_concurrently_with_grouped_output(tmp_path: Path) -> None:
repo, bin_dir = _sandbox(tmp_path)
barrier_dir = tmp_path / "barrier"
@ -159,8 +279,8 @@ def test_full_output_is_saved_to_a_log_file_in_the_git_dir(tmp_path: Path) -> No
assert "linting dashboard" in log
assert "API types" in log
assert "unstaged/untracked changes" in log
assert f"pre-commit: full log: {log_file}" in proc.stdout
assert "pre-commit: full log:" not in log
assert f"check: full log: {log_file}" in proc.stdout
assert "check: full log:" not in log
def test_unwritable_log_warns_and_falls_back_to_running_without_one(tmp_path: Path) -> None:
@ -170,7 +290,7 @@ def test_unwritable_log_warns_and_falls_back_to_running_without_one(tmp_path: Pa
assert proc.returncode == 0, proc.stdout + proc.stderr
assert "linting Python" in proc.stdout
assert "output will not be saved" in proc.stderr
assert "pre-commit: full log:" not in proc.stdout
assert "check: full log:" not in proc.stdout
failing = _run(repo, bin_dir, {"STUB_FAIL": "make-lint"})
assert failing.returncode == 1

View file

@ -7550,3 +7550,68 @@ async def test_fallback_failure_detail_from_upstream_is_bounded():
assert capture.messages, "the fallback failure path did not log at ERROR"
assert huge_message not in "".join(capture.messages)
assert max(len(message) for message in capture.messages) < 5_000
def test_stamp_or_clear_metadata_key_writes_and_clears_both_buckets():
request_kwargs = {"metadata": {}}
litellm.Router._stamp_or_clear_metadata_key(request_kwargs=request_kwargs, key="probe", value=7)
assert request_kwargs["metadata"]["probe"] == 7
stale_kwargs = {"metadata": {"probe": 7}, "litellm_metadata": {"probe": 7}}
litellm.Router._stamp_or_clear_metadata_key(request_kwargs=stale_kwargs, key="probe", value=None)
assert "probe" not in stale_kwargs["metadata"]
assert "probe" not in stale_kwargs["litellm_metadata"]
@pytest.mark.parametrize(
"complexity_router_config,expect_callback",
[
({"tiers": {"SIMPLE": "gpt-4o"}}, True),
({"tiers": {"SIMPLE": "gpt-4o"}, "deployment_affinity": False}, False),
({"tiers": {"SIMPLE": "gpt-4o"}, "deployment_affinity": False, "session_affinity": True}, True),
],
)
def test_complexity_router_registers_affinity_callback_for_deployment_pin(complexity_router_config, expect_callback):
"""The marker the complexity router stamps is inert unless a DeploymentAffinityCheck is
registered to read it, so deployment_affinity has to pull the callback in, and its default-on
means a bare config registers one. Opting out must skip the callback entirely rather than
register a filter that can never fire, including when session_affinity is on, since the two
pins are independent."""
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
DeploymentAffinityCheck,
)
router = litellm.Router(
model_list=[
{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}},
{
"model_name": "my-complexity-router",
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": complexity_router_config,
},
},
]
)
try:
registered = any(isinstance(cb, DeploymentAffinityCheck) for cb in router.optional_callbacks or [])
assert registered is expect_callback
finally:
for cb in router.optional_callbacks or []:
litellm.logging_callback_manager.remove_callback_from_all_lists(cb)
def test_ensure_deployment_affinity_callback_is_idempotent():
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
DeploymentAffinityCheck,
)
router = litellm.Router(model_list=[])
try:
router._ensure_deployment_affinity_callback()
router._ensure_deployment_affinity_callback()
affinity_callbacks = [
cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck)
]
assert len(affinity_callbacks) == 1
finally:
for cb in router.optional_callbacks or []:
litellm.logging_callback_manager.remove_callback_from_all_lists(cb)

View file

@ -27,7 +27,7 @@
"limit": 0
},
"LIT010": {
"limit": 16760
"limit": 16758
},
"LIT011": {
"limit": 5598

View file

@ -1,7 +1,9 @@
import { act, fireEvent, render, screen } from "@testing-library/react";
import type { ReactNode } from "react";
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import { act, fireEvent, render, screen, waitFor } from "@testing-library/react";
import type { ReactElement, ReactNode } from "react";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import type { RouterSettingsFormValue } from "../router_settings/RouterSettingsForm";
import { fetchAvailableModels, fetchAvailableModelsForTeam } from "@/components/llm_calls/fetch_models";
import RouterSettingsAccordion, { RouterSettingsAccordionValue } from "./RouterSettingsAccordion";
vi.mock("../networking", () => ({
@ -9,11 +11,14 @@ vi.mock("../networking", () => ({
}));
vi.mock("@/components/llm_calls/fetch_models", () => ({
fetchAvailableModels: vi.fn().mockResolvedValue([]),
fetchAvailableModels: vi.fn().mockResolvedValue([{ model_group: "global-model" }]),
fetchAvailableModelsForTeam: vi.fn().mockResolvedValue([{ model_group: "openai/*" }, { model_group: "gpt-5" }]),
}));
vi.mock("../Settings/RouterSettings/Fallbacks/FallbackSelectionForm", () => ({
FallbackSelectionForm: () => null,
FallbackSelectionForm: ({ availableModels }: { availableModels: string[] }) => (
<div data-testid="available-models">{availableModels.join(",")}</div>
),
}));
vi.mock("@tremor/react", () => ({
@ -39,9 +44,19 @@ vi.mock("../router_settings/RouterSettingsForm", () => ({
),
}));
const renderWithQueryClient = (ui: ReactElement) => {
const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } });
return render(ui, {
wrapper: ({ children }: { children: ReactNode }) => (
<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
),
});
};
describe("RouterSettingsAccordion", () => {
beforeEach(() => {
vi.useFakeTimers();
vi.clearAllMocks();
vi.useFakeTimers({ shouldAdvanceTime: true });
});
afterEach(() => {
@ -58,7 +73,7 @@ describe("RouterSettingsAccordion", () => {
it("debounces propagation and calls onChange once with the last value", async () => {
const onChange = vi.fn<(value: RouterSettingsAccordionValue) => void>();
render(<RouterSettingsAccordion accessToken="test-token" onChange={onChange} />);
renderWithQueryClient(<RouterSettingsAccordion accessToken="test-token" onChange={onChange} />);
await flushInitialPropagation(onChange);
fireEvent.click(screen.getByText("set-least-busy"));
@ -81,9 +96,51 @@ describe("RouterSettingsAccordion", () => {
expect(onChange.mock.calls[0][0].router_settings.routing_strategy).toBe("usage-based-routing");
});
it("offers the team's own models, including team-scoped BYOK ones, when a teamId is given", async () => {
renderWithQueryClient(<RouterSettingsAccordion accessToken="test-token" teamId="team-123" />);
await waitFor(() => {
expect(screen.getByTestId("available-models")).toHaveTextContent("gpt-5,openai/*");
});
expect(fetchAvailableModelsForTeam).toHaveBeenCalledWith("test-token", "team-123");
expect(fetchAvailableModels).not.toHaveBeenCalled();
});
it("falls back to the proxy-wide model listing when no teamId is given", async () => {
renderWithQueryClient(<RouterSettingsAccordion accessToken="test-token" />);
await waitFor(() => {
expect(screen.getByTestId("available-models")).toHaveTextContent("global-model");
});
expect(fetchAvailableModelsForTeam).not.toHaveBeenCalled();
});
it("ignores a stale team's model response that resolves after a newer team was selected", async () => {
const resolvers: ((models: { model_group: string }[]) => void)[] = [];
vi.mocked(fetchAvailableModelsForTeam).mockImplementation(
() => new Promise((resolve) => resolvers.push(resolve)) as Promise<{ model_group: string }[]>,
);
const { rerender } = renderWithQueryClient(<RouterSettingsAccordion accessToken="test-token" teamId="team-slow" />);
await waitFor(() => expect(resolvers).toHaveLength(1));
rerender(<RouterSettingsAccordion accessToken="test-token" teamId="team-fast" />);
await waitFor(() => expect(resolvers).toHaveLength(2));
await act(async () => {
resolvers[1]([{ model_group: "fast-team-model" }]);
resolvers[0]([{ model_group: "slow-team-model" }]);
});
await waitFor(() => {
expect(screen.getByTestId("available-models")).toHaveTextContent("fast-team-model");
});
expect(screen.getByTestId("available-models")).not.toHaveTextContent("slow-team-model");
});
it("does not call onChange when unmounted mid-wait", async () => {
const onChange = vi.fn<(value: RouterSettingsAccordionValue) => void>();
const { unmount } = render(<RouterSettingsAccordion accessToken="test-token" onChange={onChange} />);
const { unmount } = renderWithQueryClient(<RouterSettingsAccordion accessToken="test-token" onChange={onChange} />);
await flushInitialPropagation(onChange);
fireEvent.click(screen.getByText("set-least-busy"));

View file

@ -1,12 +1,13 @@
import React, { useEffect, useState, useImperativeHandle, forwardRef, useRef } from "react";
import { TabPanel, TabPanels, TabGroup, TabList, Tab } from "@tremor/react";
import { useQuery } from "@tanstack/react-query";
import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer";
import { getRouterSettingsCall } from "../networking";
import RouterSettingsForm, { RouterSettingsFormValue } from "../router_settings/RouterSettingsForm";
import { Fallbacks } from "../Settings/RouterSettings/Fallbacks/AddFallbacks";
import { FallbackSelectionForm } from "../Settings/RouterSettings/Fallbacks/FallbackSelectionForm";
import { FallbackGroup } from "../Settings/RouterSettings/Fallbacks/FallbackGroupConfig";
import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models";
import { fetchAvailableModels, fetchAvailableModelsForTeam, ModelGroup } from "@/components/llm_calls/fetch_models";
export interface RouterSettingsAccordionValue {
router_settings: {
@ -30,6 +31,7 @@ interface RouterSettingsAccordionProps {
value?: RouterSettingsAccordionValue;
onChange?: (value: RouterSettingsAccordionValue) => void;
modelData?: any;
teamId?: string | null;
}
export interface RouterSettingsAccordionRef {
@ -39,7 +41,7 @@ export interface RouterSettingsAccordionRef {
const PROPAGATE_WAIT_MS = 100;
const RouterSettingsAccordion = forwardRef<RouterSettingsAccordionRef, RouterSettingsAccordionProps>(
({ accessToken, value, onChange, modelData }, ref) => {
({ accessToken, value, onChange, modelData, teamId }, ref) => {
const [formValue, setFormValue] = useState<RouterSettingsFormValue>({
routerSettings: {},
selectedStrategy: null,
@ -47,7 +49,6 @@ const RouterSettingsAccordion = forwardRef<RouterSettingsAccordionRef, RouterSet
});
const [fallbacks, setFallbacks] = useState<Fallbacks>([]);
const [fallbackGroups, setFallbackGroups] = useState<FallbackGroup[]>([]);
const [modelInfo, setModelInfo] = useState<ModelGroup[]>([]);
const [availableRoutingStrategies, setAvailableRoutingStrategies] = useState<string[]>([]);
const [routerFieldsMetadata, setRouterFieldsMetadata] = useState<{ [key: string]: any }>({});
const [routingStrategyDescriptions, setRoutingStrategyDescriptions] = useState<{ [key: string]: string }>({});
@ -175,21 +176,11 @@ const RouterSettingsAccordion = forwardRef<RouterSettingsAccordionRef, RouterSet
});
}, [accessToken]);
// Fetch available models for fallbacks
useEffect(() => {
if (!accessToken) {
return;
}
const loadModels = async () => {
try {
const uniqueModels = await fetchAvailableModels(accessToken);
setModelInfo(uniqueModels);
} catch (error) {
console.error("Error fetching model info for fallbacks:", error);
}
};
loadModels();
}, [accessToken]);
const { data: modelInfo = [] } = useQuery<ModelGroup[]>({
queryKey: ["fallbackAvailableModels", accessToken, teamId ?? null],
queryFn: () => (teamId ? fetchAvailableModelsForTeam(accessToken, teamId) : fetchAvailableModels(accessToken)),
enabled: Boolean(accessToken),
});
// Helper function to build router_settings from current state
const buildRouterSettings = (): RouterSettingsAccordionValue["router_settings"] => {

View file

@ -0,0 +1,33 @@
import { beforeEach, describe, expect, it, vi } from "vitest";
import { modelAvailableCall } from "@/components/networking";
import { fetchAvailableModelsForTeam } from "./fetch_models";
vi.mock("@/components/networking", () => ({
modelAvailableCall: vi.fn(),
modelHubCall: vi.fn(),
}));
const modelAvailableCallMock = vi.mocked(modelAvailableCall);
describe("fetchAvailableModelsForTeam", () => {
beforeEach(() => {
vi.clearAllMocks();
});
it("requests the models scoped to the team so team-only BYOK models are included", async () => {
modelAvailableCallMock.mockResolvedValue({
data: [{ id: "all-proxy-models" }, { id: "openai/*" }, { id: "gpt-5-mini" }, { id: "openai/*" }],
});
const models = await fetchAvailableModelsForTeam("token", "team-123");
expect(modelAvailableCallMock).toHaveBeenCalledWith("token", "", "", false, "team-123");
expect(models).toEqual([{ model_group: "gpt-5-mini" }, { model_group: "openai/*" }]);
});
it("returns an empty list when the team has no models", async () => {
modelAvailableCallMock.mockResolvedValue({ data: [] });
expect(await fetchAvailableModelsForTeam("token", "team-123")).toEqual([]);
});
});

View file

@ -1,12 +1,22 @@
// fetch_models.ts
import { modelHubCall } from "@/components/networking";
import { excludeProxyWideSentinel } from "@/components/key_team_helpers/fetch_available_models_team_key";
import { modelAvailableCall, modelHubCall } from "@/components/networking";
export interface ModelGroup {
model_group: string;
mode?: string;
}
export const fetchAvailableModelsForTeam = async (accessToken: string, teamId: string): Promise<ModelGroup[]> => {
const response = await modelAvailableCall(accessToken, "", "", false, teamId);
const modelNames: string[] = (response?.data ?? []).map((model: { id: string }) => model.id);
return excludeProxyWideSentinel(Array.from(new Set(modelNames)))
.sort((a, b) => a.localeCompare(b))
.map((model) => ({ model_group: model }));
};
/**
* Fetches available models using modelHubCall and formats them for the selection dropdown.
*/

View file

@ -919,7 +919,7 @@ describe("TeamInfoView", () => {
});
};
it("prefills pairs from team metadata, hides UI-managed keys, and round-trips typed values on save", async () => {
it("should preserve metadata types and hide managed keys", async () => {
const user = userEvent.setup({ delay: null });
vi.mocked(networking.teamInfoCall).mockResolvedValue(
createMockTeamData({
@ -964,27 +964,6 @@ describe("TeamInfoView", () => {
expect(updateArg.model_tpm_limit).toEqual({ "gpt-4": 100 });
});
it("includes a newly added pair in the team update", async () => {
const user = userEvent.setup({ delay: null });
vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ models: ["gpt-4"] }));
vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any);
renderWithProviders(<TeamInfoView {...defaultProps} />);
await openSettingsEditor(user);
await user.click(screen.getByRole("button", { name: /add key-value pair/i }));
await user.type(screen.getByPlaceholderText("Key"), "cost_center");
await user.type(screen.getByPlaceholderText("Value"), "eng-1");
await user.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => {
expect(networking.teamUpdateCall).toHaveBeenCalled();
});
expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].metadata).toMatchObject({ cost_center: "eng-1" });
});
it("should keep declared keys as ordinary prefilled rows and submit the edited value", async () => {
const user = userEvent.setup({ delay: null });
vi.mocked(useTeamMetadataSchema).mockReturnValue({

View file

@ -1215,6 +1215,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
<RouterSettingsAccordion
ref={routerSettingsRef}
accessToken={accessToken || ""}
teamId={teamId}
value={info.router_settings ? { router_settings: info.router_settings } : undefined}
/>
</Form.Item>

View file

@ -31844,6 +31844,12 @@ export interface components {
* @description Default model to use if tier cannot be determined
*/
default_model?: string | null;
/**
* Deployment Affinity
* @description When True and a session_id is resolvable on the request, pin the deployment chosen inside each routed model group and reuse it whenever the session returns to that group, without pinning which group the session routes to. Independent of session_affinity, which pins the model group instead (and always carries this deployment pin with it): with session_affinity off, every turn is still classified on its own merits while a session that escalates to a stronger tier and comes back still lands on the deployment it used before, which is what keeps a provider prompt cache warm. Pins are held per model group, so switching tiers does not disturb the pin left behind in the previous group. On by default because re-shuffling a conversation across deployments of the same model discards that cache for no benefit; set False to keep every turn load-balanced across the group, which is what a deployment set with tight per-deployment rate limits wants. Inert when no session_id is resolvable, since there is nothing to key a pin on, and suppressed when plugins are configured, for the same reason session_affinity is.
* @default true
*/
deployment_affinity: boolean;
/**
* Dimension Weights
* @description Weights for each scoring dimension
@ -31901,13 +31907,13 @@ export interface components {
semantic_keyword_matching: boolean;
/**
* Session Affinity
* @description When True and a session_id is resolvable on the request, pin the model chosen on the session's first turn and reuse it for every later turn, skipping re-classification. Off by default so every turn is classified on its own merits and routed to the cheapest adequate tier. Set True to keep a multi-turn session on one model, which preserves provider prompt caches and avoids cross-model conversation-history errors.
* @description When True and a session_id is resolvable on the request, pin the model chosen on the session's first turn and reuse it for every later turn, skipping re-classification. Off by default so every turn is classified on its own merits and routed to the cheapest adequate tier. Set True to keep a multi-turn session on one model, which preserves provider prompt caches and avoids cross-model conversation-history errors. Always implies the deployment pin regardless of deployment_affinity: the session sticks to one deployment of the pinned model, since freezing the model while re-shuffling its deployments would still go cache-cold.
* @default false
*/
session_affinity: boolean;
/**
* Session Affinity Ttl Seconds
* @description TTL for the session affinity pin; refreshed on every cache hit
* @description TTL for the session affinity pin; refreshed on every cache hit. Bounds both the session_affinity model pin and the deployment_affinity deployment pin, so it measures idle time for the session's routing decisions rather than total session length
* @default 3600
*/
session_affinity_ttl_seconds: number;