diff --git a/CLAUDE.md b/CLAUDE.md index 59929143c46..f1bb46c1fd3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -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 diff --git a/Makefile b/Makefile index 493828571b7..94d8c875af5 100644 --- a/Makefile +++ b/Makefile @@ -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/ diff --git a/litellm/constants.py b/litellm/constants.py index 30d3bb1f26e..3b91f23fe39 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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 = ( diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index e273289168a..2c83406afed 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -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( diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 7abbf0c96e5..2c5f7484dac 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -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: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 1fc05ac4653..fa89df39c5f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9dc450befc2..4baa7b99a4f 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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 diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 22200567012..fbf28e223c1 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -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", diff --git a/litellm/proxy/db/autorouter_session_rollup.py b/litellm/proxy/db/autorouter_session_rollup.py index b1f074c26b7..9732b1d7402 100644 --- a/litellm/proxy/db/autorouter_session_rollup.py +++ b/litellm/proxy/db/autorouter_session_rollup.py @@ -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: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 83ae59ef050..c48fee96646 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -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", diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index f94f1302e66..9ff8bf001ec 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -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(), ) diff --git a/litellm/router.py b/litellm/router.py index 1a76eb5d59b..feaf69a44ae 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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( diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index ffcb0fc471e..1298f57e76e 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -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, diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 69609a973b0..0999af66fd8 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -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( diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index b4408d17ffb..7fb90ab89de 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -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", diff --git a/litellm/types/router.py b/litellm/types/router.py index 4280da08cbb..e166d844735 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -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) diff --git a/scripts/install_git_hooks.sh b/scripts/install_git_hooks.sh index 7ea8c3ff2e9..8f4e79e4ddd 100755 --- a/scripts/install_git_hooks.sh +++ b/scripts/install_git_hooks.sh @@ -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 diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh index af8335e0e84..afe55603466 100755 --- a/scripts/pre_commit_lint.sh +++ b/scripts/pre_commit_lint.sh @@ -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 diff --git a/tests/agent_tests/test_a2a_agent.py b/tests/agent_tests/test_a2a_agent.py index f5ad9601369..1f72ced64f1 100644 --- a/tests/agent_tests/test_a2a_agent.py +++ b/tests/agent_tests/test_a2a_agent.py @@ -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"): diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 5999e5772d1..dc69bd42171 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -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 diff --git a/tests/litellm_utils_tests/test_proxy_budget_reset.py b/tests/litellm_utils_tests/test_proxy_budget_reset.py index 44da3ea06a0..00d5380b2f4 100644 --- a/tests/litellm_utils_tests/test_proxy_budget_reset.py +++ b/tests/litellm_utils_tests/test_proxy_budget_reset.py @@ -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) diff --git a/tests/llm_responses_api_testing/base_responses_api.py b/tests/llm_responses_api_testing/base_responses_api.py index 407091a65b3..f5751aa79e8 100644 --- a/tests/llm_responses_api_testing/base_responses_api.py +++ b/tests/llm_responses_api_testing/base_responses_api.py @@ -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") diff --git a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py index 65b70f13a3b..7bc61c40ea0 100644 --- a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py +++ b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py @@ -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(), ) diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py index 2573ad5a375..82b074220fa 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py @@ -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(): diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py index b6ff3b70a4d..f39f41a6d12 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py @@ -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 diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 6fc44ef7519..60d9689dc0b 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -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.""" diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 6d6fd2e5507..22f9e6bb67a 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -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", diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 356556f3563..94b6b68855b 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -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 diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py index e29adda3328..428eb0ceafd 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -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 diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py index 0bcb0247aad..4053e6d118b 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py @@ -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) diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py index 274d3c517f6..35d98903226 100644 --- a/tests/test_litellm/test_pre_commit_lint.py +++ b/tests/test_litellm/test_pre_commit_lint.py @@ -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 diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index da2d78edb73..67fa827a8e4 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -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) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 3a670bc7345..d621e85f09b 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16760 + "limit": 16758 }, "LIT011": { "limit": 5598 diff --git a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx index a70b7602e5b..5ac3b8b2b64 100644 --- a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx @@ -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[] }) => ( +
{availableModels.join(",")}
+ ), })); 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 }) => ( + {children} + ), + }); +}; + 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(); + renderWithQueryClient(); 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(); + + 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(); + + 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(); + await waitFor(() => expect(resolvers).toHaveLength(1)); + + rerender(); + 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(); + const { unmount } = renderWithQueryClient(); await flushInitialPropagation(onChange); fireEvent.click(screen.getByText("set-least-busy")); diff --git a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx index 08b917e302f..56227abe9ea 100644 --- a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx +++ b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx @@ -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( - ({ accessToken, value, onChange, modelData }, ref) => { + ({ accessToken, value, onChange, modelData, teamId }, ref) => { const [formValue, setFormValue] = useState({ routerSettings: {}, selectedStrategy: null, @@ -47,7 +49,6 @@ const RouterSettingsAccordion = forwardRef([]); const [fallbackGroups, setFallbackGroups] = useState([]); - const [modelInfo, setModelInfo] = useState([]); const [availableRoutingStrategies, setAvailableRoutingStrategies] = useState([]); const [routerFieldsMetadata, setRouterFieldsMetadata] = useState<{ [key: string]: any }>({}); const [routingStrategyDescriptions, setRoutingStrategyDescriptions] = useState<{ [key: string]: string }>({}); @@ -175,21 +176,11 @@ const RouterSettingsAccordion = forwardRef { - 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({ + 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"] => { diff --git a/ui/litellm-dashboard/src/components/llm_calls/fetch_models.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.test.tsx new file mode 100644 index 00000000000..bd691c7f629 --- /dev/null +++ b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.test.tsx @@ -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([]); + }); +}); diff --git a/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx index 0de98330c2e..a1690b1307e 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx @@ -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 => { + 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. */ diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index 513719a2ad9..50e10285148 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -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(); - 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({ diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index e57bf59b2da..6c6fdcaedd6 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -1215,6 +1215,7 @@ const TeamInfoView: React.FC = ({ diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b90bf21e11b..3f2bfda9a8b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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;