From e75dd5663117026c3b416695b7eafaf3b5cfe1da Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 17:24:42 -0700 Subject: [PATCH 001/119] fix(proxy): expand access groups in /health scoping and allowlist health display fields /health filtered deployments by the key's literal `models` entries, so a key or UI user scoped to a model access group got 0/0 on both the live and the background-cache path, plus a misleading "missing model_info.id" warning on the cached one. The endpoint now resolves key sentinels and expands access groups the way auth does. Health entries used to copy every litellm_params field, so a deployment parameter that is not JSON-safe (a nested mapping keyed by a tuple) 500d the endpoint for every caller and admins saw internal settings nobody asked for. Entries now keep only an explicit allowlist of JSON-safe diagnostic fields; api_base and api_version stay admin-only and credentials stay out for everyone. Fixes #28206 --- litellm/proxy/health_check.py | 55 ++-- .../health_endpoints/_health_endpoints.py | 56 ++-- .../health_endpoints/test_health_endpoints.py | 286 ++++++++++++++++++ 3 files changed, 342 insertions(+), 55 deletions(-) diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 219f6f270ed..8daeb296f3e 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -32,32 +32,35 @@ from litellm.router_utils.auto_router_model_naming import ( strategy_router_dependencies, ) -ILLEGAL_DISPLAY_PARAMS: Final = [ - "messages", - "api_key", - "prompt", - "input", - "client_secret", - "azure_ad_token", - "azure_username", - "azure_password", - "vertex_credentials", - "vertex_ai_credentials", - "aws_access_key_id", - "aws_secret_access_key", - "aws_session_token", - "aws_web_identity_token", - "extra_headers", - "headers", - "exception", # internal; not JSON-serializable, never for display - "litellm_metadata", # internal tracking metadata with auth objects; not for display -] # Provider routing fields. Allowed for proxy admins so they can see which # region/version a deployment is checking; gated at the endpoint layer for # non-admin callers (see _strip_admin_only_fields_from_health_result). ADMIN_ONLY_HEALTH_DISPLAY_PARAMS: Final = ("api_base", "api_version") -MINIMAL_DISPLAY_PARAMS: Final = ["model", "mode_error"] +MINIMAL_DISPLAY_PARAMS: Final = frozenset({"model", "mode_error"}) + +HEALTH_DISPLAY_PARAMS: Final = ( + MINIMAL_DISPLAY_PARAMS + | frozenset(ADMIN_ONLY_HEALTH_DISPLAY_PARAMS) + | frozenset( + { + "custom_llm_provider", + "mode", + "base_model", + "aws_region_name", + "region_name", + "vertex_project", + "vertex_location", + "tpm", + "rpm", + "error", + "raw_request_typed_dict", + "x-ratelimit-remaining-requests", + "x-ratelimit-remaining-tokens", + "x-ms-region", + } + ) +) # Modes whose health-check probe is a chat-style completion call and # therefore accept `max_tokens`. Other modes (embedding, image_generation, @@ -143,14 +146,10 @@ def _get_random_llm_message(): def _clean_endpoint_data(endpoint_data: dict, details: bool | None = True): """ - Clean the endpoint data for display to users. + Keep only the explicitly approved, JSON-safe diagnostic fields for display to users. """ - endpoint_data.pop("litellm_logging_obj", None) - return ( - {k: v for k, v in endpoint_data.items() if k not in ILLEGAL_DISPLAY_PARAMS} - if details is not False - else {k: v for k, v in endpoint_data.items() if k in MINIMAL_DISPLAY_PARAMS} - ) + displayed: Final = HEALTH_DISPLAY_PARAMS if details is not False else MINIMAL_DISPLAY_PARAMS + return {k: v for k, v in endpoint_data.items() if k in displayed} def health_check_filter_kwargs_from_general_settings( diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 8b57bdca2fe..0810bd2b095 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -36,9 +36,13 @@ from litellm.proxy._types import ( UserAPIKeyAuth, WebhookEvent, ) +from litellm.proxy.auth.auth_checks import ( + _resolve_key_models_for_auth_check, # pyright: ignore[reportPrivateUsage] # the auth layer's sentinel resolution, reused so /health scopes exactly like a request +) from litellm.proxy.auth.auth_utils import ( _BANNED_REQUEST_BODY_PARAMS, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the request-body check ) +from litellm.proxy.auth.model_checks import get_key_models from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.db.proxy_worker_heartbeat import count_live_proxy_workers @@ -54,6 +58,7 @@ from litellm.proxy.middleware.in_flight_requests_middleware import ( get_in_flight_requests, ) from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager +from litellm.router import Router from litellm.router_utils.clientside_credential_handler import ( _ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the router path clientside_credential_keys, @@ -854,6 +859,24 @@ def _strip_admin_only_fields_from_health_result(result: dict) -> dict: return out +def _health_accessible_model_names( + user_api_key_dict: UserAPIKeyAuth, llm_router: Router | None +) -> frozenset[str] | None: + """Model names the caller may health-check, or None when the key is unrestricted.""" + granted_models: Final = _resolve_key_models_for_auth_check(user_api_key_dict) + if not granted_models or SpecialModelNames.all_proxy_models.value in granted_models: + return None + if llm_router is None: + return frozenset(granted_models) + return frozenset( + get_key_models( + user_api_key_dict=user_api_key_dict, + proxy_model_list=llm_router.get_model_names(team_id=user_api_key_dict.team_id), + model_access_groups=llm_router.get_model_access_groups(), + ) + ) + + def _resolve_targeted_model_ids(model_list: list, model: str | None, model_id: str | None) -> set | None: """ Resolve a ``/health`` ``model`` / ``model_id`` query param to the set of @@ -1080,32 +1103,11 @@ async def health_endpoint( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail={"error": "Model list not initialized"}, ) - _llm_model_list = copy.deepcopy(llm_model_list) - ### FILTER MODELS FOR ONLY THOSE USER HAS ACCESS TO ### - # Live path: scope by model_name (every deployment has one). - # Cache path: scope by model_id (the cache is keyed on model_id). - # Consequence: a deployment whose model_name the caller can access - # but which lacks model_info.id will appear in the live /health - # response but NOT in the background-cache /health response. This is - # surfaced via the "warnings" field below so operators can fix the - # missing model_info.id rather than guess at the discrepancy. - # Keys granted SpecialModelNames.all_proxy_models carry the literal - # "all-proxy-models" entry, which matches no real model_name; treat - # them as unrestricted instead of filtering the list down to nothing. - # Keys granted SpecialModelNames.all_team_models inherit the parent - # team's allowlist (same semantics as get_key_models in - # model_checks.py). Without a team_id the sentinel cannot resolve and - # stays in the list, matching nothing; denied rather than - # unrestricted, mirroring _resolve_key_models_for_auth_check. - accessible_models = list(user_api_key_dict.models) - if SpecialModelNames.all_team_models.value in accessible_models and user_api_key_dict.team_id is not None: - accessible_models = list(user_api_key_dict.team_models) - restrict_to_allowed_models: Final = ( - len(accessible_models) > 0 and SpecialModelNames.all_proxy_models.value not in accessible_models - ) - if restrict_to_allowed_models: - allowed_models: Final = set(accessible_models) - _llm_model_list = [m for m in _llm_model_list if m.get("model_name") in allowed_models] + allowed_models: Final = _health_accessible_model_names(user_api_key_dict, llm_router) + restrict_to_allowed_models: Final = allowed_models is not None + _llm_model_list: Final = [ + m for m in copy.deepcopy(llm_model_list) if allowed_models is None or m.get("model_name") in allowed_models + ] if use_background_health_checks: # The cached background result covers every model. When the # caller targets a specific model/model_id we have to narrow the @@ -1125,7 +1127,7 @@ async def health_endpoint( # intersection of "targeted" and "allowed." filter_ids: Final = targeted_ids if targeted_ids is not None else allowed_model_ids filtered: Final = _filter_health_check_results_by_model_ids(health_check_results, filter_ids) - if targeted_ids is None and not allowed_model_ids: + if targeted_ids is None and _llm_model_list and not allowed_model_ids: # Caller has accessible model_names but none of the # matching deployments expose a model_info.id, so the # cache filter (which keys on model_id) drops every diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index e3f71692c78..b14baf5f8c5 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -1,6 +1,8 @@ import asyncio import json import time +from collections.abc import Iterator, Mapping, Sequence +from contextlib import contextmanager from datetime import datetime, timedelta from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -1650,6 +1652,209 @@ async def test_health_endpoint_resolves_all_team_models_to_team_allowlist(): assert returned_names == {"model-b"}, f"all-team-models key should health-check the team's models: {returned_names}" +class _AccessGroupRouter: + """Router stand-in exposing only the two lookups /health uses to expand a key's model grants.""" + + def __init__(self, access_groups: dict[str, list[str]], model_names: list[str]) -> None: + self._access_groups = access_groups + self._model_names = model_names + + def get_model_access_groups(self, model_name=None, model_access_group=None, team_id=None): + return self._access_groups + + def get_model_names(self, team_id=None): + return self._model_names + + +_ACCESS_GROUP_MODEL_LIST = [ + { + "model_name": "bedrock-nova", + "litellm_params": {"model": "bedrock/us.amazon.nova-2-lite-v1:0"}, + "model_info": {"id": "id-bedrock", "access_groups": ["bedrock-group"]}, + }, + { + "model_name": "gpt-5.4-mini", + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "id-openai"}, + }, +] +_ACCESS_GROUP_ROUTER = _AccessGroupRouter({"bedrock-group": ["bedrock-nova"]}, ["bedrock-nova", "gpt-5.4-mini"]) +_ACCESS_GROUP_CACHED_RESULTS = { + "healthy_endpoints": [ + {"model": "bedrock/us.amazon.nova-2-lite-v1:0", "model_id": "id-bedrock"}, + {"model": "openai/gpt-5.4-mini", "model_id": "id-openai"}, + ], + "unhealthy_endpoints": [], + "healthy_count": 2, + "unhealthy_count": 0, +} + + +@contextmanager +def _proxy_health_globals( + llm_model_list: Sequence[Mapping[str, object]], + llm_router: object, + use_background_health_checks: bool = False, + health_check_results: Mapping[str, object] | None = None, +) -> Iterator[None]: + with ( + patch( # test-quality-ok: proxy module global, no injection seam + "litellm.proxy.proxy_server.llm_model_list", list(llm_model_list) + ), + patch( # test-quality-ok: proxy module global, no injection seam + "litellm.proxy.proxy_server.llm_router", llm_router + ), + patch( # test-quality-ok: proxy module global, no injection seam + "litellm.proxy.proxy_server.prisma_client", None + ), + patch( # test-quality-ok: proxy module global, no injection seam + "litellm.proxy.proxy_server.use_background_health_checks", use_background_health_checks + ), + patch( # test-quality-ok: proxy module global, no injection seam + "litellm.proxy.proxy_server.user_model", None + ), + patch( # test-quality-ok: proxy module global, no injection seam + "litellm.proxy.proxy_server.health_check_results", dict(health_check_results or {}) + ), + patch( # test-quality-ok: proxy module global, no injection seam + "litellm.proxy.proxy_server.health_check_details", True + ), + patch( # test-quality-ok: proxy module global, no injection seam + "litellm.proxy.proxy_server.health_check_concurrency", 1 + ), + ): + yield + + +@pytest.mark.asyncio +async def test_health_endpoint_expands_access_group_on_live_path(): + """ + LIT-6907 / gh-28206: a key granted a model access group carries the group + name in user_api_key_dict.models. Matching it as a literal model_name + filtered every deployment out and /health answered 0/0 for a model the + same key could call. + """ + from fastapi import Response + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + captured: dict = {} + + async def fake_perform(**kwargs): + captured["model_list"] = kwargs["model_list"] + return {"healthy_endpoints": [], "unhealthy_endpoints": [], "healthy_count": 0, "unhealthy_count": 0} + + with ( + _proxy_health_globals(_ACCESS_GROUP_MODEL_LIST, _ACCESS_GROUP_ROUTER), + patch( # test-quality-ok: the model list handed to the probe is the assertion; no injection seam + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + side_effect=fake_perform, + ), + ): + await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-group"]), + ) + + assert [m["model_name"] for m in captured["model_list"]] == ["bedrock-nova"] + + +@pytest.mark.asyncio +async def test_health_endpoint_expands_access_group_on_background_cache_path(): + """ + LIT-6907: the background-cache path scoped the cached entries through the + same literal model_name match, so an access-group key got an empty result + plus a warning blaming missing model_info.id. + """ + from fastapi import Response + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _ACCESS_GROUP_MODEL_LIST, + _ACCESS_GROUP_ROUTER, + use_background_health_checks=True, + health_check_results=_ACCESS_GROUP_CACHED_RESULTS, + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-group"]), + model=None, + model_id=None, + ) + + assert [e["model_id"] for e in result["healthy_endpoints"]] == ["id-bedrock"] + assert result["healthy_count"] == 1 + assert "warnings" not in result + + +@pytest.mark.asyncio +async def test_health_endpoint_treats_no_team_all_team_models_as_unrestricted(): + """ + A key granted "all-team-models" without a team resolves to an empty + allowlist in the auth layer, which means unrestricted. /health used to + keep the unresolved sentinel and filter every deployment out instead. + """ + from fastapi import Response + + from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + captured: dict = {} + + async def fake_perform(**kwargs): + captured["model_list"] = kwargs["model_list"] + return {"healthy_endpoints": [], "unhealthy_endpoints": [], "healthy_count": 0, "unhealthy_count": 0} + + with ( + _proxy_health_globals(_ACCESS_GROUP_MODEL_LIST, None), + patch( # test-quality-ok: the model list handed to the probe is the assertion; no injection seam + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + side_effect=fake_perform, + ), + ): + await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth( + api_key="hashed-test-key", models=[SpecialModelNames.all_team_models.value], team_id=None + ), + ) + + assert {m["model_name"] for m in captured["model_list"]} == {"bedrock-nova", "gpt-5.4-mini"} + + +@pytest.mark.asyncio +async def test_health_endpoint_omits_model_id_warning_when_no_deployment_matches(): + """ + The missing-model_info.id warning is only true when a matching deployment + exists without an id. A key whose grants match no deployment at all gets a + plain empty result, not advice to populate ids that are already there. + """ + from fastapi import Response + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _ACCESS_GROUP_MODEL_LIST, + _ACCESS_GROUP_ROUTER, + use_background_health_checks=True, + health_check_results=_ACCESS_GROUP_CACHED_RESULTS, + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["no-such-model"]), + model=None, + model_id=None, + ) + + assert result["healthy_count"] == 0 + assert result["unhealthy_count"] == 0 + assert "warnings" not in result + + @pytest.mark.asyncio async def test_health_endpoint_filters_background_cache_by_user_access(): """ @@ -2571,6 +2776,87 @@ def test_clean_endpoint_data_never_displays_credential_fields(credential_field, assert canary not in str(cleaned) +def test_clean_endpoint_data_keeps_only_json_safe_diagnostics(): + """ + LIT-6907: _clean_endpoint_data used to copy every litellm_param not on a + deny list, so a nested mapping keyed by a tuple reached jsonable_encoder + and 500'd /health. Only the explicit allowlist survives now. + """ + from fastapi.encoders import jsonable_encoder + + from litellm.proxy.health_check import _clean_endpoint_data + + cleaned = _clean_endpoint_data( + { + "model": "bedrock/us.amazon.nova-2-lite-v1:0", + "custom_llm_provider": "bedrock", + "aws_region_name": "us-east-1", + "metadata": {("us-east-1", "primary"): "canary-nested-mapping"}, + "allow_client_keepalive_override": False, + "api_key": "CANARY-API-KEY", + "x-ratelimit-remaining-requests": 99, + "raw_request_typed_dict": {"raw_request_api_base": "https://example.test"}, + }, + details=True, + ) + + assert cleaned == { + "model": "bedrock/us.amazon.nova-2-lite-v1:0", + "custom_llm_provider": "bedrock", + "aws_region_name": "us-east-1", + "x-ratelimit-remaining-requests": 99, + "raw_request_typed_dict": {"raw_request_api_base": "https://example.test"}, + } + assert jsonable_encoder(cleaned) == cleaned + + +@pytest.mark.asyncio +async def test_health_endpoint_result_survives_non_json_safe_deployment_params(): + """ + LIT-6907: the full /health path with a deployment carrying a tuple-keyed + nested mapping must produce a response FastAPI can encode, with the + approved diagnostics intact and the offending param absent. + """ + from fastapi import Response + from fastapi.encoders import jsonable_encoder + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + model_list = [ + { + "model_name": "bedrock-nova", + "litellm_params": { + "model": "bedrock/us.amazon.nova-2-lite-v1:0", + "aws_region_name": "us-east-1", + "aws_access_key_id": "CANARY-ACCESS-KEY", + "metadata": {("us-east-1", "primary"): "canary-nested-mapping"}, + }, + "model_info": {"id": "id-bedrock"}, + } + ] + + with ( + _proxy_health_globals(model_list, None), + patch( # test-quality-ok: the provider probe is faked; the assertion is the response shaping after it + "litellm.ahealth_check", AsyncMock(return_value={"x-ratelimit-remaining-requests": 99}) + ), + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-admin-key", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + encoded = jsonable_encoder(result) + assert encoded["healthy_count"] == 1 + entry = encoded["healthy_endpoints"][0] + assert entry["model_id"] == "id-bedrock" + assert entry["aws_region_name"] == "us-east-1" + assert entry["x-ratelimit-remaining-requests"] == 99 + assert "metadata" not in entry + assert "CANARY" not in str(encoded) + + class TestConfigBaseForHealthCheck: """A request that sets its own connection fields gets a base without the configuration's credentials; anything it leaves unset still comes from From a84f4d6206cfdfa8b025421b4bbdf43a55eed5ee Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 17:57:13 -0700 Subject: [PATCH 002/119] fix(proxy): keep /health inside the caller's team and admin-gate the Bedrock runtime endpoint --- litellm/proxy/health_check.py | 3 +- .../health_endpoints/_health_endpoints.py | 23 ++- .../health_endpoints/test_health_endpoints.py | 161 ++++++++++++++++-- 3 files changed, 169 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 8daeb296f3e..78bb01aa6ac 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -35,7 +35,7 @@ from litellm.router_utils.auto_router_model_naming import ( # Provider routing fields. Allowed for proxy admins so they can see which # region/version a deployment is checking; gated at the endpoint layer for # non-admin callers (see _strip_admin_only_fields_from_health_result). -ADMIN_ONLY_HEALTH_DISPLAY_PARAMS: Final = ("api_base", "api_version") +ADMIN_ONLY_HEALTH_DISPLAY_PARAMS: Final = ("api_base", "api_version", "aws_bedrock_runtime_endpoint") MINIMAL_DISPLAY_PARAMS: Final = frozenset({"model", "mode_error"}) @@ -49,6 +49,7 @@ HEALTH_DISPLAY_PARAMS: Final = ( "base_model", "aws_region_name", "region_name", + "watsonx_region_name", "vertex_project", "vertex_location", "tpm", diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 0810bd2b095..cc9ad1eb0cb 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -845,7 +845,7 @@ def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: def _strip_admin_only_fields_from_health_result(result: dict) -> dict: """ Return a copy of the /health response with provider routing fields - (``api_base``, ``api_version``) removed from each healthy/unhealthy + (``ADMIN_ONLY_HEALTH_DISPLAY_PARAMS``) removed from each healthy/unhealthy endpoint entry. Used to hide those fields from non-admin callers while still showing them which deployments they own and whether each one is healthy. Proxy admins receive the unmodified result. @@ -877,6 +877,18 @@ def _health_accessible_model_names( ) +def _caller_may_probe_deployment( + deployment: Mapping[str, object], allowed_models: frozenset[str], llm_router: Router | None, team_id: str | None +) -> bool: + """Same deployment visibility rule as request auth: another team's deployment is never in scope.""" + if llm_router is None: + return deployment.get("model_name") in allowed_models + model: Final = dict(deployment) + return any( + llm_router.should_include_deployment(model_name=name, model=model, team_id=team_id) for name in allowed_models + ) + + def _resolve_targeted_model_ids(model_list: list, model: str | None, model_id: str | None) -> set | None: """ Resolve a ``/health`` ``model`` / ``model_id`` query param to the set of @@ -1079,7 +1091,9 @@ async def health_endpoint( response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE if is_admin: return result - response.headers["Litellm-Health-Field-Notice"] = "api_base and api_version are admin-only on this endpoint" + response.headers["Litellm-Health-Field-Notice"] = ( + f"{', '.join(ADMIN_ONLY_HEALTH_DISPLAY_PARAMS)} are admin-only on this endpoint" + ) return _strip_admin_only_fields_from_health_result(result) try: @@ -1106,7 +1120,10 @@ async def health_endpoint( allowed_models: Final = _health_accessible_model_names(user_api_key_dict, llm_router) restrict_to_allowed_models: Final = allowed_models is not None _llm_model_list: Final = [ - m for m in copy.deepcopy(llm_model_list) if allowed_models is None or m.get("model_name") in allowed_models + m + for m in copy.deepcopy(llm_model_list) + if allowed_models is None + or _caller_may_probe_deployment(m, allowed_models, llm_router, user_api_key_dict.team_id) ] if use_background_health_checks: # The cached background result covers every model. When the diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index b14baf5f8c5..6e5c807afba 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -1,4 +1,5 @@ import asyncio +import copy import json import time from collections.abc import Iterator, Mapping, Sequence @@ -19,6 +20,7 @@ import litellm.proxy.health_endpoints._health_endpoints as _health_endpoints_mod from litellm.litellm_core_utils.health_check_helpers import TEST_IMAGE_BASE64 from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.router import Router from litellm.proxy.health_endpoints._health_endpoints import ( _db_health_readiness_check, _show_no_redis_warning, @@ -1652,18 +1654,8 @@ async def test_health_endpoint_resolves_all_team_models_to_team_allowlist(): assert returned_names == {"model-b"}, f"all-team-models key should health-check the team's models: {returned_names}" -class _AccessGroupRouter: - """Router stand-in exposing only the two lookups /health uses to expand a key's model grants.""" - - def __init__(self, access_groups: dict[str, list[str]], model_names: list[str]) -> None: - self._access_groups = access_groups - self._model_names = model_names - - def get_model_access_groups(self, model_name=None, model_access_group=None, team_id=None): - return self._access_groups - - def get_model_names(self, team_id=None): - return self._model_names +def _router_for(model_list: Sequence[Mapping[str, object]]) -> Router: + return Router(model_list=copy.deepcopy(list(model_list))) _ACCESS_GROUP_MODEL_LIST = [ @@ -1678,7 +1670,29 @@ _ACCESS_GROUP_MODEL_LIST = [ "model_info": {"id": "id-openai"}, }, ] -_ACCESS_GROUP_ROUTER = _AccessGroupRouter({"bedrock-group": ["bedrock-nova"]}, ["bedrock-nova", "gpt-5.4-mini"]) +_ACCESS_GROUP_ROUTER = _router_for(_ACCESS_GROUP_MODEL_LIST) +_TEAM_MODEL_LIST = [ + _ACCESS_GROUP_MODEL_LIST[0], + { + "model_name": "bedrock-nova_team-b_9f2c", + "litellm_params": {"model": "bedrock/us.amazon.nova-2-lite-v1:0"}, + "model_info": { + "id": "id-team-b", + "team_id": "team-b", + "team_public_model_name": "bedrock-nova", + "access_groups": ["bedrock-group"], + }, + }, +] +_TEAM_CACHED_RESULTS = { + "healthy_endpoints": [ + {"model": "bedrock/us.amazon.nova-2-lite-v1:0", "model_id": "id-bedrock"}, + {"model": "bedrock/us.amazon.nova-2-lite-v1:0", "model_id": "id-team-b"}, + ], + "unhealthy_endpoints": [], + "healthy_count": 2, + "unhealthy_count": 0, +} _ACCESS_GROUP_CACHED_RESULTS = { "healthy_endpoints": [ {"model": "bedrock/us.amazon.nova-2-lite-v1:0", "model_id": "id-bedrock"}, @@ -2046,7 +2060,7 @@ async def test_health_endpoint_admin_sees_routing_fields_non_admin_does_not(): # withheld so clients that previously parsed them can detect the change. assert ( non_admin_response.headers.get("Litellm-Health-Field-Notice") - == "api_base and api_version are admin-only on this endpoint" + == "api_base, api_version, aws_bedrock_runtime_endpoint are admin-only on this endpoint" ) assert "Litellm-Health-Field-Notice" not in admin_response.headers @@ -2776,6 +2790,123 @@ def test_clean_endpoint_data_never_displays_credential_fields(credential_field, assert canary not in str(cleaned) +async def _live_probed_model_ids( + model_list: Sequence[Mapping[str, object]], user_api_key_dict: UserAPIKeyAuth +) -> set[str]: + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + captured: dict = {} + + async def fake_perform(**kwargs): + captured["model_list"] = kwargs["model_list"] + return {"healthy_endpoints": [], "unhealthy_endpoints": [], "healthy_count": 0, "unhealthy_count": 0} + + with ( + _proxy_health_globals(model_list, _router_for(model_list)), + patch( # test-quality-ok: the model list handed to the probe is the assertion; no injection seam + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + side_effect=fake_perform, + ), + ): + await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict) + + return {m["model_info"]["id"] for m in captured["model_list"]} + + +@pytest.mark.asyncio +async def test_health_endpoint_hides_another_teams_deployment_behind_a_shared_access_group(): + """ + Expanding an access group must not reach past the team boundary: a + team-a key holding the group name may not probe team-b's deployment even + though that deployment sits in the same group. + """ + probed = await _live_probed_model_ids( + _TEAM_MODEL_LIST, + UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-group"], team_id="team-a"), + ) + + assert probed == {"id-bedrock"} + + +@pytest.mark.asyncio +async def test_health_endpoint_shows_a_teams_own_deployment_by_its_public_name(): + """ + A team key names its team deployment by ``team_public_model_name``, while + the proxy model list carries the internal ``__`` + name; the deployment must still be probed for its own team. + """ + probed = await _live_probed_model_ids( + _TEAM_MODEL_LIST, + UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-nova"], team_id="team-b"), + ) + + assert probed == {"id-bedrock", "id-team-b"} + + +@pytest.mark.asyncio +async def test_health_endpoint_hides_another_teams_deployment_on_background_cache_path(): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _TEAM_MODEL_LIST, + _router_for(_TEAM_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_TEAM_CACHED_RESULTS, + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-group"], team_id="team-a"), + model=None, + model_id=None, + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-bedrock"] + assert result["healthy_count"] == 1 + + +def test_health_test_connection_keeps_error_and_raw_request_through_the_allowlist(monkeypatch): + """ + The dashboard's Test Connect button reads ``result.error`` and + ``result.raw_request_typed_dict`` from /health/test_connection, so the + allowlist must keep both while dropping the probe's own params. + """ + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + app = FastAPI() + app.include_router(_health_endpoints_module.router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + client = TestClient(app) + + with ( + patch( # test-quality-ok: the endpoint reads the proxy-global DB client and 500s when it is None; it has no injection seam + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), + respx.mock(assert_all_called=True) as respx_mock, + ): + respx_mock.post(host="api.openai.com", path="/v1/chat/completions").respond( + status_code=401, json={"error": {"message": "Incorrect API key provided"}} + ) + response = client.post( + "/health/test_connection", + json={ + "mode": "chat", + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-test", "timeout": 7}, + }, + ) + + assert response.status_code == 200, response.text + body = response.json() + assert body["status"] == "error" + assert "Incorrect API key provided" in body["result"]["error"] + assert "api.openai.com" in body["result"]["raw_request_typed_dict"]["raw_request_api_base"] + assert not {"api_key", "timeout", "exception"} & set(body["result"]) + + def test_clean_endpoint_data_keeps_only_json_safe_diagnostics(): """ LIT-6907: _clean_endpoint_data used to copy every litellm_param not on a @@ -2796,6 +2927,7 @@ def test_clean_endpoint_data_keeps_only_json_safe_diagnostics(): "api_key": "CANARY-API-KEY", "x-ratelimit-remaining-requests": 99, "raw_request_typed_dict": {"raw_request_api_base": "https://example.test"}, + "aws_bedrock_runtime_endpoint": "https://vpce-bedrock.example.test", }, details=True, ) @@ -2806,6 +2938,7 @@ def test_clean_endpoint_data_keeps_only_json_safe_diagnostics(): "aws_region_name": "us-east-1", "x-ratelimit-remaining-requests": 99, "raw_request_typed_dict": {"raw_request_api_base": "https://example.test"}, + "aws_bedrock_runtime_endpoint": "https://vpce-bedrock.example.test", } assert jsonable_encoder(cleaned) == cleaned From 471f51cc4deb5aa2db57bff70c5eac32eaa3bfc3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 09:29:09 -0700 Subject: [PATCH 003/119] fix(proxy): keep another team's deployment out of /health for keys with no team --- .../health_endpoints/_health_endpoints.py | 4 +- .../health_endpoints/test_health_endpoints.py | 38 +++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index cc9ad1eb0cb..04388c773a0 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -880,7 +880,9 @@ def _health_accessible_model_names( def _caller_may_probe_deployment( deployment: Mapping[str, object], allowed_models: frozenset[str], llm_router: Router | None, team_id: str | None ) -> bool: - """Same deployment visibility rule as request auth: another team's deployment is never in scope.""" + """Same deployment visibility rule as routing: another team's deployment is never in scope, team-less callers included.""" + if not Router._deployment_usable_by_team(deployment, team_id): + return False if llm_router is None: return deployment.get("model_name") in allowed_models model: Final = dict(deployment) diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index 6e5c807afba..f4ed50fdb5f 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -2830,6 +2830,21 @@ async def test_health_endpoint_hides_another_teams_deployment_behind_a_shared_ac assert probed == {"id-bedrock"} +@pytest.mark.asyncio +async def test_health_endpoint_hides_team_deployments_from_a_key_with_no_team(): + """ + Routing never serves a team-owned deployment to a caller without a team + (``filter_team_based_models``), so a team-less access-group key must not + probe team-b's deployment with team-b's credentials either. + """ + probed = await _live_probed_model_ids( + _TEAM_MODEL_LIST, + UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-group"], team_id=None), + ) + + assert probed == {"id-bedrock"} + + @pytest.mark.asyncio async def test_health_endpoint_shows_a_teams_own_deployment_by_its_public_name(): """ @@ -2868,6 +2883,29 @@ async def test_health_endpoint_hides_another_teams_deployment_on_background_cach assert result["healthy_count"] == 1 +@pytest.mark.asyncio +async def test_health_endpoint_hides_team_deployments_from_a_key_with_no_team_on_background_cache_path(): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _TEAM_MODEL_LIST, + _router_for(_TEAM_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_TEAM_CACHED_RESULTS, + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-group"], team_id=None), + model=None, + model_id=None, + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-bedrock"] + assert result["healthy_count"] == 1 + + def test_health_test_connection_keeps_error_and_raw_request_through_the_allowlist(monkeypatch): """ The dashboard's Test Connect button reads ``result.error`` and From 69a45e81ccb49c6ee4b874ec21fd70668caff036 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 09:40:00 -0700 Subject: [PATCH 004/119] fix(proxy): refuse /health targets outside the caller's scope instead of probing the rest --- .../health_endpoints/_health_endpoints.py | 11 ++- .../health_endpoints/test_health_endpoints.py | 94 +++++++++++++++---- 2 files changed, 85 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 04388c773a0..b8c56037206 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -1127,6 +1127,16 @@ async def health_endpoint( if allowed_models is None or _caller_may_probe_deployment(m, allowed_models, llm_router, user_api_key_dict.team_id) ] + targeted_ids: Final = _resolve_targeted_model_ids(_llm_model_list, model, model_id) + if restrict_to_allowed_models and targeted_ids is not None and not targeted_ids: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": f"key not allowed to health-check model_id {model_id}" + if model_id + else f"key not allowed to health-check model {model}" + }, + ) if use_background_health_checks: # The cached background result covers every model. When the # caller targets a specific model/model_id we have to narrow the @@ -1134,7 +1144,6 @@ async def health_endpoint( # healthy_count, otherwise an unhealthy "foo" combined with any # other healthy model would still report healthy_count > 0 and # the targeted-503 path would never fire. - targeted_ids: Final = _resolve_targeted_model_ids(_llm_model_list, model, model_id) if restrict_to_allowed_models: allowed_model_ids: Final = { (m.get("model_info") or {}).get("id") diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index f4ed50fdb5f..c03284b72b6 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -1517,7 +1517,7 @@ async def test_health_endpoint_filters_model_list_by_user_access(): ): from fastapi import Response - await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict) + await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict, model=None, model_id=None) assert "model_list" in captured, "health_endpoint did not call _perform_health_check_and_save" returned_names = {m["model_name"] for m in captured["model_list"]} @@ -1580,7 +1580,7 @@ async def test_health_endpoint_keeps_full_model_list_for_all_proxy_models(): ): from fastapi import Response - await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict) + await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict, model=None, model_id=None) returned_names = {m["model_name"] for m in captured["model_list"]} assert returned_names == { @@ -1648,7 +1648,7 @@ async def test_health_endpoint_resolves_all_team_models_to_team_allowlist(): ): from fastapi import Response - await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict) + await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict, model=None, model_id=None) returned_names = {m["model_name"] for m in captured["model_list"]} assert returned_names == {"model-b"}, f"all-team-models key should health-check the team's models: {returned_names}" @@ -1769,6 +1769,8 @@ async def test_health_endpoint_expands_access_group_on_live_path(): await health_endpoint( response=Response(), user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-group"]), + model=None, + model_id=None, ) assert [m["model_name"] for m in captured["model_list"]] == ["bedrock-nova"] @@ -2149,7 +2151,7 @@ async def test_health_endpoint_blocks_cross_scope_model_id_under_background_cach cache filter was driven by an unvalidated ID and the global cache leaked id-b's entry to the caller. """ - from fastapi import Response + from fastapi import HTTPException, Response from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.health_endpoints._health_endpoints import health_endpoint @@ -2200,21 +2202,18 @@ async def test_health_endpoint_blocks_cross_scope_model_id_under_background_cach ): # Calling with model="model-b" rather than model_id="id-b" because # the model_id branch raises 404 when llm_router is None. The bug - # being verified is the same: targeted resolver must drop entries - # not in the caller's scoped model_list. With the fix, the result - # has no leaked endpoints and the targeted-503 path fires. - result = await health_endpoint( - response=response, - user_api_key_dict=user_api_key_dict, - model="model-b", - model_id=None, - ) + # being verified is the same: a target outside the caller's scoped + # model_list is refused before the cache is read. + with pytest.raises(HTTPException) as refused: + await health_endpoint( + response=response, + user_api_key_dict=user_api_key_dict, + model="model-b", + model_id=None, + ) - leaked_ids = {ep.get("model_id") for ep in result.get("healthy_endpoints", [])} - leaked_ids |= {ep.get("model_id") for ep in result.get("unhealthy_endpoints", [])} - assert "id-b" not in leaked_ids, "background cache leaked an out-of-scope deployment to a scoped caller" - assert result["healthy_count"] == 0 - assert response.status_code == 503 + assert refused.value.status_code == 403 + assert "leaky-internal.test" not in str(refused.value.detail) @pytest.mark.asyncio @@ -2810,7 +2809,7 @@ async def _live_probed_model_ids( side_effect=fake_perform, ), ): - await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict) + await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict, model=None, model_id=None) return {m["model_info"]["id"] for m in captured["model_list"]} @@ -2883,6 +2882,63 @@ async def test_health_endpoint_hides_another_teams_deployment_on_background_cach assert result["healthy_count"] == 1 +@pytest.mark.asyncio +async def test_health_endpoint_refuses_a_targeted_deployment_outside_the_callers_scope_on_live_path(): + """ + A scoped key asking for a deployment it may not see must get a 403 and no + probe at all: probing the rest of its scope instead would report another + deployment's health under the requested id and store it as such. + """ + from fastapi import HTTPException, Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + fake_perform = AsyncMock() + + with ( + _proxy_health_globals(_TEAM_MODEL_LIST, _router_for(_TEAM_MODEL_LIST)), + patch( # test-quality-ok: the probe must never run; no injection seam + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + fake_perform, + ), + pytest.raises(HTTPException) as excinfo, + ): + await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-group"], team_id=None), + model=None, + model_id="id-team-b", + ) + + assert excinfo.value.status_code == 403 + fake_perform.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_health_endpoint_refuses_a_targeted_deployment_outside_the_callers_scope_on_background_cache_path(): + from fastapi import HTTPException, Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with ( + _proxy_health_globals( + _TEAM_MODEL_LIST, + _router_for(_TEAM_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_TEAM_CACHED_RESULTS, + ), + pytest.raises(HTTPException) as excinfo, + ): + await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-group"], team_id="team-a"), + model="bedrock-nova_team-b_9f2c", + model_id=None, + ) + + assert excinfo.value.status_code == 403 + + @pytest.mark.asyncio async def test_health_endpoint_hides_team_deployments_from_a_key_with_no_team_on_background_cache_path(): from fastapi import Response From 8b3faa6ed83290828ae3520aa2d35b11f9066fe6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 09:53:53 -0700 Subject: [PATCH 005/119] fix(proxy): keep /health inside the caller's team for unrestricted non-admin keys --- .../health_endpoints/_health_endpoints.py | 16 ++++-- .../health_endpoints/test_health_endpoints.py | 53 +++++++++++++++++++ 2 files changed, 64 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index b8c56037206..bf0200acfa9 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -878,11 +878,17 @@ def _health_accessible_model_names( def _caller_may_probe_deployment( - deployment: Mapping[str, object], allowed_models: frozenset[str], llm_router: Router | None, team_id: str | None + deployment: Mapping[str, object], + allowed_models: frozenset[str] | None, + llm_router: Router | None, + team_id: str | None, + caller_is_admin: bool, ) -> bool: """Same deployment visibility rule as routing: another team's deployment is never in scope, team-less callers included.""" - if not Router._deployment_usable_by_team(deployment, team_id): + if not caller_is_admin and not Router._deployment_usable_by_team(deployment, team_id): return False + if allowed_models is None: + return True if llm_router is None: return deployment.get("model_name") in allowed_models model: Final = dict(deployment) @@ -1120,12 +1126,12 @@ async def health_endpoint( detail={"error": "Model list not initialized"}, ) allowed_models: Final = _health_accessible_model_names(user_api_key_dict, llm_router) - restrict_to_allowed_models: Final = allowed_models is not None + restrict_to_allowed_models: Final = not is_admin or allowed_models is not None _llm_model_list: Final = [ m for m in copy.deepcopy(llm_model_list) - if allowed_models is None - or _caller_may_probe_deployment(m, allowed_models, llm_router, user_api_key_dict.team_id) + if not restrict_to_allowed_models + or _caller_may_probe_deployment(m, allowed_models, llm_router, user_api_key_dict.team_id, is_admin) ] targeted_ids: Final = _resolve_targeted_model_ids(_llm_model_list, model, model_id) if restrict_to_allowed_models and targeted_ids is not None and not targeted_ids: diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index c03284b72b6..bbcae2fb68f 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -1836,6 +1836,8 @@ async def test_health_endpoint_treats_no_team_all_team_models_as_unrestricted(): user_api_key_dict=UserAPIKeyAuth( api_key="hashed-test-key", models=[SpecialModelNames.all_team_models.value], team_id=None ), + model=None, + model_id=None, ) assert {m["model_name"] for m in captured["model_list"]} == {"bedrock-nova", "gpt-5.4-mini"} @@ -2844,6 +2846,57 @@ async def test_health_endpoint_hides_team_deployments_from_a_key_with_no_team(): assert probed == {"id-bedrock"} +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_id", "expected_ids"), + [(None, {"id-bedrock"}), ("team-a", {"id-bedrock"}), ("team-b", {"id-bedrock", "id-team-b"})], +) +async def test_health_endpoint_keeps_an_unrestricted_non_admin_key_to_its_own_team(team_id, expected_ids): + """ + A key with no model restriction is still bound by routing's team rule: + it may probe global deployments and its own team's, never another team's. + """ + probed = await _live_probed_model_ids( + _TEAM_MODEL_LIST, + UserAPIKeyAuth(api_key="hashed-test-key", models=[], team_id=team_id), + ) + + assert probed == expected_ids + + +@pytest.mark.asyncio +async def test_health_endpoint_lets_a_proxy_admin_probe_every_teams_deployment(): + probed = await _live_probed_model_ids( + _TEAM_MODEL_LIST, + UserAPIKeyAuth(api_key="hashed-test-key", models=[], user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert probed == {"id-bedrock", "id-team-b"} + + +@pytest.mark.asyncio +async def test_health_endpoint_keeps_an_unrestricted_non_admin_key_to_its_own_team_on_background_cache_path(): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _TEAM_MODEL_LIST, + _router_for(_TEAM_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_TEAM_CACHED_RESULTS, + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=[], team_id="team-a"), + model=None, + model_id=None, + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-bedrock"] + assert result["healthy_count"] == 1 + + @pytest.mark.asyncio async def test_health_endpoint_shows_a_teams_own_deployment_by_its_public_name(): """ From 787f2dee0c3d9e895c14adfccd21ca7e090fea24 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 20:52:59 -0700 Subject: [PATCH 006/119] fix(health): match team public model names when targeting /health by model --- litellm/proxy/health_check.py | 9 ++- .../health_endpoints/_health_endpoints.py | 9 +-- .../health_endpoints/test_health_endpoints.py | 67 ++++++++++++++++++- .../proxy/test_health_check_functions.py | 29 ++++++++ 4 files changed, 107 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 78bb01aa6ac..8cda318b937 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -258,6 +258,13 @@ def _deployment_model(deployment: Mapping[str, object]) -> str | None: return params.get("model") if isinstance(params, Mapping) else None +def deployment_answers_to(deployment: Mapping[str, object], model_name: str) -> bool: + """True when `model_name` is the deployment's model_name or the public name a team key reaches it by.""" + info: Final = deployment.get("model_info") + public_name: Final = info.get("team_public_model_name") if isinstance(info, Mapping) else None + return model_name in (deployment.get("model_name"), public_name) + + def _narrow_to_target( model_list: Sequence[Mapping[str, object]], model: str | None, model_id: str | None ) -> tuple[Mapping[str, object], ...]: @@ -268,7 +275,7 @@ def _narrow_to_target( if model is None: return tuple(model_list) by_param: Final = tuple(x for x in model_list if _deployment_model(x) == model) - return by_param or tuple(x for x in model_list if x.get("model_name") == model) + return by_param or tuple(x for x in model_list if deployment_answers_to(x, model)) def _is_strategy_router_deployment(litellm_params: Mapping[str, object]) -> bool: diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index c66eafc724f..267b094526b 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -50,6 +50,7 @@ from litellm.proxy.health_check import ( ADMIN_ONLY_HEALTH_DISPLAY_PARAMS, _clean_endpoint_data, _update_litellm_params_for_health_check, + deployment_answers_to, health_check_filter_kwargs_from_general_settings, perform_health_check, run_with_timeout, @@ -930,9 +931,9 @@ def _resolve_targeted_model_ids(model_list: list, model: str | None, model_id: s deployment IDs the response should be scoped to. Mirrors the live-path semantics in ``perform_health_check()``: ``model`` - matches either the deployment's ``model_name`` alias or its - ``litellm_params.model`` provider string. ``model_id`` matches - ``model_info.id``. + matches the deployment's ``model_name`` alias, its ``litellm_params.model`` + provider string, or the ``model_info.team_public_model_name`` a team key + reaches it by. ``model_id`` matches ``model_info.id``. Both query params are validated against the supplied ``model_list``. Callers pass an already-scoped list (filtered to the caller's allowed @@ -956,7 +957,7 @@ def _resolve_targeted_model_ids(model_list: list, model: str | None, model_id: s continue if model: litellm_model = (m.get("litellm_params") or {}).get("model") - if m.get("model_name") == model or litellm_model == model: + if litellm_model == model or deployment_answers_to(m, model): target_ids.add(deployment_id) return target_ids diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index 022cf1e5505..e0cf6d97141 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -2802,7 +2802,7 @@ def test_clean_endpoint_data_never_displays_credential_fields(credential_field, async def _live_probed_model_ids( - model_list: Sequence[Mapping[str, object]], user_api_key_dict: UserAPIKeyAuth + model_list: Sequence[Mapping[str, object]], user_api_key_dict: UserAPIKeyAuth, model: str | None = None ) -> set[str]: from fastapi import Response @@ -2821,7 +2821,7 @@ async def _live_probed_model_ids( side_effect=fake_perform, ), ): - await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict, model=None, model_id=None) + await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict, model=model, model_id=None) return {m["model_info"]["id"] for m in captured["model_list"]} @@ -3025,6 +3025,69 @@ async def test_health_endpoint_hides_team_deployments_from_a_key_with_no_team_on assert result["healthy_count"] == 1 +_TEAM_ONLY_MODEL_LIST = [_TEAM_MODEL_LIST[1]] + + +@pytest.mark.asyncio +async def test_health_endpoint_probes_a_team_only_deployment_by_its_public_name_on_live_path(): + """ + A team key targets its deployment by ``team_public_model_name``; when that + name resolves to nothing but the team deployment, the probe must run rather + than 403 as if the key were out of scope. + """ + probed = await _live_probed_model_ids( + _TEAM_ONLY_MODEL_LIST, + UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-nova"], team_id="team-b"), + model="bedrock-nova", + ) + + assert probed == {"id-team-b"} + + +@pytest.mark.asyncio +async def test_health_endpoint_returns_a_team_only_deployment_by_its_public_name_on_background_cache_path(): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _TEAM_ONLY_MODEL_LIST, + _router_for(_TEAM_ONLY_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_TEAM_CACHED_RESULTS, + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-nova"], team_id="team-b"), + model="bedrock-nova", + model_id=None, + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-team-b"] + + +@pytest.mark.asyncio +async def test_health_endpoint_targets_both_deployments_behind_a_shared_public_name_on_background_cache_path(): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _TEAM_MODEL_LIST, + _router_for(_TEAM_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_TEAM_CACHED_RESULTS, + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-nova"], team_id="team-b"), + model="bedrock-nova", + model_id=None, + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-bedrock", "id-team-b"] + + def test_health_test_connection_keeps_error_and_raw_request_through_the_allowlist(monkeypatch): """ The dashboard's Test Connect button reads ``result.error`` and diff --git a/tests/test_litellm/proxy/test_health_check_functions.py b/tests/test_litellm/proxy/test_health_check_functions.py index fdae11d517a..137eda223d1 100644 --- a/tests/test_litellm/proxy/test_health_check_functions.py +++ b/tests/test_litellm/proxy/test_health_check_functions.py @@ -623,6 +623,35 @@ async def test_perform_health_check_and_save_forwards_skip_disabled_background_f assert call_kwargs["health_check_skip_disabled_background_models"] is True +@pytest.mark.asyncio +async def test_perform_health_check_narrows_to_a_team_deployment_by_its_public_name(): + """``/health?model=`` must probe the team deployment, not an empty list.""" + from litellm.proxy.health_check import perform_health_check + + team_deployment = { + "model_name": "bedrock-nova_team-b_9f2c", + "litellm_params": {"model": "bedrock/us.amazon.nova-2-lite-v1:0"}, + "model_info": {"id": "id-team-b", "team_id": "team-b", "team_public_model_name": "bedrock-nova"}, + } + other_deployment = { + "model_name": "gpt-5.4-mini", + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "id-openai"}, + } + probe = AsyncMock(return_value=([{"model": "bedrock/us.amazon.nova-2-lite-v1:0", "model_id": "id-team-b"}], [], {})) + + with patch( # test-quality-ok: the deployments handed to the probe are the assertion; no injection seam + "litellm.proxy.health_check._perform_health_check", probe + ): + healthy, unhealthy, _ = await perform_health_check( + model_list=[team_deployment, other_deployment], model="bedrock-nova" + ) + + assert [m["model_info"]["id"] for m in probe.call_args.args[0]] == ["id-team-b"] + assert [ep["model_id"] for ep in healthy] == ["id-team-b"] + assert unhealthy == [] + + def test_parse_background_health_check_model_groups_unset_returns_none(): from litellm.proxy.health_check import parse_background_health_check_model_groups From 90ac77e58e32eeabb12b8df28144bf2ebae28dee Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 5 Sep 2026 17:25:02 -0700 Subject: [PATCH 007/119] test(e2e): add JWT issuer harness and first live JWT auth tests Add tests/e2e/jwt_issuer.py, a test-only RS256 issuer that serves a JWKS document and signs arbitrary claims over an open loopback-only /token endpoint, so e2e tests can mint tokens without ever holding a signing key. One key per process keeps the proxy's cached JWKS valid for the whole run. Add the first five live JWT auth tests in tests/e2e/other: a valid token for an existing team is accepted and its spend row carries the claimed team and user, a tampered signature and an expired token are refused with 401, a token naming a team that does not exist is refused with 403, and a plain sk- virtual key keeps working with enable_jwt_auth on. Harness unit tests cover the issuer itself. Team and user create/delete land on the shared ProxyClient (warn-only teardown, /user/delete typed as the int it returns) instead of a fourth per-suite copy. Document the issuer command, the E2E_JWT_ISSUER_PORT convention (default 4190), JWT_PUBLIC_KEY_URL and the litellm_jwtauth config block in tests/e2e/CONTRIBUTING.md, and register the new cells in other.yaml. --- tests/e2e/CLAUDE.md | 2 +- tests/e2e/CONTRIBUTING.md | 21 ++- tests/e2e/coverage_registry/other.yaml | 9 +- tests/e2e/e2e_config.py | 4 + tests/e2e/e2e_http.py | 21 +++ tests/e2e/jwt_issuer.py | 230 +++++++++++++++++++++++++ tests/e2e/models.py | 12 ++ tests/e2e/other/other_client.py | 34 +++- tests/e2e/other/test_jwt_auth_e2e.py | 129 ++++++++++++++ tests/e2e/proxy_client.py | 41 +++++ tests/e2e/test_jwt_issuer.py | 141 +++++++++++++++ 11 files changed, 633 insertions(+), 11 deletions(-) create mode 100644 tests/e2e/jwt_issuer.py create mode 100644 tests/e2e/other/test_jwt_auth_e2e.py create mode 100644 tests/e2e/test_jwt_issuer.py diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 89c04208d65..4985da9e95d 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -19,7 +19,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (fallbacks, cooldowns) - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What remains here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`) and markerless harness unit tests for the Locust/session-anomaly aggregation logic -- `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite +- `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (RS256 tokens minted by the test-only issuer in `jwt_issuer.py`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at; see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke - `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index b270feb820e..ad414bb6647 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -27,14 +27,31 @@ The suites run against a live proxy, so bring one up first by running the litell 2. Bring up a Postgres and a Redis for the proxy to use. The repo-root `docker-compose.yml` already defines a Postgres on `5432`; a `docker run -p 6379:6379 redis:7` covers Redis. Point `DATABASE_URL` / `REDIS_HOST` / `REDIS_PORT` at whatever you run. Tests that read Redis directly default to the deployed shape (TLS + cluster mode) whenever `REDIS_HOST` is set, so for a local standalone Redis also set `REDIS_CLUSTER=false` and `REDIS_SSL=false` (plus `REDIS_PASSWORD` when your Redis requires auth) -3. Start the litellm proxy locally against your config and confirm it is live: +3. Start the test-only JWT issuer the `other/` suite mints tokens from, then the litellm proxy against your config, and confirm both are live. The issuer is a fake identity provider (`jwt_issuer.py`) with an open mint endpoint, so it binds loopback only; it generates one RSA key per process and serves it as a JWKS, and the proxy caches that JWKS for `public_key_ttl` (600s) without refetching on an unknown `kid`, so restart the proxy whenever you restart the issuer: ```bash set -a && source .env && set +a - litellm --config .yml --port 4000 + uv run python tests/e2e/jwt_issuer.py & + curl -fs http://127.0.0.1:4190/.well-known/jwks.json + JWT_PUBLIC_KEY_URL=http://127.0.0.1:4190/.well-known/jwks.json litellm --config .yml --port 4000 curl -fs http://localhost:4000/health/liveliness ``` + The issuer's port is `E2E_JWT_ISSUER_PORT` (default `4190`, the port in the URLs above), read by both the issuer process and the tests, so set it in one place if you change it. JWT auth is an enterprise feature, so the proxy also needs `LITELLM_LICENSE` in its environment, and its config needs the JWT block below. `enable_jwt_auth` only routes bearer tokens with three dot-separated segments into the JWT path, so `sk-` virtual keys and the master key keep working for every other suite. `proxy_batch_write_at` is lowered so the JWT spend-attribution test sees its row well inside the poll deadline: + + ```yaml + general_settings: + proxy_batch_write_at: 5 + enable_jwt_auth: true + litellm_jwtauth: + user_id_jwt_field: sub + user_email_jwt_field: email + team_ids_jwt_field: groups + user_id_upsert: true + ``` + + Leave `JWT_AUDIENCE` and `JWT_ISSUER` unset unless you also put matching `aud` / `iss` claims in the tokens the tests mint; the issuer sets `iss` to its own base URL + 4. Run a suite against it; the harness reads `LITELLM_PROXY_URL` (default `http://localhost:4000`): ```bash diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index 814ebae2e0b..8b9cb4d2adf 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -9,9 +9,12 @@ - {id: other.auth.llm_chat.not_bearer_scheme_denied, module: other, tier: P0, area: auth, assertions: [not_bearer_scheme_denied], source: "vendor testing strategy §11.1 / LIT-4778", rationale: "NotBearer scheme on chat is 401/403"} - {id: other.auth.realtime.missing_header_denied, module: other, tier: P1, area: auth, assertions: [missing_header_denied], source: "vendor testing strategy §9.19 / LIT-4778", rationale: "Realtime client-secret and calls routes reject requests without Authorization"} - {id: other.config.responses.metadata_redis_ttl_bounded, module: other, tier: P0, area: config, assertions: [ttl_bounded], source: "responses + redis cache", rationale: "Responses store+metadata must not leave TTL-unbounded Redis entries (LIT-1201)"} -- {id: other.auth.jwt.valid_token_allows, module: other, tier: P0, area: auth, assertions: [valid_token_allows], source: "handle_jwt.py:77-150", rationale: "Valid JWT with correct issuer + claims grants access"} -- {id: other.auth.jwt.expired_denied, module: other, tier: P0, area: auth, assertions: [expired_denied], source: "handle_jwt.py:125-135", rationale: "Expired JWT rejected even with valid signature"} -- {id: other.auth.jwt.invalid_signature_denied, module: other, tier: P0, area: auth, assertions: [invalid_signature_denied], source: "handle_jwt.py:145-150", rationale: "Bad/missing signature fails verification"} +- {id: other.auth.jwt.valid_token_allows, module: other, tier: P0, area: auth, assertions: [valid_token_allows], source: "handle_jwt.py:1217-1256 auth_jwt / user_api_key_auth.py:1365-1377", rationale: "An RS256 JWT signed by the configured JWKS whose groups claim names an existing team is accepted on /chat/completions"} +- {id: other.auth.jwt.spend_attributed_to_claims, module: other, tier: P0, area: auth, assertions: [spend_attributed_to_claims], source: "handle_jwt.py:2224 auth_builder / user_api_key_auth.py:1438-1474", rationale: "The spend log row for a JWT-authenticated call carries the team_id from the groups claim and the user_id from sub, not a virtual key's identity. Single-group claim only: the proxy picks the team from a set, so attribution over several groups is unordered"} +- {id: other.auth.jwt.expired_denied, module: other, tier: P0, area: auth, assertions: [expired_denied], source: "handle_jwt.py:1244-1250", rationale: "Expired JWT rejected 401 (Token Expired) even with valid signature; leeway is 0"} +- {id: other.auth.jwt.invalid_signature_denied, module: other, tier: P0, area: auth, assertions: [invalid_signature_denied], source: "handle_jwt.py:1158-1166 _decode_jwt_with_public_key", rationale: "A genuine token whose signature bytes were altered fails verification with 401"} +- {id: other.auth.jwt.unknown_team_denied, module: other, tier: P0, area: auth, assertions: [unknown_team_denied], source: "handle_jwt.py:1549-1632 find_team_with_model_access", rationale: "A verified JWT whose groups claim resolves to no existing team is denied with 403 naming the unresolved team, never silently admitted without a team. The proxy words it as a model-access denial, the same body an existing team without model access gets"} +- {id: other.auth.jwt.virtual_key_unaffected, module: other, tier: P0, area: auth, assertions: [virtual_key_unaffected], source: "handle_jwt.py:213 is_jwt / user_api_key_auth.py:1332-1333", rationale: "enable_jwt_auth only routes three-segment bearer tokens into the JWT branch, so sk- virtual keys keep working on the same proxy"} - {id: other.auth.model_access_group.wildcard_bare_name_allowed, module: other, tier: P0, area: auth, assertions: [wildcard_bare_name_allowed], source: "auth_checks.py:3232 / LIT-5813", fail_before_fix: proven, rationale: "A grant of a group holding a wildcard deployment covers the bare model names callers actually send, not only the provider-prefixed spelling"} - {id: other.auth.model_access_group.member_allowed, module: other, tier: P0, area: auth, assertions: [member_allowed], source: "auth_checks.py:3232", rationale: "A key whose allow-list is a model access group can call the deployments in that group"} - {id: other.auth.model_access_group.non_member_denied, module: other, tier: P0, area: auth, assertions: [non_member_denied], source: "auth_checks.py:3232", rationale: "That same grant reaches nothing outside the group, including provider models the group's wildcard does not cover"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 09d17b9a0db..0e7f4c0ad2b 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -15,6 +15,7 @@ from typing import Final from dotenv import load_dotenv from fixture_mode import deterministic_marker, parse_fixture_mode +from jwt_issuer import jwt_issuer_url from provider_edge import provider_edge_api_base # Local runs keep provider / DataDog keys in tests/e2e/.env (see CONTRIBUTING.md). @@ -43,6 +44,9 @@ UI_BASE_URL = os.environ.get("E2E_UI_BASE_URL", PROXY_BASE_URL).rstrip("/") CHEAP_ANTHROPIC_MODEL = os.environ.get("E2E_CHEAP_ANTHROPIC_MODEL", "claude-haiku-4-5") CHEAP_OPENAI_MODEL = os.environ.get("E2E_CHEAP_OPENAI_MODEL", "gpt-5.5") +# Test-only JWT issuer (jwt_issuer.py); the port is E2E_JWT_ISSUER_PORT (see CONTRIBUTING.md). +JWT_ISSUER_URL = jwt_issuer_url() + LINEAR_MCP_URL = os.environ.get("E2E_LINEAR_MCP_URL", "https://mcp.linear.app/mcp") LINEAR_STORAGE_STATE = os.environ.get("E2E_LINEAR_STORAGE_STATE", "") diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 9d5f1658e91..4c65c53b653 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -386,6 +386,27 @@ def get_external[R: BaseModel]( return _classify(resp, response_type) +def post_external[R: BaseModel]( + url: str, + *, + json: BaseModel, + response_type: type[R], + timeout: float = 30.0, +) -> Result[R]: + """POST an absolute URL outside the proxy (e.g. the e2e JWT issuer's mint + endpoint). Like get_external: no proxy base url, no proxy auth, and the same + tagged-union classification as every other call.""" + try: + resp = requests.post( + url, + json=json.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return _classify(resp, response_type) + + def delete[R: BaseModel]( url: URL, *, diff --git a/tests/e2e/jwt_issuer.py b/tests/e2e/jwt_issuer.py new file mode 100644 index 00000000000..4006fcaa314 --- /dev/null +++ b/tests/e2e/jwt_issuer.py @@ -0,0 +1,230 @@ +"""Test-only fake identity provider for the e2e JWT suite. Never deploy it. + +Run it next to the proxy (`uv run python tests/e2e/jwt_issuer.py`). On start it +generates one RSA signing key and keeps it for the life of the process, serving +the public half at `GET /.well-known/jwks.json` and signing whatever claims are +POSTed to `/token`. The proxy's `JWT_PUBLIC_KEY_URL` points at the JWKS URL, and +tests mint RS256 tokens by POSTing claims, so the private key never leaves this +process and no test holds it. + +One key per process, rather than per pytest run, is what survives the proxy's +JWKS cache: the proxy caches the JWKS for `litellm_jwtauth.public_key_ttl` +(600s by default) and does not refetch on an unknown `kid`, so a key rotated +every run would be rejected until the cache expired. Restart the proxy whenever +you restart the issuer. + +The mint endpoint takes no credential: anyone who can reach it gets a token the +proxy trusts. It therefore binds 127.0.0.1 only, must never be exposed beyond +loopback, and must only ever be trusted by a proxy under test. +""" + +from __future__ import annotations + +import logging +import os +import threading +import time +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Final, Literal + +import jwt +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from jwt.utils import to_base64url_uint +from pydantic import BaseModel, RootModel, ValidationError + +JWT_ISSUER_PORT_ENV: Final = "E2E_JWT_ISSUER_PORT" +DEFAULT_JWT_ISSUER_PORT: Final = 4190 +LOOPBACK_HOST: Final = "127.0.0.1" +JWKS_PATH: Final = "/.well-known/jwks.json" +TOKEN_PATH: Final = "/token" +DEFAULT_TOKEN_LIFETIME_SECONDS: Final = 300 + +ClaimValue = str | int | float | bool | None | list[str] + + +def jwt_issuer_port() -> int: + raw: Final = os.environ.get(JWT_ISSUER_PORT_ENV, "").strip() + return int(raw) if raw else DEFAULT_JWT_ISSUER_PORT + + +def jwt_issuer_url() -> str: + return f"http://{LOOPBACK_HOST}:{jwt_issuer_port()}" + + +class RsaJwk(BaseModel): + kty: Literal["RSA"] = "RSA" + alg: Literal["RS256"] = "RS256" + use: Literal["sig"] = "sig" + kid: str + n: str + e: str + + +class JwksDocument(BaseModel): + keys: tuple[RsaJwk, ...] + + +class TokenRequest(RootModel[Mapping[str, ClaimValue]]): + """The JSON body of POST /token: the claims to sign, verbatim. Nested objects + are not supported; every value is a scalar or a list of strings.""" + + +class MintedToken(BaseModel): + token: str + + +class IssuerError(BaseModel): + error: str + + +@dataclass(frozen=True, slots=True) +class SigningKey: + kid: str + private_pem: str + jwk: RsaJwk + + +def generate_signing_key(kid: str | None = None) -> SigningKey: + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + numbers: Final = private_key.public_key().public_numbers() + resolved_kid: Final = kid if kid is not None else uuid.uuid4().hex + return SigningKey( + kid=resolved_kid, + private_pem=private_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ).decode(), + jwk=RsaJwk( + kid=resolved_kid, + n=to_base64url_uint(numbers.n).decode(), + e=to_base64url_uint(numbers.e).decode(), + ), + ) + + +def mint( + key: SigningKey, + claims: Mapping[str, ClaimValue], + *, + issuer: str, + now: int, + lifetime_seconds: int = DEFAULT_TOKEN_LIFETIME_SECONDS, +) -> str: + """Sign `claims` as a compact RS256 JWT carrying `key.kid` in its header. + `iss`, `iat`, and `exp` are filled in when absent and left alone when the + caller sets them, so a test can mint an already-expired token.""" + payload: Final[dict[str, ClaimValue]] = { + "iss": issuer, + "iat": now, + "exp": now + lifetime_seconds, + **claims, + } + return jwt.encode(payload, key.private_pem, algorithm="RS256", headers={"kid": key.kid}) + + +class _IssuerServer(ThreadingHTTPServer): + daemon_threads = True + + def __init__(self, bind: tuple[str, int], *, key: SigningKey, clock: Callable[[], int]) -> None: + super().__init__(bind, _IssuerHandler) + self.key: Final = key + self.clock: Final = clock + + @property + def url(self) -> str: + host, port = self.server_address[0], self.server_address[1] + return f"http://{host}:{port}" + + +class _IssuerHandler(BaseHTTPRequestHandler): + def _issuer(self) -> _IssuerServer: + issuer: Final = self.server + assert isinstance(issuer, _IssuerServer) + return issuer + + def do_GET(self) -> None: + if self.path != JWKS_PATH: + self._send(404, IssuerError(error=f"unknown path {self.path}; the JWKS is at {JWKS_PATH}")) + return + self._send(200, JwksDocument(keys=(self._issuer().key.jwk,))) + + def do_POST(self) -> None: + if self.path != TOKEN_PATH: + self._send(404, IssuerError(error=f"unknown path {self.path}; mint tokens at {TOKEN_PATH}")) + return + length: Final = int(self.headers.get("Content-Length", "0")) + try: + request: Final = TokenRequest.model_validate_json(self.rfile.read(length)) + except ValidationError as exc: + self._send(400, IssuerError(error=f"claims must be a JSON object of scalar or string-list values: {exc}")) + return + issuer: Final = self._issuer() + token: Final = mint(issuer.key, request.root, issuer=issuer.url, now=issuer.clock()) + self._send(200, MintedToken(token=token)) + + def _send(self, status: int, body: BaseModel) -> None: + payload: Final = body.model_dump_json().encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + +@dataclass(frozen=True, slots=True) +class RunningIssuer: + url: str + key: SigningKey + server: _IssuerServer + + @property + def jwks_url(self) -> str: + return f"{self.url}{JWKS_PATH}" + + def shutdown(self) -> None: + self.server.shutdown() + self.server.server_close() + + +def _wall_clock() -> int: + return int(time.time()) + + +def start_jwt_issuer( + *, + port: int = 0, + key: SigningKey | None = None, + clock: Callable[[], int] = _wall_clock, +) -> RunningIssuer: + """Serve the issuer on loopback in a daemon thread. `port=0` takes an + OS-assigned port for in-process tests; the CLI passes the documented one.""" + server: Final = _IssuerServer((LOOPBACK_HOST, port), key=key or generate_signing_key(), clock=clock) + thread: Final = threading.Thread(target=server.serve_forever, name="e2e-jwt-issuer", daemon=True) + thread.start() + return RunningIssuer(url=server.url, key=server.key, server=server) + + +def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(message)s") + running: Final = start_jwt_issuer(port=jwt_issuer_port()) + logging.getLogger(__name__).info( + "e2e jwt issuer listening on %s jwks=%s mint=POST %s%s kid=%s (test-only, loopback, ctrl-c to stop)", + running.url, + running.jwks_url, + running.url, + TOKEN_PATH, + running.key.kid, + ) + try: + threading.Event().wait() + except KeyboardInterrupt: + running.shutdown() + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 2c6c0e9bbd4..17f5418bd73 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1078,6 +1078,18 @@ class UserListResponse(BaseModel): total: int +class JwtClaimsBody(BaseModel): + """Claims POSTed to the e2e JWT issuer's /token: exactly what the + `litellm_jwtauth` block in CONTRIBUTING.md reads (sub -> user_id, email -> + user_email, groups -> team ids), plus an explicit `exp` for the expired case; + the issuer fills in iss/iat/exp when they are left unset.""" + + sub: str + email: str + groups: Sequence[str] + exp: int | None = None + + class OrgNewBody(BaseModel): organization_alias: str models: list[str] = [] diff --git a/tests/e2e/other/other_client.py b/tests/e2e/other/other_client.py index 1aa83ac42c7..b51f3338f52 100644 --- a/tests/e2e/other/other_client.py +++ b/tests/e2e/other/other_client.py @@ -1,19 +1,27 @@ """Client for the `other` holding-pen suite: the auth gate (master key vs an -invalid key on an admin route) and the process-lifecycle health probes -(liveness, public readiness, authenticated readiness diagnostics). +invalid key on an admin route), JWT auth against the test-only issuer +(jwt_issuer.py), and the process-lifecycle health probes (liveness, public +readiness, authenticated readiness diagnostics). Holds the shared ProxyClient so `resources` / `scoped_key` still clean up, and adds only the routes these behaviors need. The health probes deliberately send no auth header (public routes), so they go through the transport with an empty -headers model rather than a bearer. +headers model rather than a bearer. Tokens are minted by POSTing claims to the +issuer, so no test ever holds a signing key. """ from __future__ import annotations from dataclasses import dataclass +from typing import Final -from e2e_http import NoBody, ProbeResult, Result +import pytest + +from e2e_config import JWT_ISSUER_URL +from e2e_http import NetworkError, NoBody, ProbeResult, Result, Success, post_external +from jwt_issuer import TOKEN_PATH, MintedToken from models import ( + JwtClaimsBody, ReadinessDetailsResponse, ReadinessResponse, UserListParams, @@ -25,6 +33,7 @@ from proxy_client import ProxyClient @dataclass(frozen=True, slots=True) class OtherClient: proxy: ProxyClient + jwt_issuer_url: str def liveness(self) -> ProbeResult: """GET /health/liveliness. Unauthenticated; the probe returns status + @@ -68,6 +77,21 @@ class OtherClient: response_type=UserListResponse, ) + def mint_jwt(self, claims: JwtClaimsBody) -> str: + """Have the test-only issuer sign `claims` into a compact RS256 JWT. A + missing issuer is a hard failure naming the start command, not a skip.""" + result: Final = post_external(f"{self.jwt_issuer_url}{TOKEN_PATH}", json=claims, response_type=MintedToken) + match result: + case Success(data=minted): + return minted.token + case NetworkError(message=message): + pytest.fail( + f"No live JWT issuer at {self.jwt_issuer_url}: {message}. Start it next to the proxy with " + "`uv run python tests/e2e/jwt_issuer.py` (see CONTRIBUTING.md)" + ) + case _: + raise AssertionError(result) + def build_client(proxy: ProxyClient) -> OtherClient: - return OtherClient(proxy=proxy) + return OtherClient(proxy=proxy, jwt_issuer_url=JWT_ISSUER_URL) diff --git a/tests/e2e/other/test_jwt_auth_e2e.py b/tests/e2e/other/test_jwt_auth_e2e.py new file mode 100644 index 00000000000..49a9350a3a2 --- /dev/null +++ b/tests/e2e/other/test_jwt_auth_e2e.py @@ -0,0 +1,129 @@ +"""Live e2e: RS256 JWTs minted by the test-only issuer (jwt_issuer.py) against a +proxy running with `enable_jwt_auth: true` and the `litellm_jwtauth` block from +CONTRIBUTING.md (sub -> user_id, email -> user_email, groups -> team ids, +user_id_upsert). + +Every case mints through the issuer, so the tests never hold a signing key: the +bad-signature case corrupts a genuine signature, the expired case asks the +issuer for a token whose `exp` is already in the past. Those identities get +their own freshly created team so a rejection can only be blamed on the token, +while the unknown-team case names a team that was never created. An acceptance +is proven twice, at the boundary (200 from a real provider) and in the spend log +the proxy attributes to the claims. The last case keeps a plain `sk-` virtual +key working on the same proxy, guarding against the flag turning JWT on for +everyone. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Final + +import pytest + +from e2e_config import CHEAP_OPENAI_MODEL, unique_marker +from e2e_http import UnauthorizedError, UnknownApiError, unwrap +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, JwtClaimsBody, TeamNewBody +from other_client import OtherClient + +pytestmark = pytest.mark.e2e + + +@dataclass(frozen=True, slots=True) +class JwtIdentity: + user_id: str + team_id: str + + def claims(self, *, exp: int | None = None) -> JwtClaimsBody: + return JwtClaimsBody(sub=self.user_id, email=f"{self.user_id}@example.com", groups=(self.team_id,), exp=exp) + + +@pytest.fixture +def identity(client: OtherClient, resources: ResourceManager) -> JwtIdentity: + marker: Final = unique_marker() + team_id: Final = client.proxy.create_team( + TeamNewBody(team_alias=f"e2e-jwt-{marker}", team_id=f"e2e-jwt-team-{marker}") + ) + resources.defer(lambda: client.proxy.delete_team(team_id)) + user_id: Final = f"e2e-jwt-user-{marker}" + resources.defer(lambda: client.proxy.delete_user(user_id)) + return JwtIdentity(user_id=user_id, team_id=team_id) + + +def _ping() -> ChatBody: + return ChatBody( + model=CHEAP_OPENAI_MODEL, + messages=[ChatMessage(role="user", content=f"Reply with the single word pong. {unique_marker()}")], + max_tokens=16, + ) + + +def _corrupt_signature(token: str) -> str: + header, payload, signature = token.split(".") + flipped: Final = "A" if signature[10] != "A" else "B" + return f"{header}.{payload}.{signature[:10]}{flipped}{signature[11:]}" + + +class TestJwtAuth: + @pytest.mark.covers("other.auth.jwt.valid_token_allows", "other.auth.jwt.spend_attributed_to_claims") + def test_valid_token_for_an_existing_team_is_accepted_and_attributed( + self, client: OtherClient, identity: JwtIdentity + ) -> None: + token: Final = client.mint_jwt(identity.claims()) + + response: Final = unwrap(client.proxy.chat(token, _ping())) + assert response.id is not None and response.choices, ( + f"chat under a valid JWT returned no completion: {response}" + ) + + rows: Final = client.proxy.poll_logs_for_request_id(response.id) + assert rows, f"no spend log row for request {response.id} within the poll deadline" + row: Final = rows[0] + assert row.team_id == identity.team_id, ( + f"spend row must carry the team from the JWT groups claim {identity.team_id!r}, got {row.team_id!r}" + ) + assert row.user == identity.user_id, ( + f"spend row must carry the user from the JWT sub claim {identity.user_id!r}, got {row.user!r}" + ) + + @pytest.mark.covers("other.auth.jwt.invalid_signature_denied") + def test_tampered_signature_is_rejected(self, client: OtherClient, identity: JwtIdentity) -> None: + tampered: Final = _corrupt_signature(client.mint_jwt(identity.claims())) + + result: Final = client.proxy.chat(tampered, _ping()) + assert isinstance(result, UnauthorizedError), ( + f"a JWT whose signature does not verify must be rejected with 401, got {result}" + ) + assert "signature verification failed" in result.body.lower(), ( + f"the 401 must come from signature verification, not another auth failure, got {result.body[:300]}" + ) + + @pytest.mark.covers("other.auth.jwt.expired_denied") + def test_expired_token_is_rejected(self, client: OtherClient, identity: JwtIdentity) -> None: + expired: Final = client.mint_jwt(identity.claims(exp=1)) + + result: Final = client.proxy.chat(expired, _ping()) + assert isinstance(result, UnauthorizedError), ( + f"an expired JWT must be rejected with 401 even though its signature verifies, got {result}" + ) + assert "expired" in result.body.lower(), f"the 401 must say the token expired, got {result.body[:300]}" + + @pytest.mark.covers("other.auth.jwt.unknown_team_denied") + def test_token_naming_a_team_that_does_not_exist_is_rejected(self, client: OtherClient) -> None: + marker: Final = unique_marker() + never_created: Final = JwtIdentity(user_id=f"e2e-jwt-user-{marker}", team_id=f"e2e-jwt-missing-team-{marker}") + token: Final = client.mint_jwt(never_created.claims()) + + result: Final = client.proxy.chat(token, _ping()) + assert isinstance(result, UnknownApiError) and result.status_code == 403, ( + f"a valid JWT whose groups name no existing team must be rejected with 403, got {result}" + ) + assert never_created.team_id in result.body, ( + f"the 403 must name the team it could not resolve ({never_created.team_id}), got {result.body[:300]}" + ) + + @pytest.mark.covers("other.auth.jwt.virtual_key_unaffected") + def test_plain_virtual_key_still_works_with_jwt_auth_enabled(self, client: OtherClient, scoped_key: str) -> None: + response: Final = unwrap(client.proxy.chat(scoped_key, _ping())) + assert response.choices, f"an sk- key must keep working on a proxy with enable_jwt_auth, got {response}" diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index cdc20e5299a..5292f56b324 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -21,6 +21,7 @@ from e2e_http import ( Result, StreamingResponse, Success, + UnknownApiError, is_ok, unwrap, ) @@ -65,6 +66,11 @@ from models import ( SpendLogsPage, SpendLogsPageParams, SpendLogsParams, + TeamDeleteBody, + TeamNewBody, + TeamNewResponse, + UserDeleteBody, + UserDeleteResponse, ) from e2e_config import ( CONTROL_PLANE_BASE_URL, @@ -409,6 +415,41 @@ class ProxyClient: if not is_ok(result): warnings.warn(f"delete_credential({credential_name!r}) failed: {result}", stacklevel=2) + def create_team(self, body: TeamNewBody) -> str: + return unwrap( + self.transport.post( + "/team/new", + headers=self.transport.master, + json=body, + response_type=TeamNewResponse, + ) + ).team_id + + def delete_team(self, team_id: str) -> None: + result = self.transport.post( + "/team/delete", + headers=self.transport.master, + json=TeamDeleteBody(team_ids=[team_id]), + response_type=NoBody, + ) + if not is_ok(result): + warnings.warn(f"delete_team({team_id!r}) failed: {result}", stacklevel=2) + + def delete_user(self, user_id: str) -> None: + """Best-effort teardown; a 404 is not a leak, since JWT tests defer this for + a user the proxy only upserts after a successful auth.""" + result = self.transport.post( + "/user/delete", + headers=self.transport.master, + json=UserDeleteBody(user_ids=[user_id]), + response_type=UserDeleteResponse, + ) + match result: + case Success() | UnknownApiError(status_code=404): + return + case _: + warnings.warn(f"delete_user({user_id!r}) failed: {result}", stacklevel=2) + # ---- LLM calls ------------------------------------------------------ def chat(self, key: str, body: ChatBody) -> Result[ChatResponse]: diff --git a/tests/e2e/test_jwt_issuer.py b/tests/e2e/test_jwt_issuer.py new file mode 100644 index 00000000000..7acaaf13949 --- /dev/null +++ b/tests/e2e/test_jwt_issuer.py @@ -0,0 +1,141 @@ +"""Harness coverage for the test-only JWT issuer (jwt_issuer.py). + +No proxy and no ``e2e`` marker. The issuer is booted in-process on an +OS-assigned port with a fixed clock and driven over HTTP through +``e2e_http.post_external`` / ``get_external``, the same transport the live +suite uses, so what is pinned here is the contract the live JWT tests lean on: +a token minted at ``/token`` verifies against the key served at +``/.well-known/jwks.json`` under the ``kid`` in its header, ``iss``/``iat``/``exp`` +are filled in only when the caller left them out, and malformed claim bodies or +unknown paths are refused instead of signed. +""" + +from __future__ import annotations + +import time +from collections.abc import Iterator +from typing import Final + +import jwt +import pytest +from pydantic import BaseModel, RootModel + +from e2e_http import UnknownApiError, get_external, post_external, unwrap +from jwt_issuer import ( + DEFAULT_TOKEN_LIFETIME_SECONDS, + JWKS_PATH, + JWT_ISSUER_PORT_ENV, + TOKEN_PATH, + JwksDocument, + MintedToken, + RunningIssuer, + jwt_issuer_port, + start_jwt_issuer, +) + +FROZEN_NOW: Final = int(time.time()) + + +class _Claims(BaseModel): + sub: str + groups: tuple[str, ...] = () + exp: int | None = None + + +class _DecodedClaims(BaseModel): + sub: str + iss: str + iat: int + exp: int + groups: tuple[str, ...] = () + + +class _NotAnObject(RootModel[tuple[str, ...]]): + pass + + +@pytest.fixture(scope="module") +def issuer() -> Iterator[RunningIssuer]: + running: Final = start_jwt_issuer(clock=lambda: FROZEN_NOW) + yield running + running.shutdown() + + +def _mint(issuer: RunningIssuer, claims: _Claims) -> str: + return unwrap(post_external(f"{issuer.url}{TOKEN_PATH}", json=claims, response_type=MintedToken)).token + + +def _served_jwks(issuer: RunningIssuer) -> JwksDocument: + return unwrap(get_external(issuer.jwks_url, response_type=JwksDocument)) + + +def _decode(token: str, jwks: JwksDocument, *, verify_exp: bool = True) -> _DecodedClaims: + key: Final = jwt.PyJWK.from_json(jwks.keys[0].model_dump_json()) + decoded: Final = jwt.decode(token, key, algorithms=["RS256"], options={"verify_exp": verify_exp}) + return _DecodedClaims.model_validate(decoded) + + +class TestJwtIssuer: + def test_minted_token_verifies_against_the_served_jwks(self, issuer: RunningIssuer) -> None: + token: Final = _mint(issuer, _Claims(sub="alice", groups=("team-a",))) + jwks: Final = _served_jwks(issuer) + + assert len(jwks.keys) == 1 + assert jwt.get_unverified_header(token)["kid"] == jwks.keys[0].kid + claims: Final = _decode(token, jwks) + assert claims.sub == "alice" + assert claims.groups == ("team-a",) + assert claims.iss == issuer.url + assert claims.iat == FROZEN_NOW + assert claims.exp == FROZEN_NOW + DEFAULT_TOKEN_LIFETIME_SECONDS + + def test_a_token_signed_by_another_key_does_not_verify(self, issuer: RunningIssuer) -> None: + other: Final = start_jwt_issuer(clock=lambda: FROZEN_NOW) + try: + foreign_token: Final = _mint(other, _Claims(sub="alice")) + finally: + other.shutdown() + + with pytest.raises(jwt.InvalidSignatureError): + _decode(foreign_token, _served_jwks(issuer)) + + def test_an_explicit_exp_is_signed_as_given(self, issuer: RunningIssuer) -> None: + expired_at: Final = FROZEN_NOW - 60 + token: Final = _mint(issuer, _Claims(sub="alice", exp=expired_at)) + jwks: Final = _served_jwks(issuer) + + assert _decode(token, jwks, verify_exp=False).exp == expired_at + with pytest.raises(jwt.ExpiredSignatureError): + _ = _decode(token, jwks) + + def test_non_object_claims_are_refused(self, issuer: RunningIssuer) -> None: + result: Final = post_external( + f"{issuer.url}{TOKEN_PATH}", json=_NotAnObject(("not", "claims")), response_type=MintedToken + ) + assert isinstance(result, UnknownApiError) + assert result.status_code == 400 + + @pytest.mark.parametrize( + ("method", "path"), + [("GET", TOKEN_PATH), ("POST", JWKS_PATH), ("GET", "/token/anything")], + ids=["get-token", "post-jwks", "get-other"], + ) + def test_unknown_routes_are_404(self, issuer: RunningIssuer, method: str, path: str) -> None: + url: Final = f"{issuer.url}{path}" + result: Final = ( + get_external(url, response_type=MintedToken) + if method == "GET" + else post_external(url, json=_Claims(sub="alice"), response_type=MintedToken) + ) + assert isinstance(result, UnknownApiError) + assert result.status_code == 404 + + +class TestIssuerPort: + def test_defaults_to_the_documented_port(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv(JWT_ISSUER_PORT_ENV, raising=False) + assert jwt_issuer_port() == 4190, "CONTRIBUTING.md hardcodes 4190 in JWT_PUBLIC_KEY_URL" + + def test_env_override_wins(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(JWT_ISSUER_PORT_ENV, " 4321 ") + assert jwt_issuer_port() == 4321 From 8c5b29519e6609d743913bc8d3c85070fc5cdf16 Mon Sep 17 00:00:00 2001 From: Jon Walton Date: Wed, 9 Sep 2026 18:15:53 +0800 Subject: [PATCH 008/119] fix(proxy): emit internal user budget alerts --- .../SlackAlerting/slack_alerting.py | 2 +- litellm/proxy/auth/auth_checks.py | 13 ++ .../SlackAlerting/test_slack_alerting.py | 15 ++ .../proxy/auth/test_auth_checks.py | 131 +++++++++++++++++- 4 files changed, 158 insertions(+), 3 deletions(-) diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index dc41c7dadc8..caac8e888fd 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -673,7 +673,7 @@ class SlackAlerting(CustomBatchLogger): Create a standard message for a budget alert """ _all_fields_as_dict: Final[dict[str, object]] = user_info.model_dump(exclude_none=True) - _all_fields_as_dict.pop("token") + _all_fields_as_dict.pop("token", None) msg = "" for k, v in _all_fields_as_dict.items(): if isinstance(v, Litellm_EntityType): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index dc693317de0..d26474778d9 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1020,6 +1020,19 @@ async def common_checks( fallback_spend=user_object.spend or 0.0, max_budget=user_budget, ) + call_info: Final = CallInfo( + spend=user_spend, + max_budget=user_budget, + user_id=user_object.user_id, + user_email=user_object.user_email, + event_group=Litellm_EntityType.USER, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="user_budget", + user_info=call_info, + ) + ) if math.isfinite(user_budget) and user_spend >= user_budget: raise litellm.BudgetExceededError( current_cost=user_spend, diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py index 55e2dcdc270..44dda57dd27 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py @@ -45,6 +45,21 @@ class TestSlackAlerting(unittest.TestCase): result = self.slack_alerting._get_percent_of_max_budget_left(user_info) self.assertEqual(result, -0.2) + def test_get_user_info_str_omits_absent_token_for_user_alert(self): + user_info = CallInfo( + spend=85.0, + max_budget=100.0, + user_id="user-1", + user_email="person@example.com", + event_group=Litellm_EntityType.USER, + ) + + result = self.slack_alerting._get_user_info_str(user_info) + + self.assertIn("*user_id:* `user-1`", result) + self.assertIn("*user_email:* `person@example.com`", result) + self.assertNotIn("*token:*", result) + def test_get_event_and_event_message_max_budget(self): # Initial setup with no event event = None diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 2284a05b2e9..1fc2e0b3dc1 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -1,7 +1,7 @@ import asyncio import json from types import SimpleNamespace -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Final, Literal, Optional from unittest.mock import AsyncMock, MagicMock, patch if TYPE_CHECKING: @@ -29,6 +29,7 @@ from litellm.proxy._types import ( ProxyException, SSOUserDefinedValues, UserAPIKeyAuth, + WebhookEvent, ) from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, @@ -52,6 +53,7 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache +from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.constants import ( DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, END_USER_RESTRICTED_REGISTRY_MAX_SIZE, @@ -5375,6 +5377,127 @@ async def test_common_checks_personal_user_budget_blocks_in_gather(): assert "User=u1" in str(over.value) +async def _run_internal_user_budget_alert( + *, + spend: float, +) -> tuple[AsyncMock, litellm.BudgetExceededError | None]: + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + user: Final = LiteLLM_UserTable( + user_id="user-1", + user_email="person@example.com", + spend=0.0, + max_budget=100.0, + ) + token: Final = UserAPIKeyAuth(token="hashed-key-1", user_id="user-1") + slack_alerting: Final = SlackAlerting(alerting=["webhook"]) + send_alert: Final = AsyncMock() + alert_finished: Final = asyncio.Event() + + async def _get_spend( + counter_key: str, + fallback_spend: float, + max_budget: float | None = None, + **kwargs: object, + ) -> float: + assert counter_key == "spend:user:user-1" + assert fallback_spend == 0.0 + assert max_budget == 100.0 + return spend + + async def _budget_alerts( + *, + type: Literal["user_budget"], + user_info: CallInfo, + ) -> None: + try: + await slack_alerting.budget_alerts(type=type, user_info=user_info) + finally: + alert_finished.set() + + proxy_logging_obj: Final = MagicMock(budget_alerts=_budget_alerts) + + async def _check() -> bool: + return await common_checks( + request_body={"messages": [{"role": "user", "content": "hi"}]}, + team_object=None, + user_object=user, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=proxy_logging_obj, + valid_token=token, + request=MagicMock(spec=Request), + ) + + async def _check_for_error() -> litellm.BudgetExceededError | None: + if spend < 100.0: + assert await _check() is True + return None + + with pytest.raises(litellm.BudgetExceededError) as raised: + await _check() + return raised.value + + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: common_checks has no database seam + patch("litellm.proxy.proxy_server.get_current_spend", _get_spend), # test-quality-ok: common_checks imports it locally + patch.object(slack_alerting, "send_alert", send_alert), + ): + error: Final = await _check_for_error() + await asyncio.wait_for(alert_finished.wait(), timeout=1.0) + + return send_alert, error + + +@pytest.mark.asyncio +async def test_common_checks_internal_user_budget_below_threshold_does_not_emit_alert(): + send_alert, error = await _run_internal_user_budget_alert(spend=84.0) + + assert error is None + send_alert.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_common_checks_internal_user_budget_emits_user_threshold_event(): + send_alert, error = await _run_internal_user_budget_alert(spend=85.0) + + assert error is None + send_alert.assert_awaited_once() + event: Final = send_alert.await_args.kwargs["user_info"] + assert isinstance(event, WebhookEvent) + assert event.event == "threshold_crossed" + assert event.event_group == Litellm_EntityType.USER + assert event.user_id == "user-1" + assert event.user_email == "person@example.com" + assert event.spend == 85.0 + assert event.max_budget == 100.0 + assert event.token is None + assert event.key_alias is None + assert event.team_id is None + assert event.organization_id is None + + +@pytest.mark.asyncio +async def test_common_checks_internal_user_budget_emits_crossed_event_and_rejects(): + send_alert, error = await _run_internal_user_budget_alert(spend=100.0) + + assert error is not None + assert error.current_cost == 100.0 + assert error.max_budget == 100.0 + send_alert.assert_awaited_once() + event: Final = send_alert.await_args.kwargs["user_info"] + assert isinstance(event, WebhookEvent) + assert event.event == "budget_crossed" + assert event.event_group == Litellm_EntityType.USER + assert event.user_id == "user-1" + assert event.user_email == "person@example.com" + + @pytest.mark.asyncio async def test_common_checks_personal_user_budget_skipped_for_team_key(): """A user's personal max_budget does not apply to a team-scoped key. @@ -5398,6 +5521,9 @@ async def test_common_checks_personal_user_budget_skipped_for_team_key(): async def _no_membership(*args, **kwargs): return None + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + with ( patch("litellm.proxy.proxy_server.prisma_client", None), patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), @@ -5412,11 +5538,12 @@ async def test_common_checks_personal_user_budget_skipped_for_team_key(): general_settings={}, route="/chat/completions", llm_router=None, - proxy_logging_obj=MagicMock(), + proxy_logging_obj=proxy_logging_obj, valid_token=token, request=MagicMock(spec=Request), ) assert result is True + proxy_logging_obj.budget_alerts.assert_not_awaited() @pytest.mark.asyncio From 6d6659a81bd1b6d92c5e8c43383a37b95711aa68 Mon Sep 17 00:00:00 2001 From: Jon Walton Date: Wed, 9 Sep 2026 18:39:36 +0800 Subject: [PATCH 009/119] test(proxy): harden user budget alert coverage --- .../proxy/auth/test_auth_checks.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 1fc2e0b3dc1..f9e25e9beaf 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5356,6 +5356,9 @@ async def test_common_checks_personal_user_budget_blocks_in_gather(): async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): return 999.0 if counter_key == "spend:user:u1" else 0.0 + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + with ( patch("litellm.proxy.proxy_server.prisma_client", None), patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), @@ -5370,10 +5373,11 @@ async def test_common_checks_personal_user_budget_blocks_in_gather(): general_settings={}, route="/chat/completions", llm_router=None, - proxy_logging_obj=MagicMock(), + proxy_logging_obj=proxy_logging_obj, valid_token=token, request=MagicMock(spec=Request), ) + await asyncio.sleep(0) assert "User=u1" in str(over.value) @@ -5412,6 +5416,7 @@ async def _run_internal_user_budget_alert( type: Literal["user_budget"], user_info: CallInfo, ) -> None: + assert type == "user_budget" try: await slack_alerting.budget_alerts(type=type, user_info=user_info) finally: @@ -5568,6 +5573,9 @@ async def test_common_checks_personal_user_budget_enforced_on_team_key_when_flag async def _no_membership(*args, **kwargs): return None + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + with ( patch("litellm.proxy.proxy_server.prisma_client", None), patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), @@ -5583,10 +5591,11 @@ async def test_common_checks_personal_user_budget_enforced_on_team_key_when_flag general_settings={"apply_user_budget_to_team_keys": True}, route="/chat/completions", llm_router=None, - proxy_logging_obj=MagicMock(), + proxy_logging_obj=proxy_logging_obj, valid_token=token, request=MagicMock(spec=Request), ) + await asyncio.sleep(0) assert "ExceededBudget: User=u1" in str(exc_info.value) @@ -5603,6 +5612,9 @@ async def test_common_checks_personal_user_budget_still_enforced_on_personal_key async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): return 999.0 if counter_key == "spend:user:u1" else 0.0 + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + with ( patch("litellm.proxy.proxy_server.prisma_client", None), patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), @@ -5617,10 +5629,11 @@ async def test_common_checks_personal_user_budget_still_enforced_on_personal_key general_settings={"apply_user_budget_to_team_keys": True}, route="/chat/completions", llm_router=None, - proxy_logging_obj=MagicMock(), + proxy_logging_obj=proxy_logging_obj, valid_token=token, request=MagicMock(spec=Request), ) + await asyncio.sleep(0) @pytest.mark.parametrize( From 058ff8c63c222d52392814b0fda84d5b030a84ec Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 9 Sep 2026 15:33:44 -0700 Subject: [PATCH 010/119] test(e2e): keep answering while every Redis command times out Add tests/e2e/router/test_redis_timeout_e2e.py against a proxy booted from tests/e2e/gateway/redis_timeout_ci_config.yml: a real Redis with socket_timeout 0.001, so every command times out and the circuit breaker opens, plus a primary deployment that always fails and falls back to a healthy one, so every request carries retry breadcrumbs into cost tracking. The test drives twenty chat requests through the proxy and asserts each answers within ten seconds, the last third is no slower than the first, /health/liveliness stays fast, and every request still reaches the spend log. Gate it behind the redis_timeout marker and E2E_REDIS_TIMEOUT, exclude it from the per-PR e2e-changed selector, register the reliability.circuit_breaker.redis_timeout.stays_responsive cell, and run it as its own job in the weekly load anomaly workflow with a Postgres and Valkey service. Against a v1.100.0 proxy the run wedges the worker: requests time out and liveliness stops answering (LIT-6780). Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_014ZDULyJPp17ZFiJenRxs2T --- .github/e2e-stack/select_tests.py | 1 + .github/workflows/weekly_load_anomaly.yml | 84 +++++++++++++++++ tests/e2e/CLAUDE.md | 2 +- tests/e2e/CONTRIBUTING.md | 2 +- tests/e2e/conftest.py | 9 +- tests/e2e/coverage_registry/reliability.yaml | 1 + tests/e2e/e2e_config.py | 5 +- tests/e2e/gateway/redis_timeout_ci_config.yml | 29 ++++++ tests/e2e/pytest.ini | 1 + tests/e2e/router/conftest.py | 21 +++-- tests/e2e/router/test_redis_timeout_e2e.py | 93 +++++++++++++++++++ 11 files changed, 234 insertions(+), 14 deletions(-) create mode 100644 tests/e2e/gateway/redis_timeout_ci_config.yml create mode 100644 tests/e2e/router/test_redis_timeout_e2e.py diff --git a/.github/e2e-stack/select_tests.py b/.github/e2e-stack/select_tests.py index a62358f81ff..10f26bd3b8a 100644 --- a/.github/e2e-stack/select_tests.py +++ b/.github/e2e-stack/select_tests.py @@ -8,6 +8,7 @@ UNSUPPORTED: Final = re.compile( r"|^tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e\.py$" r"|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$" r"|^tests/e2e/guardrails/test_presidio_masking_e2e\.py$" + r"|^tests/e2e/router/test_redis_timeout_e2e\.py$" ) HARNESS: Final = re.compile( r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$" diff --git a/.github/workflows/weekly_load_anomaly.yml b/.github/workflows/weekly_load_anomaly.yml index 3e1fca89645..6e564cf7e78 100644 --- a/.github/workflows/weekly_load_anomaly.yml +++ b/.github/workflows/weekly_load_anomaly.yml @@ -83,3 +83,87 @@ jobs: - name: Show proxy log on failure if: failure() run: tail -n 300 proxy.log + + redis-timeout-e2e: + if: github.event_name != 'schedule' || github.repository == 'BerriAI/litellm' + runs-on: ubuntu-latest + timeout-minutes: 30 + services: + postgres: + image: postgres:16.6 + env: + POSTGRES_USER: llmproxy + POSTGRES_PASSWORD: dbpassword9090 + POSTGRES_DB: litellm + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U llmproxy" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + valkey: + image: valkey/valkey:8.1.4@sha256:81db6d39e1bba3b3ff32bd3a1b19a6d69690f94a3954ec131277b9a26b95b3aa + ports: + - 6379:6379 + options: >- + --health-cmd "valkey-cli ping" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + env: + DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm + LITELLM_MASTER_KEY: sk-redis-timeout-e2e + LITELLM_LOG: WARNING + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Set up uv + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Cache the Rust build + uses: ./.github/actions/cache-cargo-build + + - name: Install dependencies + run: | + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy + + - name: Cache Prisma binaries + uses: ./.github/actions/cache-prisma-binaries + + - name: Generate Prisma client + run: | + uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma + + - name: Start the proxy with a Redis that times out on every command + run: | + nohup uv run --no-sync litellm --config tests/e2e/gateway/redis_timeout_ci_config.yml --port 4000 > proxy.log 2>&1 & + for _ in $(seq 1 90); do + if curl -fs http://localhost:4000/health/liveliness > /dev/null; then + exit 0 + fi + sleep 2 + done + echo "proxy never became live" + tail -n 100 proxy.log + exit 1 + + - name: Run the Redis timeout e2e test + env: + E2E_REDIS_TIMEOUT: "1" + LITELLM_PROXY_URL: http://localhost:4000 + run: | + uv run --no-sync pytest tests/e2e/router/test_redis_timeout_e2e.py -v --tb=short -rA + + - name: Show proxy log on failure + if: failure() + run: tail -n 300 proxy.log diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 89c04208d65..53342f66b67 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -17,7 +17,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below); plus the gateway-managed OAuth (authorization_code) path exercised through `/chat/completions`, the one behavior Datadog's static-header auth cannot reach, seeding the per-user upstream token via the interactive authorize dance driven with the mcp SDK's own OAuth client (headless-browser consent from a saved session) and asserting the completion lists and executes the server's tools with the stored per-user token - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection -- `router/` - routing and reliability behavior (fallbacks, cooldowns) +- `router/` - routing and reliability behavior (fallbacks, cooldowns). Also holds `test_redis_timeout_e2e.py`, which needs a proxy booted from `gateway/redis_timeout_ci_config.yml` (real Redis, `socket_timeout: 0.001`, so every command times out); it is marked `redis_timeout`, deselected unless `E2E_REDIS_TIMEOUT` is set, excluded from the per-PR check, and driven by `.github/workflows/weekly_load_anomaly.yml` - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What remains here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`) and markerless harness unit tests for the Locust/session-anomaly aggregation logic - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 1183096b81e..ccd27df6f46 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -65,7 +65,7 @@ A couple of logging destinations are configured on the proxy rather than by the ### The pull request check -Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. A change to the harness itself, meaning a root-level `tests/e2e/*.py` file or `pytest.ini`, `tests/e2e/gateway/`, `.github/e2e-stack/`, or the workflow, also runs the `access_control` suite as a canary, because those files have no test of their own that exercises the stack. `.github/e2e-stack/select_tests.py` applies both rules. The stack config at `tests/e2e/gateway/stage_mirror_ci_config.yml` must declare every model the selected suites use; a missing one shows up as a failed test id in the public log. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Jaeger, and TLS cluster-mode Valkey. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, deleted-file, and application-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories, `batches/test_managed_files_enforcement_e2e.py`, `llm_translation/realtime/test_realtime_pipecat_audio_e2e.py`, and `guardrails/test_presidio_masking_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack: the pipecat audio suite skips itself at import time unless the NLTK `punkt_tab` data is installed, and the presidio suite fails without the analyzer and anonymizer services this stack does not start +Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. A change to the harness itself, meaning a root-level `tests/e2e/*.py` file or `pytest.ini`, `tests/e2e/gateway/`, `.github/e2e-stack/`, or the workflow, also runs the `access_control` suite as a canary, because those files have no test of their own that exercises the stack. `.github/e2e-stack/select_tests.py` applies both rules. The stack config at `tests/e2e/gateway/stage_mirror_ci_config.yml` must declare every model the selected suites use; a missing one shows up as a failed test id in the public log. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Jaeger, and TLS cluster-mode Valkey. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, deleted-file, and application-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories, `batches/test_managed_files_enforcement_e2e.py`, `llm_translation/realtime/test_realtime_pipecat_audio_e2e.py`, `guardrails/test_presidio_masking_e2e.py`, and `router/test_redis_timeout_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack: the pipecat audio suite skips itself at import time unless the NLTK `punkt_tab` data is installed, and the presidio suite fails without the analyzer and anonymizer services this stack does not start, and the Redis timeout test needs a proxy whose Redis times out on every command (`gateway/redis_timeout_ci_config.yml`), which the weekly load workflow boots Every selected file must execute at least one passing test in each pass, and any test failure, collection error, or entirely skipped or deselected file fails the check. A file whose tests are all marked skip therefore cannot pass this check, so unskip at least one of them, or add the file to `UNSUPPORTED` in `select_tests.py` with the reason, before changing one. A failed pass stops the run. The public log prints pytest's one-line summary for each pass, including the rerun count, and names each failed or errored test as `classname::name`, so a retried network error or a failing test is visible without the raw output. The final `e2e-changed-tests` job succeeds only when no supported test files changed or the approved run completed all three passes. Fork PRs with selected tests fail this gate until a maintainer brings the reviewed change onto a same-repository branch diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index e1b987cbfd9..ff3d2dcc6bd 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -20,16 +20,14 @@ from datetime import datetime, timezone import pytest import requests - from e2e_config import CONTROL_PLANE_BASE_URL, FIXTURE_DIR, FIXTURE_MODE_RAW, PROXY_BASE_URL from e2e_db import RESET_OPT_IN_ENV, reset_spend_logs, run_spend_log_cleanup from fixture_mode import fixture_mode_collection_error, fixture_report_lines -from provider_edge import replay_leftover_error from junit_properties import attach_result_properties from lifecycle import ProxyClientProvider, ResourceManager +from provider_edge import replay_leftover_error from proxy_client import ProxyClient, build_proxy_client - _E2E_TEST_RAN = pytest.StashKey[bool]() _CALL_PASSED = pytest.StashKey[bool]() @@ -60,6 +58,11 @@ def pytest_configure(config: pytest.Config) -> None: "markers", "managed_files: needs a proxy running with require_managed_files enabled; deselected unless E2E_MANAGED_FILES_STACK is set", ) + config.addinivalue_line( + "markers", + "redis_timeout: needs a proxy booted from gateway/redis_timeout_ci_config.yml whose Redis times out on every " + "command; deselected unless E2E_REDIS_TIMEOUT is set", + ) def pytest_sessionstart(session: pytest.Session) -> None: diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index 6b69677d490..d62f0105931 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -31,6 +31,7 @@ - {id: reliability.cache.exact.returns_cached, module: reliability, tier: P1, behavior: cache, variant: exact, assertions: [returns_cached], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/caching.py", rationale: "Response cache returns cached on exact match"} - {id: reliability.cache.prompt_caching_model_select.returns_cached, module: reliability, tier: P1, behavior: cache, variant: prompt_caching_model_select, assertions: [returns_cached], exercised_on: [chat_completions], source: "router_utils/prompt_caching_cache.py", rationale: "Selects model supporting prompt caching for cacheable prefix"} - {id: reliability.circuit_breaker.redis.trips_then_recovers, module: reliability, tier: P0, behavior: circuit_breaker, variant: redis, assertions: [trips_then_recovers], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/redis_cache.py:99", rationale: "Redis breaker CLOSED->OPEN->HALF_OPEN; guards all cache/rate-limit ops"} +- {id: reliability.circuit_breaker.redis_timeout.stays_responsive, module: reliability, tier: P0, behavior: circuit_breaker, variant: redis_timeout, assertions: [stays_responsive], exercised_on: [chat_completions], source: "litellm/proxy/hooks/proxy_track_cost_callback.py:386", fail_before_fix: proven, rationale: "With every Redis command timing out and every request retrying then falling back, per-request latency stays flat, liveliness stays fast, and spend rows still land; on v1.100.0 the failed-tracking alert body doubled per request until the worker OOMed (LIT-6780)"} - {id: reliability.timeout.request_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: request_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions, messages], source: "litellm/router.py:545-551", rationale: "Per-request timeout raises Timeout"} - {id: reliability.timeout.stream_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: stream_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions], source: "litellm/router.py:551", rationale: "Streaming chunk-delivery timeout"} - {id: reliability.perf.throughput.under_slo, module: reliability, tier: P1, behavior: perf, variant: throughput, assertions: [under_slo], exercised_on: [chat_completions, messages], source: grammar, rationale: "Throughput SLO under load"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 691335ffdd5..6cb59ea4e90 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -13,7 +13,6 @@ from pathlib import Path from typing import Final from dotenv import load_dotenv - from fixture_mode import deterministic_marker, parse_fixture_mode from provider_edge import provider_edge_api_base @@ -144,6 +143,7 @@ LOAD_MIN_CONCURRENCY_EFFICIENCY = float(os.environ.get("E2E_LOAD_MIN_CONCURRENCY WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY" MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK" +REDIS_TIMEOUT_OPT_IN_ENV = "E2E_REDIS_TIMEOUT" ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3")) @@ -181,8 +181,7 @@ def datadog_mcp_url(*, toolsets: str = "core") -> str: site = ( os.environ.get("DD_SITE", DD_SITE) or "datadoghq.com" ).strip().removeprefix("https://").removeprefix("http://").rstrip("/") - if site.startswith("app."): - site = site[len("app.") :] + site = site.removeprefix("app.") host = "mcp.datadoghq.com" if site in ("", "datadoghq.com") else f"mcp.{site}" base = f"https://{host}/v1/mcp" return f"{base}?toolsets={toolsets}" if toolsets else base diff --git a/tests/e2e/gateway/redis_timeout_ci_config.yml b/tests/e2e/gateway/redis_timeout_ci_config.yml new file mode 100644 index 00000000000..69ab2ee14dc --- /dev/null +++ b/tests/e2e/gateway/redis_timeout_ci_config.yml @@ -0,0 +1,29 @@ +general_settings: + master_key: os.environ/LITELLM_MASTER_KEY + store_model_in_db: true + +litellm_settings: + cache: true + cache_params: + type: redis + host: 127.0.0.1 + port: 6379 + socket_timeout: 0.001 + +router_settings: + num_retries: 1 + fallbacks: + - redis-timeout-primary: + - redis-timeout-backup + +model_list: + - model_name: redis-timeout-primary + litellm_params: + model: openai/gpt-5-mini + api_key: sk-redis-timeout-primary-not-used + mock_response: "litellm.InternalServerError" + - model_name: redis-timeout-backup + litellm_params: + model: openai/gpt-5-mini + api_key: sk-redis-timeout-backup-not-used + mock_response: "ok" diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index c3f8865f218..f69d5d058f3 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -9,3 +9,4 @@ markers = load: heavy throughput/load test; collected last so it never perturbs latency-sensitive suites weekly: real-provider anomaly load test that spends real money; deselected unless E2E_WEEKLY_ANOMALY is set managed_files: needs a proxy running with require_managed_files enabled; deselected unless E2E_MANAGED_FILES_STACK is set + redis_timeout: needs a proxy booted from gateway/redis_timeout_ci_config.yml whose Redis times out on every command; deselected unless E2E_REDIS_TIMEOUT is set diff --git a/tests/e2e/router/conftest.py b/tests/e2e/router/conftest.py index 98501f9bd7c..57fde729f41 100644 --- a/tests/e2e/router/conftest.py +++ b/tests/e2e/router/conftest.py @@ -10,13 +10,12 @@ proxy does not already list it (compose has it in static config; stage does not) from __future__ import annotations +import os from collections.abc import Iterator import pytest -from requests import RequestException - from complexity_router_client import ComplexityRouterClient, build_client -from proxy_client import ProxyClient +from e2e_config import REDIS_TIMEOUT_OPT_IN_ENV from e2e_http import NoBody, Success from lifecycle import ResourceManager from models import ( @@ -26,6 +25,8 @@ from models import ( LiteLLMParamsBody, ModelsListResponse, ) +from proxy_client import ProxyClient +from requests import RequestException ROUTER_MODEL = "complexity-smart-router" ROUTER_PARAMS = LiteLLMParamsBody( @@ -45,6 +46,16 @@ ROUTER_PARAMS = LiteLLMParamsBody( ROUTER_KEY_MODELS = [ROUTER_MODEL, "gpt-5.5", "claude-haiku-4-5"] +def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None: + if os.environ.get(REDIS_TIMEOUT_OPT_IN_ENV): + return + deselected = [item for item in items if item.get_closest_marker("redis_timeout") is not None] + if not deselected: + return + config.hook.pytest_deselected(items=deselected) + items[:] = [item for item in items if item.get_closest_marker("redis_timeout") is None] + + @pytest.fixture(scope="session") def client(proxy: ProxyClient) -> ComplexityRouterClient: return build_client(proxy) @@ -120,8 +131,6 @@ def _ensure_complexity_smart_router( # pyright: ignore[reportUnusedFunction] # @pytest.fixture def complexity_key(resources: ResourceManager, client: ComplexityRouterClient) -> str: """Per-test key allowed to call the complexity router and its tier backends.""" - key = client.proxy.generate_key( - KeyGenerateBody(models=ROUTER_KEY_MODELS, user_id="e2e-complexity-router") - ) + key = client.proxy.generate_key(KeyGenerateBody(models=ROUTER_KEY_MODELS, user_id="e2e-complexity-router")) resources.defer(lambda: client.proxy.delete_key(key)) return key diff --git a/tests/e2e/router/test_redis_timeout_e2e.py b/tests/e2e/router/test_redis_timeout_e2e.py new file mode 100644 index 00000000000..eeb5224e051 --- /dev/null +++ b/tests/e2e/router/test_redis_timeout_e2e.py @@ -0,0 +1,93 @@ +"""Live e2e: the proxy keeps answering while every Redis command times out. + +Runs only against a proxy booted from tests/e2e/gateway/redis_timeout_ci_config.yml, which +points cache_params at a real Redis with socket_timeout 0.001 so every command times out and +the circuit breaker opens. Each request fails its primary deployment, retries, falls back to the +backup and succeeds, so it carries retry breadcrumbs; its cost tracking then fails on the spend +counter increment and stringifies the request metadata into a failed-tracking alert. On v1.100.0 +that string doubled per request until the worker hung (LIT-6780). Deselected unless +E2E_REDIS_TIMEOUT is set, since it needs that dedicated proxy. +""" + +from __future__ import annotations + +import time +from typing import Final + +import pytest +from complexity_router_client import ComplexityRouterClient +from e2e_config import unique_marker +from e2e_http import NoBody, Success +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, ChatResponse, KeyGenerateBody + +pytestmark = [pytest.mark.e2e, pytest.mark.redis_timeout] + +PRIMARY_MODEL: Final = "redis-timeout-primary" +BACKUP_MODEL: Final = "redis-timeout-backup" +REQUESTS: Final = 20 +MAX_SECONDS_PER_REQUEST: Final = 10.0 +MAX_LATENCY_GROWTH_RATIO: Final = 3.0 +MAX_LIVELINESS_SECONDS: Final = 2.0 + + +class TestRedisTimeout: + @pytest.mark.covers( + "reliability.circuit_breaker.redis_timeout.stays_responsive", + exercised_on=["chat_completions"], + ) + def test_retries_under_redis_timeouts_keep_answering( + self, client: ComplexityRouterClient, resources: ResourceManager + ) -> None: + proxy = client.proxy + key = proxy.generate_key( + KeyGenerateBody(models=[PRIMARY_MODEL, BACKUP_MODEL], key_alias=f"e2e-redis-timeout-{unique_marker()}") + ) + resources.defer(lambda: proxy.delete_key(key)) + + latencies: list[float] = [] + for request_number in range(1, REQUESTS + 1): + started = time.monotonic() + result = proxy.transport.post( + "/chat/completions", + headers=proxy.transport.bearer(key), + json=ChatBody( + model=PRIMARY_MODEL, + messages=[ChatMessage(role="user", content=f"redis timeout {unique_marker()} {request_number}")], + max_tokens=5, + ), + response_type=ChatResponse, + timeout=MAX_SECONDS_PER_REQUEST, + ) + elapsed = time.monotonic() - started + assert isinstance(result, Success), ( + f"request {request_number} failed after {elapsed:.1f}s with Redis timing out: {result}; " + f"earlier requests took {[round(seconds, 2) for seconds in latencies]}" + ) + assert result.data.choices, f"request {request_number}: fallback to {BACKUP_MODEL} returned no choices" + assert elapsed < MAX_SECONDS_PER_REQUEST, ( + f"request {request_number} took {elapsed:.1f}s with Redis timing out; " + f"earlier requests took {[round(seconds, 2) for seconds in latencies]}" + ) + latencies.append(elapsed) + + third = REQUESTS // 3 + early = sum(latencies[:third]) / third + late = sum(latencies[-third:]) / third + assert late <= max(early * MAX_LATENCY_GROWTH_RATIO, 0.5), ( + f"per-request latency grew from {early:.2f}s to {late:.2f}s across {REQUESTS} requests " + "while Redis timed out; the proxy is paying more for each failed request" + ) + + started = time.monotonic() + probe = proxy.transport.probe("/health/liveliness", params=NoBody()) + liveliness_seconds = time.monotonic() - started + assert probe.healthy, f"/health/liveliness returned {probe.status_code} after the Redis timeout loop" + assert liveliness_seconds < MAX_LIVELINESS_SECONDS, ( + f"/health/liveliness took {liveliness_seconds:.1f}s after the loop; the worker is stalled" + ) + + rows = proxy.poll_logs_for_key(key, min_rows=REQUESTS) + assert len(rows) >= REQUESTS, ( + f"only {len(rows)} of {REQUESTS} requests reached the spend log; a Redis outage must not lose spend rows" + ) From d9e822d2562ddd716d8f6a5e86d3d8fedcb30df1 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 9 Sep 2026 15:51:11 -0700 Subject: [PATCH 011/119] ci: run the Redis timeout e2e test from its own weekly workflow It is a functional e2e test, not a load test, so give it its own workflow instead of a job inside the load anomaly run. It keeps the Saturday 12:00 UTC cadence and manual dispatch, and boots the timeout-config proxy with Postgres and Valkey services exactly as before. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_014ZDULyJPp17ZFiJenRxs2T --- .github/workflows/test-e2e-redis-timeout.yml | 94 ++++++++++++++++++++ .github/workflows/weekly_load_anomaly.yml | 84 ----------------- tests/e2e/CLAUDE.md | 2 +- tests/e2e/CONTRIBUTING.md | 2 +- 4 files changed, 96 insertions(+), 86 deletions(-) create mode 100644 .github/workflows/test-e2e-redis-timeout.yml diff --git a/.github/workflows/test-e2e-redis-timeout.yml b/.github/workflows/test-e2e-redis-timeout.yml new file mode 100644 index 00000000000..a956ed79590 --- /dev/null +++ b/.github/workflows/test-e2e-redis-timeout.yml @@ -0,0 +1,94 @@ +name: "Weekly Redis Timeout E2E" + +on: + schedule: + - cron: "0 12 * * 6" + workflow_dispatch: + +permissions: + contents: read + +jobs: + redis-timeout-e2e: + if: github.event_name != 'schedule' || github.repository == 'BerriAI/litellm' + runs-on: ubuntu-latest + timeout-minutes: 30 + services: + postgres: + image: postgres:16.6 + env: + POSTGRES_USER: llmproxy + POSTGRES_PASSWORD: dbpassword9090 + POSTGRES_DB: litellm + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U llmproxy" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + valkey: + image: valkey/valkey:8.1.4@sha256:81db6d39e1bba3b3ff32bd3a1b19a6d69690f94a3954ec131277b9a26b95b3aa + ports: + - 6379:6379 + options: >- + --health-cmd "valkey-cli ping" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + env: + DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm + LITELLM_MASTER_KEY: sk-redis-timeout-e2e + LITELLM_LOG: WARNING + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Set up uv + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Cache the Rust build + uses: ./.github/actions/cache-cargo-build + + - name: Install dependencies + run: | + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy + + - name: Cache Prisma binaries + uses: ./.github/actions/cache-prisma-binaries + + - name: Generate Prisma client + run: | + uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma + + - name: Start the proxy with a Redis that times out on every command + run: | + nohup uv run --no-sync litellm --config tests/e2e/gateway/redis_timeout_ci_config.yml --port 4000 > proxy.log 2>&1 & + for _ in $(seq 1 90); do + if curl -fs http://localhost:4000/health/liveliness > /dev/null; then + exit 0 + fi + sleep 2 + done + echo "proxy never became live" + tail -n 100 proxy.log + exit 1 + + - name: Run the Redis timeout e2e test + env: + E2E_REDIS_TIMEOUT: "1" + LITELLM_PROXY_URL: http://localhost:4000 + run: | + uv run --no-sync pytest tests/e2e/router/test_redis_timeout_e2e.py -v --tb=short -rA + + - name: Show proxy log on failure + if: failure() + run: tail -n 300 proxy.log diff --git a/.github/workflows/weekly_load_anomaly.yml b/.github/workflows/weekly_load_anomaly.yml index 6e564cf7e78..3e1fca89645 100644 --- a/.github/workflows/weekly_load_anomaly.yml +++ b/.github/workflows/weekly_load_anomaly.yml @@ -83,87 +83,3 @@ jobs: - name: Show proxy log on failure if: failure() run: tail -n 300 proxy.log - - redis-timeout-e2e: - if: github.event_name != 'schedule' || github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - timeout-minutes: 30 - services: - postgres: - image: postgres:16.6 - env: - POSTGRES_USER: llmproxy - POSTGRES_PASSWORD: dbpassword9090 - POSTGRES_DB: litellm - ports: - - 5432:5432 - options: >- - --health-cmd "pg_isready -U llmproxy" - --health-interval 5s - --health-timeout 5s - --health-retries 10 - valkey: - image: valkey/valkey:8.1.4@sha256:81db6d39e1bba3b3ff32bd3a1b19a6d69690f94a3954ec131277b9a26b95b3aa - ports: - - 6379:6379 - options: >- - --health-cmd "valkey-cli ping" - --health-interval 5s - --health-timeout 5s - --health-retries 10 - env: - DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm - LITELLM_MASTER_KEY: sk-redis-timeout-e2e - LITELLM_LOG: WARNING - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - - name: Cache the Rust build - uses: ./.github/actions/cache-cargo-build - - - name: Install dependencies - run: | - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy - - - name: Cache Prisma binaries - uses: ./.github/actions/cache-prisma-binaries - - - name: Generate Prisma client - run: | - uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma - - - name: Start the proxy with a Redis that times out on every command - run: | - nohup uv run --no-sync litellm --config tests/e2e/gateway/redis_timeout_ci_config.yml --port 4000 > proxy.log 2>&1 & - for _ in $(seq 1 90); do - if curl -fs http://localhost:4000/health/liveliness > /dev/null; then - exit 0 - fi - sleep 2 - done - echo "proxy never became live" - tail -n 100 proxy.log - exit 1 - - - name: Run the Redis timeout e2e test - env: - E2E_REDIS_TIMEOUT: "1" - LITELLM_PROXY_URL: http://localhost:4000 - run: | - uv run --no-sync pytest tests/e2e/router/test_redis_timeout_e2e.py -v --tb=short -rA - - - name: Show proxy log on failure - if: failure() - run: tail -n 300 proxy.log diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 53342f66b67..76e8550cb9a 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -17,7 +17,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below); plus the gateway-managed OAuth (authorization_code) path exercised through `/chat/completions`, the one behavior Datadog's static-header auth cannot reach, seeding the per-user upstream token via the interactive authorize dance driven with the mcp SDK's own OAuth client (headless-browser consent from a saved session) and asserting the completion lists and executes the server's tools with the stored per-user token - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection -- `router/` - routing and reliability behavior (fallbacks, cooldowns). Also holds `test_redis_timeout_e2e.py`, which needs a proxy booted from `gateway/redis_timeout_ci_config.yml` (real Redis, `socket_timeout: 0.001`, so every command times out); it is marked `redis_timeout`, deselected unless `E2E_REDIS_TIMEOUT` is set, excluded from the per-PR check, and driven by `.github/workflows/weekly_load_anomaly.yml` +- `router/` - routing and reliability behavior (fallbacks, cooldowns). Also holds `test_redis_timeout_e2e.py`, which needs a proxy booted from `gateway/redis_timeout_ci_config.yml` (real Redis, `socket_timeout: 0.001`, so every command times out); it is marked `redis_timeout`, deselected unless `E2E_REDIS_TIMEOUT` is set, excluded from the per-PR check, and driven weekly by `.github/workflows/test-e2e-redis-timeout.yml` - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What remains here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`) and markerless harness unit tests for the Locust/session-anomaly aggregation logic - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index ccd27df6f46..7a859965c80 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -65,7 +65,7 @@ A couple of logging destinations are configured on the proxy rather than by the ### The pull request check -Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. A change to the harness itself, meaning a root-level `tests/e2e/*.py` file or `pytest.ini`, `tests/e2e/gateway/`, `.github/e2e-stack/`, or the workflow, also runs the `access_control` suite as a canary, because those files have no test of their own that exercises the stack. `.github/e2e-stack/select_tests.py` applies both rules. The stack config at `tests/e2e/gateway/stage_mirror_ci_config.yml` must declare every model the selected suites use; a missing one shows up as a failed test id in the public log. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Jaeger, and TLS cluster-mode Valkey. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, deleted-file, and application-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories, `batches/test_managed_files_enforcement_e2e.py`, `llm_translation/realtime/test_realtime_pipecat_audio_e2e.py`, `guardrails/test_presidio_masking_e2e.py`, and `router/test_redis_timeout_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack: the pipecat audio suite skips itself at import time unless the NLTK `punkt_tab` data is installed, and the presidio suite fails without the analyzer and anonymizer services this stack does not start, and the Redis timeout test needs a proxy whose Redis times out on every command (`gateway/redis_timeout_ci_config.yml`), which the weekly load workflow boots +Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. A change to the harness itself, meaning a root-level `tests/e2e/*.py` file or `pytest.ini`, `tests/e2e/gateway/`, `.github/e2e-stack/`, or the workflow, also runs the `access_control` suite as a canary, because those files have no test of their own that exercises the stack. `.github/e2e-stack/select_tests.py` applies both rules. The stack config at `tests/e2e/gateway/stage_mirror_ci_config.yml` must declare every model the selected suites use; a missing one shows up as a failed test id in the public log. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Jaeger, and TLS cluster-mode Valkey. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, deleted-file, and application-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories, `batches/test_managed_files_enforcement_e2e.py`, `llm_translation/realtime/test_realtime_pipecat_audio_e2e.py`, `guardrails/test_presidio_masking_e2e.py`, and `router/test_redis_timeout_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack: the pipecat audio suite skips itself at import time unless the NLTK `punkt_tab` data is installed, and the presidio suite fails without the analyzer and anonymizer services this stack does not start, and the Redis timeout test needs a proxy whose Redis times out on every command (`gateway/redis_timeout_ci_config.yml`), which `.github/workflows/test-e2e-redis-timeout.yml` boots on a weekly schedule Every selected file must execute at least one passing test in each pass, and any test failure, collection error, or entirely skipped or deselected file fails the check. A file whose tests are all marked skip therefore cannot pass this check, so unskip at least one of them, or add the file to `UNSUPPORTED` in `select_tests.py` with the reason, before changing one. A failed pass stops the run. The public log prints pytest's one-line summary for each pass, including the rerun count, and names each failed or errored test as `classname::name`, so a retried network error or a failing test is visible without the raw output. The final `e2e-changed-tests` job succeeds only when no supported test files changed or the approved run completed all three passes. Fork PRs with selected tests fail this gate until a maintainer brings the reviewed change onto a same-repository branch From 02351da51f14d644660c4a5583968389752af46b Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 9 Sep 2026 15:51:35 -0700 Subject: [PATCH 012/119] test(e2e): register the Redis timeout cell at tier P1 Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_014ZDULyJPp17ZFiJenRxs2T --- tests/e2e/coverage_registry/reliability.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index d62f0105931..38427b647b9 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -31,7 +31,7 @@ - {id: reliability.cache.exact.returns_cached, module: reliability, tier: P1, behavior: cache, variant: exact, assertions: [returns_cached], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/caching.py", rationale: "Response cache returns cached on exact match"} - {id: reliability.cache.prompt_caching_model_select.returns_cached, module: reliability, tier: P1, behavior: cache, variant: prompt_caching_model_select, assertions: [returns_cached], exercised_on: [chat_completions], source: "router_utils/prompt_caching_cache.py", rationale: "Selects model supporting prompt caching for cacheable prefix"} - {id: reliability.circuit_breaker.redis.trips_then_recovers, module: reliability, tier: P0, behavior: circuit_breaker, variant: redis, assertions: [trips_then_recovers], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/redis_cache.py:99", rationale: "Redis breaker CLOSED->OPEN->HALF_OPEN; guards all cache/rate-limit ops"} -- {id: reliability.circuit_breaker.redis_timeout.stays_responsive, module: reliability, tier: P0, behavior: circuit_breaker, variant: redis_timeout, assertions: [stays_responsive], exercised_on: [chat_completions], source: "litellm/proxy/hooks/proxy_track_cost_callback.py:386", fail_before_fix: proven, rationale: "With every Redis command timing out and every request retrying then falling back, per-request latency stays flat, liveliness stays fast, and spend rows still land; on v1.100.0 the failed-tracking alert body doubled per request until the worker OOMed (LIT-6780)"} +- {id: reliability.circuit_breaker.redis_timeout.stays_responsive, module: reliability, tier: P1, behavior: circuit_breaker, variant: redis_timeout, assertions: [stays_responsive], exercised_on: [chat_completions], source: "litellm/proxy/hooks/proxy_track_cost_callback.py:386", fail_before_fix: proven, rationale: "With every Redis command timing out and every request retrying then falling back, per-request latency stays flat, liveliness stays fast, and spend rows still land; on v1.100.0 the failed-tracking alert body doubled per request until the worker OOMed (LIT-6780)"} - {id: reliability.timeout.request_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: request_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions, messages], source: "litellm/router.py:545-551", rationale: "Per-request timeout raises Timeout"} - {id: reliability.timeout.stream_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: stream_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions], source: "litellm/router.py:551", rationale: "Streaming chunk-delivery timeout"} - {id: reliability.perf.throughput.under_slo, module: reliability, tier: P1, behavior: perf, variant: throughput, assertions: [under_slo], exercised_on: [chat_completions, messages], source: grammar, rationale: "Throughput SLO under load"} From 939ac04a932baae63f2f8b9440ee43ed70570763 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 9 Sep 2026 16:03:36 -0700 Subject: [PATCH 013/119] test(e2e): cover /v1/responses in the Redis timeout test and fail the primary for real The Responses path returns a mock for any mock_response string, so the InternalServerError sentinel never failed there. Point the primary deployment's api_base at a closed port instead, which fails every endpoint the same way, then parametrize the test over /chat/completions and /v1/responses and register both on the coverage cell. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_014ZDULyJPp17ZFiJenRxs2T --- tests/e2e/coverage_registry/reliability.yaml | 2 +- tests/e2e/gateway/redis_timeout_ci_config.yml | 2 +- tests/e2e/router/test_redis_timeout_e2e.py | 105 +++++++++++++----- 3 files changed, 81 insertions(+), 28 deletions(-) diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index 38427b647b9..3b378b56870 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -31,7 +31,7 @@ - {id: reliability.cache.exact.returns_cached, module: reliability, tier: P1, behavior: cache, variant: exact, assertions: [returns_cached], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/caching.py", rationale: "Response cache returns cached on exact match"} - {id: reliability.cache.prompt_caching_model_select.returns_cached, module: reliability, tier: P1, behavior: cache, variant: prompt_caching_model_select, assertions: [returns_cached], exercised_on: [chat_completions], source: "router_utils/prompt_caching_cache.py", rationale: "Selects model supporting prompt caching for cacheable prefix"} - {id: reliability.circuit_breaker.redis.trips_then_recovers, module: reliability, tier: P0, behavior: circuit_breaker, variant: redis, assertions: [trips_then_recovers], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/redis_cache.py:99", rationale: "Redis breaker CLOSED->OPEN->HALF_OPEN; guards all cache/rate-limit ops"} -- {id: reliability.circuit_breaker.redis_timeout.stays_responsive, module: reliability, tier: P1, behavior: circuit_breaker, variant: redis_timeout, assertions: [stays_responsive], exercised_on: [chat_completions], source: "litellm/proxy/hooks/proxy_track_cost_callback.py:386", fail_before_fix: proven, rationale: "With every Redis command timing out and every request retrying then falling back, per-request latency stays flat, liveliness stays fast, and spend rows still land; on v1.100.0 the failed-tracking alert body doubled per request until the worker OOMed (LIT-6780)"} +- {id: reliability.circuit_breaker.redis_timeout.stays_responsive, module: reliability, tier: P1, behavior: circuit_breaker, variant: redis_timeout, assertions: [stays_responsive], exercised_on: [chat_completions, responses], source: "litellm/proxy/hooks/proxy_track_cost_callback.py:386", fail_before_fix: proven, rationale: "With every Redis command timing out and every request retrying then falling back, per-request latency stays flat, liveliness stays fast, and spend rows still land; on v1.100.0 the failed-tracking alert body doubled per request until the worker OOMed (LIT-6780)"} - {id: reliability.timeout.request_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: request_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions, messages], source: "litellm/router.py:545-551", rationale: "Per-request timeout raises Timeout"} - {id: reliability.timeout.stream_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: stream_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions], source: "litellm/router.py:551", rationale: "Streaming chunk-delivery timeout"} - {id: reliability.perf.throughput.under_slo, module: reliability, tier: P1, behavior: perf, variant: throughput, assertions: [under_slo], exercised_on: [chat_completions, messages], source: grammar, rationale: "Throughput SLO under load"} diff --git a/tests/e2e/gateway/redis_timeout_ci_config.yml b/tests/e2e/gateway/redis_timeout_ci_config.yml index 69ab2ee14dc..735285adb36 100644 --- a/tests/e2e/gateway/redis_timeout_ci_config.yml +++ b/tests/e2e/gateway/redis_timeout_ci_config.yml @@ -21,7 +21,7 @@ model_list: litellm_params: model: openai/gpt-5-mini api_key: sk-redis-timeout-primary-not-used - mock_response: "litellm.InternalServerError" + api_base: http://127.0.0.1:1 - model_name: redis-timeout-backup litellm_params: model: openai/gpt-5-mini diff --git a/tests/e2e/router/test_redis_timeout_e2e.py b/tests/e2e/router/test_redis_timeout_e2e.py index eeb5224e051..d10f32b6482 100644 --- a/tests/e2e/router/test_redis_timeout_e2e.py +++ b/tests/e2e/router/test_redis_timeout_e2e.py @@ -1,25 +1,29 @@ """Live e2e: the proxy keeps answering while every Redis command times out. Runs only against a proxy booted from tests/e2e/gateway/redis_timeout_ci_config.yml, which -points cache_params at a real Redis with socket_timeout 0.001 so every command times out and -the circuit breaker opens. Each request fails its primary deployment, retries, falls back to the -backup and succeeds, so it carries retry breadcrumbs; its cost tracking then fails on the spend -counter increment and stringifies the request metadata into a failed-tracking alert. On v1.100.0 -that string doubled per request until the worker hung (LIT-6780). Deselected unless -E2E_REDIS_TIMEOUT is set, since it needs that dedicated proxy. +points cache_params at a real Redis with socket_timeout 0.001 so commands time out and the +circuit breaker opens. Each request fails its primary deployment, whose api_base is a closed +port, retries, falls back to the backup and succeeds, so it carries retry breadcrumbs; its cost +tracking then fails on the spend counter increment and stringifies the request metadata into a +failed-tracking alert. On v1.100.0 that string doubled per request until the worker hung +(LIT-6780). Deselected unless E2E_REDIS_TIMEOUT is set, since it needs that dedicated proxy. """ from __future__ import annotations import time +from collections.abc import Callable +from dataclasses import dataclass from typing import Final import pytest from complexity_router_client import ComplexityRouterClient from e2e_config import unique_marker -from e2e_http import NoBody, Success +from e2e_http import NoBody, Result, Success from lifecycle import ResourceManager from models import ChatBody, ChatMessage, ChatResponse, KeyGenerateBody +from proxy_client import ProxyClient +from pydantic import BaseModel pytestmark = [pytest.mark.e2e, pytest.mark.redis_timeout] @@ -31,42 +35,90 @@ MAX_LATENCY_GROWTH_RATIO: Final = 3.0 MAX_LIVELINESS_SECONDS: Final = 2.0 +class ResponsesBody(BaseModel): + model: str + input: str + max_output_tokens: int = 5 + + +class ResponsesObject(BaseModel): + id: str | None = None + status: str | None = None + output: list[object] = [] + + +@dataclass(frozen=True, slots=True) +class Endpoint: + name: str + send: Callable[[ProxyClient, str, str], Result[BaseModel]] + served: Callable[[BaseModel], bool] + + +def _send_chat(proxy: ProxyClient, key: str, marker: str) -> Result[BaseModel]: + return proxy.transport.post( + "/chat/completions", + headers=proxy.transport.bearer(key), + json=ChatBody(model=PRIMARY_MODEL, messages=[ChatMessage(role="user", content=marker)], max_tokens=5), + response_type=ChatResponse, + timeout=MAX_SECONDS_PER_REQUEST, + ) + + +def _send_responses(proxy: ProxyClient, key: str, marker: str) -> Result[BaseModel]: + return proxy.transport.post( + "/v1/responses", + headers=proxy.transport.bearer(key), + json=ResponsesBody(model=PRIMARY_MODEL, input=marker), + response_type=ResponsesObject, + timeout=MAX_SECONDS_PER_REQUEST, + ) + + +ENDPOINTS: Final = ( + Endpoint( + name="chat_completions", + send=_send_chat, + served=lambda data: isinstance(data, ChatResponse) and bool(data.choices), + ), + Endpoint( + name="responses", + send=_send_responses, + served=lambda data: isinstance(data, ResponsesObject) and bool(data.output), + ), +) + + class TestRedisTimeout: + @pytest.mark.parametrize("endpoint", ENDPOINTS, ids=[endpoint.name for endpoint in ENDPOINTS]) @pytest.mark.covers( "reliability.circuit_breaker.redis_timeout.stays_responsive", - exercised_on=["chat_completions"], + exercised_on=["chat_completions", "responses"], ) def test_retries_under_redis_timeouts_keep_answering( - self, client: ComplexityRouterClient, resources: ResourceManager + self, client: ComplexityRouterClient, resources: ResourceManager, endpoint: Endpoint ) -> None: proxy = client.proxy key = proxy.generate_key( - KeyGenerateBody(models=[PRIMARY_MODEL, BACKUP_MODEL], key_alias=f"e2e-redis-timeout-{unique_marker()}") + KeyGenerateBody( + models=[PRIMARY_MODEL, BACKUP_MODEL], key_alias=f"e2e-redis-timeout-{endpoint.name}-{unique_marker()}" + ) ) resources.defer(lambda: proxy.delete_key(key)) latencies: list[float] = [] for request_number in range(1, REQUESTS + 1): started = time.monotonic() - result = proxy.transport.post( - "/chat/completions", - headers=proxy.transport.bearer(key), - json=ChatBody( - model=PRIMARY_MODEL, - messages=[ChatMessage(role="user", content=f"redis timeout {unique_marker()} {request_number}")], - max_tokens=5, - ), - response_type=ChatResponse, - timeout=MAX_SECONDS_PER_REQUEST, - ) + result = endpoint.send(proxy, key, f"redis timeout {unique_marker()} {request_number}") elapsed = time.monotonic() - started assert isinstance(result, Success), ( - f"request {request_number} failed after {elapsed:.1f}s with Redis timing out: {result}; " + f"{endpoint.name} request {request_number} failed after {elapsed:.1f}s with Redis timing out: {result}; " f"earlier requests took {[round(seconds, 2) for seconds in latencies]}" ) - assert result.data.choices, f"request {request_number}: fallback to {BACKUP_MODEL} returned no choices" + assert endpoint.served(result.data), ( + f"{endpoint.name} request {request_number}: fallback to {BACKUP_MODEL} returned no output" + ) assert elapsed < MAX_SECONDS_PER_REQUEST, ( - f"request {request_number} took {elapsed:.1f}s with Redis timing out; " + f"{endpoint.name} request {request_number} took {elapsed:.1f}s with Redis timing out; " f"earlier requests took {[round(seconds, 2) for seconds in latencies]}" ) latencies.append(elapsed) @@ -75,7 +127,7 @@ class TestRedisTimeout: early = sum(latencies[:third]) / third late = sum(latencies[-third:]) / third assert late <= max(early * MAX_LATENCY_GROWTH_RATIO, 0.5), ( - f"per-request latency grew from {early:.2f}s to {late:.2f}s across {REQUESTS} requests " + f"{endpoint.name} per-request latency grew from {early:.2f}s to {late:.2f}s across {REQUESTS} requests " "while Redis timed out; the proxy is paying more for each failed request" ) @@ -89,5 +141,6 @@ class TestRedisTimeout: rows = proxy.poll_logs_for_key(key, min_rows=REQUESTS) assert len(rows) >= REQUESTS, ( - f"only {len(rows)} of {REQUESTS} requests reached the spend log; a Redis outage must not lose spend rows" + f"only {len(rows)} of {REQUESTS} {endpoint.name} requests reached the spend log; " + "a Redis outage must not lose spend rows" ) From 2d7e5999b78f726c97a7142231de566d788d4fc0 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 9 Sep 2026 16:26:03 -0700 Subject: [PATCH 014/119] test(e2e): pause Redis writes for the run and prove the timeouts from /metrics A loopback Redis answers many commands inside the 1 ms socket timeout, so nothing guaranteed the failure path ran. The test now holds the proxy's Redis in CLIENT PAUSE WRITE for its duration, so every write the proxy sends, the spend counter increment included, outlives the timeout, and lifts the pause in teardown. Reads stay live so the control connection can do that. Enable the prometheus callback in the gateway config and assert from /metrics that the proxy counted at least the breaker's five timeouts and that, during each case, it saw fresh timeouts, a breaker transition, or an open breaker rejecting every call. The open breaker is the state a customer's worker sits in, and cost tracking fails on every request either way. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_014ZDULyJPp17ZFiJenRxs2T --- .github/workflows/test-e2e-redis-timeout.yml | 2 + tests/e2e/gateway/redis_timeout_ci_config.yml | 2 + tests/e2e/router/test_redis_timeout_e2e.py | 60 +++++++++++++++++-- 3 files changed, 60 insertions(+), 4 deletions(-) diff --git a/.github/workflows/test-e2e-redis-timeout.yml b/.github/workflows/test-e2e-redis-timeout.yml index a956ed79590..9c318a1346c 100644 --- a/.github/workflows/test-e2e-redis-timeout.yml +++ b/.github/workflows/test-e2e-redis-timeout.yml @@ -86,6 +86,8 @@ jobs: env: E2E_REDIS_TIMEOUT: "1" LITELLM_PROXY_URL: http://localhost:4000 + REDIS_HOST: 127.0.0.1 + REDIS_PORT: "6379" run: | uv run --no-sync pytest tests/e2e/router/test_redis_timeout_e2e.py -v --tb=short -rA diff --git a/tests/e2e/gateway/redis_timeout_ci_config.yml b/tests/e2e/gateway/redis_timeout_ci_config.yml index 735285adb36..375a51f6127 100644 --- a/tests/e2e/gateway/redis_timeout_ci_config.yml +++ b/tests/e2e/gateway/redis_timeout_ci_config.yml @@ -3,6 +3,8 @@ general_settings: store_model_in_db: true litellm_settings: + callbacks: ["prometheus"] + require_auth_for_metrics_endpoint: false cache: true cache_params: type: redis diff --git a/tests/e2e/router/test_redis_timeout_e2e.py b/tests/e2e/router/test_redis_timeout_e2e.py index d10f32b6482..e99bd2fbe80 100644 --- a/tests/e2e/router/test_redis_timeout_e2e.py +++ b/tests/e2e/router/test_redis_timeout_e2e.py @@ -1,8 +1,10 @@ """Live e2e: the proxy keeps answering while every Redis command times out. Runs only against a proxy booted from tests/e2e/gateway/redis_timeout_ci_config.yml, which -points cache_params at a real Redis with socket_timeout 0.001 so commands time out and the -circuit breaker opens. Each request fails its primary deployment, whose api_base is a closed +points cache_params at a real Redis with socket_timeout 0.001. The test holds that Redis in +CLIENT PAUSE WRITE for its duration, so every write the proxy sends, the spend counter increment +included, hangs past the timeout, and it proves the degradation was real from the breaker metrics on /metrics: fresh timeouts, a breaker transition, +or an already-open breaker rejecting every call, which is the state a customer's worker sits in. Each request fails its primary deployment, whose api_base is a closed port, retries, falls back to the backup and succeeds, so it carries retry breadcrumbs; its cost tracking then fails on the spend counter increment and stringifies the request metadata into a failed-tracking alert. On v1.100.0 that string doubled per request until the worker hung @@ -11,12 +13,15 @@ failed-tracking alert. On v1.100.0 that string doubled per request until the wor from __future__ import annotations +import os +import re import time -from collections.abc import Callable +from collections.abc import Callable, Iterator from dataclasses import dataclass from typing import Final import pytest +import redis from complexity_router_client import ComplexityRouterClient from e2e_config import unique_marker from e2e_http import NoBody, Result, Success @@ -33,6 +38,11 @@ REQUESTS: Final = 20 MAX_SECONDS_PER_REQUEST: Final = 10.0 MAX_LATENCY_GROWTH_RATIO: Final = 3.0 MAX_LIVELINESS_SECONDS: Final = 2.0 +REDIS_PAUSE_MS: Final = 600_000 +BREAKER_FAILURE_THRESHOLD: Final = 5 +TIMEOUT_FAILURES_RE: Final = re.compile( + r'^litellm_redis_circuit_breaker_failures_total\{failure_class="timeout"\} ([0-9.e+]+)$', re.M +) class ResponsesBody(BaseModel): @@ -86,6 +96,32 @@ ENDPOINTS: Final = ( served=lambda data: isinstance(data, ResponsesObject) and bool(data.output), ), ) +BREAKER_OPEN_RE: Final = re.compile(r'^litellm_redis_circuit_breaker_state\{state="open"\} ([0-9.e+]+)$', re.M) +BREAKER_TRANSITIONS_RE: Final = re.compile( + r'^litellm_redis_circuit_breaker_transitions_total\{state="[a-z_]+"\} ([0-9.e+]+)$', re.M +) + + +@pytest.fixture +def paused_redis() -> Iterator[None]: + """Hold the proxy's Redis in CLIENT PAUSE WRITE so every write it sends outlives the 1 ms socket + timeout. A loopback Redis otherwise answers many commands inside that budget. Reads stay live so + this control connection can lift the pause in teardown.""" + host = os.environ.get("REDIS_HOST") + port = os.environ.get("REDIS_PORT") + assert host and port, "REDIS_HOST and REDIS_PORT must name the Redis the proxy under test uses" + control = redis.Redis(host=host, port=int(port), socket_timeout=5) + control.client_pause(REDIS_PAUSE_MS, all=False) # pyright: ignore[reportUnknownMemberType] # redis-py stubs return Any + try: + yield + finally: + control.client_unpause() # pyright: ignore[reportUnknownMemberType] # redis-py stubs return Any + control.close() + + +def _metric(proxy: ProxyClient, pattern: re.Pattern[str]) -> float: + body = proxy.probe("/metrics", params=NoBody()).body + return sum(float(match.group(1)) for match in pattern.finditer(body)) class TestRedisTimeout: @@ -95,9 +131,11 @@ class TestRedisTimeout: exercised_on=["chat_completions", "responses"], ) def test_retries_under_redis_timeouts_keep_answering( - self, client: ComplexityRouterClient, resources: ResourceManager, endpoint: Endpoint + self, client: ComplexityRouterClient, resources: ResourceManager, endpoint: Endpoint, paused_redis: None ) -> None: proxy = client.proxy + timeouts_before = _metric(proxy, TIMEOUT_FAILURES_RE) + transitions_before = _metric(proxy, BREAKER_TRANSITIONS_RE) key = proxy.generate_key( KeyGenerateBody( models=[PRIMARY_MODEL, BACKUP_MODEL], key_alias=f"e2e-redis-timeout-{endpoint.name}-{unique_marker()}" @@ -139,6 +177,20 @@ class TestRedisTimeout: f"/health/liveliness took {liveliness_seconds:.1f}s after the loop; the worker is stalled" ) + timeouts_total = _metric(proxy, TIMEOUT_FAILURES_RE) + timeouts = timeouts_total - timeouts_before + transitions = _metric(proxy, BREAKER_TRANSITIONS_RE) - transitions_before + breaker_open = _metric(proxy, BREAKER_OPEN_RE) >= 1 + assert timeouts_total >= BREAKER_FAILURE_THRESHOLD, ( + f"the proxy counted only {timeouts_total:.0f} Redis timeouts in its lifetime; the write-paused Redis " + "never made its spend counter writes time out, so this run proved nothing" + ) + assert timeouts >= REQUESTS or transitions >= 1 or breaker_open, ( + f"during {REQUESTS} {endpoint.name} requests the breaker counted {timeouts:.0f} new timeouts, " + f"{transitions:.0f} state transitions, and ended {'open' if breaker_open else 'closed'}; " + "Redis was healthy for this case, so it proved nothing" + ) + rows = proxy.poll_logs_for_key(key, min_rows=REQUESTS) assert len(rows) >= REQUESTS, ( f"only {len(rows)} of {REQUESTS} {endpoint.name} requests reached the spend log; " From 2e4461ec458756c940f3a8863ea805bc44c33907 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 9 Sep 2026 16:33:18 -0700 Subject: [PATCH 015/119] test(e2e): register the Redis timeout deployments through /model/new The e2e directive has every test create its deployments through the management API and delete them on teardown. Drop the static model_list from the gateway config; the test now registers the closed-port primary and the mock backup itself, and the fallback map stays in router_settings where proxy-level config belongs. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_014ZDULyJPp17ZFiJenRxs2T --- tests/e2e/gateway/redis_timeout_ci_config.yml | 12 ----------- tests/e2e/router/test_redis_timeout_e2e.py | 20 ++++++++++++++++--- 2 files changed, 17 insertions(+), 15 deletions(-) diff --git a/tests/e2e/gateway/redis_timeout_ci_config.yml b/tests/e2e/gateway/redis_timeout_ci_config.yml index 375a51f6127..9bcf65162d6 100644 --- a/tests/e2e/gateway/redis_timeout_ci_config.yml +++ b/tests/e2e/gateway/redis_timeout_ci_config.yml @@ -17,15 +17,3 @@ router_settings: fallbacks: - redis-timeout-primary: - redis-timeout-backup - -model_list: - - model_name: redis-timeout-primary - litellm_params: - model: openai/gpt-5-mini - api_key: sk-redis-timeout-primary-not-used - api_base: http://127.0.0.1:1 - - model_name: redis-timeout-backup - litellm_params: - model: openai/gpt-5-mini - api_key: sk-redis-timeout-backup-not-used - mock_response: "ok" diff --git a/tests/e2e/router/test_redis_timeout_e2e.py b/tests/e2e/router/test_redis_timeout_e2e.py index e99bd2fbe80..58d10f1d8ea 100644 --- a/tests/e2e/router/test_redis_timeout_e2e.py +++ b/tests/e2e/router/test_redis_timeout_e2e.py @@ -4,8 +4,8 @@ Runs only against a proxy booted from tests/e2e/gateway/redis_timeout_ci_config. points cache_params at a real Redis with socket_timeout 0.001. The test holds that Redis in CLIENT PAUSE WRITE for its duration, so every write the proxy sends, the spend counter increment included, hangs past the timeout, and it proves the degradation was real from the breaker metrics on /metrics: fresh timeouts, a breaker transition, -or an already-open breaker rejecting every call, which is the state a customer's worker sits in. Each request fails its primary deployment, whose api_base is a closed -port, retries, falls back to the backup and succeeds, so it carries retry breadcrumbs; its cost +or an already-open breaker rejecting every call, which is the state a customer's worker sits in. The test registers two deployments through /model/new: a primary whose api_base is a closed port +and a backup that answers with a mock. Each request fails the primary, retries, falls back and succeeds, so it carries retry breadcrumbs; its cost tracking then fails on the spend counter increment and stringifies the request metadata into a failed-tracking alert. On v1.100.0 that string doubled per request until the worker hung (LIT-6780). Deselected unless E2E_REDIS_TIMEOUT is set, since it needs that dedicated proxy. @@ -26,7 +26,7 @@ from complexity_router_client import ComplexityRouterClient from e2e_config import unique_marker from e2e_http import NoBody, Result, Success from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, ChatResponse, KeyGenerateBody +from models import ChatBody, ChatMessage, ChatResponse, KeyGenerateBody, LiteLLMParamsBody from proxy_client import ProxyClient from pydantic import BaseModel @@ -34,6 +34,8 @@ pytestmark = [pytest.mark.e2e, pytest.mark.redis_timeout] PRIMARY_MODEL: Final = "redis-timeout-primary" BACKUP_MODEL: Final = "redis-timeout-backup" +BACKING_MODEL: Final = "openai/gpt-5-mini" +CLOSED_PORT_API_BASE: Final = "http://127.0.0.1:1" REQUESTS: Final = 20 MAX_SECONDS_PER_REQUEST: Final = 10.0 MAX_LATENCY_GROWTH_RATIO: Final = 3.0 @@ -134,6 +136,18 @@ class TestRedisTimeout: self, client: ComplexityRouterClient, resources: ResourceManager, endpoint: Endpoint, paused_redis: None ) -> None: proxy = client.proxy + primary_id = proxy.create_model( + PRIMARY_MODEL, + LiteLLMParamsBody( + model=BACKING_MODEL, api_key="sk-redis-timeout-primary-not-used", api_base=CLOSED_PORT_API_BASE + ), + ) + resources.defer(lambda: proxy.delete_model(primary_id)) + backup_id = proxy.create_model( + BACKUP_MODEL, + LiteLLMParamsBody(model=BACKING_MODEL, api_key="sk-redis-timeout-backup-not-used", mock_response="ok"), + ) + resources.defer(lambda: proxy.delete_model(backup_id)) timeouts_before = _metric(proxy, TIMEOUT_FAILURES_RE) transitions_before = _metric(proxy, BREAKER_TRANSITIONS_RE) key = proxy.generate_key( From 2085a37b82f100edc9652e6859713b97e0fb7cdd Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 9 Sep 2026 16:56:07 -0700 Subject: [PATCH 016/119] test(e2e): issue the JWT suite's tokens from a real Keycloak realm The suite used to mint its own RS256 tokens from a stand-in issuer, which could only ever prove the proxy agreed with the tests: the claims were whatever the tests chose to sign. Every JWT bug worth catching lives in the shape of what an identity provider really emits, so the suite now runs against Keycloak (realm in idp_realm.json), provisions a group and a user per test through its admin API, and signs in through the direct-access grant. That changes what the tokens look like: sub is Keycloak's opaque user uuid rather than a friendly name, groups arrives from a protocol mapper, aud is the IdP's own audience, and the JWKS carries an encryption key beside the signing key so the proxy has to select on kid. The expiry case now takes a one-second token from a second client in the realm and waits for it to lapse instead of forging a stale exp. The proxy config the suite needs is unchanged. CI runs it against a Keycloak deployed beside the ephemeral stack, which lives in the releaser repo. Claude-Session: https://claude.ai/code/session_01EX13mWex6RaBo9PYnkAtFW --- tests/e2e/CLAUDE.md | 2 +- tests/e2e/CONTRIBUTING.md | 17 +- tests/e2e/coverage_registry/other.yaml | 6 +- tests/e2e/e2e_config.py | 4 - tests/e2e/e2e_http.py | 72 ++++++-- tests/e2e/idp.py | 226 ++++++++++++++++++++++++ tests/e2e/idp_realm.json | 56 ++++++ tests/e2e/jwt_issuer.py | 230 ------------------------- tests/e2e/models.py | 12 -- tests/e2e/other/other_client.py | 43 ++--- tests/e2e/other/test_jwt_auth_e2e.py | 93 +++++----- tests/e2e/test_idp.py | 79 +++++++++ tests/e2e/test_jwt_issuer.py | 141 --------------- 13 files changed, 502 insertions(+), 479 deletions(-) create mode 100644 tests/e2e/idp.py create mode 100644 tests/e2e/idp_realm.json delete mode 100644 tests/e2e/jwt_issuer.py create mode 100644 tests/e2e/test_idp.py delete mode 100644 tests/e2e/test_jwt_issuer.py diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 4985da9e95d..c954c03dfb7 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -19,7 +19,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (fallbacks, cooldowns) - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What remains here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`) and markerless harness unit tests for the Locust/session-anomaly aggregation logic -- `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (RS256 tokens minted by the test-only issuer in `jwt_issuer.py`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at; see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite +- `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at; see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke - `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index ad414bb6647..4802feb10c4 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -27,17 +27,20 @@ The suites run against a live proxy, so bring one up first by running the litell 2. Bring up a Postgres and a Redis for the proxy to use. The repo-root `docker-compose.yml` already defines a Postgres on `5432`; a `docker run -p 6379:6379 redis:7` covers Redis. Point `DATABASE_URL` / `REDIS_HOST` / `REDIS_PORT` at whatever you run. Tests that read Redis directly default to the deployed shape (TLS + cluster mode) whenever `REDIS_HOST` is set, so for a local standalone Redis also set `REDIS_CLUSTER=false` and `REDIS_SSL=false` (plus `REDIS_PASSWORD` when your Redis requires auth) -3. Start the test-only JWT issuer the `other/` suite mints tokens from, then the litellm proxy against your config, and confirm both are live. The issuer is a fake identity provider (`jwt_issuer.py`) with an open mint endpoint, so it binds loopback only; it generates one RSA key per process and serves it as a JWKS, and the proxy caches that JWKS for `public_key_ttl` (600s) without refetching on an unknown `kid`, so restart the proxy whenever you restart the issuer: +3. Start the identity provider the `other/` suite authenticates against, then the litellm proxy against your config, and confirm both are live. It is a real Keycloak, running the realm in `tests/e2e/idp_realm.json`, and the proxy trusts it because `JWT_PUBLIC_KEY_URL` points at that realm's JWKS. The proxy caches the JWKS for `public_key_ttl` (600s) and does not refetch on an unknown `kid`, so restart the proxy whenever you recreate the container: ```bash set -a && source .env && set +a - uv run python tests/e2e/jwt_issuer.py & - curl -fs http://127.0.0.1:4190/.well-known/jwks.json - JWT_PUBLIC_KEY_URL=http://127.0.0.1:4190/.well-known/jwks.json litellm --config .yml --port 4000 + docker run -d --name litellm-e2e-idp -p 8480:8080 \ + -e KC_BOOTSTRAP_ADMIN_USERNAME=admin -e KC_BOOTSTRAP_ADMIN_PASSWORD=admin \ + -v "$PWD/tests/e2e/idp_realm.json:/opt/keycloak/data/import/realm.json:ro" \ + quay.io/keycloak/keycloak:26.7.3 start-dev --import-realm + curl -fs --retry 30 --retry-delay 2 --retry-all-errors http://127.0.0.1:8480/realms/litellm-e2e/.well-known/openid-configuration + JWT_PUBLIC_KEY_URL=http://127.0.0.1:8480/realms/litellm-e2e/protocol/openid-connect/certs litellm --config .yml --port 4000 curl -fs http://localhost:4000/health/liveliness ``` - The issuer's port is `E2E_JWT_ISSUER_PORT` (default `4190`, the port in the URLs above), read by both the issuer process and the tests, so set it in one place if you change it. JWT auth is an enterprise feature, so the proxy also needs `LITELLM_LICENSE` in its environment, and its config needs the JWT block below. `enable_jwt_auth` only routes bearer tokens with three dot-separated segments into the JWT path, so `sk-` virtual keys and the master key keep working for every other suite. `proxy_batch_write_at` is lowered so the JWT spend-attribution test sees its row well inside the poll deadline: + The tests reach Keycloak at `E2E_KEYCLOAK_URL` (default `http://127.0.0.1:8480`) and provision their identities through its admin API, so they also need `E2E_KEYCLOAK_ADMIN_USER` and `E2E_KEYCLOAK_ADMIN_PASSWORD` (`admin` / `admin` for the throwaway container above; the deployed stacks take theirs from a secret). JWT auth is an enterprise feature, so the proxy needs `LITELLM_LICENSE` in its environment, and its config needs the JWT block below. `enable_jwt_auth` only routes bearer tokens with three dot-separated segments into the JWT path, so `sk-` virtual keys and the master key keep working for every other suite. `proxy_batch_write_at` is lowered so the JWT spend-attribution test sees its row well inside the poll deadline: ```yaml general_settings: @@ -50,7 +53,9 @@ The suites run against a live proxy, so bring one up first by running the litell user_id_upsert: true ``` - Leave `JWT_AUDIENCE` and `JWT_ISSUER` unset unless you also put matching `aud` / `iss` claims in the tokens the tests mint; the issuer sets `iss` to its own base URL + Leave `JWT_AUDIENCE` and `JWT_ISSUER` unset. Keycloak's access tokens carry `aud: account`, its own audience rather than the proxy's, and their `iss` is whatever base URL the token was requested through, so pinning either one only makes sense once the deployment fixes Keycloak's hostname + + CI runs this suite against a Keycloak deployed beside the ephemeral stack, and that deployment (the JWT config block, `JWT_PUBLIC_KEY_URL`, and the admin credential handed to the run pod) lives in the project-releaser repo, not here. A stack without it fails the JWT tests rather than skipping them 4. Run a suite against it; the harness reads `LITELLM_PROXY_URL` (default `http://localhost:4000`): diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index 8b9cb4d2adf..ff25155924b 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -9,9 +9,9 @@ - {id: other.auth.llm_chat.not_bearer_scheme_denied, module: other, tier: P0, area: auth, assertions: [not_bearer_scheme_denied], source: "vendor testing strategy §11.1 / LIT-4778", rationale: "NotBearer scheme on chat is 401/403"} - {id: other.auth.realtime.missing_header_denied, module: other, tier: P1, area: auth, assertions: [missing_header_denied], source: "vendor testing strategy §9.19 / LIT-4778", rationale: "Realtime client-secret and calls routes reject requests without Authorization"} - {id: other.config.responses.metadata_redis_ttl_bounded, module: other, tier: P0, area: config, assertions: [ttl_bounded], source: "responses + redis cache", rationale: "Responses store+metadata must not leave TTL-unbounded Redis entries (LIT-1201)"} -- {id: other.auth.jwt.valid_token_allows, module: other, tier: P0, area: auth, assertions: [valid_token_allows], source: "handle_jwt.py:1217-1256 auth_jwt / user_api_key_auth.py:1365-1377", rationale: "An RS256 JWT signed by the configured JWKS whose groups claim names an existing team is accepted on /chat/completions"} -- {id: other.auth.jwt.spend_attributed_to_claims, module: other, tier: P0, area: auth, assertions: [spend_attributed_to_claims], source: "handle_jwt.py:2224 auth_builder / user_api_key_auth.py:1438-1474", rationale: "The spend log row for a JWT-authenticated call carries the team_id from the groups claim and the user_id from sub, not a virtual key's identity. Single-group claim only: the proxy picks the team from a set, so attribution over several groups is unordered"} -- {id: other.auth.jwt.expired_denied, module: other, tier: P0, area: auth, assertions: [expired_denied], source: "handle_jwt.py:1244-1250", rationale: "Expired JWT rejected 401 (Token Expired) even with valid signature; leeway is 0"} +- {id: other.auth.jwt.valid_token_allows, module: other, tier: P0, area: auth, assertions: [valid_token_allows], source: "handle_jwt.py:1217-1256 auth_jwt / user_api_key_auth.py:1365-1377", rationale: "An access token issued by the configured IdP whose groups claim names an existing team is accepted on /chat/completions"} +- {id: other.auth.jwt.spend_attributed_to_claims, module: other, tier: P0, area: auth, assertions: [spend_attributed_to_claims], source: "handle_jwt.py:2224 auth_builder / user_api_key_auth.py:1438-1474", rationale: "The spend log row for a JWT-authenticated call carries the team_id from the groups claim and the user_id from sub, which for a real IdP is an opaque uuid, not a virtual key's identity. Single-group claim only: the proxy picks the team from a set, so attribution over several groups is unordered"} +- {id: other.auth.jwt.expired_denied, module: other, tier: P0, area: auth, assertions: [expired_denied], source: "handle_jwt.py:1244-1250", rationale: "A token the IdP issued with a one-second lifespan is rejected 401 (Token Expired) once it lapses, even though its signature still verifies; leeway is 0"} - {id: other.auth.jwt.invalid_signature_denied, module: other, tier: P0, area: auth, assertions: [invalid_signature_denied], source: "handle_jwt.py:1158-1166 _decode_jwt_with_public_key", rationale: "A genuine token whose signature bytes were altered fails verification with 401"} - {id: other.auth.jwt.unknown_team_denied, module: other, tier: P0, area: auth, assertions: [unknown_team_denied], source: "handle_jwt.py:1549-1632 find_team_with_model_access", rationale: "A verified JWT whose groups claim resolves to no existing team is denied with 403 naming the unresolved team, never silently admitted without a team. The proxy words it as a model-access denial, the same body an existing team without model access gets"} - {id: other.auth.jwt.virtual_key_unaffected, module: other, tier: P0, area: auth, assertions: [virtual_key_unaffected], source: "handle_jwt.py:213 is_jwt / user_api_key_auth.py:1332-1333", rationale: "enable_jwt_auth only routes three-segment bearer tokens into the JWT branch, so sk- virtual keys keep working on the same proxy"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 0e7f4c0ad2b..09d17b9a0db 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -15,7 +15,6 @@ from typing import Final from dotenv import load_dotenv from fixture_mode import deterministic_marker, parse_fixture_mode -from jwt_issuer import jwt_issuer_url from provider_edge import provider_edge_api_base # Local runs keep provider / DataDog keys in tests/e2e/.env (see CONTRIBUTING.md). @@ -44,9 +43,6 @@ UI_BASE_URL = os.environ.get("E2E_UI_BASE_URL", PROXY_BASE_URL).rstrip("/") CHEAP_ANTHROPIC_MODEL = os.environ.get("E2E_CHEAP_ANTHROPIC_MODEL", "claude-haiku-4-5") CHEAP_OPENAI_MODEL = os.environ.get("E2E_CHEAP_OPENAI_MODEL", "gpt-5.5") -# Test-only JWT issuer (jwt_issuer.py); the port is E2E_JWT_ISSUER_PORT (see CONTRIBUTING.md). -JWT_ISSUER_URL = jwt_issuer_url() - LINEAR_MCP_URL = os.environ.get("E2E_LINEAR_MCP_URL", "https://mcp.linear.app/mcp") LINEAR_STORAGE_STATE = os.environ.get("E2E_LINEAR_STORAGE_STATE", "") diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 4c65c53b653..d6924aff5c0 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -125,6 +125,20 @@ class ProbeResult(BaseModel): return 200 <= self.status_code < 500 and self.status_code != 404 +class ExternalWrite(BaseModel): + """Outcome of a write to a non-proxy API (an identity provider's admin API) + that answers with a status and, on create, a Location header naming the new + resource rather than a JSON body.""" + + status_code: int + location: str = "" + body: str = "" + + @property + def ok(self) -> bool: + return 200 <= self.status_code < 300 + + class StreamingResponse(BaseModel): """Raw outcome for calls whose body is provider-native or streamed: status, the x-litellm-call-id header, the x-litellm-response-cost header (StandardLogging @@ -252,16 +266,17 @@ def assert_auth_denied(result: StreamingResponse, context: str) -> None: f"{context}: expected 401/403, got {result.status_code}: {result.body[:300]}" ) -def _headers(headers: BaseModel) -> dict[str, str]: - dumped: dict[str, object] = headers.model_dump(by_alias=True, exclude_none=True) +def _flat(model: BaseModel) -> dict[str, str]: + dumped: dict[str, object] = model.model_dump(by_alias=True, exclude_none=True) return {key: str(value) for key, value in dumped.items()} +def _headers(headers: BaseModel) -> dict[str, str]: + return _flat(headers) + + def _params(params: BaseModel | None) -> dict[str, str]: - if params is None: - return {} - dumped: dict[str, object] = params.model_dump(by_alias=True, exclude_none=True) - return {key: str(value) for key, value in dumped.items()} + return _flat(params) if params is not None else {} TRANSIENT_STATUSES: frozenset[int] = frozenset({529}) @@ -386,20 +401,20 @@ def get_external[R: BaseModel]( return _classify(resp, response_type) -def post_external[R: BaseModel]( +def post_form_external[R: BaseModel]( url: str, *, - json: BaseModel, + form: BaseModel, response_type: type[R], timeout: float = 30.0, ) -> Result[R]: - """POST an absolute URL outside the proxy (e.g. the e2e JWT issuer's mint - endpoint). Like get_external: no proxy base url, no proxy auth, and the same - tagged-union classification as every other call.""" + """POST an absolute URL outside the proxy as `application/x-www-form-urlencoded`, + the encoding OAuth 2 token endpoints take. Like get_external: no proxy base url, + no proxy auth, and the same tagged-union classification as every other call.""" try: resp = requests.post( url, - json=json.model_dump(by_alias=True, exclude_none=True), + data=_flat(form), timeout=timeout, ) except requests.RequestException as exc: @@ -407,6 +422,39 @@ def post_external[R: BaseModel]( return _classify(resp, response_type) +def post_json_external( + url: str, + *, + headers: BaseModel, + json: BaseModel, + timeout: float = 30.0, +) -> ExternalWrite: + """POST an absolute URL outside the proxy under its own bearer, for an API that + answers a create with a status and a Location header rather than a JSON body.""" + try: + resp = requests.post( + url, + headers=_headers(headers), + json=json.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + except requests.RequestException as exc: + return ExternalWrite(status_code=-1, body=str(exc)) + return ExternalWrite( + status_code=resp.status_code, + location=resp.headers.get("Location", ""), + body=resp.text, + ) + + +def delete_external(url: str, *, headers: BaseModel, timeout: float = 30.0) -> ExternalWrite: + try: + resp = requests.delete(url, headers=_headers(headers), timeout=timeout) + except requests.RequestException as exc: + return ExternalWrite(status_code=-1, body=str(exc)) + return ExternalWrite(status_code=resp.status_code, body=resp.text) + + def delete[R: BaseModel]( url: URL, *, diff --git a/tests/e2e/idp.py b/tests/e2e/idp.py new file mode 100644 index 00000000000..50cf44b4159 --- /dev/null +++ b/tests/e2e/idp.py @@ -0,0 +1,226 @@ +"""The identity provider the JWT suite authenticates against: a real Keycloak +realm, imported from `idp_realm.json`. + +A real IdP rather than a hand-rolled signer because every JWT bug this suite +exists to catch lives in the shape of what an IdP actually emits: `sub` is an +opaque uuid and not a friendly name, group membership arrives as a claim built +by a protocol mapper, the JWKS carries a signing key next to an encryption key +so the proxy has to select on `kid`, and `aud` is the IdP's own audience rather +than the proxy's. A stand-in issuer that mints exactly the claims the tests +assert on can only prove the proxy agrees with the tests. + +Tests never hold a signing key. They provision an identity through Keycloak's +admin API (a group named after the litellm team, a user in it with a password +generated for that test alone), then ask Keycloak for an access token through +the direct-access grant, the same way a CLI or service account signs in. The +proxy's `JWT_PUBLIC_KEY_URL` points at this realm's JWKS, so the token the +tests carry is trusted for exactly one reason: Keycloak signed it. + +The realm declares two clients. `litellm-e2e-tests` mints ordinary tokens; the +`litellm-e2e-shortlived` client sets `access.token.lifespan` to one second, so +the expiry test lets a genuine token expire instead of forging a stale `exp`. + +Connection details come from the environment (`E2E_KEYCLOAK_URL` and the admin +credential). A missing or unreachable IdP is a hard failure naming the start +command, never a skip, so a stack deployed without it turns the run red. +""" + +from __future__ import annotations + +import os +import secrets +from dataclasses import dataclass +from typing import Final, Literal + +import pytest +from pydantic import BaseModel, Field + +from e2e_http import ( + AuthHeaders, + ExternalWrite, + NetworkError, + Result, + Success, + delete_external, + post_form_external, + post_json_external, +) + +KEYCLOAK_URL_ENV: Final = "E2E_KEYCLOAK_URL" +KEYCLOAK_REALM_ENV: Final = "E2E_KEYCLOAK_REALM" +KEYCLOAK_ADMIN_USER_ENV: Final = "E2E_KEYCLOAK_ADMIN_USER" +KEYCLOAK_ADMIN_PASSWORD_ENV: Final = "E2E_KEYCLOAK_ADMIN_PASSWORD" + +DEFAULT_KEYCLOAK_URL: Final = "http://127.0.0.1:8480" +DEFAULT_REALM: Final = "litellm-e2e" +TESTS_CLIENT_ID: Final = "litellm-e2e-tests" +SHORT_LIVED_CLIENT_ID: Final = "litellm-e2e-shortlived" +SHORT_LIVED_TOKEN_SECONDS: Final = 1 + +_START_HINT: Final = ( + "Start it with the `docker run ... quay.io/keycloak/keycloak` command in tests/e2e/CONTRIBUTING.md, " + f"and point {KEYCLOAK_URL_ENV} / {KEYCLOAK_ADMIN_USER_ENV} / {KEYCLOAK_ADMIN_PASSWORD_ENV} at it" +) + + +class TokenGrantForm(BaseModel): + """The direct-access (password) grant an OAuth 2 token endpoint takes, form encoded.""" + + grant_type: Literal["password"] = "password" + client_id: str + username: str + password: str + + +class TokenResponse(BaseModel): + access_token: str + + +class GroupCreateBody(BaseModel): + name: str + + +class PasswordCredential(BaseModel): + type: Literal["password"] = "password" + value: str + temporary: bool = False + + +class UserCreateBody(BaseModel): + """Keycloak's admin representation of a new user. `firstName` / `lastName` and + an empty `requiredActions` matter: a realm's default VERIFY_PROFILE action + otherwise leaves the account "not fully set up" and every grant fails.""" + + username: str + email: str + email_verified: bool = Field(default=True, alias="emailVerified") + first_name: str = Field(default="E2E", alias="firstName") + last_name: str = Field(default="Tester", alias="lastName") + enabled: bool = True + groups: tuple[str, ...] + credentials: tuple[PasswordCredential, ...] + required_actions: tuple[str, ...] = Field(default=(), alias="requiredActions") + + +def created_id(write: ExternalWrite, context: str) -> str: + """The new resource's id, which Keycloak returns only as the last segment of + the Location header on a 201.""" + if not write.ok: + pytest.fail(f"Keycloak refused to create {context}: HTTP {write.status_code} {write.body[:300]}") + return write.location.rsplit("/", 1)[-1] + + +@dataclass(frozen=True, slots=True) +class Identity: + """One provisioned IdP user: the `sub` the proxy will see, the credential the + test signs in with, and the group whose name the litellm team carries.""" + + user_id: str + username: str + password: str + group: str + group_id: str + + +@dataclass(frozen=True, slots=True) +class Keycloak: + base_url: str + realm: str + admin_username: str + admin_password: str + + @property + def issuer(self) -> str: + return f"{self.base_url}/realms/{self.realm}" + + @property + def jwks_url(self) -> str: + return f"{self.issuer}/protocol/openid-connect/certs" + + def token_url(self, realm: str) -> str: + return f"{self.base_url}/realms/{realm}/protocol/openid-connect/token" + + def _admin_url(self, path: str) -> str: + return f"{self.base_url}/admin/realms/{self.realm}{path}" + + def _admin_headers(self) -> AuthHeaders: + """A fresh admin token per call: the master realm's tokens are short lived, + and a cached one would expire in the middle of a slow test.""" + form: Final = TokenGrantForm(client_id="admin-cli", username=self.admin_username, password=self.admin_password) + result: Final = post_form_external(self.token_url("master"), form=form, response_type=TokenResponse) + return AuthHeaders(authorization=f"Bearer {self._token(result, 'the Keycloak admin credential')}") + + def _token(self, result: Result[TokenResponse], context: str) -> str: + match result: + case Success(data=granted): + return granted.access_token + case NetworkError(message=message): + return pytest.fail(f"No live Keycloak at {self.base_url} for {context}: {message}. {_START_HINT}") + case _: + return pytest.fail(f"Keycloak refused {context}: {result}") + + def create_group(self, name: str) -> str: + return created_id( + post_json_external( + self._admin_url("/groups"), headers=self._admin_headers(), json=GroupCreateBody(name=name) + ), + f"group {name}", + ) + + def create_user(self, *, username: str, email: str, password: str, group: str) -> str: + return created_id( + post_json_external( + self._admin_url("/users"), + headers=self._admin_headers(), + json=UserCreateBody( + username=username, + email=email, + groups=(group,), + credentials=(PasswordCredential(value=password),), + ), + ), + f"user {username}", + ) + + def delete_user(self, user_id: str) -> None: + delete_external(self._admin_url(f"/users/{user_id}"), headers=self._admin_headers()) + + def delete_group(self, group_id: str) -> None: + delete_external(self._admin_url(f"/groups/{group_id}"), headers=self._admin_headers()) + + def provision(self, *, marker: str, group: str) -> Identity: + """Create `group` and a user in it, credentialed with a password generated + for this test alone, and hand back the identity a token can be minted for.""" + group_id: Final = self.create_group(group) + username: Final = f"e2e-jwt-user-{marker}" + password: Final = secrets.token_urlsafe(24) + user_id: Final = self.create_user( + username=username, email=f"{username}@example.com", password=password, group=group + ) + return Identity(user_id=user_id, username=username, password=password, group=group, group_id=group_id) + + def access_token(self, identity: Identity, *, client_id: str = TESTS_CLIENT_ID) -> str: + """Sign `identity` in through the direct-access grant and hand back the + access token Keycloak signed, exactly as it came off the wire.""" + result: Final = post_form_external( + self.token_url(self.realm), + form=TokenGrantForm(client_id=client_id, username=identity.username, password=identity.password), + response_type=TokenResponse, + ) + return self._token(result, f"a token for {identity.username}") + + +def keycloak_from_env() -> Keycloak: + admin_username: Final = os.environ.get(KEYCLOAK_ADMIN_USER_ENV, "").strip() + admin_password: Final = os.environ.get(KEYCLOAK_ADMIN_PASSWORD_ENV, "").strip() + if not admin_username or not admin_password: + pytest.fail( + f"The JWT suite needs {KEYCLOAK_ADMIN_USER_ENV} and {KEYCLOAK_ADMIN_PASSWORD_ENV} to provision " + f"identities in its Keycloak realm, and neither may be empty. {_START_HINT}" + ) + return Keycloak( + base_url=os.environ.get(KEYCLOAK_URL_ENV, DEFAULT_KEYCLOAK_URL).rstrip("/"), + realm=os.environ.get(KEYCLOAK_REALM_ENV, "").strip() or DEFAULT_REALM, + admin_username=admin_username, + admin_password=admin_password, + ) diff --git a/tests/e2e/idp_realm.json b/tests/e2e/idp_realm.json new file mode 100644 index 00000000000..2a8bd89b428 --- /dev/null +++ b/tests/e2e/idp_realm.json @@ -0,0 +1,56 @@ +{ + "realm": "litellm-e2e", + "enabled": true, + "sslRequired": "none", + "registrationAllowed": false, + "accessTokenLifespan": 300, + "clients": [ + { + "clientId": "litellm-e2e-tests", + "enabled": true, + "publicClient": true, + "standardFlowEnabled": false, + "directAccessGrantsEnabled": true, + "protocolMappers": [ + { + "name": "groups", + "protocol": "openid-connect", + "protocolMapper": "oidc-group-membership-mapper", + "consentRequired": false, + "config": { + "claim.name": "groups", + "full.path": "false", + "access.token.claim": "true", + "id.token.claim": "true", + "userinfo.token.claim": "true" + } + } + ] + }, + { + "clientId": "litellm-e2e-shortlived", + "enabled": true, + "publicClient": true, + "standardFlowEnabled": false, + "directAccessGrantsEnabled": true, + "attributes": { + "access.token.lifespan": "1" + }, + "protocolMappers": [ + { + "name": "groups", + "protocol": "openid-connect", + "protocolMapper": "oidc-group-membership-mapper", + "consentRequired": false, + "config": { + "claim.name": "groups", + "full.path": "false", + "access.token.claim": "true", + "id.token.claim": "true", + "userinfo.token.claim": "true" + } + } + ] + } + ] +} diff --git a/tests/e2e/jwt_issuer.py b/tests/e2e/jwt_issuer.py deleted file mode 100644 index 4006fcaa314..00000000000 --- a/tests/e2e/jwt_issuer.py +++ /dev/null @@ -1,230 +0,0 @@ -"""Test-only fake identity provider for the e2e JWT suite. Never deploy it. - -Run it next to the proxy (`uv run python tests/e2e/jwt_issuer.py`). On start it -generates one RSA signing key and keeps it for the life of the process, serving -the public half at `GET /.well-known/jwks.json` and signing whatever claims are -POSTed to `/token`. The proxy's `JWT_PUBLIC_KEY_URL` points at the JWKS URL, and -tests mint RS256 tokens by POSTing claims, so the private key never leaves this -process and no test holds it. - -One key per process, rather than per pytest run, is what survives the proxy's -JWKS cache: the proxy caches the JWKS for `litellm_jwtauth.public_key_ttl` -(600s by default) and does not refetch on an unknown `kid`, so a key rotated -every run would be rejected until the cache expired. Restart the proxy whenever -you restart the issuer. - -The mint endpoint takes no credential: anyone who can reach it gets a token the -proxy trusts. It therefore binds 127.0.0.1 only, must never be exposed beyond -loopback, and must only ever be trusted by a proxy under test. -""" - -from __future__ import annotations - -import logging -import os -import threading -import time -import uuid -from collections.abc import Callable, Mapping -from dataclasses import dataclass -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from typing import Final, Literal - -import jwt -from cryptography.hazmat.primitives import serialization -from cryptography.hazmat.primitives.asymmetric import rsa -from jwt.utils import to_base64url_uint -from pydantic import BaseModel, RootModel, ValidationError - -JWT_ISSUER_PORT_ENV: Final = "E2E_JWT_ISSUER_PORT" -DEFAULT_JWT_ISSUER_PORT: Final = 4190 -LOOPBACK_HOST: Final = "127.0.0.1" -JWKS_PATH: Final = "/.well-known/jwks.json" -TOKEN_PATH: Final = "/token" -DEFAULT_TOKEN_LIFETIME_SECONDS: Final = 300 - -ClaimValue = str | int | float | bool | None | list[str] - - -def jwt_issuer_port() -> int: - raw: Final = os.environ.get(JWT_ISSUER_PORT_ENV, "").strip() - return int(raw) if raw else DEFAULT_JWT_ISSUER_PORT - - -def jwt_issuer_url() -> str: - return f"http://{LOOPBACK_HOST}:{jwt_issuer_port()}" - - -class RsaJwk(BaseModel): - kty: Literal["RSA"] = "RSA" - alg: Literal["RS256"] = "RS256" - use: Literal["sig"] = "sig" - kid: str - n: str - e: str - - -class JwksDocument(BaseModel): - keys: tuple[RsaJwk, ...] - - -class TokenRequest(RootModel[Mapping[str, ClaimValue]]): - """The JSON body of POST /token: the claims to sign, verbatim. Nested objects - are not supported; every value is a scalar or a list of strings.""" - - -class MintedToken(BaseModel): - token: str - - -class IssuerError(BaseModel): - error: str - - -@dataclass(frozen=True, slots=True) -class SigningKey: - kid: str - private_pem: str - jwk: RsaJwk - - -def generate_signing_key(kid: str | None = None) -> SigningKey: - private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) - numbers: Final = private_key.public_key().public_numbers() - resolved_kid: Final = kid if kid is not None else uuid.uuid4().hex - return SigningKey( - kid=resolved_kid, - private_pem=private_key.private_bytes( - encoding=serialization.Encoding.PEM, - format=serialization.PrivateFormat.PKCS8, - encryption_algorithm=serialization.NoEncryption(), - ).decode(), - jwk=RsaJwk( - kid=resolved_kid, - n=to_base64url_uint(numbers.n).decode(), - e=to_base64url_uint(numbers.e).decode(), - ), - ) - - -def mint( - key: SigningKey, - claims: Mapping[str, ClaimValue], - *, - issuer: str, - now: int, - lifetime_seconds: int = DEFAULT_TOKEN_LIFETIME_SECONDS, -) -> str: - """Sign `claims` as a compact RS256 JWT carrying `key.kid` in its header. - `iss`, `iat`, and `exp` are filled in when absent and left alone when the - caller sets them, so a test can mint an already-expired token.""" - payload: Final[dict[str, ClaimValue]] = { - "iss": issuer, - "iat": now, - "exp": now + lifetime_seconds, - **claims, - } - return jwt.encode(payload, key.private_pem, algorithm="RS256", headers={"kid": key.kid}) - - -class _IssuerServer(ThreadingHTTPServer): - daemon_threads = True - - def __init__(self, bind: tuple[str, int], *, key: SigningKey, clock: Callable[[], int]) -> None: - super().__init__(bind, _IssuerHandler) - self.key: Final = key - self.clock: Final = clock - - @property - def url(self) -> str: - host, port = self.server_address[0], self.server_address[1] - return f"http://{host}:{port}" - - -class _IssuerHandler(BaseHTTPRequestHandler): - def _issuer(self) -> _IssuerServer: - issuer: Final = self.server - assert isinstance(issuer, _IssuerServer) - return issuer - - def do_GET(self) -> None: - if self.path != JWKS_PATH: - self._send(404, IssuerError(error=f"unknown path {self.path}; the JWKS is at {JWKS_PATH}")) - return - self._send(200, JwksDocument(keys=(self._issuer().key.jwk,))) - - def do_POST(self) -> None: - if self.path != TOKEN_PATH: - self._send(404, IssuerError(error=f"unknown path {self.path}; mint tokens at {TOKEN_PATH}")) - return - length: Final = int(self.headers.get("Content-Length", "0")) - try: - request: Final = TokenRequest.model_validate_json(self.rfile.read(length)) - except ValidationError as exc: - self._send(400, IssuerError(error=f"claims must be a JSON object of scalar or string-list values: {exc}")) - return - issuer: Final = self._issuer() - token: Final = mint(issuer.key, request.root, issuer=issuer.url, now=issuer.clock()) - self._send(200, MintedToken(token=token)) - - def _send(self, status: int, body: BaseModel) -> None: - payload: Final = body.model_dump_json().encode() - self.send_response(status) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(payload))) - self.end_headers() - self.wfile.write(payload) - - -@dataclass(frozen=True, slots=True) -class RunningIssuer: - url: str - key: SigningKey - server: _IssuerServer - - @property - def jwks_url(self) -> str: - return f"{self.url}{JWKS_PATH}" - - def shutdown(self) -> None: - self.server.shutdown() - self.server.server_close() - - -def _wall_clock() -> int: - return int(time.time()) - - -def start_jwt_issuer( - *, - port: int = 0, - key: SigningKey | None = None, - clock: Callable[[], int] = _wall_clock, -) -> RunningIssuer: - """Serve the issuer on loopback in a daemon thread. `port=0` takes an - OS-assigned port for in-process tests; the CLI passes the documented one.""" - server: Final = _IssuerServer((LOOPBACK_HOST, port), key=key or generate_signing_key(), clock=clock) - thread: Final = threading.Thread(target=server.serve_forever, name="e2e-jwt-issuer", daemon=True) - thread.start() - return RunningIssuer(url=server.url, key=server.key, server=server) - - -def main() -> None: - logging.basicConfig(level=logging.INFO, format="%(message)s") - running: Final = start_jwt_issuer(port=jwt_issuer_port()) - logging.getLogger(__name__).info( - "e2e jwt issuer listening on %s jwks=%s mint=POST %s%s kid=%s (test-only, loopback, ctrl-c to stop)", - running.url, - running.jwks_url, - running.url, - TOKEN_PATH, - running.key.kid, - ) - try: - threading.Event().wait() - except KeyboardInterrupt: - running.shutdown() - - -if __name__ == "__main__": - main() diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 17f5418bd73..2c6c0e9bbd4 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1078,18 +1078,6 @@ class UserListResponse(BaseModel): total: int -class JwtClaimsBody(BaseModel): - """Claims POSTed to the e2e JWT issuer's /token: exactly what the - `litellm_jwtauth` block in CONTRIBUTING.md reads (sub -> user_id, email -> - user_email, groups -> team ids), plus an explicit `exp` for the expired case; - the issuer fills in iss/iat/exp when they are left unset.""" - - sub: str - email: str - groups: Sequence[str] - exp: int | None = None - - class OrgNewBody(BaseModel): organization_alias: str models: list[str] = [] diff --git a/tests/e2e/other/other_client.py b/tests/e2e/other/other_client.py index b51f3338f52..4313bbe4068 100644 --- a/tests/e2e/other/other_client.py +++ b/tests/e2e/other/other_client.py @@ -1,27 +1,23 @@ """Client for the `other` holding-pen suite: the auth gate (master key vs an -invalid key on an admin route), JWT auth against the test-only issuer -(jwt_issuer.py), and the process-lifecycle health probes (liveness, public -readiness, authenticated readiness diagnostics). +invalid key on an admin route), JWT auth against the suite's Keycloak realm +(idp.py), and the process-lifecycle health probes (liveness, public readiness, +authenticated readiness diagnostics). Holds the shared ProxyClient so `resources` / `scoped_key` still clean up, and adds only the routes these behaviors need. The health probes deliberately send no auth header (public routes), so they go through the transport with an empty -headers model rather than a bearer. Tokens are minted by POSTing claims to the -issuer, so no test ever holds a signing key. +headers model rather than a bearer. JWT tests reach the identity provider +through `idp`, which provisions identities and mints tokens through Keycloak's +own endpoints, so no test ever holds a signing key. """ from __future__ import annotations from dataclasses import dataclass -from typing import Final -import pytest - -from e2e_config import JWT_ISSUER_URL -from e2e_http import NetworkError, NoBody, ProbeResult, Result, Success, post_external -from jwt_issuer import TOKEN_PATH, MintedToken +from e2e_http import NoBody, ProbeResult, Result +from idp import Keycloak, keycloak_from_env from models import ( - JwtClaimsBody, ReadinessDetailsResponse, ReadinessResponse, UserListParams, @@ -33,7 +29,11 @@ from proxy_client import ProxyClient @dataclass(frozen=True, slots=True) class OtherClient: proxy: ProxyClient - jwt_issuer_url: str + + @property + def idp(self) -> Keycloak: + """Resolved per use, so the suite's non-JWT tests never need the IdP env.""" + return keycloak_from_env() def liveness(self) -> ProbeResult: """GET /health/liveliness. Unauthenticated; the probe returns status + @@ -77,21 +77,6 @@ class OtherClient: response_type=UserListResponse, ) - def mint_jwt(self, claims: JwtClaimsBody) -> str: - """Have the test-only issuer sign `claims` into a compact RS256 JWT. A - missing issuer is a hard failure naming the start command, not a skip.""" - result: Final = post_external(f"{self.jwt_issuer_url}{TOKEN_PATH}", json=claims, response_type=MintedToken) - match result: - case Success(data=minted): - return minted.token - case NetworkError(message=message): - pytest.fail( - f"No live JWT issuer at {self.jwt_issuer_url}: {message}. Start it next to the proxy with " - "`uv run python tests/e2e/jwt_issuer.py` (see CONTRIBUTING.md)" - ) - case _: - raise AssertionError(result) - def build_client(proxy: ProxyClient) -> OtherClient: - return OtherClient(proxy=proxy, jwt_issuer_url=JWT_ISSUER_URL) + return OtherClient(proxy=proxy) diff --git a/tests/e2e/other/test_jwt_auth_e2e.py b/tests/e2e/other/test_jwt_auth_e2e.py index 49a9350a3a2..25a5e7d4b9a 100644 --- a/tests/e2e/other/test_jwt_auth_e2e.py +++ b/tests/e2e/other/test_jwt_auth_e2e.py @@ -1,54 +1,63 @@ -"""Live e2e: RS256 JWTs minted by the test-only issuer (jwt_issuer.py) against a +"""Live e2e: access tokens issued by a real Keycloak realm (idp.py) against a proxy running with `enable_jwt_auth: true` and the `litellm_jwtauth` block from CONTRIBUTING.md (sub -> user_id, email -> user_email, groups -> team ids, user_id_upsert). -Every case mints through the issuer, so the tests never hold a signing key: the -bad-signature case corrupts a genuine signature, the expired case asks the -issuer for a token whose `exp` is already in the past. Those identities get -their own freshly created team so a rejection can only be blamed on the token, -while the unknown-team case names a team that was never created. An acceptance -is proven twice, at the boundary (200 from a real provider) and in the spend log -the proxy attributes to the claims. The last case keeps a plain `sk-` virtual -key working on the same proxy, guarding against the flag turning JWT on for -everyone. +Every identity is provisioned in Keycloak for the test that uses it: a group +named after the litellm team, and a user in that group whose password exists +only for the length of the test. Tokens then come from Keycloak's direct-access +grant, so no test ever holds a signing key and the claims the proxy reads are +the ones an IdP really emits (`sub` is Keycloak's user uuid, `groups` comes off +a protocol mapper, `aud` is Keycloak's own audience). + +The rejection cases stay honest about where the rejection has to come from: the +bad-signature case corrupts a genuine signature, and the expiry case takes its +token from the realm's one-second client and waits for it to lapse rather than +forging a stale `exp`. Those identities get their own freshly created team, so a +rejection can only be blamed on the token, while the unknown-team case names a +group no litellm team was ever created for. An acceptance is proven twice, at +the boundary (200 from a real provider) and in the spend log the proxy +attributes to the claims. The last case keeps a plain `sk-` virtual key working +on the same proxy, guarding against the flag turning JWT on for everyone. """ from __future__ import annotations -from dataclasses import dataclass +import time from typing import Final import pytest from e2e_config import CHEAP_OPENAI_MODEL, unique_marker from e2e_http import UnauthorizedError, UnknownApiError, unwrap +from idp import SHORT_LIVED_CLIENT_ID, SHORT_LIVED_TOKEN_SECONDS, Identity from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, JwtClaimsBody, TeamNewBody +from models import ChatBody, ChatMessage, TeamNewBody from other_client import OtherClient pytestmark = pytest.mark.e2e -@dataclass(frozen=True, slots=True) -class JwtIdentity: - user_id: str - team_id: str - - def claims(self, *, exp: int | None = None) -> JwtClaimsBody: - return JwtClaimsBody(sub=self.user_id, email=f"{self.user_id}@example.com", groups=(self.team_id,), exp=exp) +def _provision(client: OtherClient, resources: ResourceManager, *, marker: str) -> Identity: + """A Keycloak group and a user in it, torn down with the test. The group name + is what the token's `groups` claim carries, which is what the proxy resolves + as a litellm team id.""" + identity: Final = client.idp.provision(marker=marker, group=f"e2e-jwt-team-{marker}") + resources.defer(lambda: client.idp.delete_user(identity.user_id)) + resources.defer(lambda: client.idp.delete_group(identity.group_id)) + resources.defer(lambda: client.proxy.delete_user(identity.user_id)) + return identity @pytest.fixture -def identity(client: OtherClient, resources: ResourceManager) -> JwtIdentity: +def identity(client: OtherClient, resources: ResourceManager) -> Identity: + """An IdP identity whose group is also a real litellm team, so anything the + proxy rejects is about the token and never about an unresolvable team.""" marker: Final = unique_marker() - team_id: Final = client.proxy.create_team( - TeamNewBody(team_alias=f"e2e-jwt-{marker}", team_id=f"e2e-jwt-team-{marker}") - ) + provisioned: Final = _provision(client, resources, marker=marker) + team_id: Final = client.proxy.create_team(TeamNewBody(team_alias=f"e2e-jwt-{marker}", team_id=provisioned.group)) resources.defer(lambda: client.proxy.delete_team(team_id)) - user_id: Final = f"e2e-jwt-user-{marker}" - resources.defer(lambda: client.proxy.delete_user(user_id)) - return JwtIdentity(user_id=user_id, team_id=team_id) + return provisioned def _ping() -> ChatBody: @@ -68,9 +77,9 @@ def _corrupt_signature(token: str) -> str: class TestJwtAuth: @pytest.mark.covers("other.auth.jwt.valid_token_allows", "other.auth.jwt.spend_attributed_to_claims") def test_valid_token_for_an_existing_team_is_accepted_and_attributed( - self, client: OtherClient, identity: JwtIdentity + self, client: OtherClient, identity: Identity ) -> None: - token: Final = client.mint_jwt(identity.claims()) + token: Final = client.idp.access_token(identity) response: Final = unwrap(client.proxy.chat(token, _ping())) assert response.id is not None and response.choices, ( @@ -80,16 +89,16 @@ class TestJwtAuth: rows: Final = client.proxy.poll_logs_for_request_id(response.id) assert rows, f"no spend log row for request {response.id} within the poll deadline" row: Final = rows[0] - assert row.team_id == identity.team_id, ( - f"spend row must carry the team from the JWT groups claim {identity.team_id!r}, got {row.team_id!r}" + assert row.team_id == identity.group, ( + f"spend row must carry the team from the JWT groups claim {identity.group!r}, got {row.team_id!r}" ) assert row.user == identity.user_id, ( f"spend row must carry the user from the JWT sub claim {identity.user_id!r}, got {row.user!r}" ) @pytest.mark.covers("other.auth.jwt.invalid_signature_denied") - def test_tampered_signature_is_rejected(self, client: OtherClient, identity: JwtIdentity) -> None: - tampered: Final = _corrupt_signature(client.mint_jwt(identity.claims())) + def test_tampered_signature_is_rejected(self, client: OtherClient, identity: Identity) -> None: + tampered: Final = _corrupt_signature(client.idp.access_token(identity)) result: Final = client.proxy.chat(tampered, _ping()) assert isinstance(result, UnauthorizedError), ( @@ -100,27 +109,29 @@ class TestJwtAuth: ) @pytest.mark.covers("other.auth.jwt.expired_denied") - def test_expired_token_is_rejected(self, client: OtherClient, identity: JwtIdentity) -> None: - expired: Final = client.mint_jwt(identity.claims(exp=1)) + def test_expired_token_is_rejected(self, client: OtherClient, identity: Identity) -> None: + expiring: Final = client.idp.access_token(identity, client_id=SHORT_LIVED_CLIENT_ID) + time.sleep(SHORT_LIVED_TOKEN_SECONDS + 1) - result: Final = client.proxy.chat(expired, _ping()) + result: Final = client.proxy.chat(expiring, _ping()) assert isinstance(result, UnauthorizedError), ( f"an expired JWT must be rejected with 401 even though its signature verifies, got {result}" ) assert "expired" in result.body.lower(), f"the 401 must say the token expired, got {result.body[:300]}" @pytest.mark.covers("other.auth.jwt.unknown_team_denied") - def test_token_naming_a_team_that_does_not_exist_is_rejected(self, client: OtherClient) -> None: - marker: Final = unique_marker() - never_created: Final = JwtIdentity(user_id=f"e2e-jwt-user-{marker}", team_id=f"e2e-jwt-missing-team-{marker}") - token: Final = client.mint_jwt(never_created.claims()) + def test_token_naming_a_team_that_does_not_exist_is_rejected( + self, client: OtherClient, resources: ResourceManager + ) -> None: + stranger: Final = _provision(client, resources, marker=unique_marker()) + token: Final = client.idp.access_token(stranger) result: Final = client.proxy.chat(token, _ping()) assert isinstance(result, UnknownApiError) and result.status_code == 403, ( f"a valid JWT whose groups name no existing team must be rejected with 403, got {result}" ) - assert never_created.team_id in result.body, ( - f"the 403 must name the team it could not resolve ({never_created.team_id}), got {result.body[:300]}" + assert stranger.group in result.body, ( + f"the 403 must name the team it could not resolve ({stranger.group}), got {result.body[:300]}" ) @pytest.mark.covers("other.auth.jwt.virtual_key_unaffected") diff --git a/tests/e2e/test_idp.py b/tests/e2e/test_idp.py new file mode 100644 index 00000000000..93be2806eb5 --- /dev/null +++ b/tests/e2e/test_idp.py @@ -0,0 +1,79 @@ +"""Harness coverage for idp.py: the pure parts of the Keycloak client, which are +the ones a wrong value in silently mistargets. No proxy and no IdP needed, so +these carry no `e2e` marker and run everywhere.""" + +from __future__ import annotations + +from typing import Final + +import pytest + +from e2e_http import ExternalWrite +from idp import ( + KEYCLOAK_ADMIN_PASSWORD_ENV, + KEYCLOAK_ADMIN_USER_ENV, + KEYCLOAK_REALM_ENV, + KEYCLOAK_URL_ENV, + Keycloak, + PasswordCredential, + created_id, + UserCreateBody, + keycloak_from_env, +) + +_REALM: Final = Keycloak( + base_url="http://keycloak:8080", realm="litellm-e2e", admin_username="admin", admin_password="pw" +) + + +def test_realm_urls_match_keycloaks_own_layout() -> None: + assert _REALM.issuer == "http://keycloak:8080/realms/litellm-e2e" + assert _REALM.jwks_url == "http://keycloak:8080/realms/litellm-e2e/protocol/openid-connect/certs" + assert _REALM.token_url("master") == "http://keycloak:8080/realms/master/protocol/openid-connect/token" + + +def test_created_id_is_the_last_segment_of_the_location_header() -> None: + created: Final = ExternalWrite( + status_code=201, location="http://keycloak:8080/admin/realms/litellm-e2e/groups/abc-123" + ) + assert created_id(created, "a group") == "abc-123" + + +def test_a_refused_create_fails_the_test_with_the_idps_own_words() -> None: + with pytest.raises(BaseException, match=r"409.*already exists"): + created_id(ExternalWrite(status_code=409, body="Group already exists"), "a group") + + +def test_new_users_are_born_fully_set_up() -> None: + """A user without a profile or with a pending required action authenticates + nowhere: Keycloak answers every grant with "Account is not fully set up".""" + body: Final = UserCreateBody( + username="e2e", email="e2e@example.com", groups=("team",), credentials=(PasswordCredential(value="pw"),) + ).model_dump(by_alias=True) + + assert body["requiredActions"] == () + assert body["firstName"] and body["lastName"] and body["emailVerified"] is True + assert body["credentials"][0]["temporary"] is False + + +def test_connection_details_come_from_the_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(KEYCLOAK_URL_ENV, "http://keycloak.litellm.svc.cluster.local:8080/") + monkeypatch.setenv(KEYCLOAK_REALM_ENV, "other-realm") + monkeypatch.setenv(KEYCLOAK_ADMIN_USER_ENV, "admin") + monkeypatch.setenv(KEYCLOAK_ADMIN_PASSWORD_ENV, "pw") + + resolved: Final = keycloak_from_env() + + assert resolved.issuer == "http://keycloak.litellm.svc.cluster.local:8080/realms/other-realm" + assert resolved.admin_username == "admin" and resolved.admin_password == "pw" + + +@pytest.mark.parametrize("blank", ["", " "]) +def test_a_missing_admin_credential_fails_loudly_instead_of_skipping( + monkeypatch: pytest.MonkeyPatch, blank: str +) -> None: + monkeypatch.setenv(KEYCLOAK_ADMIN_USER_ENV, "admin") + monkeypatch.setenv(KEYCLOAK_ADMIN_PASSWORD_ENV, blank) + + with pytest.raises(BaseException, match=KEYCLOAK_ADMIN_PASSWORD_ENV): + keycloak_from_env() diff --git a/tests/e2e/test_jwt_issuer.py b/tests/e2e/test_jwt_issuer.py deleted file mode 100644 index 7acaaf13949..00000000000 --- a/tests/e2e/test_jwt_issuer.py +++ /dev/null @@ -1,141 +0,0 @@ -"""Harness coverage for the test-only JWT issuer (jwt_issuer.py). - -No proxy and no ``e2e`` marker. The issuer is booted in-process on an -OS-assigned port with a fixed clock and driven over HTTP through -``e2e_http.post_external`` / ``get_external``, the same transport the live -suite uses, so what is pinned here is the contract the live JWT tests lean on: -a token minted at ``/token`` verifies against the key served at -``/.well-known/jwks.json`` under the ``kid`` in its header, ``iss``/``iat``/``exp`` -are filled in only when the caller left them out, and malformed claim bodies or -unknown paths are refused instead of signed. -""" - -from __future__ import annotations - -import time -from collections.abc import Iterator -from typing import Final - -import jwt -import pytest -from pydantic import BaseModel, RootModel - -from e2e_http import UnknownApiError, get_external, post_external, unwrap -from jwt_issuer import ( - DEFAULT_TOKEN_LIFETIME_SECONDS, - JWKS_PATH, - JWT_ISSUER_PORT_ENV, - TOKEN_PATH, - JwksDocument, - MintedToken, - RunningIssuer, - jwt_issuer_port, - start_jwt_issuer, -) - -FROZEN_NOW: Final = int(time.time()) - - -class _Claims(BaseModel): - sub: str - groups: tuple[str, ...] = () - exp: int | None = None - - -class _DecodedClaims(BaseModel): - sub: str - iss: str - iat: int - exp: int - groups: tuple[str, ...] = () - - -class _NotAnObject(RootModel[tuple[str, ...]]): - pass - - -@pytest.fixture(scope="module") -def issuer() -> Iterator[RunningIssuer]: - running: Final = start_jwt_issuer(clock=lambda: FROZEN_NOW) - yield running - running.shutdown() - - -def _mint(issuer: RunningIssuer, claims: _Claims) -> str: - return unwrap(post_external(f"{issuer.url}{TOKEN_PATH}", json=claims, response_type=MintedToken)).token - - -def _served_jwks(issuer: RunningIssuer) -> JwksDocument: - return unwrap(get_external(issuer.jwks_url, response_type=JwksDocument)) - - -def _decode(token: str, jwks: JwksDocument, *, verify_exp: bool = True) -> _DecodedClaims: - key: Final = jwt.PyJWK.from_json(jwks.keys[0].model_dump_json()) - decoded: Final = jwt.decode(token, key, algorithms=["RS256"], options={"verify_exp": verify_exp}) - return _DecodedClaims.model_validate(decoded) - - -class TestJwtIssuer: - def test_minted_token_verifies_against_the_served_jwks(self, issuer: RunningIssuer) -> None: - token: Final = _mint(issuer, _Claims(sub="alice", groups=("team-a",))) - jwks: Final = _served_jwks(issuer) - - assert len(jwks.keys) == 1 - assert jwt.get_unverified_header(token)["kid"] == jwks.keys[0].kid - claims: Final = _decode(token, jwks) - assert claims.sub == "alice" - assert claims.groups == ("team-a",) - assert claims.iss == issuer.url - assert claims.iat == FROZEN_NOW - assert claims.exp == FROZEN_NOW + DEFAULT_TOKEN_LIFETIME_SECONDS - - def test_a_token_signed_by_another_key_does_not_verify(self, issuer: RunningIssuer) -> None: - other: Final = start_jwt_issuer(clock=lambda: FROZEN_NOW) - try: - foreign_token: Final = _mint(other, _Claims(sub="alice")) - finally: - other.shutdown() - - with pytest.raises(jwt.InvalidSignatureError): - _decode(foreign_token, _served_jwks(issuer)) - - def test_an_explicit_exp_is_signed_as_given(self, issuer: RunningIssuer) -> None: - expired_at: Final = FROZEN_NOW - 60 - token: Final = _mint(issuer, _Claims(sub="alice", exp=expired_at)) - jwks: Final = _served_jwks(issuer) - - assert _decode(token, jwks, verify_exp=False).exp == expired_at - with pytest.raises(jwt.ExpiredSignatureError): - _ = _decode(token, jwks) - - def test_non_object_claims_are_refused(self, issuer: RunningIssuer) -> None: - result: Final = post_external( - f"{issuer.url}{TOKEN_PATH}", json=_NotAnObject(("not", "claims")), response_type=MintedToken - ) - assert isinstance(result, UnknownApiError) - assert result.status_code == 400 - - @pytest.mark.parametrize( - ("method", "path"), - [("GET", TOKEN_PATH), ("POST", JWKS_PATH), ("GET", "/token/anything")], - ids=["get-token", "post-jwks", "get-other"], - ) - def test_unknown_routes_are_404(self, issuer: RunningIssuer, method: str, path: str) -> None: - url: Final = f"{issuer.url}{path}" - result: Final = ( - get_external(url, response_type=MintedToken) - if method == "GET" - else post_external(url, json=_Claims(sub="alice"), response_type=MintedToken) - ) - assert isinstance(result, UnknownApiError) - assert result.status_code == 404 - - -class TestIssuerPort: - def test_defaults_to_the_documented_port(self, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.delenv(JWT_ISSUER_PORT_ENV, raising=False) - assert jwt_issuer_port() == 4190, "CONTRIBUTING.md hardcodes 4190 in JWT_PUBLIC_KEY_URL" - - def test_env_override_wins(self, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv(JWT_ISSUER_PORT_ENV, " 4321 ") - assert jwt_issuer_port() == 4321 From 1213d7d39929500883fee4c3823f8dfbc073fdb8 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 9 Sep 2026 17:17:33 -0700 Subject: [PATCH 017/119] test(e2e): cover /embeddings and assert memory and fallbacks in the Redis timeout test Add an embeddings case with its own closed-port primary and mock backup (the fallback map in the gateway config gains the pair; LiteLLMParamsBody.mock_response accepts the list an embedding mock needs). Assert from /metrics that the proxy's resident memory grows by no more than 200 MB across each case where the process collector reports it (Linux), that the router counted a successful fallback for every request, and that every spend row is a success. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_014ZDULyJPp17ZFiJenRxs2T --- tests/e2e/coverage_registry/reliability.yaml | 2 +- tests/e2e/gateway/redis_timeout_ci_config.yml | 2 + tests/e2e/models.py | 2 +- tests/e2e/router/test_redis_timeout_e2e.py | 142 +++++++++++++----- 4 files changed, 110 insertions(+), 38 deletions(-) diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index 3b378b56870..a42e4b2779a 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -31,7 +31,7 @@ - {id: reliability.cache.exact.returns_cached, module: reliability, tier: P1, behavior: cache, variant: exact, assertions: [returns_cached], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/caching.py", rationale: "Response cache returns cached on exact match"} - {id: reliability.cache.prompt_caching_model_select.returns_cached, module: reliability, tier: P1, behavior: cache, variant: prompt_caching_model_select, assertions: [returns_cached], exercised_on: [chat_completions], source: "router_utils/prompt_caching_cache.py", rationale: "Selects model supporting prompt caching for cacheable prefix"} - {id: reliability.circuit_breaker.redis.trips_then_recovers, module: reliability, tier: P0, behavior: circuit_breaker, variant: redis, assertions: [trips_then_recovers], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/redis_cache.py:99", rationale: "Redis breaker CLOSED->OPEN->HALF_OPEN; guards all cache/rate-limit ops"} -- {id: reliability.circuit_breaker.redis_timeout.stays_responsive, module: reliability, tier: P1, behavior: circuit_breaker, variant: redis_timeout, assertions: [stays_responsive], exercised_on: [chat_completions, responses], source: "litellm/proxy/hooks/proxy_track_cost_callback.py:386", fail_before_fix: proven, rationale: "With every Redis command timing out and every request retrying then falling back, per-request latency stays flat, liveliness stays fast, and spend rows still land; on v1.100.0 the failed-tracking alert body doubled per request until the worker OOMed (LIT-6780)"} +- {id: reliability.circuit_breaker.redis_timeout.stays_responsive, module: reliability, tier: P1, behavior: circuit_breaker, variant: redis_timeout, assertions: [stays_responsive], exercised_on: [chat_completions, responses, embeddings], source: "litellm/proxy/hooks/proxy_track_cost_callback.py:386", fail_before_fix: proven, rationale: "With every Redis command timing out and every request retrying then falling back, per-request latency stays flat, liveliness stays fast, and spend rows still land; on v1.100.0 the failed-tracking alert body doubled per request until the worker OOMed (LIT-6780)"} - {id: reliability.timeout.request_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: request_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions, messages], source: "litellm/router.py:545-551", rationale: "Per-request timeout raises Timeout"} - {id: reliability.timeout.stream_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: stream_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions], source: "litellm/router.py:551", rationale: "Streaming chunk-delivery timeout"} - {id: reliability.perf.throughput.under_slo, module: reliability, tier: P1, behavior: perf, variant: throughput, assertions: [under_slo], exercised_on: [chat_completions, messages], source: grammar, rationale: "Throughput SLO under load"} diff --git a/tests/e2e/gateway/redis_timeout_ci_config.yml b/tests/e2e/gateway/redis_timeout_ci_config.yml index 9bcf65162d6..fd283dbc7df 100644 --- a/tests/e2e/gateway/redis_timeout_ci_config.yml +++ b/tests/e2e/gateway/redis_timeout_ci_config.yml @@ -17,3 +17,5 @@ router_settings: fallbacks: - redis-timeout-primary: - redis-timeout-backup + - redis-timeout-embed-primary: + - redis-timeout-embed-backup diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 62810e6cfd9..f3c111bdf06 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -920,7 +920,7 @@ class LiteLLMParamsBody(BaseModel): auto_router_default_model: str | None = None auto_router_embedding_model: str | None = None tags: list[str] | None = None - mock_response: str | None = None + mock_response: str | list[float] | None = None timeout: float | None = None tpm: int | None = None weight: int | None = None diff --git a/tests/e2e/router/test_redis_timeout_e2e.py b/tests/e2e/router/test_redis_timeout_e2e.py index 58d10f1d8ea..3f1984c0696 100644 --- a/tests/e2e/router/test_redis_timeout_e2e.py +++ b/tests/e2e/router/test_redis_timeout_e2e.py @@ -3,12 +3,12 @@ Runs only against a proxy booted from tests/e2e/gateway/redis_timeout_ci_config.yml, which points cache_params at a real Redis with socket_timeout 0.001. The test holds that Redis in CLIENT PAUSE WRITE for its duration, so every write the proxy sends, the spend counter increment -included, hangs past the timeout, and it proves the degradation was real from the breaker metrics on /metrics: fresh timeouts, a breaker transition, -or an already-open breaker rejecting every call, which is the state a customer's worker sits in. The test registers two deployments through /model/new: a primary whose api_base is a closed port -and a backup that answers with a mock. Each request fails the primary, retries, falls back and succeeds, so it carries retry breadcrumbs; its cost -tracking then fails on the spend counter increment and stringifies the request metadata into a -failed-tracking alert. On v1.100.0 that string doubled per request until the worker hung -(LIT-6780). Deselected unless E2E_REDIS_TIMEOUT is set, since it needs that dedicated proxy. +included, hangs past the timeout. For each endpoint it registers two deployments through +/model/new: a primary whose api_base is a closed port and a backup that answers with a mock. Each request fails the primary, +retries, falls back and succeeds, so it carries retry breadcrumbs; its cost tracking then fails on +the spend counter increment and stringifies the request metadata into a failed-tracking alert. On +v1.100.0 that string doubled per request until the worker hung (LIT-6780). Deselected unless +E2E_REDIS_TIMEOUT is set, since it needs that dedicated proxy. """ from __future__ import annotations @@ -26,25 +26,29 @@ from complexity_router_client import ComplexityRouterClient from e2e_config import unique_marker from e2e_http import NoBody, Result, Success from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, ChatResponse, KeyGenerateBody, LiteLLMParamsBody +from models import ChatBody, ChatMessage, ChatResponse, EmbedBody, KeyGenerateBody, LiteLLMParamsBody from proxy_client import ProxyClient from pydantic import BaseModel pytestmark = [pytest.mark.e2e, pytest.mark.redis_timeout] -PRIMARY_MODEL: Final = "redis-timeout-primary" -BACKUP_MODEL: Final = "redis-timeout-backup" -BACKING_MODEL: Final = "openai/gpt-5-mini" -CLOSED_PORT_API_BASE: Final = "http://127.0.0.1:1" REQUESTS: Final = 20 MAX_SECONDS_PER_REQUEST: Final = 10.0 MAX_LATENCY_GROWTH_RATIO: Final = 3.0 MAX_LIVELINESS_SECONDS: Final = 2.0 +MAX_RSS_GROWTH_BYTES: Final = 200 * 1024 * 1024 REDIS_PAUSE_MS: Final = 600_000 BREAKER_FAILURE_THRESHOLD: Final = 5 +CLOSED_PORT_API_BASE: Final = "http://127.0.0.1:1" TIMEOUT_FAILURES_RE: Final = re.compile( r'^litellm_redis_circuit_breaker_failures_total\{failure_class="timeout"\} ([0-9.e+]+)$', re.M ) +BREAKER_OPEN_RE: Final = re.compile(r'^litellm_redis_circuit_breaker_state\{state="open"\} ([0-9.e+]+)$', re.M) +BREAKER_TRANSITIONS_RE: Final = re.compile( + r'^litellm_redis_circuit_breaker_transitions_total\{state="[a-z_]+"\} ([0-9.e+]+)$', re.M +) +FALLBACKS_RE: Final = re.compile(r"^litellm_deployment_successful_fallbacks_total\{[^}]*\} ([0-9.e+]+)$", re.M) +RSS_RE: Final = re.compile(r"^process_resident_memory_bytes ([0-9.e+]+)$", re.M) class ResponsesBody(BaseModel): @@ -59,48 +63,95 @@ class ResponsesObject(BaseModel): output: list[object] = [] +class EmbeddingsObject(BaseModel): + model: str | None = None + data: list[object] = [] + + @dataclass(frozen=True, slots=True) class Endpoint: + """One endpoint's deployments and request shape.""" + name: str - send: Callable[[ProxyClient, str, str], Result[BaseModel]] + primary: str + backup: str + primary_params: LiteLLMParamsBody + backup_params: LiteLLMParamsBody + send: Callable[[ProxyClient, str, str, str], Result[BaseModel]] served: Callable[[BaseModel], bool] -def _send_chat(proxy: ProxyClient, key: str, marker: str) -> Result[BaseModel]: +def _send_chat(proxy: ProxyClient, key: str, model: str, marker: str) -> Result[BaseModel]: return proxy.transport.post( "/chat/completions", headers=proxy.transport.bearer(key), - json=ChatBody(model=PRIMARY_MODEL, messages=[ChatMessage(role="user", content=marker)], max_tokens=5), + json=ChatBody(model=model, messages=[ChatMessage(role="user", content=marker)], max_tokens=5), response_type=ChatResponse, timeout=MAX_SECONDS_PER_REQUEST, ) -def _send_responses(proxy: ProxyClient, key: str, marker: str) -> Result[BaseModel]: +def _send_responses(proxy: ProxyClient, key: str, model: str, marker: str) -> Result[BaseModel]: return proxy.transport.post( "/v1/responses", headers=proxy.transport.bearer(key), - json=ResponsesBody(model=PRIMARY_MODEL, input=marker), + json=ResponsesBody(model=model, input=marker), response_type=ResponsesObject, timeout=MAX_SECONDS_PER_REQUEST, ) +def _send_embeddings(proxy: ProxyClient, key: str, model: str, marker: str) -> Result[BaseModel]: + return proxy.transport.post( + "/embeddings", + headers=proxy.transport.bearer(key), + json=EmbedBody(model=model, input=marker), + response_type=EmbeddingsObject, + timeout=MAX_SECONDS_PER_REQUEST, + ) + + +_CHAT_PRIMARY: Final = LiteLLMParamsBody( + model="openai/gpt-5-mini", api_key="sk-redis-timeout-primary-not-used", api_base=CLOSED_PORT_API_BASE +) +_CHAT_BACKUP: Final = LiteLLMParamsBody( + model="openai/gpt-5-nano", api_key="sk-redis-timeout-backup-not-used", mock_response="ok" +) +_EMBED_PRIMARY: Final = LiteLLMParamsBody( + model="openai/text-embedding-3-small", api_key="sk-redis-timeout-primary-not-used", api_base=CLOSED_PORT_API_BASE +) +_EMBED_BACKUP: Final = LiteLLMParamsBody( + model="openai/text-embedding-3-large", api_key="sk-redis-timeout-backup-not-used", mock_response=[0.1, 0.2, 0.3] +) + ENDPOINTS: Final = ( Endpoint( name="chat_completions", + primary="redis-timeout-primary", + backup="redis-timeout-backup", + primary_params=_CHAT_PRIMARY, + backup_params=_CHAT_BACKUP, send=_send_chat, served=lambda data: isinstance(data, ChatResponse) and bool(data.choices), ), Endpoint( name="responses", + primary="redis-timeout-primary", + backup="redis-timeout-backup", + primary_params=_CHAT_PRIMARY, + backup_params=_CHAT_BACKUP, send=_send_responses, served=lambda data: isinstance(data, ResponsesObject) and bool(data.output), ), -) -BREAKER_OPEN_RE: Final = re.compile(r'^litellm_redis_circuit_breaker_state\{state="open"\} ([0-9.e+]+)$', re.M) -BREAKER_TRANSITIONS_RE: Final = re.compile( - r'^litellm_redis_circuit_breaker_transitions_total\{state="[a-z_]+"\} ([0-9.e+]+)$', re.M + Endpoint( + name="embeddings", + primary="redis-timeout-embed-primary", + backup="redis-timeout-embed-backup", + primary_params=_EMBED_PRIMARY, + backup_params=_EMBED_BACKUP, + send=_send_embeddings, + served=lambda data: isinstance(data, EmbeddingsObject) and bool(data.data), + ), ) @@ -126,48 +177,51 @@ def _metric(proxy: ProxyClient, pattern: re.Pattern[str]) -> float: return sum(float(match.group(1)) for match in pattern.finditer(body)) +def _rss_bytes(proxy: ProxyClient) -> float | None: + """The proxy's resident memory from the Prometheus process collector, which reads /proc and + so reports on Linux only; None where the metric is absent.""" + body = proxy.probe("/metrics", params=NoBody()).body + match = RSS_RE.search(body) + return float(match.group(1)) if match else None + + class TestRedisTimeout: @pytest.mark.parametrize("endpoint", ENDPOINTS, ids=[endpoint.name for endpoint in ENDPOINTS]) @pytest.mark.covers( "reliability.circuit_breaker.redis_timeout.stays_responsive", - exercised_on=["chat_completions", "responses"], + exercised_on=["chat_completions", "responses", "embeddings"], ) def test_retries_under_redis_timeouts_keep_answering( self, client: ComplexityRouterClient, resources: ResourceManager, endpoint: Endpoint, paused_redis: None ) -> None: proxy = client.proxy - primary_id = proxy.create_model( - PRIMARY_MODEL, - LiteLLMParamsBody( - model=BACKING_MODEL, api_key="sk-redis-timeout-primary-not-used", api_base=CLOSED_PORT_API_BASE - ), - ) + primary_id = proxy.create_model(endpoint.primary, endpoint.primary_params) resources.defer(lambda: proxy.delete_model(primary_id)) - backup_id = proxy.create_model( - BACKUP_MODEL, - LiteLLMParamsBody(model=BACKING_MODEL, api_key="sk-redis-timeout-backup-not-used", mock_response="ok"), - ) + backup_id = proxy.create_model(endpoint.backup, endpoint.backup_params) resources.defer(lambda: proxy.delete_model(backup_id)) - timeouts_before = _metric(proxy, TIMEOUT_FAILURES_RE) - transitions_before = _metric(proxy, BREAKER_TRANSITIONS_RE) key = proxy.generate_key( KeyGenerateBody( - models=[PRIMARY_MODEL, BACKUP_MODEL], key_alias=f"e2e-redis-timeout-{endpoint.name}-{unique_marker()}" + models=[endpoint.primary, endpoint.backup], + key_alias=f"e2e-redis-timeout-{endpoint.name}-{unique_marker()}", ) ) resources.defer(lambda: proxy.delete_key(key)) + timeouts_before = _metric(proxy, TIMEOUT_FAILURES_RE) + transitions_before = _metric(proxy, BREAKER_TRANSITIONS_RE) + fallbacks_before = _metric(proxy, FALLBACKS_RE) + rss_before = _rss_bytes(proxy) latencies: list[float] = [] for request_number in range(1, REQUESTS + 1): started = time.monotonic() - result = endpoint.send(proxy, key, f"redis timeout {unique_marker()} {request_number}") + result = endpoint.send(proxy, key, endpoint.primary, f"redis timeout {unique_marker()} {request_number}") elapsed = time.monotonic() - started assert isinstance(result, Success), ( f"{endpoint.name} request {request_number} failed after {elapsed:.1f}s with Redis timing out: {result}; " f"earlier requests took {[round(seconds, 2) for seconds in latencies]}" ) assert endpoint.served(result.data), ( - f"{endpoint.name} request {request_number}: fallback to {BACKUP_MODEL} returned no output" + f"{endpoint.name} request {request_number}: fallback to {endpoint.backup} returned no output" ) assert elapsed < MAX_SECONDS_PER_REQUEST, ( f"{endpoint.name} request {request_number} took {elapsed:.1f}s with Redis timing out; " @@ -191,6 +245,20 @@ class TestRedisTimeout: f"/health/liveliness took {liveliness_seconds:.1f}s after the loop; the worker is stalled" ) + if rss_before is not None: + rss_after = _rss_bytes(proxy) + assert rss_after is not None + assert rss_after - rss_before <= MAX_RSS_GROWTH_BYTES, ( + f"proxy RSS grew {(rss_after - rss_before) / 2**20:.0f} MB across {REQUESTS} {endpoint.name} requests " + "with Redis timing out; on v1.100.0 this path grew by gigabytes" + ) + + fallbacks = _metric(proxy, FALLBACKS_RE) - fallbacks_before + assert fallbacks >= REQUESTS, ( + f"only {fallbacks:.0f} successful fallbacks were counted across {REQUESTS} {endpoint.name} requests; " + "the closed-port primary did not fail every request, so the retry path was not exercised" + ) + timeouts_total = _metric(proxy, TIMEOUT_FAILURES_RE) timeouts = timeouts_total - timeouts_before transitions = _metric(proxy, BREAKER_TRANSITIONS_RE) - transitions_before @@ -210,3 +278,5 @@ class TestRedisTimeout: f"only {len(rows)} of {REQUESTS} {endpoint.name} requests reached the spend log; " "a Redis outage must not lose spend rows" ) + failed_rows = [row.status for row in rows if row.status not in (None, "success")] + assert not failed_rows, f"{len(failed_rows)} {endpoint.name} spend rows are not successes: {failed_rows[:3]}" From 8b965880efd8e28321c1db74850fbb2b6f59242c Mon Sep 17 00:00:00 2001 From: dclark Date: Thu, 10 Sep 2026 11:29:37 +0100 Subject: [PATCH 018/119] fix(proxy): prevent spend counter double counting --- litellm/proxy/db/spend_counter_reseed.py | 24 ++++++++-- .../proxy/db/test_spend_counter_reseed.py | 48 +++++++++++++++++++ 2 files changed, 69 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index e35b1c8c82b..c6060a47f84 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -245,7 +245,16 @@ class SpendCounterReseed: value=current_value, ) else: - await spend_counter_cache.async_increment_cache(key=counter_key, value=db_spend, refresh_ttl=True) + # Repair/reservations can populate the counter during the DB read. + # Seed a floor without adding the database balance again. + # No await between read/compare/write: atomic within this worker. + cached = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + current_value = float(db_spend) + if cached is not None: + current_value = max(current_value, float(cached)) + spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, value=current_value + ) except Exception: verbose_proxy_logger.exception( "SpendCounterReseed.coalesced: failed to warm counter %s", @@ -438,11 +447,20 @@ class SpendCounterReseed: value=current_value, ) else: - await spend_counter_cache.async_increment_cache(key=counter_key, value=window_spend) + # Repair/reservations can populate the counter during the DB read. + # Seed a floor without adding the database balance again. + # No await between read/compare/write: atomic within this worker. + cached = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + current_value = float(window_spend) + if cached is not None: + current_value = max(current_value, float(cached)) + spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, value=current_value + ) except Exception: verbose_proxy_logger.exception( "SpendCounterReseed.coalesced_window: failed to warm counter %s", counter_key, ) raise - return window_spend + return current_value diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index 0361d5cfe8f..a0cf9128cc9 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -15,6 +15,7 @@ from typing import Final import pytest from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import PROXY_DB_LOOKUP_MAX_CONCURRENCY from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed @@ -270,6 +271,53 @@ async def test_coalesced_window_seeds_a_cold_counter_from_the_row(): assert prisma.db.litellm_spendlogs.call_count == 0 +@pytest.mark.asyncio +@pytest.mark.parametrize("window", [False, True], ids=["primary", "window"]) +@pytest.mark.parametrize("concurrent_spend", [989.01459411, 995.0, 900.0]) +async def test_cold_reseed_does_not_add_database_spend_to_concurrent_cache( + monkeypatch: pytest.MonkeyPatch, + window: bool, + concurrent_spend: float, +): + """A repair/reservation write may populate the counter while DB read runs. + + Reseeding must establish the larger value, not increment the concurrent + value by the same authoritative spend a second time. + """ + cache = DualCache(in_memory_cache=InMemoryCache()) + counter_key = ( + "spend:team:team-1:window:1d" if window else "spend:user:user-1" + ) + db_spend = 989.01459411 + + async def read_db(*args, **kwargs): + cache.in_memory_cache.set_cache(key=counter_key, value=concurrent_spend) + return db_spend + + if window: + monkeypatch.setattr(SpendCounterReseed, "window_from_db", staticmethod(read_db)) + result = await SpendCounterReseed.coalesced_window( + prisma_client=None, + spend_counter_cache=cache, + counter_key=counter_key, + entity_type="Team", + entity_id="team-1", + window_duration="1d", + window_start=WINDOW_START, + ) + else: + monkeypatch.setattr(SpendCounterReseed, "from_db", staticmethod(read_db)) + result = await SpendCounterReseed.coalesced( + prisma_client=None, + spend_counter_cache=cache, + counter_key=counter_key, + ) + + expected = max(db_spend, concurrent_spend) + assert cache.in_memory_cache.get_cache(key=counter_key) == expected + assert result == expected + + @pytest.mark.asyncio async def test_end_user_from_db_reads_the_end_user_row_by_user_id(): prisma: Final = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="customer-42", spend=0.0)) From a39c35878569f74c50de3295320457dd7bd6adb2 Mon Sep 17 00:00:00 2001 From: dclark Date: Thu, 10 Sep 2026 13:44:08 +0100 Subject: [PATCH 019/119] fix(proxy): preserve local spend adjustments during reseeding --- litellm/proxy/db/spend_counter_reseed.py | 37 +++--- litellm/proxy/proxy_server.py | 12 +- .../proxy/db/test_spend_counter_reseed.py | 117 ++++++++++++------ 3 files changed, 102 insertions(+), 64 deletions(-) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index c6060a47f84..0131a67db8b 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -104,6 +104,15 @@ class SpendCounterReseed: SpendCounterReseed._locks.popitem(last=False) return lock + @staticmethod + async def increment_in_memory(spend_counter_cache: "DualCache", counter_key: str, increment: float) -> float | None: + """Apply local deltas after an in-flight reseed establishes the spend balance.""" + lock: Final = await SpendCounterReseed._get_lock(counter_key) + async with lock: + return await spend_counter_cache.async_increment_cache( + key=counter_key, value=increment, local_only=True, refresh_ttl=True + ) + @staticmethod async def from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> float | None: """ @@ -245,16 +254,10 @@ class SpendCounterReseed: value=current_value, ) else: - # Repair/reservations can populate the counter during the DB read. - # Seed a floor without adding the database balance again. - # No await between read/compare/write: atomic within this worker. - cached = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) - current_value = float(db_spend) - if cached is not None: - current_value = max(current_value, float(cached)) - spend_counter_cache.in_memory_cache.set_cache( - key=counter_key, value=current_value - ) + cached_spend: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + seeded_spend: Final = max(db_spend, float(cached_spend)) if cached_spend is not None else db_spend + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=seeded_spend) + return seeded_spend except Exception: verbose_proxy_logger.exception( "SpendCounterReseed.coalesced: failed to warm counter %s", @@ -447,16 +450,12 @@ class SpendCounterReseed: value=current_value, ) else: - # Repair/reservations can populate the counter during the DB read. - # Seed a floor without adding the database balance again. - # No await between read/compare/write: atomic within this worker. - cached = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) - current_value = float(window_spend) - if cached is not None: - current_value = max(current_value, float(cached)) - spend_counter_cache.in_memory_cache.set_cache( - key=counter_key, value=current_value + cached_spend: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + seeded_spend: Final = ( + max(window_spend, float(cached_spend)) if cached_spend is not None else window_spend ) + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=seeded_spend) + return seeded_spend except Exception: verbose_proxy_logger.exception( "SpendCounterReseed.coalesced_window: failed to warm counter %s", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9269fd48e6c..e8608661c48 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3330,10 +3330,8 @@ async def _increment_spend_counter_cache(counter_key: str, increment: float): ) return current_value - return await spend_counter_cache.async_increment_cache( - key=counter_key, - value=increment, - refresh_ttl=True, + return await SpendCounterReseed.increment_in_memory( + spend_counter_cache=spend_counter_cache, counter_key=counter_key, increment=increment ) @@ -3356,10 +3354,8 @@ async def _apply_spend_counter_increments(pending: Sequence[_PendingSpendIncreme redis_cache: Final = spend_counter_cache.redis_cache if redis_cache is None: for item in pending: - await spend_counter_cache.async_increment_cache( - key=item.counter_key, - value=item.increment, - refresh_ttl=True, + await SpendCounterReseed.increment_in_memory( + spend_counter_cache=spend_counter_cache, counter_key=item.counter_key, increment=item.increment ) return ttl: Final = redis_cache.get_ttl() diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index a0cf9128cc9..69ec174903b 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -8,6 +8,7 @@ allowed to run: only when the row is missing or belongs to an older window. from __future__ import annotations import asyncio +from collections.abc import Mapping from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final @@ -79,6 +80,39 @@ def _row(window_start: datetime, spend: float) -> SimpleNamespace: return SimpleNamespace(window_start=window_start, spend=spend) +class _PausedSpendTable: + def __init__(self, spend: float) -> None: + self.spend: Final = spend + self.read_started: Final = asyncio.Event() + self.resume_read: Final = asyncio.Event() + + async def find_unique(self, where: Mapping[str, object]) -> SimpleNamespace: + self.read_started.set() + await self.resume_read.wait() + return _row(WINDOW_START, self.spend) + + +async def _reseed_with_paused_table( + table: _PausedSpendTable, cache: DualCache, counter_key: str, window: bool +) -> float | None: + prisma: Final = SimpleNamespace(db=SimpleNamespace(litellm_usertable=table, litellm_budgetwindowspend=table)) + if window: + return await SpendCounterReseed.coalesced_window( + prisma_client=prisma, + spend_counter_cache=cache, + counter_key=counter_key, + entity_type="Team", + entity_id="team-1", + window_duration="1d", + window_start=WINDOW_START, + ) + return await SpendCounterReseed.coalesced( + prisma_client=prisma, + spend_counter_cache=cache, + counter_key=counter_key, + ) + + @pytest.mark.asyncio async def test_window_from_table_reads_row_by_primary_key(): """The lookup must use the table's own entity_type values ("key"), not the @@ -275,49 +309,59 @@ async def test_coalesced_window_seeds_a_cold_counter_from_the_row(): @pytest.mark.parametrize("window", [False, True], ids=["primary", "window"]) @pytest.mark.parametrize("concurrent_spend", [989.01459411, 995.0, 900.0]) async def test_cold_reseed_does_not_add_database_spend_to_concurrent_cache( - monkeypatch: pytest.MonkeyPatch, window: bool, concurrent_spend: float, -): - """A repair/reservation write may populate the counter while DB read runs. +) -> None: + cache: Final = DualCache(in_memory_cache=InMemoryCache()) + counter_key: Final = "spend:team:team-1:window:1d" if window else "spend:user:user-1" + db_spend: Final = 989.01459411 + table: Final = _PausedSpendTable(db_spend) + reseed_task: Final = asyncio.create_task(_reseed_with_paused_table(table, cache, counter_key, window)) - Reseeding must establish the larger value, not increment the concurrent - value by the same authoritative spend a second time. - """ - cache = DualCache(in_memory_cache=InMemoryCache()) - counter_key = ( - "spend:team:team-1:window:1d" if window else "spend:user:user-1" - ) - db_spend = 989.01459411 + await asyncio.wait_for(table.read_started.wait(), timeout=5) + cache.in_memory_cache.set_cache(key=counter_key, value=concurrent_spend) + table.resume_read.set() + result: Final = await asyncio.wait_for(reseed_task, timeout=5) - async def read_db(*args, **kwargs): - cache.in_memory_cache.set_cache(key=counter_key, value=concurrent_spend) - return db_spend - - if window: - monkeypatch.setattr(SpendCounterReseed, "window_from_db", staticmethod(read_db)) - result = await SpendCounterReseed.coalesced_window( - prisma_client=None, - spend_counter_cache=cache, - counter_key=counter_key, - entity_type="Team", - entity_id="team-1", - window_duration="1d", - window_start=WINDOW_START, - ) - else: - monkeypatch.setattr(SpendCounterReseed, "from_db", staticmethod(read_db)) - result = await SpendCounterReseed.coalesced( - prisma_client=None, - spend_counter_cache=cache, - counter_key=counter_key, - ) - - expected = max(db_spend, concurrent_spend) + expected: Final = max(db_spend, concurrent_spend) assert cache.in_memory_cache.get_cache(key=counter_key) == expected assert result == expected +@pytest.mark.asyncio +@pytest.mark.parametrize("window", [False, True], ids=["primary", "window"]) +@pytest.mark.parametrize("batch", [False, True], ids=["single_increment", "batch_increment"]) +@pytest.mark.parametrize("increment", [5.0, -5.0], ids=["charge", "refund"]) +async def test_cold_reseed_preserves_concurrent_local_increment( + monkeypatch: pytest.MonkeyPatch, window: bool, batch: bool, increment: float +) -> None: + from litellm.proxy import proxy_server + + cache: Final = DualCache(in_memory_cache=InMemoryCache()) + counter_key: Final = ( + f"spend:team:concurrent-{batch}-{increment}:window:1d" + if window + else f"spend:user:concurrent-{batch}-{increment}" + ) + table: Final = _PausedSpendTable(100.0) + monkeypatch.setattr(proxy_server, "spend_counter_cache", cache) + reseed_task: Final = asyncio.create_task(_reseed_with_paused_table(table, cache, counter_key, window)) + await asyncio.wait_for(table.read_started.wait(), timeout=5) + + increment_task: Final = asyncio.create_task( + proxy_server._apply_spend_counter_increments( + pending=(proxy_server._PendingSpendIncrement(counter_key=counter_key, increment=increment),) + ) + if batch + else proxy_server._increment_spend_counter_cache(counter_key=counter_key, increment=increment) + ) + await asyncio.sleep(0) + table.resume_read.set() + await asyncio.wait_for(asyncio.gather(reseed_task, increment_task), timeout=5) + + assert cache.in_memory_cache.get_cache(key=counter_key) == 100.0 + increment + + @pytest.mark.asyncio async def test_end_user_from_db_reads_the_end_user_row_by_user_id(): prisma: Final = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="customer-42", spend=0.0)) @@ -352,8 +396,7 @@ async def test_end_user_from_db_ignores_other_counter_kinds_without_touching_the @pytest.mark.asyncio async def test_end_user_from_db_returns_none_without_a_row_a_client_or_on_db_error(): assert ( - await SpendCounterReseed.end_user_from_db(prisma_client=None, counter_key="spend:end_user:customer-42") - is None + await SpendCounterReseed.end_user_from_db(prisma_client=None, counter_key="spend:end_user:customer-42") is None ) assert ( await SpendCounterReseed.end_user_from_db( From 8deb465346317baf550acb948ff2c72bdf8b9315 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 13:18:52 -0700 Subject: [PATCH 020/119] fix(realtime): dial Azure's GA realtime upstream for GA clients Azure realtime defaulted to the beta upstream whenever realtime_protocol was not configured, so a GA client's session.update (session.type, output_modalities, nested audio) was forwarded unchanged to /openai/realtime and Azure rejected it with "Unknown parameter: 'session.type'" on gpt-realtime and gpt-realtime-1.5. The unset default now follows the client the way the OpenAI handler already does: a client that sends OpenAI-Beta: realtime=v1 keeps the beta upstream, any other client gets /openai/v1/realtime. An explicit realtime_protocol in litellm_params or LITELLM_AZURE_REALTIME_PROTOCOL still wins. --- litellm/realtime_api/main.py | 22 ++++++-- tests/test_litellm/realtime_api/test_main.py | 59 +++++++++++++++++++- 2 files changed, 74 insertions(+), 7 deletions(-) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index b824a5928c6..5310efba1c4 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -14,6 +14,7 @@ from litellm.constants import ( request_timeout, ) from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.litellm_core_utils.realtime_streaming import client_sent_openai_beta_realtime_header from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.xai.common_utils import XAIModelInfo @@ -413,14 +414,14 @@ async def _arealtime( api_version = api_version or litellm_params.api_version or "2024-10-01-preview" - realtime_protocol = ( + configured_realtime_protocol: Final = ( kwargs.get("realtime_protocol") or litellm_params.get("realtime_protocol") or os.environ.get("LITELLM_AZURE_REALTIME_PROTOCOL") ) - if realtime_protocol is None and (query_params or {}).get("intent") == "transcription": - realtime_protocol = "GA" - realtime_protocol = realtime_protocol or "beta" + realtime_protocol: Final = _azure_realtime_protocol_for_client( + configured_realtime_protocol, query_params=query_params, websocket=websocket + ) resolved_azure_ad_token: Final = ( None if api_key else get_azure_ad_token(GenericLiteLLMParams(**kwargs, azure_ad_token=azure_ad_token)) ) @@ -576,6 +577,19 @@ def _is_transcription_only_realtime_model(model: str, custom_llm_provider: str) _TRANSCRIPTION_QUERY_PARAMS: Final[RealtimeQueryParams] = {"intent": "transcription"} +def _azure_realtime_protocol_for_client( + configured_protocol: object, + *, + query_params: RealtimeQueryParams | None, + websocket: "WebSocket", +) -> str: + if isinstance(configured_protocol, str) and configured_protocol: + return configured_protocol + if (query_params or {}).get("intent") == "transcription": + return "GA" + return "beta" if client_sent_openai_beta_realtime_header(websocket) else "GA" + + def _azure_realtime_health_protocol( model: str, realtime_protocol: str | None, model_params: Mapping[str, object] ) -> tuple[str, RealtimeQueryParams | None]: diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py index 761e87ac764..8a9abe819e5 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/test_litellm/realtime_api/test_main.py @@ -327,7 +327,60 @@ async def test_arealtime_azure_ai_on_a_foundry_host_connects_to_the_azure_openai api_key="fake-key", litellm_logging_obj=FakeLogging(), ) - assert connect.url == ( - "wss://my-project.services.ai.azure.com/openai/realtime" - "?api-version=2024-10-01-preview&deployment=gpt-realtime-mini" + assert connect.url == "wss://my-project.services.ai.azure.com/openai/v1/realtime?model=gpt-realtime-mini" + + +class _ClientWebSocketWithHeaders: + def __init__(self, headers: tuple[tuple[bytes, bytes], ...]) -> None: + self.scope: Final = {"headers": headers} + + +_GA_CLIENT: Final = _ClientWebSocketWithHeaders(headers=()) +_BETA_CLIENT: Final = _ClientWebSocketWithHeaders(headers=((b"openai-beta", b"realtime=v1"),)) + + +async def _azure_backend_url_dialed_for(websocket: _ClientWebSocketWithHeaders, **kwargs: object) -> str | None: + connect: Final = _ConnectThatStopsAfterCapturingTheUrl() + with patch("websockets.connect", connect): + await realtime_main._arealtime.__wrapped__( + model="azure/gpt-realtime", + websocket=websocket, + api_base="https://my-endpoint.openai.azure.com", + api_key="fake-key", + litellm_logging_obj=FakeLogging(), + **kwargs, + ) + return connect.url + + +@pytest.mark.asyncio +async def test_arealtime_azure_ga_client_without_beta_header_dials_the_ga_upstream(monkeypatch): + monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) + assert ( + await _azure_backend_url_dialed_for(_GA_CLIENT) + == "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=gpt-realtime" + ) + + +@pytest.mark.asyncio +async def test_arealtime_azure_beta_header_client_keeps_the_beta_upstream(monkeypatch): + monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) + assert await _azure_backend_url_dialed_for(_BETA_CLIENT) == ( + "wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=gpt-realtime" + ) + + +@pytest.mark.asyncio +async def test_arealtime_azure_explicit_beta_protocol_wins_over_a_ga_client(monkeypatch): + monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) + assert await _azure_backend_url_dialed_for(_GA_CLIENT, realtime_protocol="beta") == ( + "wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=gpt-realtime" + ) + + +@pytest.mark.asyncio +async def test_arealtime_azure_env_beta_protocol_wins_over_a_ga_client(monkeypatch): + monkeypatch.setenv("LITELLM_AZURE_REALTIME_PROTOCOL", "beta") + assert await _azure_backend_url_dialed_for(_GA_CLIENT) == ( + "wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=gpt-realtime" ) From 8fe2094a5579e048c5d79f374f4a6ad7d7261326 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 13:46:13 -0700 Subject: [PATCH 021/119] refactor(realtime): move the Azure protocol picker into the Azure realtime handler --- litellm/llms/azure/realtime/handler.py | 23 +++++++++++++++++++++-- litellm/realtime_api/main.py | 18 ++---------------- 2 files changed, 23 insertions(+), 18 deletions(-) diff --git a/litellm/llms/azure/realtime/handler.py b/litellm/llms/azure/realtime/handler.py index e9913f0108d..88813c21cd9 100644 --- a/litellm/llms/azure/realtime/handler.py +++ b/litellm/llms/azure/realtime/handler.py @@ -6,17 +6,23 @@ This requires websockets, and is currently only supported on LiteLLM Proxy. from collections.abc import Mapping from types import MappingProxyType -from typing import Any, Final, Protocol, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast from litellm._logging import _redact_string, verbose_proxy_logger from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.types.realtime import RealtimeQueryParams from ....litellm_core_utils.litellm_logging import Logging as LiteLLMLogging -from ....litellm_core_utils.realtime_streaming import RealTimeStreaming +from ....litellm_core_utils.realtime_streaming import ( + RealTimeStreaming, + client_sent_openai_beta_realtime_header, +) from ....llms.custom_httpx.http_handler import get_shared_realtime_ssl_context from ..azure import AzureChatCompletion +if TYPE_CHECKING: + from fastapi import WebSocket + # BACKEND_WS_URL = "ws://localhost:8080/v1/realtime?model=gpt-4o-realtime-preview-2024-10-01" @@ -31,6 +37,19 @@ async def forward_messages(client_ws: Any, backend_ws: Any): pass +def azure_realtime_protocol_for_client( + configured_protocol: object, + *, + query_params: RealtimeQueryParams | None, + websocket: "WebSocket", +) -> str: + if isinstance(configured_protocol, str) and configured_protocol: + return configured_protocol + if (query_params or {}).get("intent") == "transcription": + return "GA" + return "beta" if client_sent_openai_beta_realtime_header(websocket) else "GA" + + class _ProxyClientWebSocket(Protocol): """Client-facing websocket handle: this path only closes it after a failed handshake.""" diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 5310efba1c4..d42e7e18b75 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -14,7 +14,6 @@ from litellm.constants import ( request_timeout, ) from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider -from litellm.litellm_core_utils.realtime_streaming import client_sent_openai_beta_realtime_header from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.xai.common_utils import XAIModelInfo @@ -34,7 +33,7 @@ from litellm.utils import ProviderConfigManager from ..litellm_core_utils.get_litellm_params import get_litellm_params from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from ..llms.azure.common_utils import get_azure_ad_token -from ..llms.azure.realtime.handler import AzureOpenAIRealtime +from ..llms.azure.realtime.handler import AzureOpenAIRealtime, azure_realtime_protocol_for_client from ..llms.bedrock.realtime.handler import BedrockRealtime from ..llms.custom_httpx.http_handler import get_shared_realtime_ssl_context from ..llms.openai.realtime.handler import OpenAIRealtime @@ -419,7 +418,7 @@ async def _arealtime( or litellm_params.get("realtime_protocol") or os.environ.get("LITELLM_AZURE_REALTIME_PROTOCOL") ) - realtime_protocol: Final = _azure_realtime_protocol_for_client( + realtime_protocol: Final = azure_realtime_protocol_for_client( configured_realtime_protocol, query_params=query_params, websocket=websocket ) resolved_azure_ad_token: Final = ( @@ -577,19 +576,6 @@ def _is_transcription_only_realtime_model(model: str, custom_llm_provider: str) _TRANSCRIPTION_QUERY_PARAMS: Final[RealtimeQueryParams] = {"intent": "transcription"} -def _azure_realtime_protocol_for_client( - configured_protocol: object, - *, - query_params: RealtimeQueryParams | None, - websocket: "WebSocket", -) -> str: - if isinstance(configured_protocol, str) and configured_protocol: - return configured_protocol - if (query_params or {}).get("intent") == "transcription": - return "GA" - return "beta" if client_sent_openai_beta_realtime_header(websocket) else "GA" - - def _azure_realtime_health_protocol( model: str, realtime_protocol: str | None, model_params: Mapping[str, object] ) -> tuple[str, RealtimeQueryParams | None]: From 4a3950cf678fa1b65acfc468c61ae3dcbb67b472 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 13:57:48 -0700 Subject: [PATCH 022/119] refactor(realtime): type the Azure protocol picker with the streaming module's websocket protocol --- litellm/litellm_core_utils/realtime_streaming.py | 8 ++++---- litellm/llms/azure/realtime/handler.py | 8 +++----- 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 75046f2cf87..4923bdda305 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -90,12 +90,12 @@ class _ResponseDoneBody(TypedDict, total=False): output: ReadOnly[Sequence[Mapping[str, object]]] -class _ScopedWebSocket(Protocol): +class ScopedWebSocket(Protocol): @property def scope(self) -> _ASGIScope: ... -class _ClientWebSocket(_ScopedWebSocket, Protocol): +class _ClientWebSocket(ScopedWebSocket, Protocol): async def send_text(self, data: str) -> None: ... async def receive_text(self) -> str: ... async def close(self, code: int = 1000, reason: str | None = None) -> None: ... @@ -1149,7 +1149,7 @@ class RealTimeStreaming: ) @staticmethod - def _detect_beta_header(websocket: _ScopedWebSocket) -> bool: + def _detect_beta_header(websocket: ScopedWebSocket) -> bool: """Return True if the client sent 'OpenAI-Beta: realtime=v1'. Checks the raw ASGI scope headers so it works for both FastAPI WebSocket @@ -1584,6 +1584,6 @@ class RealTimeStreaming: verbose_logger.debug("Could not relay the upstream close to the client: %s", e) -def client_sent_openai_beta_realtime_header(websocket: _ScopedWebSocket) -> bool: +def client_sent_openai_beta_realtime_header(websocket: ScopedWebSocket) -> bool: """True when the client WebSocket includes ``OpenAI-Beta: realtime=v1``.""" return RealTimeStreaming._detect_beta_header(websocket) diff --git a/litellm/llms/azure/realtime/handler.py b/litellm/llms/azure/realtime/handler.py index 88813c21cd9..146915dd6fd 100644 --- a/litellm/llms/azure/realtime/handler.py +++ b/litellm/llms/azure/realtime/handler.py @@ -6,7 +6,7 @@ This requires websockets, and is currently only supported on LiteLLM Proxy. from collections.abc import Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Protocol, cast +from typing import Any, Final, Protocol, cast from litellm._logging import _redact_string, verbose_proxy_logger from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES @@ -15,14 +15,12 @@ from litellm.types.realtime import RealtimeQueryParams from ....litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from ....litellm_core_utils.realtime_streaming import ( RealTimeStreaming, + ScopedWebSocket, client_sent_openai_beta_realtime_header, ) from ....llms.custom_httpx.http_handler import get_shared_realtime_ssl_context from ..azure import AzureChatCompletion -if TYPE_CHECKING: - from fastapi import WebSocket - # BACKEND_WS_URL = "ws://localhost:8080/v1/realtime?model=gpt-4o-realtime-preview-2024-10-01" @@ -41,7 +39,7 @@ def azure_realtime_protocol_for_client( configured_protocol: object, *, query_params: RealtimeQueryParams | None, - websocket: "WebSocket", + websocket: ScopedWebSocket, ) -> str: if isinstance(configured_protocol, str) and configured_protocol: return configured_protocol From be7dce30a3ab832bdd104b2d2f15bace5d622b4d Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Thu, 10 Sep 2026 15:16:24 -0700 Subject: [PATCH 023/119] ci(e2e): pin the Redis timeout workflow's Postgres image by digest Co-Authored-By: Claude Code --- .github/workflows/test-e2e-redis-timeout.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test-e2e-redis-timeout.yml b/.github/workflows/test-e2e-redis-timeout.yml index 9c318a1346c..501af0f240f 100644 --- a/.github/workflows/test-e2e-redis-timeout.yml +++ b/.github/workflows/test-e2e-redis-timeout.yml @@ -15,7 +15,7 @@ jobs: timeout-minutes: 30 services: postgres: - image: postgres:16.6 + image: postgres:16.6@sha256:557fea37a744d5f4c8faab304b0a90858b53ab119735a88c131fd19dab802f36 env: POSTGRES_USER: llmproxy POSTGRES_PASSWORD: dbpassword9090 From dcf8228a9d94e35da7d6d3066b842d385e96b701 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 17:18:20 -0700 Subject: [PATCH 024/119] test(e2e): memory regression test for failing requests on the release gate --- tests/e2e/CLAUDE.md | 4 +- tests/e2e/coverage_registry/reliability.yaml | 1 + tests/e2e/e2e_config.py | 8 + tests/e2e/models.py | 31 ++- tests/e2e/proxy_client.py | 16 ++ tests/e2e/router/reliability_support.py | 18 +- .../e2e/router/test_reliability_memory_e2e.py | 233 ++++++++++++++++++ 7 files changed, 306 insertions(+), 5 deletions(-) create mode 100644 tests/e2e/router/test_reliability_memory_e2e.py diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 34fbe9d9247..f61c76c2d0a 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -17,7 +17,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below); plus the gateway-managed OAuth (authorization_code) path exercised through `/chat/completions`, the one behavior Datadog's static-header auth cannot reach, seeding the per-user upstream token via the interactive authorize dance driven with the mcp SDK's own OAuth client (headless-browser consent from a saved session) and asserting the completion lists and executes the server's tools with the stored per-user token - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection -- `router/` - routing and reliability behavior (fallbacks, cooldowns) +- `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory regression test (`test_reliability_memory_e2e.py`: a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a fixed budget nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak) - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What remains here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`) and markerless harness unit tests for the Locust/session-anomaly aggregation logic - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests @@ -163,7 +163,7 @@ reliability... behavior : fallback | retry | cooldown | timeout | routing | cache | circuit_breaker | perf variant : 5xx | context_window | content_policy | 429 | timeout simple_shuffle | usage_based | latency_based | cost_based | least_busy - latency | throughput | session_anomaly (perf only; SLO/threshold assertion, not binary) + latency | throughput | session_anomaly | memory (perf only; SLO/threshold assertion, not binary) assertion : routes_to_fallback | succeeds_within_retries | picks_under_tpm | returns_cached | trips_then_recovers | under_slo e.g. reliability.fallback.context_window.routes_to_fallback exercised_on=[chat_completions] diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index 6b69677d490..63c6505569e 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -34,4 +34,5 @@ - {id: reliability.timeout.request_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: request_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions, messages], source: "litellm/router.py:545-551", rationale: "Per-request timeout raises Timeout"} - {id: reliability.timeout.stream_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: stream_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions], source: "litellm/router.py:551", rationale: "Streaming chunk-delivery timeout"} - {id: reliability.perf.throughput.under_slo, module: reliability, tier: P1, behavior: perf, variant: throughput, assertions: [under_slo], exercised_on: [chat_completions, messages], source: grammar, rationale: "Throughput SLO under load"} +- {id: reliability.perf.memory.under_slo, module: reliability, tier: P1, behavior: perf, variant: memory, assertions: [under_slo], exercised_on: [chat_completions], source: grammar, rationale: "Proxy RSS and the stored request snapshot stay within fixed budgets across a few hundred failing requests with retries and fallbacks, the v1.100.0 retry-breadcrumb leak shape (MAT-335)"} - {id: reliability.perf.session_anomaly.under_slo, module: reliability, tier: P1, behavior: perf, variant: session_anomaly, assertions: [under_slo], exercised_on: [messages], source: grammar, rationale: "Weekly Claude Code-shaped multi-turn session load against real providers; ceilings on error rate, warm-turn cache read/write, p95 turn time, and gateway-recorded spend (LIT-4562)"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 691335ffdd5..98c3e1cc575 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -160,6 +160,14 @@ ANOMALY_MAX_KEY_SPEND_USD = float( ANOMALY_SPEND_SETTLE_SECONDS = float( os.environ.get("E2E_ANOMALY_SPEND_SETTLE_SECONDS", "75") ) +MEMORY_REQUESTS_PER_PHASE = int(os.environ.get("E2E_MEMORY_REQUESTS_PER_PHASE", "300")) +MEMORY_RETRIES_PER_REQUEST = int(os.environ.get("E2E_MEMORY_RETRIES_PER_REQUEST", "2")) +MEMORY_TRANSCRIPT_TURNS = int(os.environ.get("E2E_MEMORY_TRANSCRIPT_TURNS", "40")) +MEMORY_CONCURRENCY = int(os.environ.get("E2E_MEMORY_CONCURRENCY", "4")) +MEMORY_RSS_SETTLE_SAMPLES = int(os.environ.get("E2E_MEMORY_RSS_SETTLE_SAMPLES", "15")) +MEMORY_RSS_SAMPLE_INTERVAL_SECONDS = float(os.environ.get("E2E_MEMORY_RSS_SAMPLE_INTERVAL_SECONDS", "1")) +MEMORY_RSS_BUDGET_MB = float(os.environ.get("E2E_MEMORY_RSS_BUDGET_MB", "48")) +MEMORY_STORED_REQUEST_BUDGET_KB = float(os.environ.get("E2E_MEMORY_STORED_REQUEST_BUDGET_KB", "64")) def ws_base_url() -> str: diff --git a/tests/e2e/models.py b/tests/e2e/models.py index faf8557498b..c7904966653 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -11,7 +11,16 @@ from datetime import datetime from typing import Final, Literal from e2e_http import PartialBody -from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_serializer, model_validator +from pydantic import ( + AliasChoices, + BaseModel, + ConfigDict, + Field, + JsonValue, + RootModel, + model_serializer, + model_validator, +) # ---------- keys ---------- @@ -695,6 +704,7 @@ class SpendLogRow(BaseModel): total_tokens: int | None = None request_tags: list[str] | None = None metadata: SpendLogMetadata | None = None + proxy_server_request: JsonValue = None class SpendLogs(RootModel[list[SpendLogRow]]): @@ -926,6 +936,7 @@ class LiteLLMParamsBody(BaseModel): timeout: float | None = None tpm: int | None = None weight: int | None = None + cooldown_time: float | None = None ModelMode = Literal["batch", "realtime", "image_generation"] @@ -1260,6 +1271,24 @@ class TagListResponse(RootModel[list[TagListEntry]]): # ---------- health / lifecycle ---------- +class ProcessMemory(BaseModel): + """The `memory` block of GET /debug/memory/summary: the serving worker's resident + set in MB, or `error` when the proxy has no psutil to read it with.""" + + ram_usage_mb: float | None = None + system_memory_percent: float | None = None + error: str | None = None + + +class MemorySummaryResponse(BaseModel): + """GET /debug/memory/summary (master key). One worker's resident memory, keyed by + its pid so readings behind a load balancer can be told apart per pod.""" + + worker_pid: int + status: str + memory: ProcessMemory + + class ReadinessResponse(BaseModel): """GET /health/readiness (public probe). The low-detail payload a load balancer sees: `status` plus the resolved DB state (`connected`, diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 1bac5116a9d..b7865667cd5 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -63,6 +63,7 @@ from models import ( ModelNewBody, ModelNewResponse, ModelsListParams, + MemorySummaryResponse, ModelsListResponse, ModelUpdateBody, OcrBody, @@ -467,6 +468,21 @@ class ProxyClient: ) ).info + def memory_summary_everywhere(self) -> Mapping[str, Result[MemorySummaryResponse]]: + """GET /debug/memory/summary under the master key on every replica in + PROXY_REPLICA_URLS (the data-plane URL alone when the stack exports no + per-gateway addresses). Each read reports the pid of the worker that answered, + so a single address in front of several pods still tells its readings apart.""" + return { + url: transport.get( + "/debug/memory/summary", + headers=transport.master, + params=NoBody(), + response_type=MemorySummaryResponse, + ) + for url, transport in self.replicas.items() + } + def read_back_everywhere[R: BaseModel]( self, path: str, diff --git a/tests/e2e/router/reliability_support.py b/tests/e2e/router/reliability_support.py index 1efcb1a045b..f7dfeb0ef23 100644 --- a/tests/e2e/router/reliability_support.py +++ b/tests/e2e/router/reliability_support.py @@ -11,6 +11,8 @@ body, so a single long-lived proxy serves every reliability behavior. from __future__ import annotations +from collections.abc import Sequence + from pydantic import ValidationError from proxy_client import ProxyClient @@ -49,6 +51,16 @@ def create_bad_base_deployment(proxy: ProxyClient, name: str) -> str: ) +def create_never_benched_refusing_deployment(proxy: ProxyClient, name: str) -> str: + """Register a deployment that refuses every call at the socket and opts out of the + stack's cooldown policy (cooldown_time 0), so the router keeps retrying it for the + whole run instead of benching it after allowed_fails and skipping the retry loop.""" + return proxy.create_model( + name, + LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, api_base="http://127.0.0.1:9/v1", cooldown_time=0), + ) + + def create_timeout_deployment(proxy: ProxyClient, name: str) -> str: """Register a deployment with a 1ms deadline the real backend always exceeds.""" return proxy.create_model(name, LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, timeout=0.001)) @@ -109,15 +121,17 @@ def chat_override( override: RouterSettingsOverride | None = None, stream: bool = False, cache: dict[str, bool] | None = {"no-cache": True}, + history: Sequence[ChatMessage] = (), ) -> StreamingResponse: """POST /chat/completions with an optional per-request router_settings_override, - returning the raw outcome so tests read status, body, and reliability headers.""" + returning the raw outcome so tests read status, body, and reliability headers. + `history` is the conversation sent ahead of the user turn carrying `content`.""" return proxy.transport.send( "/chat/completions", headers=proxy.transport.bearer(key), json=ReliabilityChatBody( model=model, - messages=[ChatMessage(role="user", content=content)], + messages=[*history, ChatMessage(role="user", content=content)], max_tokens=512, stream=stream, router_settings_override=override, diff --git a/tests/e2e/router/test_reliability_memory_e2e.py b/tests/e2e/router/test_reliability_memory_e2e.py new file mode 100644 index 00000000000..48dee01c07f --- /dev/null +++ b/tests/e2e/router/test_reliability_memory_e2e.py @@ -0,0 +1,233 @@ +"""Live e2e: a few hundred requests that fail before any provider answers must not +grow the proxy's resident memory past a fixed budget once the proxy is warm. + +The regression this guards shipped in v1.100.0: every retry breadcrumb copied the +whole request and the copies nested into one router-global list, so a proxy under +retry-heavy failing traffic grew until it was OOM-killed. The traffic here has that +shape: a model group whose deployments refuse at the socket (an unreachable base +URL) with cooldown_time 0 so the router keeps retrying them, per-request retries, +and a fallback group that refuses the same way, each request carrying a long chat +transcript so every whole-request copy costs hundreds of containers instead of a +handful. Under the stack's cooldown policy a +deployment that fails a handful of times in a row is benched (a bad-credential 401 +included), the router answers "No deployments available" without retrying, and the +retry loop that leaks stops running; cooldown_time 0 keeps it running. + +Two identical phases run back to back. The first is the warmup that grows the +proxy's caches and allocator arenas to their steady state, the second is the one +the budget applies to, so a healthy proxy shows the second phase adding roughly +nothing while a leaking one adds a fixed amount per request. RSS is read through +/debug/memory/summary on every configured replica; a burst of failing calls leaves +a transient bulge of garbage that gc reclaims within seconds, so each checkpoint +samples for a settle window and keeps the lowest reading per worker, and the growth +is judged per worker (by pid) so a stack serving one address from several pods +compares each pod with itself. + +RSS alone is a coarse gauge: on the release stack (spend logs storing prompts, +json logs, prometheus and otel callbacks) the same v1.100.0 breadcrumbs grew RSS +by only about 15 MB per 300 failing requests, while every failing request's stored +request snapshot carried a copy of the request per failed attempt, over 100 KB on +the first call and a couple of MB once the copies nested, against tens of KB with +the fix. So the first check sends one failing request before the phases, reads its +spend log back through /spend/logs, and holds the stored request body to a fixed +size budget: the deterministic catch for a breadcrumb that copies the whole +request. It runs before the phases because the leaking writer drops its own rows +under the phases' traffic (a queue budget hit, a recursion limit on the nested +copies), which would turn the size check into a missing-row check. +""" + +from __future__ import annotations + +import json +import time +from collections.abc import Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final + +import pytest + +from complexity_router_client import ComplexityRouterClient +from e2e_config import ( + MEMORY_CONCURRENCY, + MEMORY_REQUESTS_PER_PHASE, + MEMORY_RETRIES_PER_REQUEST, + MEMORY_RSS_BUDGET_MB, + MEMORY_RSS_SAMPLE_INTERVAL_SECONDS, + MEMORY_RSS_SETTLE_SAMPLES, + MEMORY_STORED_REQUEST_BUDGET_KB, + MEMORY_TRANSCRIPT_TURNS, + unique_marker, +) +from e2e_http import unwrap +from lifecycle import ResourceManager +from models import ChatMessage, RouterSettingsOverride, SpendLogRow +from proxy_client import ProxyClient +from reliability_support import chat_override, create_never_benched_refusing_deployment + +pytestmark = pytest.mark.e2e + +DEPLOYMENTS_PER_GROUP: Final = 2 + + +@dataclass(frozen=True, slots=True) +class FailedCall: + status_code: int + seconds: float + body_head: str + call_id: str | None + + +@dataclass(frozen=True, slots=True) +class RssReading: + replica: str + worker_pid: int + ram_usage_mb: float + + +@dataclass(frozen=True, slots=True) +class WorkerGrowth: + warm: RssReading + after: RssReading + + @property + def growth_mb(self) -> float: + return self.after.ram_usage_mb - self.warm.ram_usage_mb + + +def _register_refusing_group(proxy: ProxyClient, resources: ResourceManager, name: str) -> None: + for model_id in tuple(create_never_benched_refusing_deployment(proxy, name) for _ in range(DEPLOYMENTS_PER_GROUP)): + resources.defer(lambda model_id=model_id: proxy.delete_model(model_id)) + + +def _transcript(turns: int) -> tuple[ChatMessage, ...]: + return tuple( + ChatMessage(role=role, content=f"turn {turn} {role}") + for turn in range(turns) + for role in ("user", "assistant") + ) + + +TRANSCRIPT: Final = _transcript(MEMORY_TRANSCRIPT_TURNS) + + +def _fail_once(proxy: ProxyClient, key: str, model: str, override: RouterSettingsOverride) -> FailedCall: + started: Final = time.perf_counter() + resp: Final = chat_override( + proxy, key, model, f"memory regression {unique_marker()}", override=override, history=TRANSCRIPT + ) + return FailedCall(resp.status_code, time.perf_counter() - started, resp.body[:300], resp.call_id) + + +def _fail_many(proxy: ProxyClient, key: str, model: str, override: RouterSettingsOverride) -> tuple[FailedCall, ...]: + with ThreadPoolExecutor(max_workers=MEMORY_CONCURRENCY) as pool: + futures: Final = tuple( + pool.submit(_fail_once, proxy, key, model, override) for _ in range(MEMORY_REQUESTS_PER_PHASE) + ) + return tuple(future.result() for future in futures) + + +def _read_rss_everywhere_after_pause(proxy: ProxyClient) -> tuple[RssReading, ...]: + time.sleep(MEMORY_RSS_SAMPLE_INTERVAL_SECONDS) + return tuple( + RssReading(replica, body.worker_pid, body.memory.ram_usage_mb) + for replica, result in proxy.memory_summary_everywhere().items() + for body in (unwrap(result),) + if body.memory.ram_usage_mb is not None + ) + + +def _settled_rss_per_worker(proxy: ProxyClient) -> Mapping[int, RssReading]: + readings: Final = tuple( + reading for _ in range(MEMORY_RSS_SETTLE_SAMPLES) for reading in _read_rss_everywhere_after_pause(proxy) + ) + assert readings, "no /debug/memory/summary read carried ram_usage_mb, so the proxy cannot report its RSS" + return MappingProxyType( + { + pid: min((reading for reading in readings if reading.worker_pid == pid), key=lambda r: r.ram_usage_mb) + for pid in {reading.worker_pid for reading in readings} + } + ) + + +def _heaviest_worker_growth(warm: Mapping[int, RssReading], after: Mapping[int, RssReading]) -> WorkerGrowth: + growths: Final = tuple(WorkerGrowth(warm[pid], after[pid]) for pid in warm.keys() & after.keys()) + assert growths, ( + f"no worker answered /debug/memory/summary at both checkpoints (warm pids {sorted(warm)}, " + f"after pids {sorted(after)}), so no worker can be compared with itself" + ) + return max(growths, key=lambda growth: growth.growth_mb) + + +def _assert_every_call_failed_through_fallback(calls: Sequence[FailedCall], fallback: str) -> None: + served: Final = tuple(call for call in calls if call.status_code == 200) + assert not served, ( + f"{len(served)} of {len(calls)} calls came back 200, so they reached a provider and never " + f"exercised the retry loop: {served[0].body_head}" + ) + without_fallback: Final = tuple(call for call in calls if fallback not in call.body_head) + assert not without_fallback, ( + f"{len(without_fallback)} of {len(calls)} failures never named the fallback group {fallback}, " + f"so the request did not run through retries into the fallback: {without_fallback[0].body_head}" + ) + + +def _stored_request_kb(proxy: ProxyClient, call: FailedCall) -> float: + assert call.call_id, ( + f"the failing call carried no x-litellm-call-id header, so its spend log cannot be read back: {call.body_head}" + ) + rows: Final[Sequence[SpendLogRow]] = proxy.poll_logs_for_request_id(call.call_id) + assert rows, ( + f"no spend log row appeared for failing call {call.call_id} within the poll window: either the stack " + "writes no spend logs or its writer dropped the row, which the v1.100.0 one did once the stored " + "request outgrew the writer's queue budget" + ) + snapshot: Final = rows[0].proxy_server_request + assert snapshot, ( + f"spend log {call.call_id} stored no request body, so the stack is not running with " + "general_settings.store_prompts_in_spend_logs and the stored-request check would pass vacuously" + ) + return len(json.dumps(snapshot).encode()) / 1024 + + +@pytest.mark.covers("reliability.perf.memory.under_slo") +def test_failing_requests_do_not_grow_rss_or_stored_request( + client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str +) -> None: + marker: Final = unique_marker() + primary: Final = f"reliability-memory-{marker}" + fallback: Final = f"reliability-memory-fb-{marker}" + _register_refusing_group(client.proxy, resources, primary) + _register_refusing_group(client.proxy, resources, fallback) + override: Final = RouterSettingsOverride( + num_retries=MEMORY_RETRIES_PER_REQUEST, fallbacks=[{primary: [fallback]}] + ) + + probe: Final = _fail_once(client.proxy, scoped_key, primary, override) + _assert_every_call_failed_through_fallback((probe,), fallback) + stored_kb: Final = _stored_request_kb(client.proxy, probe) + assert stored_kb <= MEMORY_STORED_REQUEST_BUDGET_KB, ( + f"the spend log of one failing request stored a {stored_kb:.0f} KB request body, past the " + f"{MEMORY_STORED_REQUEST_BUDGET_KB:.0f} KB budget for a {len(TRANSCRIPT)}-message transcript with " + f"{MEMORY_RETRIES_PER_REQUEST} retries and a fallback; the retry breadcrumbs are copying the whole " + f"request into the stored snapshot the way the v1.100.0 ones did" + ) + + warmup: Final = _fail_many(client.proxy, scoped_key, primary, override) + _assert_every_call_failed_through_fallback(warmup, fallback) + warm: Final = _settled_rss_per_worker(client.proxy) + + measured: Final = _fail_many(client.proxy, scoped_key, primary, override) + _assert_every_call_failed_through_fallback(measured, fallback) + after: Final = _settled_rss_per_worker(client.proxy) + + heaviest: Final = _heaviest_worker_growth(warm, after) + assert heaviest.growth_mb <= MEMORY_RSS_BUDGET_MB, ( + f"proxy RSS grew {heaviest.growth_mb:.1f} MB over a second batch of {MEMORY_REQUESTS_PER_PHASE} failing " + f"requests ({MEMORY_RETRIES_PER_REQUEST} retries each plus a fallback) after an identical warmup batch, " + f"past the {MEMORY_RSS_BUDGET_MB:.0f} MB budget: worker pid {heaviest.warm.worker_pid} at " + f"{heaviest.warm.replica} settled at {heaviest.warm.ram_usage_mb:.1f} MB warm and " + f"{heaviest.after.ram_usage_mb:.1f} MB after; failing requests are leaking memory the way the " + f"v1.100.0 retry breadcrumbs did" + ) From 93e2393c5374ab03bead6b7597aebdabc370be4f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 17:29:36 -0700 Subject: [PATCH 025/119] test(e2e): compare RSS per replica and pid so same-pid pods are not merged --- .../e2e/router/test_reliability_memory_e2e.py | 24 ++++++++++++------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/tests/e2e/router/test_reliability_memory_e2e.py b/tests/e2e/router/test_reliability_memory_e2e.py index 48dee01c07f..6a975722801 100644 --- a/tests/e2e/router/test_reliability_memory_e2e.py +++ b/tests/e2e/router/test_reliability_memory_e2e.py @@ -20,8 +20,8 @@ nothing while a leaking one adds a fixed amount per request. RSS is read through /debug/memory/summary on every configured replica; a burst of failing calls leaves a transient bulge of garbage that gc reclaims within seconds, so each checkpoint samples for a settle window and keeps the lowest reading per worker, and the growth -is judged per worker (by pid) so a stack serving one address from several pods -compares each pod with itself. +is judged per worker (by replica address and pid, since pods in their own pid +namespaces report the same pids) so each worker is compared with itself. RSS alone is a coarse gauge: on the release stack (spend logs storing prompts, json logs, prometheus and otel callbacks) the same v1.100.0 breadcrumbs grew RSS @@ -85,6 +85,10 @@ class RssReading: worker_pid: int ram_usage_mb: float + @property + def worker(self) -> tuple[str, int]: + return (self.replica, self.worker_pid) + @dataclass(frozen=True, slots=True) class WorkerGrowth: @@ -138,24 +142,26 @@ def _read_rss_everywhere_after_pause(proxy: ProxyClient) -> tuple[RssReading, .. ) -def _settled_rss_per_worker(proxy: ProxyClient) -> Mapping[int, RssReading]: +def _settled_rss_per_worker(proxy: ProxyClient) -> Mapping[tuple[str, int], RssReading]: readings: Final = tuple( reading for _ in range(MEMORY_RSS_SETTLE_SAMPLES) for reading in _read_rss_everywhere_after_pause(proxy) ) assert readings, "no /debug/memory/summary read carried ram_usage_mb, so the proxy cannot report its RSS" return MappingProxyType( { - pid: min((reading for reading in readings if reading.worker_pid == pid), key=lambda r: r.ram_usage_mb) - for pid in {reading.worker_pid for reading in readings} + worker: min((reading for reading in readings if reading.worker == worker), key=lambda r: r.ram_usage_mb) + for worker in {reading.worker for reading in readings} } ) -def _heaviest_worker_growth(warm: Mapping[int, RssReading], after: Mapping[int, RssReading]) -> WorkerGrowth: - growths: Final = tuple(WorkerGrowth(warm[pid], after[pid]) for pid in warm.keys() & after.keys()) +def _heaviest_worker_growth( + warm: Mapping[tuple[str, int], RssReading], after: Mapping[tuple[str, int], RssReading] +) -> WorkerGrowth: + growths: Final = tuple(WorkerGrowth(warm[worker], after[worker]) for worker in warm.keys() & after.keys()) assert growths, ( - f"no worker answered /debug/memory/summary at both checkpoints (warm pids {sorted(warm)}, " - f"after pids {sorted(after)}), so no worker can be compared with itself" + f"no worker answered /debug/memory/summary at both checkpoints (warm workers {sorted(warm)}, " + f"after workers {sorted(after)}), so no worker can be compared with itself" ) return max(growths, key=lambda growth: growth.growth_mb) From 857b9ad7d344b2cb6d76f0f0b4072b464ad3b297 Mon Sep 17 00:00:00 2001 From: ryan Date: Fri, 11 Sep 2026 01:08:20 +0000 Subject: [PATCH 026/119] fix(ui): jump straight to the last Request Logs page instead of advancing one page Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../spend_management_endpoints.py | 11 +++-- .../test_spend_management_endpoints.py | 42 +++++++++++++++++++ .../view_logs/RequestLogsPanel.test.tsx | 38 +++++++++++++++++ .../components/view_logs/RequestLogsPanel.tsx | 7 ++-- 4 files changed, 91 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index da79328fa59..4a5995167f7 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2944,7 +2944,10 @@ async def _ui_session_grouped_spend_logs( next ``page_size`` sessions ordered by ``(MAX(startTime), session_key, api_key)``, resumed from the ``session_cursor`` keyset ``'||'`` instead of an OFFSET, so - page depth does not degrade the query plan. Each session is represented + page depth does not degrade the query plan. A request for ``page > 1`` + without a cursor (the UI jumping straight to the last page, or back to a + page it never walked through) falls back to ``OFFSET (page - 1) * + page_size``, bounded by the capped total. Each session is represented by its newest non-MCP row, enriched by ``_build_ui_spend_logs_response`` exactly like the flat listing, and the response carries ``next_session_cursor`` / ``has_more`` while ``total`` counts sessions @@ -2963,6 +2966,8 @@ async def _ui_session_grouped_spend_logs( ) cursor_params: Final[tuple[object, ...]] = cursor if cursor else () limit_index: Final = next_param_index + len(cursor_params) + offset_params: Final[tuple[int, ...]] = ((page - 1) * page_size,) if cursor is None and page > 1 else () + offset_clause: Final = f"OFFSET ${limit_index + 1}" if offset_params else "" page_query: Final = f""" SELECT {_SESSION_KEY_EXPR} AS session_key, @@ -2973,10 +2978,10 @@ async def _ui_session_grouped_spend_logs( GROUP BY {_SESSION_GROUP_KEY_SQL} {having_clause} ORDER BY MAX("startTime") {direction}, {_SESSION_KEY_EXPR} {direction}, api_key {direction} - LIMIT ${limit_index} + LIMIT ${limit_index} {offset_clause} """ page_rows: Final[Sequence[_SessionPageRow]] = await _query_raw( - prisma_client, page_query, *sql_params, *cursor_params, page_size + 1 + prisma_client, page_query, *sql_params, *cursor_params, page_size + 1, *offset_params ) has_more: Final = len(page_rows) > page_size diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 671a8ae63fc..5052cf3c085 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -6742,6 +6742,48 @@ async def test_ui_view_spend_logs_group_by_session_cursor_page(client, monkeypat app.dependency_overrides.pop(ps.user_api_key_auth, None) +@pytest.mark.asyncio +async def test_ui_view_spend_logs_group_by_session_jumps_to_page_without_cursor(client, monkeypatch): + """page > 1 with no session_cursor (the UI's last-page jump) skips (page - 1) * page_size sessions by OFFSET.""" + page_rows = [_session_page_row("sess-3", "2026-08-29 06:00:00")] + reps = [_session_representative_row("req-3", "sess-3")] + mock_prisma = _session_grouped_mock_prisma(page_rows, 60, reps) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", + lambda user_api_key_dict: True, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "start_date": start_date, + "end_date": end_date, + "group_by_session": "true", + "page": 3, + "page_size": 25, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + assert data["page"] == 3 + assert data["has_more"] is False + assert [row["request_id"] for row in data["data"]] == ["req-3"] + + page_query_call = mock_prisma.db.query_raw.await_args_list[0] + page_query_sql = page_query_call.args[0] + assert "HAVING" not in page_query_sql + assert "OFFSET" in page_query_sql + assert page_query_call.args[-2:] == (26, 50), "LIMIT page_size + 1 then OFFSET (page - 1) * page_size" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_ui_view_spend_logs_group_by_session_offset_for_non_starttime_sort( client, monkeypatch diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx index 295446186b1..e760d0b8dc3 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx @@ -258,6 +258,44 @@ describe("RequestLogsPanel", () => { }); }); + it("jumps straight to the last page without a cursor when the last-page button is clicked", async () => { + const firstPage = Array.from({ length: 25 }, (_, index) => logEntry({ request_id: `req-${index}` })); + const lastPage = Array.from({ length: 10 }, (_, index) => logEntry({ request_id: `req-last-${index}` })); + vi.mocked(uiSpendLogsCall).mockImplementation(async ({ page }) => + page === 3 + ? { + data: lastPage, + total: 60, + page: 3, + page_size: 25, + total_pages: 3, + next_session_cursor: null, + has_more: false, + } + : { + data: firstPage, + total: 60, + page: 1, + page_size: 25, + total_pages: 3, + next_session_cursor: "2026-07-07 09:50:13|key-1|sess-1", + has_more: true, + }, + ); + renderPanel(); + + await waitFor(() => expect(row("req-0")).not.toBeNull()); + expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 1 of 3"); + fireEvent.click(screen.getByTestId("pagination-last")); + + await waitFor(() => expect(row("req-last-0")).not.toBeNull()); + expect(lastCall()?.page).toBe(3); + expect(lastCall()?.params?.session_cursor).toBeUndefined(); + expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 3 of 3"); + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 51-60 of 60"); + expect(vi.mocked(uiSpendLogsCall).mock.calls.filter(([options]) => options.page === 2)).toHaveLength(0); + }); + it("drops the cursor and returns to the first page when a filter changes", async () => { const firstPage = Array.from({ length: 50 }, (_, index) => logEntry({ request_id: `req-${index}` })); vi.mocked(uiSpendLogsCall).mockResolvedValue({ diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx index 6e984297bf2..4f39bb3b79b 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx @@ -209,15 +209,14 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, setPagination({ ...requested, pageIndex: 0 }); return; } - if (requested.pageIndex <= pagination.pageIndex) { + if (requested.pageIndex !== pagination.pageIndex + 1) { setPagination(requested); return; } const nextCursor = filteredLogs.next_session_cursor; if (!nextCursor || logsQuery.isPlaceholderData) return; - const nextPageIndex = pagination.pageIndex + 1; - setSessionCursors((previous) => ({ ...previous, [nextPageIndex]: nextCursor })); - setPagination({ ...requested, pageIndex: nextPageIndex }); + setSessionCursors((previous) => ({ ...previous, [requested.pageIndex]: nextCursor })); + setPagination(requested); }, [usesSessionCursor, pagination, filteredLogs.next_session_cursor, logsQuery.isPlaceholderData], ); From e4f958d7cf17c31cffeded2ddb46030c7d2599e8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 18:12:40 -0700 Subject: [PATCH 027/119] fix(gateway): serve /debug/memory/summary on the data plane so the memory regression test can read each gateway's RSS --- gateway/routes/allowlist.py | 4 +++- .../proxy/test_component_allowlists.py | 14 ++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/gateway/routes/allowlist.py b/gateway/routes/allowlist.py index 92b73867e67..3733072a948 100644 --- a/gateway/routes/allowlist.py +++ b/gateway/routes/allowlist.py @@ -3,7 +3,8 @@ The gateway exposes the LLM data-plane surface: chat/completions, embeddings, audio, batches, files, fine-tuning, rerank, ocr, rag, video, search, image, responses, vector stores, passthrough providers, realtime websockets, MCP -tool-call endpoints, and operational endpoints (/health, /metrics). +tool-call endpoints, and operational endpoints (/health, /metrics, and the +/debug/memory/summary read of the serving worker's RSS). Any path not listed here is dropped from the gateway process so management/UI endpoints don't ride on the same pods. @@ -121,6 +122,7 @@ GATEWAY_EXACT_PATHS: frozenset[str] = frozenset( "/docs/oauth2-redirect", "/redoc", "/test", + "/debug/memory/summary", } ) diff --git a/tests/test_litellm/proxy/test_component_allowlists.py b/tests/test_litellm/proxy/test_component_allowlists.py index 0fdb43d60da..3073908fa54 100644 --- a/tests/test_litellm/proxy/test_component_allowlists.py +++ b/tests/test_litellm/proxy/test_component_allowlists.py @@ -196,6 +196,20 @@ def test_gateway_drops_ui_and_swagger_mounts(): f"Mount {path} must not be served by the gateway" +def test_gateway_keeps_memory_summary_and_trims_the_other_debug_routes(): + """The gateway serves /debug/memory/summary, since the RSS that matters is the + serving worker's and the memory regression e2e test reads it on every gateway + replica; the heavier and mutating /debug/memory routes stay on the backend.""" + debug_memory_routes = { + getattr(r, "path"): r for r in app.router.routes if str(getattr(r, "path", "")).startswith("/debug/memory/") + } + assert {"/debug/memory/summary", "/debug/memory/details", "/debug/memory/gc/configure"} <= set(debug_memory_routes) + assert _is_gateway_route(debug_memory_routes["/debug/memory/summary"]), \ + "/debug/memory/summary must survive the gateway route trim" + for path in ("/debug/memory/details", "/debug/memory/gc/configure"): + assert not _is_gateway_route(debug_memory_routes[path]), f"{path} must not be served by the gateway" + + def test_every_app_mount_is_assigned_to_a_component(): """Every Mount on the proxy app must be consciously assigned to a component. From a9bd86b3711d2ca88c8b60be3e7d7b5677cebcfb Mon Sep 17 00:00:00 2001 From: ryan Date: Fri, 11 Sep 2026 01:37:43 +0000 Subject: [PATCH 028/119] feat(ui): search, sort and role filter for the team member table Rebuild the shared member table on DataTable so admins can search members by name, email or user id, sort by name, email, role, budget and spend, and filter by role. /team/info now returns each member's user_alias so the table can show a human-readable name Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 5 + .../management_endpoints/team_endpoints.py | 60 ++-- .../test_team_endpoints.py | 110 +++--- .../common_components/MemberTable.test.tsx | 215 ++++++++++++ .../common_components/MemberTable.tsx | 329 +++++++++++++----- .../src/components/networking.tsx | 1 + .../organization/organization_view.tsx | 19 +- .../shared/DataTable/DataTable.test.tsx | 22 ++ .../components/shared/DataTable/DataTable.tsx | 10 + .../src/components/team/TeamMemberTab.tsx | 18 +- 10 files changed, 595 insertions(+), 194 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 2cbff128635..22591eab651 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4358,7 +4358,12 @@ class TeamAccessGroupModelGrant(LiteLLMPydanticObjectBase): agent_ids: tuple[str, ...] = () +class TeamInfoMember(Member): + user_alias: str | None = None + + class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): + members_with_roles: tuple[TeamInfoMember, ...] = () team_member_budget_table: LiteLLM_BudgetTableFull | None = None # Resources inherited from access groups (separate from direct assignments) access_group_models: list[str] | None = None diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index c050368b3fe..e9e37540dd8 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -61,6 +61,7 @@ from litellm.proxy._types import ( SpecialProxyStrings, TeamAccessGroupModelGrant, TeamAddMemberResponse, + TeamInfoMember, TeamInfoResponseObject, TeamInfoResponseObjectTeamTable, TeamListResponseObject, @@ -4292,37 +4293,35 @@ async def _add_team_member_budget_table( return team_info_response_object -async def _hydrate_member_emails( +async def _hydrate_member_user_details( prisma_client: PrismaClient, members: Sequence[Member], -) -> tuple[Member, ...]: - """Fill in ``user_email`` for roster entries that were stored without one. - - ``members_with_roles`` is a denormalized snapshot written at add-time, so an entry - stored with ``user_email=None`` keeps that null even once the user row has an email. - Look the missing ones up in ``LiteLLM_UserTable`` (one indexed query) and fill them - in. A stored email is never overwritten - the snapshot stays the source of truth - wherever it has a value. - """ - missing_user_ids: Final = frozenset(m.user_id for m in members if not m.user_email and m.user_id is not None) - if not missing_user_ids: - return tuple(members) - - user_rows: Final[Sequence[prisma_models.LiteLLM_UserTable]] = await _user_db(prisma_client).find_many( - where={ # mutable-ok: Prisma query filters are dict-shaped - "user_id": { # mutable-ok: Prisma query filters are dict-shaped - "in": sorted(missing_user_ids) +) -> tuple[TeamInfoMember, ...]: + """Attach ``user_alias`` and fill in a missing ``user_email`` from ``LiteLLM_UserTable`` in one query.""" + user_ids: Final = frozenset(m.user_id for m in members if m.user_id is not None) + user_rows: Final[Sequence[prisma_models.LiteLLM_UserTable]] = ( + await _user_db(prisma_client).find_many( + where={ # mutable-ok: Prisma query filters are dict-shaped + "user_id": { # mutable-ok: Prisma query filters are dict-shaped + "in": sorted(user_ids) + } } - } + ) + if user_ids + else () ) - email_by_user_id: Final = MappingProxyType({u.user_id: u.user_email for u in user_rows if u.user_email}) + user_by_id: Final = MappingProxyType({u.user_id: u for u in user_rows}) - return tuple( - m.model_copy(update={"user_email": email_by_user_id[m.user_id]}) # mutable-ok: pydantic update payload - if not m.user_email and m.user_id is not None and m.user_id in email_by_user_id - else m - for m in members - ) + def hydrate(m: Member) -> TeamInfoMember: + user_row: Final = user_by_id.get(m.user_id) if m.user_id is not None else None + return TeamInfoMember( + role=m.role, + user_id=m.user_id, + user_email=m.user_email or (user_row.user_email if user_row is not None else None), + user_alias=user_row.user_alias if user_row is not None else None, + ) + + return tuple(hydrate(m) for m in members) async def _resolve_team_access_group_resources( @@ -4462,17 +4461,12 @@ async def team_info( # Resolve resources inherited from access groups resolved_team_info: Final = await _resolve_team_access_group_resources(_team_info) - # Fill in emails the add-time roster snapshot never captured - hydrated_members: Final = await _hydrate_member_emails( + hydrated_members: Final = await _hydrate_member_user_details( prisma_client=prisma_client, members=resolved_team_info.members_with_roles, ) hydrated_team_info: Final = resolved_team_info.model_copy( - update={ # mutable-ok: pydantic update payload - # list(), not the tuple: model_copy skips validation, so the field has - # to be handed the list[Member] the response model declares. - "members_with_roles": list(hydrated_members) # mutable-ok: declared list[Member] - } + update={"members_with_roles": hydrated_members} # mutable-ok: pydantic update payload ) response_object: Final = TeamInfoResponseObject( diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 2f6561046b1..2fb496d6231 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -11220,13 +11220,14 @@ async def test_team_info_returns_model_aliases(): @pytest.mark.asyncio -async def test_team_info_hydrates_member_emails_from_the_user_table(): - """/team/info must fill in emails missing from the members_with_roles snapshot. +async def test_team_info_hydrates_member_names_and_emails_from_the_user_table(): + """/team/info must attach each member's display name and fill in emails missing + from the members_with_roles snapshot. - members_with_roles is written at add-time, so a member added by user_id alone - carries user_email=None forever. Without this join the Admin UI's member table - shows "-" for a user that has an email on their user row. A stored email is left - exactly as-is. + members_with_roles is written at add-time, so it never carries user_alias and a + member added by user_id alone carries user_email=None forever. Without this join + the Admin UI's member table can only show emails. A stored email is left exactly + as-is. """ from fastapi import Request @@ -11246,13 +11247,8 @@ async def test_team_info_hydrates_member_emails_from_the_user_table(): find_many = AsyncMock( return_value=[ - LiteLLM_UserTable( - user_id="no-email-on-roster", - user_email="real@example.com", - max_budget=None, - spend=0.0, - models=[], - ) + _user_row("no-email-on-roster", "real@example.com", "Real Person"), + _user_row("already-stored", "current@example.com", "Stored Person"), ] ) @@ -11270,12 +11266,12 @@ async def test_team_info_hydrates_member_emails_from_the_user_table(): ) members = response["team_info"].members_with_roles - assert [(m.user_id, m.user_email) for m in members] == [ - ("no-email-on-roster", "real@example.com"), - ("already-stored", "stored@example.com"), + assert [(m.user_id, m.user_email, m.user_alias) for m in members] == [ + ("no-email-on-roster", "real@example.com", "Real Person"), + ("already-stored", "stored@example.com", "Stored Person"), ] - # only the member actually missing an email is looked up - assert find_many.await_args.kwargs["where"] == {"user_id": {"in": ["no-email-on-roster"]}} + find_many.assert_awaited_once() + assert find_many.await_args.kwargs["where"] == {"user_id": {"in": ["already-stored", "no-email-on-roster"]}} @pytest.mark.asyncio @@ -12472,89 +12468,93 @@ async def test_resolve_existing_member_user_ids_skips_the_query_when_no_user_ids repo.return_value.table.find_many.assert_not_awaited() -def _user_row(user_id: str, user_email: str | None) -> LiteLLM_UserTable: +def _user_row(user_id: str, user_email: str | None, user_alias: str | None = None) -> LiteLLM_UserTable: return LiteLLM_UserTable( - user_id=user_id, user_email=user_email, max_budget=None, spend=0.0, models=[] + user_id=user_id, user_email=user_email, user_alias=user_alias, max_budget=None, spend=0.0, models=[] ) @pytest.mark.asyncio -async def test_hydrate_member_emails_fills_in_emails_the_roster_snapshot_never_captured(): - """A member added by user_id alone has user_email=None on the stored roster entry. - - /team/info has to fill it in from the user row, or the UI renders "-" for a user - that plainly has an email. +async def test_hydrate_member_user_details_attaches_alias_and_fills_in_missing_email(): + """The stored roster never carries a display name, and a member added by user_id + alone has user_email=None. /team/info has to fill both in from the user row so the + UI can show and search by a human-readable name instead of only an email. """ - from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_user_details - find_many = AsyncMock(return_value=[_user_row("by-id", "found@example.com")]) + find_many = AsyncMock(return_value=[_user_row("by-id", "found@example.com", "Found Person")]) with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: repo.return_value.table.find_many = find_many - hydrated = await _hydrate_member_emails( + hydrated = await _hydrate_member_user_details( prisma_client=MagicMock(), members=[Member(user_id="by-id", role="admin")], ) - assert [(m.user_id, m.user_email, m.role) for m in hydrated] == [("by-id", "found@example.com", "admin")] + assert [(m.user_id, m.user_email, m.user_alias, m.role) for m in hydrated] == [ + ("by-id", "found@example.com", "Found Person", "admin") + ] find_many.assert_awaited_once() assert find_many.await_args.kwargs["where"] == {"user_id": {"in": ["by-id"]}} @pytest.mark.asyncio -async def test_hydrate_member_emails_never_overwrites_a_stored_email(): - """The snapshot wins wherever it has a value - hydration only fills blanks. - - Overwriting would be a real behavior change to /team/info; filling a null is not. - """ - from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails - - find_many = AsyncMock(return_value=[_user_row("has-email", "current@example.com")]) +async def test_hydrate_member_user_details_never_overwrites_a_stored_email(): + """The snapshot wins wherever it has a value - hydration only fills blanks.""" + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_user_details with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: - repo.return_value.table.find_many = find_many + repo.return_value.table.find_many = AsyncMock( + return_value=[_user_row("has-email", "current@example.com", "Current Name")] + ) - hydrated = await _hydrate_member_emails( + hydrated = await _hydrate_member_user_details( prisma_client=MagicMock(), members=[Member(user_id="has-email", user_email="stored@example.com", role="user")], ) - assert hydrated[0].user_email == "stored@example.com" - # nothing was missing, so no round-trip either - find_many.assert_not_awaited() + assert (hydrated[0].user_email, hydrated[0].user_alias) == ("stored@example.com", "Current Name") @pytest.mark.asyncio -async def test_hydrate_member_emails_leaves_members_alone_when_the_user_row_has_no_email(): - """A user row with no email leaves the member as-is rather than inventing one.""" - from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails +async def test_hydrate_member_user_details_leaves_blanks_when_the_user_row_is_bare_or_missing(): + """A user row with no email or alias, or no user row at all, must not invent values.""" + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_user_details with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: - repo.return_value.table.find_many = AsyncMock(return_value=[_user_row("no-email", None)]) + repo.return_value.table.find_many = AsyncMock(return_value=[_user_row("bare", None)]) - hydrated = await _hydrate_member_emails( + hydrated = await _hydrate_member_user_details( prisma_client=MagicMock(), - members=[Member(user_id="no-email", role="user"), Member(user_email="e@example.com", role="user")], + members=[ + Member(user_id="bare", role="user"), + Member(user_id="deleted", user_email="gone@example.com", role="user"), + Member(user_email="e@example.com", role="user"), + ], ) - assert [m.user_email for m in hydrated] == [None, "e@example.com"] + assert [(m.user_id, m.user_email, m.user_alias) for m in hydrated] == [ + ("bare", None, None), + ("deleted", "gone@example.com", None), + (None, "e@example.com", None), + ] @pytest.mark.asyncio -async def test_hydrate_member_emails_skips_the_query_when_every_member_has_one(): - """No blanks means /team/info pays for no extra query.""" - from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails +async def test_hydrate_member_user_details_skips_the_query_when_no_member_has_a_user_id(): + """Email-only roster entries give nothing to look up, so /team/info pays for no query.""" + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_user_details with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: repo.return_value.table.find_many = AsyncMock() - hydrated = await _hydrate_member_emails( + hydrated = await _hydrate_member_user_details( prisma_client=MagicMock(), - members=[Member(user_id="a", user_email="a@example.com", role="user")], + members=[Member(user_email="a@example.com", role="user")], ) - assert hydrated[0].user_email == "a@example.com" + assert [(m.user_email, m.user_alias) for m in hydrated] == [("a@example.com", None)] repo.return_value.table.find_many.assert_not_awaited() diff --git a/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx b/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx new file mode 100644 index 00000000000..a16f11eab60 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx @@ -0,0 +1,215 @@ +import { fireEvent, screen, waitFor, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import React from "react"; +import { describe, expect, it, vi } from "vitest"; + +import type { Member } from "@/components/networking"; + +import { renderWithProviders } from "../../../tests/test-utils"; +import MemberTable, { MemberTableColumn, memberRoleOptions } from "./MemberTable"; + +const MEMBERS: Member[] = [ + { user_id: "u-zed", user_email: "zed@example.com", user_alias: "Zed Ortiz", role: "user" }, + { user_id: "u-nameless", user_email: "mystery@example.com", user_alias: null, role: "user" }, + { user_id: "u-amy", user_email: "amy@example.com", user_alias: "amy chen", role: "admin" }, + { user_id: "u-bob", user_email: null, user_alias: "Bob Lee", role: "user" }, +]; + +const BUDGETS: Record = { "u-zed": 50, "u-nameless": null, "u-amy": 1000, "u-bob": 5 }; + +const budgetColumn: MemberTableColumn = { + title: "Budget", + key: "budget", + sortValue: (member) => BUDGETS[member.user_id ?? ""] ?? null, + render: (member) => {BUDGETS[member.user_id ?? ""] ?? "Unlimited"}, +}; + +const renderTable = (overrides: Partial> = {}) => { + const props = { + members: MEMBERS, + canEdit: true, + onEdit: vi.fn(), + onDelete: vi.fn(), + extraColumns: [budgetColumn], + ...overrides, + }; + renderWithProviders(); + return props; +}; + +const rowIds = (): (string | null)[] => + Array.from(document.querySelectorAll("tbody tr[data-row-id]")).map((row) => row.getAttribute("data-row-id")); + +const search = (value: string) => fireEvent.change(screen.getByTestId("datatable-search"), { target: { value } }); + +describe("MemberTable display", () => { + it("shows each member's name, falling back to a dash when there is none", () => { + renderTable(); + + expect(within(screen.getByRole("row", { name: /zed@example\.com/ })).getByText("Zed Ortiz")).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: /^name/i })).toBeInTheDocument(); + const cells = within(screen.getByRole("row", { name: /mystery@example\.com/ })).getAllByRole("cell"); + expect(cells[0]).toHaveTextContent("-"); + }); + + it("orders members by name with nameless members last by default", () => { + renderTable(); + + expect(rowIds()).toEqual(["u-amy", "u-bob", "u-zed", "u-nameless"]); + }); + + it("reports the full member count", () => { + renderTable(); + + expect(screen.getByText("4 Members")).toBeInTheDocument(); + }); +}); + +describe("MemberTable search", () => { + it("matches on name case-insensitively", async () => { + renderTable(); + + search("ZED"); + + await waitFor(() => expect(rowIds()).toEqual(["u-zed"])); + }); + + it("matches on email", async () => { + renderTable(); + + search("mystery@"); + + await waitFor(() => expect(rowIds()).toEqual(["u-nameless"])); + }); + + it("matches on user id", async () => { + renderTable(); + + search("u-bob"); + + await waitFor(() => expect(rowIds()).toEqual(["u-bob"])); + }); + + it("still searches names when the first member has no name", async () => { + renderTable({ members: [MEMBERS[1], MEMBERS[0]] }); + + search("ortiz"); + + await waitFor(() => expect(rowIds()).toEqual(["u-zed"])); + }); + + it("does not match on role", async () => { + renderTable(); + + search("admin"); + + await waitFor(() => expect(rowIds()).toEqual([])); + expect(screen.getByText("No members match your search or filters")).toBeInTheDocument(); + }); +}); + +describe("MemberTable sorting", () => { + it("sorts by email with missing emails last in both directions", async () => { + const user = userEvent.setup(); + renderTable(); + + await user.click(screen.getByTestId("sort-header-user_email")); + expect(rowIds()).toEqual(["u-amy", "u-nameless", "u-zed", "u-bob"]); + + await user.click(screen.getByTestId("sort-header-user_email")); + expect(rowIds()).toEqual(["u-zed", "u-nameless", "u-amy", "u-bob"]); + }); + + it("sorts by role", async () => { + const user = userEvent.setup(); + renderTable(); + + await user.click(screen.getByTestId("sort-header-role")); + expect(rowIds()[0]).toBe("u-amy"); + + await user.click(screen.getByTestId("sort-header-role")); + expect(rowIds()[3]).toBe("u-amy"); + }); + + it("sorts an extra column numerically by its sort value with blanks last", async () => { + const user = userEvent.setup(); + renderTable(); + + await user.click(screen.getByTestId("sort-header-budget")); + expect(rowIds()).toEqual(["u-bob", "u-zed", "u-amy", "u-nameless"]); + + await user.click(screen.getByTestId("sort-header-budget")); + expect(rowIds()).toEqual(["u-amy", "u-zed", "u-bob", "u-nameless"]); + }); + + it("flips name order on the second click", async () => { + const user = userEvent.setup(); + renderTable(); + + await user.click(screen.getByTestId("sort-header-user_alias")); + expect(rowIds()).toEqual(["u-zed", "u-bob", "u-amy", "u-nameless"]); + }); + + it("leaves extra columns without a sort value unsortable", () => { + renderTable({ + extraColumns: [{ title: "Rate Limits", key: "rate_limits", render: () => No Limits }], + }); + + expect(screen.getByRole("columnheader", { name: "Rate Limits" })).toBeInTheDocument(); + expect(screen.queryByTestId("sort-header-rate_limits")).not.toBeInTheDocument(); + }); +}); + +describe("MemberTable role filter", () => { + it("shows only members with the chosen role and clears on reset", async () => { + const user = userEvent.setup(); + renderTable({ roleColumnTitle: "Team Role" }); + + await user.click(screen.getByTestId("datatable-filters-trigger")); + await user.click(screen.getByTestId("filter-role")); + await user.click(await screen.findByRole("option", { name: "admin" })); + await user.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => expect(rowIds()).toEqual(["u-amy"])); + expect(screen.getByTestId("filter-chip-role")).toHaveTextContent("Team Role"); + + await user.click(screen.getByTestId("datatable-filters-trigger")); + await user.click(screen.getByTestId("filter-drawer-reset")); + + await waitFor(() => expect(rowIds()).toHaveLength(4)); + }); + + it("offers the roles present in the roster", () => { + expect(memberRoleOptions(MEMBERS)).toEqual(["admin", "user"]); + expect(memberRoleOptions([{ user_id: "x", role: "" }])).toEqual([]); + }); +}); + +describe("MemberTable actions", () => { + it("passes the clicked member to onEdit and onDelete", async () => { + const user = userEvent.setup(); + const { onEdit, onDelete } = renderTable(); + const row = screen.getByRole("row", { name: /amy@example\.com/ }); + + await user.click(within(row).getByTestId("edit-member")); + await user.click(within(row).getByTestId("delete-member")); + + expect(onEdit).toHaveBeenCalledWith(MEMBERS[2]); + expect(onDelete).toHaveBeenCalledWith(MEMBERS[2]); + }); + + it("hides delete for members the caller excludes", () => { + renderTable({ showDeleteForMember: (member) => member.role !== "admin" }); + + expect(screen.getAllByTestId("delete-member")).toHaveLength(3); + expect( + within(screen.getByRole("row", { name: /amy@example\.com/ })).queryByTestId("delete-member"), + ).not.toBeInTheDocument(); + }); + + it("shows the empty text when there are no members at all", () => { + renderTable({ members: [], emptyText: "No members found" }); + + expect(screen.getByText("No members found")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx b/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx index 8a0b7671078..e54994c1ab7 100644 --- a/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx +++ b/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx @@ -1,17 +1,29 @@ -import { SimpleTooltip } from "@/components/ui/tooltip"; +import type { ColumnDef, ColumnFiltersState } from "@tanstack/react-table"; +import { Crown, Info, User, UserPlus } from "lucide-react"; +import React, { useState } from "react"; + import { Member } from "@/components/networking"; +import { + DataTable, + DataTableFilterDrawer, + DataTableFilterField, + DataTableSortHeader, + DataTableToolbar, +} from "@/components/shared/DataTable"; import { StatusBadge } from "@/components/shared/table_cells"; import { Button } from "@/components/ui/button"; -import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; -import { Crown, Info, User, UserPlus } from "lucide-react"; -import React from "react"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { SimpleTooltip } from "@/components/ui/tooltip"; + import TableIconActionButton from "./IconActionButton/TableIconActionButtons/TableIconActionButton"; +export type MemberTableSortValue = string | number | null | undefined; + export interface MemberTableColumn { title: React.ReactNode; - key: React.Key; - dataIndex?: keyof Member; - render?: (value: Member[keyof Member], member: Member, index: number) => React.ReactNode; + key: string; + render: (member: Member) => React.ReactNode; + sortValue?: (member: Member) => MemberTableSortValue; } export interface MemberTableProps { @@ -27,12 +39,152 @@ export interface MemberTableProps { emptyText?: string; } -const extraColumnCell = (column: MemberTableColumn, member: Member, index: number): React.ReactNode => { - const value = column.dataIndex ? member[column.dataIndex] : undefined; - return column.render ? column.render(value, member, index) : value; +const ALL_ROLES = "all"; + +export const memberRowId = (member: Member): string => member.user_id ?? member.user_email ?? JSON.stringify(member); + +export const memberRoleOptions = (members: readonly Member[]): string[] => + Array.from(new Set(members.map((member) => member.role).filter((role) => role !== ""))).sort(); + +const isAdminRole = (role: string): boolean => { + const normalized = role.toLowerCase(); + return normalized === "admin" || normalized === "org_admin"; }; -const STICKY_ACTIONS_CLASS = "sticky right-0 w-[120px] bg-background"; +function RoleHeaderTitle({ title, tooltip }: { title: string; tooltip?: string }) { + if (tooltip === undefined) return <>{title}; + return ( + + {title} + + + + + ); +} + +const ACTIONS_COLUMN_WIDTH = 120; + +interface MemberColumnDeps { + canEdit: boolean; + onEdit: (member: Member) => void; + onDelete: (member: Member) => void; + roleColumnTitle: string; + roleTooltip?: string; + extraColumns: MemberTableColumn[]; + showDeleteForMember?: (member: Member) => boolean; +} + +const extraColumnDef = (column: MemberTableColumn): ColumnDef => { + const { sortValue } = column; + if (sortValue === undefined) { + return { + id: column.key, + header: () => {column.title}, + enableSorting: false, + enableGlobalFilter: false, + cell: ({ row }) => column.render(row.original), + }; + } + return { + id: column.key, + accessorFn: (member) => sortValue(member) ?? undefined, + header: ({ column: tableColumn }) => , + sortDescFirst: false, + sortUndefined: "last", + enableGlobalFilter: false, + cell: ({ row }) => column.render(row.original), + }; +}; + +const buildColumns = ({ + canEdit, + onEdit, + onDelete, + roleColumnTitle, + roleTooltip, + extraColumns, + showDeleteForMember, +}: MemberColumnDeps): ColumnDef[] => [ + { + id: "user_alias", + accessorFn: (member) => member.user_alias || undefined, + header: ({ column }) => , + sortingFn: "text", + sortUndefined: "last", + enableGlobalFilter: true, + meta: { title: "Name" }, + cell: ({ row }) => row.original.user_alias || -, + }, + { + id: "user_email", + accessorFn: (member) => member.user_email || undefined, + header: ({ column }) => , + sortingFn: "text", + sortUndefined: "last", + enableGlobalFilter: true, + meta: { title: "User Email" }, + cell: ({ row }) => row.original.user_email || "-", + }, + { + id: "user_id", + accessorFn: (member) => member.user_id ?? undefined, + header: "User ID", + enableSorting: false, + enableGlobalFilter: true, + cell: ({ row }) => + row.original.user_id === "default_user_id" ? ( + + ) : ( + row.original.user_id || "-" + ), + }, + { + id: "role", + accessorFn: (member) => member.role, + header: ({ column }) => ( + } /> + ), + sortingFn: "text", + filterFn: "equalsString", + enableGlobalFilter: false, + meta: { title: roleColumnTitle }, + cell: ({ row }) => ( + + {isAdminRole(row.original.role) ? : } + {row.original.role || "-"} + + ), + }, + ...extraColumns.map(extraColumnDef), + { + id: "actions", + header: "Actions", + size: ACTIONS_COLUMN_WIDTH, + enableSorting: false, + enableGlobalFilter: false, + meta: { pinned: "right" }, + cell: ({ row }) => + canEdit ? ( + + onEdit(row.original)} + /> + {(!showDeleteForMember || showDeleteForMember(row.original)) && ( + onDelete(row.original)} + /> + )} + + ) : null, + }, +]; export default function MemberTable({ members, @@ -46,90 +198,89 @@ export default function MemberTable({ showDeleteForMember, emptyText, }: MemberTableProps) { + const [globalFilter, setGlobalFilter] = useState(""); + const [columnFilters, setColumnFilters] = useState([]); + const [filtersOpen, setFiltersOpen] = useState(false); + + const columnDeps: MemberColumnDeps = { + canEdit, + onEdit, + onDelete, + roleColumnTitle, + roleTooltip, + extraColumns, + showDeleteForMember, + }; + const columns = buildColumns(columnDeps); + const roleFilterItems = [ + { value: ALL_ROLES, label: "All Roles" }, + ...memberRoleOptions(members).map((role) => ({ value: role, label: role })), + ]; + + const isNarrowed = globalFilter !== "" || columnFilters.length > 0; + return (
{members.length} Member{members.length !== 1 ? "s" : ""} - - - - User Email - User ID - - {roleTooltip ? ( - - {roleColumnTitle} - - - - - ) : ( - roleColumnTitle + + {isNarrowed ? "No members match your search or filters" : emptyText ?? "No data"} + + } + toolbar={(table) => ( + <> + setFiltersOpen(true)} + showViewOptions={false} + /> + + {({ get, set }) => ( + + + )} - - {extraColumns.map((column) => ( - {column.title} - ))} - Actions - - - - {members.length === 0 ? ( - - - {emptyText ?? "No data"} - - - ) : ( - members.map((member, memberIndex) => ( - - {member.user_email || "-"} - - {member.user_id === "default_user_id" ? ( - - ) : ( - member.user_id || "-" - )} - - - - {member.role?.toLowerCase() === "admin" || member.role?.toLowerCase() === "org_admin" ? ( - - ) : ( - - )} - {member.role || "-"} - - - {extraColumns.map((column) => ( - {extraColumnCell(column, member, memberIndex)} - ))} - - {canEdit ? ( - - onEdit(member)} - /> - {(!showDeleteForMember || showDeleteForMember(member)) && ( - onDelete(member)} - /> - )} - - ) : null} - - - )) - )} - -
+ + + )} + /> {onAddMember && canEdit && ( @@ -71,6 +75,14 @@ export function IdCell({ ); + const idElement = linked ? ( + + {value} + + ) : ( + unlinkedElement + ); + const withTooltip = ; if (!copyable) { @@ -94,3 +106,21 @@ export function IdCell({ ); } + +interface IdLinkProps extends React.ComponentPropsWithoutRef<"a"> { + href: string; + dataTestId?: string; +} + +const IdLink = React.forwardRef(function IdLink( + { href, dataTestId, children, ...props }, + ref, +) { + const handleClick = useEntityLinkClick(href); + + return ( + + {children} + + ); +}); From 06964e5603ecce13b0e4aabcbcab5d1382c4c6e5 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 11 Sep 2026 17:48:57 -0700 Subject: [PATCH 096/119] feat(ui): link the User ID, Created By and Deleted By cells on Deleted Keys (#40750) * feat(ui): link the User ID, Created By and Deleted By cells on Deleted Keys All three columns rendered as plain text, so auditing a deleted key meant copying an id into the Users page search box. Route them through IdentityCell with userDetailHref, which keeps the proxy admin placeholder unlinked. User Email and Team Alias stay as they are: the deleted key table has no column for either, so the API never populates them. Claude-Session: https://claude.ai/code/session_01NfwfQhamRNnSqgXMUjf3h4 * test(ui): mount a router mock for the Deleted Keys page test The page test renders the table, and the newly linked cells call useRouter, which throws without an App Router mounted. Matches how the other 35 test files in the suite stub next/navigation. Claude-Session: https://claude.ai/code/session_01NfwfQhamRNnSqgXMUjf3h4 --- .../DeletedKeysPage/DeletedKeysPage.test.tsx | 2 ++ .../DeletedKeysTable.test.tsx | 18 +++++++++++++++++ .../DeletedKeysTableColumns.tsx | 20 +++++++++++++++---- 3 files changed, 36 insertions(+), 4 deletions(-) diff --git a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx index ee2f42ad85b..31cd407a5e6 100644 --- a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx +++ b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx @@ -5,6 +5,8 @@ import { renderWithProviders } from "../../../tests/test-utils"; import DeletedKeysPage from "./DeletedKeysPage"; import { useDeletedKeys, DeletedKeyResponse } from "@/app/(dashboard)/hooks/keys/useKeys"; +vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) })); + vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ useDeletedKeys: vi.fn(), })); diff --git a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTable.test.tsx b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTable.test.tsx index 7e30ef2c135..dd8007be264 100644 --- a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTable.test.tsx @@ -5,6 +5,8 @@ import { renderWithProviders } from "../../../../tests/test-utils"; import { DeletedKeysTable } from "./DeletedKeysTable"; import { DeletedKeyResponse } from "@/app/(dashboard)/hooks/keys/useKeys"; +vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) })); + const makeDeletedKey = (overrides: Partial = {}): DeletedKeyResponse => ({ token: "sk-1234567890abcdef", @@ -86,3 +88,19 @@ it("should show the empty state when there are no deleted keys", () => { expect(screen.getByText("No deleted keys found")).toBeInTheDocument(); }); + +it("links the owner, creator and deleter cells to their user detail pages", () => { + renderWithProviders(); + + expect(screen.getByRole("link", { name: "user-1" })).toHaveAttribute("href", "/ui/users?user=user-1"); + expect(screen.getByRole("link", { name: "creator-1" })).toHaveAttribute("href", "/ui/users?user=creator-1"); + expect(screen.getByRole("link", { name: "deleter-1" })).toHaveAttribute("href", "/ui/users?user=deleter-1"); +}); + +it("leaves the default_user_id placeholder unlinked", () => { + const placeholderKey = makeDeletedKey({ user_id: "default_user_id", created_by: "default_user_id" }); + renderWithProviders(); + + expect(screen.getAllByText("default_user_id")).toHaveLength(2); + expect(screen.queryByRole("link", { name: "default_user_id" })).not.toBeInTheDocument(); +}); diff --git a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTableColumns.tsx b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTableColumns.tsx index aa7d6380cc3..32185f867b4 100644 --- a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysTable/DeletedKeysTableColumns.tsx @@ -3,8 +3,9 @@ import { ColumnDef } from "@tanstack/react-table"; import { DataTableSortHeader } from "@/components/shared/DataTable"; -import { DateCell, IdCell, MoneyCell } from "@/components/shared/table_cells"; +import { DateCell, IdCell, IdentityCell, MoneyCell } from "@/components/shared/table_cells"; import { DeletedKeyResponse } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { userDetailHref } from "@/utils/entityLinks"; function TruncatedTextCell({ value }: { value: string | null | undefined }) { if (!value) { @@ -17,6 +18,17 @@ function TruncatedTextCell({ value }: { value: string | null | undefined }) { ); } +function UserLinkCell({ userId }: { userId: string | null | undefined }) { + if (!userId) { + return -; + } + return ( + + + + ); +} + export const getDeletedKeysTableColumns = (): ColumnDef[] => [ { id: "token", @@ -89,7 +101,7 @@ export const getDeletedKeysTableColumns = (): ColumnDef[] => header: "User ID", size: 120, enableSorting: false, - cell: ({ row }) => , + cell: ({ row }) => , }, { id: "created_at", @@ -107,7 +119,7 @@ export const getDeletedKeysTableColumns = (): ColumnDef[] => header: "Created By", size: 120, enableSorting: false, - cell: ({ row }) => , + cell: ({ row }) => , }, { id: "deleted_at", @@ -125,6 +137,6 @@ export const getDeletedKeysTableColumns = (): ColumnDef[] => header: "Deleted By", size: 120, enableSorting: false, - cell: ({ row }) => , + cell: ({ row }) => , }, ]; From c53f72c764f9c1ae9015df667e350b130ebb8a9f Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 11 Sep 2026 17:49:01 -0700 Subject: [PATCH 097/119] feat(ui): link the Created By cell on the Prompts page (#40753) The column was plain muted text, so finding out who owns a prompt meant copying the id into the Users page search box. Route it through IdentityCell with userDetailHref, which keeps the proxy admin placeholder unlinked. Claude-Session: https://claude.ai/code/session_01NfwfQhamRNnSqgXMUjf3h4 --- .../prompts/_components/PromptTable.test.tsx | 11 +++++++++++ .../prompts/_components/PromptTableColumns.tsx | 12 ++++++++++-- 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.test.tsx index edbb897fb2f..009abf93d1f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.test.tsx @@ -10,6 +10,8 @@ vi.mock("@/components/networking", () => ({ modelHubCall: vi.fn().mockResolvedValue({ data: [] }), })); +vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) })); + const mockPrompts: PromptSpec[] = [ { prompt_id: "prompt-newer", @@ -53,6 +55,15 @@ describe("PromptTable", () => { } }); + it("links the Created By cell to the creator's detail page, leaving the placeholder unlinked", () => { + const prompts = [mockPrompts[0], { ...mockPrompts[1], created_by: "default_user_id" }]; + render(); + + expect(screen.getByRole("link", { name: "user-1" })).toHaveAttribute("href", "/ui/users?user=user-1"); + expect(screen.getByText("default_user_id")).toBeInTheDocument(); + expect(screen.queryByRole("link", { name: "default_user_id" })).not.toBeInTheDocument(); + }); + it("should display the empty state when data is empty", () => { render(); expect(screen.getByText("No prompts yet")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTableColumns.tsx index f927a6d1486..db0392bbd6e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTableColumns.tsx @@ -17,6 +17,7 @@ import { } from "@/components/ui/dropdown-menu"; import { cn } from "@/lib/cva.config"; import { copyToClipboard } from "@/utils/dataUtils"; +import { userDetailHref } from "@/utils/entityLinks"; import { extractModel, getProviderFromModelHub, ModelGroupInfo } from "./prompt_utils"; @@ -191,9 +192,16 @@ export const getPromptTableColumns = ({ enableSorting: false, cell: ({ row }) => { const createdBy = row.original.created_by; + if (!createdBy) { + return -; + } return ( - - {createdBy || "-"} + + ); }, From b27d2cce77cd1d0ba9662ccb5b2fe9640471f2c6 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 11 Sep 2026 17:49:07 -0700 Subject: [PATCH 098/119] feat(ui): link the Organization and Deleted By cells on Deleted Teams (#40751) * feat(ui): link the Organization and Deleted By cells on Deleted Teams Both columns rendered as plain text, so tracing a deleted team back to its org or to whoever removed it meant copying an id into another page's search box. Route them through IdentityCell with orgDetailHref and userDetailHref. Team ID stays unlinked because the team itself is gone. Claude-Session: https://claude.ai/code/session_01NfwfQhamRNnSqgXMUjf3h4 * test(ui): mount a router mock for the Deleted Teams page test The page test renders the table, and the newly linked cells call useRouter, which throws without an App Router mounted. Matches how the other 35 test files in the suite stub next/navigation. Claude-Session: https://claude.ai/code/session_01NfwfQhamRNnSqgXMUjf3h4 --- .../DeletedTeamsPage.test.tsx | 2 ++ .../DeletedTeamsTable.test.tsx | 20 +++++++++++++ .../DeletedTeamsTableColumns.tsx | 30 ++++++++++++------- 3 files changed, 41 insertions(+), 11 deletions(-) diff --git a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx index 952d8764463..d2d25f5e47e 100644 --- a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx +++ b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx @@ -5,6 +5,8 @@ import { renderWithProviders } from "../../../tests/test-utils"; import DeletedTeamsPage from "./DeletedTeamsPage"; import { useDeletedTeams, DeletedTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; +vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) })); + vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ useDeletedTeams: vi.fn(), })); diff --git a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.test.tsx b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.test.tsx index e166f6b0d1b..837fb583f30 100644 --- a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.test.tsx +++ b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.test.tsx @@ -4,6 +4,8 @@ import { renderWithProviders } from "../../../../tests/test-utils"; import { DeletedTeamsTable } from "./DeletedTeamsTable"; import { DeletedTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; +vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) })); + const makeDeletedTeam = (overrides: Partial = {}): DeletedTeam => ({ team_id: "team-1", team_alias: "Test Team", @@ -81,3 +83,21 @@ it("renders the shared pagination footer with the server row count", () => { expect(screen.getByTestId("pagination-prev")).toBeEnabled(); expect(screen.getByTestId("pagination-next")).toBeDisabled(); }); + +it("links the organization and deleted by cells, leaving the deleted team id unlinked", () => { + renderWithProviders( + , + ); + + expect(screen.getByRole("link", { name: "org-1" })).toHaveAttribute("href", "/ui/organizations?org=org-1"); + expect(screen.getByRole("link", { name: "user-1" })).toHaveAttribute("href", "/ui/users?user=user-1"); + expect(screen.queryByRole("link", { name: "team-1" })).not.toBeInTheDocument(); +}); + +it("leaves the default_user_id placeholder unlinked in the deleted by cell", () => { + const team = makeDeletedTeam({ deleted_by: "default_user_id", organization_id: null }); + renderWithProviders(); + + expect(screen.getByText("default_user_id")).toBeInTheDocument(); + expect(screen.queryByRole("link", { name: "default_user_id" })).not.toBeInTheDocument(); +}); diff --git a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTableColumns.tsx b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTableColumns.tsx index e36077fd2c3..172f0417027 100644 --- a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTableColumns.tsx @@ -3,8 +3,20 @@ import { ColumnDef } from "@tanstack/react-table"; import { DataTableSortHeader } from "@/components/shared/DataTable"; -import { DateCell, IdCell, ModelsCell, MoneyCell } from "@/components/shared/table_cells"; +import { DateCell, IdCell, IdentityCell, ModelsCell, MoneyCell } from "@/components/shared/table_cells"; import { DeletedTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { orgDetailHref, userDetailHref } from "@/utils/entityLinks"; + +function EntityCell({ value, href }: { value: string | null | undefined; href: string | undefined }) { + if (!value) { + return -; + } + return ( + + + + ); +} export const getDeletedTeamsTableColumns = (): ColumnDef[] => [ { @@ -78,7 +90,10 @@ export const getDeletedTeamsTableColumns = (): ColumnDef[] => [ header: "Organization", size: 150, enableSorting: false, - cell: ({ row }) => , + cell: ({ row }) => { + const orgId = row.original.organization_id; + return ; + }, }, { id: "deleted_at", @@ -97,15 +112,8 @@ export const getDeletedTeamsTableColumns = (): ColumnDef[] => [ size: 120, enableSorting: false, cell: ({ row }) => { - const value = row.original.deleted_by; - if (!value) { - return -; - } - return ( - - {value} - - ); + const deletedBy = row.original.deleted_by; + return ; }, }, ]; From 115535c3add1f79bf8b5056663c1456a0be5e1cf Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 12 Sep 2026 00:49:26 +0000 Subject: [PATCH 099/119] test(model_prices): cover Fireworks DeepSeek V4.1 costs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_fireworks_serverless_model_costs.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/test_litellm/test_fireworks_serverless_model_costs.py b/tests/test_litellm/test_fireworks_serverless_model_costs.py index b8d36b996df..e3404e477b7 100644 --- a/tests/test_litellm/test_fireworks_serverless_model_costs.py +++ b/tests/test_litellm/test_fireworks_serverless_model_costs.py @@ -14,6 +14,8 @@ import os import pytest +from litellm import completion_cost +from litellm.types.utils import Choices, Message, ModelResponse, Usage from litellm.utils import get_model_info @@ -56,6 +58,20 @@ def test_bare_fireworks_ids_resolve_through_prefixed_entries(): assert info["max_output_tokens"] == expected["max_output_tokens"] +def test_deepseek_v4p1_flash_twin_costs(): + for model in ( + "fireworks_ai/deepseek-v4p1-flash", + "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash", + ): + response = ModelResponse( + model=model, + choices=[Choices(index=0, message=Message(role="assistant", content="ok"))], + usage=Usage(prompt_tokens=1000, completion_tokens=1000, total_tokens=2000), + ) + cost = completion_cost(completion_response=response, model=model) + assert cost == pytest.approx(8.8e-04) + + TWIN_PINNED_PRICES = { "deepseek-v4-flash-0731": { "input_cost_per_token": 2.2e-07, From 05d2c316f5d389979ed60ba26cd0c2ef71ec3692 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Fri, 11 Sep 2026 17:53:51 -0700 Subject: [PATCH 100/119] refactor(mcp): reuse in-memory discovery storage --- litellm/caching/in_memory_cache.py | 13 ++-- .../mcp_server/mcp_server_manager.py | 67 +++++++++---------- .../caching/test_in_memory_cache.py | 24 +++++++ .../mcp_server/test_mcp_server_manager.py | 51 +++++++++++++- 4 files changed, 112 insertions(+), 43 deletions(-) diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 4058a3d72dd..56c9147e066 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -13,6 +13,7 @@ import json import sys import threading import time +from collections.abc import Callable from typing import TYPE_CHECKING, Any, Final if TYPE_CHECKING: @@ -34,6 +35,7 @@ class InMemoryCache(BaseCache): default_ttl: int | None = 600, # default ttl is 10 minutes. At maximum litellm rate limiting logic requires objects to be in memory for 1 minute max_size_per_item: int | None = 1024, # 1MB = 1024KB + clock: Callable[[], float] | None = None, ): """ max_size_in_memory [int]: Maximum number of items in cache. done to prevent memory leaks. Use 200 items as a default @@ -49,6 +51,7 @@ class InMemoryCache(BaseCache): self.ttl_dict: dict = {} self.expiration_heap: list[tuple[float, str]] = [] self._increment_lock = threading.Lock() + self._clock = clock if clock is not None else lambda: time.time() def check_value_size(self, value: Any): """ @@ -91,7 +94,7 @@ class InMemoryCache(BaseCache): """ Check if a specific key is expired """ - return key in self.ttl_dict and time.time() > self.ttl_dict[key] + return key in self.ttl_dict and self._clock() > self.ttl_dict[key] def _remove_key(self, key: str) -> None: """ @@ -113,7 +116,7 @@ class InMemoryCache(BaseCache): - 3. the size of in-memory cache is bounded """ - current_time: Final = time.time() + current_time: Final = self._clock() # Step 1: Remove expired or outdated items while self.expiration_heap: @@ -147,7 +150,7 @@ class InMemoryCache(BaseCache): Check if ttl is set for a key """ ttl_time: Final = self.ttl_dict.get(key) - if ttl_time is None or float(ttl_time) < time.time(): # if ttl is not set, allow override + if ttl_time is None or float(ttl_time) < self._clock(): # if ttl is not set, allow override return True else: return False @@ -167,10 +170,10 @@ class InMemoryCache(BaseCache): self.cache_dict[key] = value if self.allow_ttl_override(key): # if ttl is not set, set it to default ttl if "ttl" in kwargs and kwargs["ttl"] is not None: - self.ttl_dict[key] = time.time() + float(kwargs["ttl"]) + self.ttl_dict[key] = self._clock() + float(kwargs["ttl"]) heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key)) else: - self.ttl_dict[key] = time.time() + self.default_ttl + self.ttl_dict[key] = self._clock() + self.default_ttl heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key)) async def async_set_cache(self, key, value, **kwargs): diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index ca81e8191f7..35dfeb9ee74 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -51,6 +51,7 @@ from typing_extensions import ReadOnly import litellm from litellm._logging import verbose_logger +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( MCP_CLIENT_TIMEOUT, MCP_HEALTH_CHECK_TIMEOUT, @@ -195,7 +196,6 @@ if TYPE_CHECKING: from mcp.shared.context import RequestContext from mcp.types import CreateMessageRequestParams - from litellm.caching.caching import InMemoryCache from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.mcp_server.mcp_toolset import MCPToolset @@ -1686,21 +1686,29 @@ _DISCOVERY_CACHE_LIMIT: Final = 1024 @dataclass(frozen=True, slots=True) class _DiscoveryEntry(Generic[_DiscoveryItem]): - expires_at: float items: tuple[_DiscoveryItem, ...] class _DiscoveryCache(Generic[_DiscoveryItem]): def __init__(self, ttl: float, clock: Callable[[], float]) -> None: self._ttl = ttl - self._clock = clock - self._entries: Mapping[_DiscoveryKey, _DiscoveryEntry[_DiscoveryItem]] = MappingProxyType({}) - self._pending: Mapping[_DiscoveryKey, asyncio.Task[list[_DiscoveryItem]]] = MappingProxyType({}) - self._waiters: Mapping[asyncio.Task[list[_DiscoveryItem]], int] = MappingProxyType({}) + self._entries = InMemoryCache(max_size_in_memory=_DISCOVERY_CACHE_LIMIT, clock=clock) + self._pending: dict[ + _DiscoveryKey, asyncio.Task[list[_DiscoveryItem]] + ] = {} # mutable-ok: constant-time fetch registration + self._waiters: dict[asyncio.Task[list[_DiscoveryItem]], int] = {} # mutable-ok: constant-time waiter accounting def invalidate(self, server_id: str) -> None: - self._entries = MappingProxyType({key: entry for key, entry in self._entries.items() if key[0] != server_id}) - self._pending = MappingProxyType({key: task for key, task in self._pending.items() if key[0] != server_id}) + prefix: Final = f"[{json.dumps(server_id)}," + keys: Final = cast( # cast-ok: private cache contains only JSON string keys + "tuple[str, ...]", tuple(self._entries.cache_dict) + ) + for entry_key in keys: + if entry_key.startswith(prefix): + self._entries.delete_cache(entry_key) + for key in tuple(self._pending): + if key[0] == server_id: + self._pending.pop(key) @staticmethod def _observe_completion(task: asyncio.Task[list[_DiscoveryItem]]) -> None: @@ -1712,8 +1720,10 @@ class _DiscoveryCache(Generic[_DiscoveryItem]): ) -> tuple[_DiscoveryItem, ...]: if self._ttl <= 0: return tuple(await fetch()) - entry: Final = self._entries.get(key) - if entry is not None and entry.expires_at > self._clock(): + entry: Final = cast( # cast-ok: private cache contains only entries for this item type + "_DiscoveryEntry[_DiscoveryItem] | None", self._entries.get_cache(json.dumps(key)) + ) + if entry is not None: return tuple(item.model_copy(deep=True) for item in entry.items) pending: Final = self._pending.get(key) if pending is not None: @@ -1721,28 +1731,24 @@ class _DiscoveryCache(Generic[_DiscoveryItem]): if len(self._pending) >= _DISCOVERY_CACHE_LIMIT: return tuple(await fetch()) task: Final = asyncio.create_task(self._fetch(key, fetch)) - self._pending = MappingProxyType({**self._pending, key: task}) + self._pending[key] = task task.add_done_callback(self._observe_completion) return await self._await_fetch(key, task) async def _await_fetch( self, key: _DiscoveryKey, task: asyncio.Task[list[_DiscoveryItem]] ) -> tuple[_DiscoveryItem, ...]: - self._waiters = MappingProxyType({**self._waiters, task: self._waiters.get(task, 0) + 1}) + self._waiters[task] = self._waiters.get(task, 0) + 1 try: return tuple(item.model_copy(deep=True) for item in await asyncio.shield(task)) finally: remaining: Final = self._waiters[task] - 1 if remaining: - self._waiters = MappingProxyType({**self._waiters, task: remaining}) + self._waiters[task] = remaining else: - self._waiters = MappingProxyType( - {pending: count for pending, count in self._waiters.items() if pending is not task} - ) + self._waiters.pop(task) if self._pending.get(key) is task: - self._pending = MappingProxyType( - {entry_key: pending for entry_key, pending in self._pending.items() if entry_key != key} - ) + self._pending.pop(key) if not task.done(): task.cancel() @@ -1752,28 +1758,15 @@ class _DiscoveryCache(Generic[_DiscoveryItem]): try: items: Final = await fetch() if self._pending.get(key) is asyncio.current_task(): - now: Final = self._clock() - live_entries: Final = tuple( - (entry_key, entry) for entry_key, entry in self._entries.items() if entry.expires_at > now - ) - self._entries = MappingProxyType( - { - entry_key: entry - for entry_key, entry in ( - *live_entries[-(_DISCOVERY_CACHE_LIMIT - 1) :], - ( - key, - _DiscoveryEntry(now + self._ttl, tuple(item.model_copy(deep=True) for item in items)), - ), - ) - } + self._entries.set_cache( + json.dumps(key), + _DiscoveryEntry(tuple(item.model_copy(deep=True) for item in items)), + ttl=self._ttl, ) return items finally: if self._pending.get(key) is asyncio.current_task(): - self._pending = MappingProxyType( - {entry_key: task for entry_key, task in self._pending.items() if entry_key != key} - ) + self._pending.pop(key) def _mcp_discovery_cache_ttl() -> float: diff --git a/tests/test_litellm/caching/test_in_memory_cache.py b/tests/test_litellm/caching/test_in_memory_cache.py index 85e8308ae91..40ad4f0c6f0 100644 --- a/tests/test_litellm/caching/test_in_memory_cache.py +++ b/tests/test_litellm/caching/test_in_memory_cache.py @@ -250,3 +250,27 @@ def test_in_memory_cache_prunes_expired_heap_entries_below_capacity(): assert len(in_memory_cache.cache_dict) == 5 assert len(in_memory_cache.ttl_dict) == 5 assert len(in_memory_cache.expiration_heap) == 5 + + +def test_in_memory_cache_injected_clock_controls_expiry_and_eviction() -> None: + class Clock: + now = 0.0 + + def __call__(self) -> float: + return self.now + + clock = Clock() + cache = InMemoryCache(max_size_in_memory=2, default_ttl=60, clock=clock) + cache.set_cache("first", "original", ttl=10) + clock.now = 9.0 + cache.set_cache("second", "survivor") + assert cache.get_cache("first") == "original" + clock.now = 10.001 + assert cache.get_cache("first") is None + cache.set_cache("third", "replacement") + assert cache.get_cache("second") == "survivor" + clock.now = 69.001 + cache.set_cache("fourth", "new") + assert cache.get_cache("second") is None + assert cache.get_cache("third") == "replacement" + assert cache.get_cache("fourth") == "new" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index f1343bd143d..6630e184a16 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -13095,7 +13095,7 @@ async def test_discovery_cache_reuses_raw_results_and_expires(kind: str) -> None clock.now = 59.999 assert (await operation(server, None))[0].name == "discovery-example" assert upstream.initializes == 1 - clock.now = 60.0 + clock.now = 60.001 assert (await operation(server, None))[0].name == "discovery-example" assert upstream.initializes == 2 @@ -13383,3 +13383,52 @@ async def test_discovery_resolves_stored_oauth_for_the_requesting_user() -> None assert store.calls == (("requesting-user", "discovery"), ("requesting-user", "discovery")) assert upstream.initializes == 1 assert ("prompts/list", "Bearer stored-token") in upstream.requests + + +@pytest.mark.asyncio +async def test_discovery_cache_evicts_results_at_capacity() -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache + + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock()) + + async def original() -> list[Prompt]: + return [Prompt(name="original")] + + async def refetched() -> list[Prompt]: + return [Prompt(name="refetched")] + + for index in range(1025): + assert (await cache.get((f"server-{index:04}", None), original))[0].name == "original" + assert (await cache.get(("server-1024", None), refetched))[0].name == "original" + assert (await cache.get(("server-0000", None), refetched))[0].name == "refetched" + + +@pytest.mark.asyncio +async def test_discovery_cache_invalidation_preserves_other_servers_and_pending_fetches() -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache + + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock()) + entered: Final = asyncio.Event() + release: Final = asyncio.Event() + + async def original() -> list[Prompt]: + return [Prompt(name="original")] + + async def blocked() -> list[Prompt]: + entered.set() + await release.wait() + return [Prompt(name="pending")] + + async def refetched() -> list[Prompt]: + return [Prompt(name="refetched")] + + assert (await cache.get(("server", None), original))[0].name == "original" + assert (await cache.get(("server-extra", None), original))[0].name == "original" + task: Final = asyncio.create_task(cache.get(("other", None), blocked)) + await asyncio.wait_for(entered.wait(), timeout=5) + cache.invalidate("server") + release.set() + assert (await asyncio.wait_for(task, timeout=5))[0].name == "pending" + assert (await cache.get(("other", None), refetched))[0].name == "pending" + assert (await cache.get(("server-extra", None), refetched))[0].name == "original" + assert (await cache.get(("server", None), refetched))[0].name == "refetched" From 6c07876dcfe9ce74401526c3a29657848ffb4afe Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 12 Sep 2026 01:00:11 +0000 Subject: [PATCH 101/119] test(model_prices): use local cost map for Fireworks cost coverage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_fireworks_serverless_model_costs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/test_fireworks_serverless_model_costs.py b/tests/test_litellm/test_fireworks_serverless_model_costs.py index e3404e477b7..1303f46e8fa 100644 --- a/tests/test_litellm/test_fireworks_serverless_model_costs.py +++ b/tests/test_litellm/test_fireworks_serverless_model_costs.py @@ -58,7 +58,7 @@ def test_bare_fireworks_ids_resolve_through_prefixed_entries(): assert info["max_output_tokens"] == expected["max_output_tokens"] -def test_deepseek_v4p1_flash_twin_costs(): +def test_deepseek_v4p1_flash_twin_costs(local_model_cost_map): for model in ( "fireworks_ai/deepseek-v4p1-flash", "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash", From 9d31de2f20d444d02674c8ec3b8f25a2cfa3257e Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Fri, 11 Sep 2026 18:08:27 -0700 Subject: [PATCH 102/119] fix(mcp): bound discovery cache result bytes --- .../mcp_server/mcp_server_manager.py | 34 ++++++++++--------- .../mcp_server/test_mcp_server_manager.py | 25 ++++++++++---- 2 files changed, 37 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 35dfeb9ee74..9211d8ab003 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -46,7 +46,7 @@ from mcp.types import ( ResourceTemplate, ) from mcp.types import Tool as MCPTool -from pydantic import AnyUrl, BaseModel +from pydantic import AnyUrl, BaseModel, TypeAdapter from typing_extensions import ReadOnly import litellm @@ -1684,15 +1684,13 @@ _DiscoveryKey: TypeAlias = tuple[str, str | None] _DISCOVERY_CACHE_LIMIT: Final = 1024 -@dataclass(frozen=True, slots=True) -class _DiscoveryEntry(Generic[_DiscoveryItem]): - items: tuple[_DiscoveryItem, ...] - - class _DiscoveryCache(Generic[_DiscoveryItem]): - def __init__(self, ttl: float, clock: Callable[[], float]) -> None: + def __init__( + self, ttl: float, clock: Callable[[], float], adapter: TypeAdapter[tuple[_DiscoveryItem, ...]] + ) -> None: self._ttl = ttl - self._entries = InMemoryCache(max_size_in_memory=_DISCOVERY_CACHE_LIMIT, clock=clock) + self._adapter = adapter + self._entries = InMemoryCache(max_size_in_memory=_DISCOVERY_CACHE_LIMIT, max_size_per_item=64, clock=clock) self._pending: dict[ _DiscoveryKey, asyncio.Task[list[_DiscoveryItem]] ] = {} # mutable-ok: constant-time fetch registration @@ -1720,11 +1718,9 @@ class _DiscoveryCache(Generic[_DiscoveryItem]): ) -> tuple[_DiscoveryItem, ...]: if self._ttl <= 0: return tuple(await fetch()) - entry: Final = cast( # cast-ok: private cache contains only entries for this item type - "_DiscoveryEntry[_DiscoveryItem] | None", self._entries.get_cache(json.dumps(key)) - ) + entry: Final[object] = self._entries.get_cache(json.dumps(key)) if entry is not None: - return tuple(item.model_copy(deep=True) for item in entry.items) + return self._adapter.validate_python(entry) pending: Final = self._pending.get(key) if pending is not None: return await self._await_fetch(key, pending) @@ -1760,7 +1756,7 @@ class _DiscoveryCache(Generic[_DiscoveryItem]): if self._pending.get(key) is asyncio.current_task(): self._entries.set_cache( json.dumps(key), - _DiscoveryEntry(tuple(item.model_copy(deep=True) for item in items)), + self._adapter.dump_json(tuple(items)), ttl=self._ttl, ) return items @@ -1909,9 +1905,15 @@ class MCPServerManager: token_exchanger=build_token_exchanger(), ) discovery_ttl: Final = _mcp_discovery_cache_ttl() - self._prompt_discovery_cache = _DiscoveryCache[Prompt](discovery_ttl, discovery_clock) - self._resource_discovery_cache = _DiscoveryCache[Resource](discovery_ttl, discovery_clock) - self._template_discovery_cache = _DiscoveryCache[ResourceTemplate](discovery_ttl, discovery_clock) + self._prompt_discovery_cache = _DiscoveryCache[Prompt]( + discovery_ttl, discovery_clock, TypeAdapter(tuple[Prompt, ...]) + ) + self._resource_discovery_cache = _DiscoveryCache[Resource]( + discovery_ttl, discovery_clock, TypeAdapter(tuple[Resource, ...]) + ) + self._template_discovery_cache = _DiscoveryCache[ResourceTemplate]( + discovery_ttl, discovery_clock, TypeAdapter(tuple[ResourceTemplate, ...]) + ) self.registry: dict[str, MCPServer] = {} self._openapi_health_probes: Callable[[str], _OpenAPIHealthProbe] = lru_cache(maxsize=128)(_OpenAPIHealthProbe) self.config_mcp_servers: dict[str, MCPServer] = {} diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 6630e184a16..d2987c5112e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -30,7 +30,7 @@ from mcp.types import ( TextResourceContents, ) from mcp.types import Tool as MCPTool -from pydantic import AnyUrl +from pydantic import AnyUrl, TypeAdapter from litellm.constants import MCP_METADATA_TIMEOUT from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( @@ -13224,7 +13224,7 @@ def test_discovery_cache_keys_isolate_user_dependent_auth(auth_type: MCPAuth) -> async def test_discovery_cache_retries_cancelled_fetches() -> None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache - cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock()) + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock(), TypeAdapter(tuple[Prompt, ...])) async def cancelled() -> list[Prompt]: raise asyncio.CancelledError() @@ -13241,7 +13241,7 @@ async def test_discovery_cache_retries_cancelled_fetches() -> None: async def test_discovery_cache_cancels_fetch_when_last_waiter_leaves() -> None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache - cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock()) + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock(), TypeAdapter(tuple[Prompt, ...])) entered: Final = asyncio.Event() stopped: Final = asyncio.Event() release: Final = asyncio.Event() @@ -13270,7 +13270,7 @@ async def test_discovery_cache_cancels_fetch_when_last_waiter_leaves() -> None: async def test_discovery_cache_bounds_detached_fetches_without_dropping_results() -> None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache - cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock()) + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock(), TypeAdapter(tuple[Prompt, ...])) entered: Final[asyncio.Queue[None]] = asyncio.Queue() release: Final = asyncio.Event() @@ -13389,7 +13389,7 @@ async def test_discovery_resolves_stored_oauth_for_the_requesting_user() -> None async def test_discovery_cache_evicts_results_at_capacity() -> None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache - cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock()) + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock(), TypeAdapter(tuple[Prompt, ...])) async def original() -> list[Prompt]: return [Prompt(name="original")] @@ -13407,7 +13407,7 @@ async def test_discovery_cache_evicts_results_at_capacity() -> None: async def test_discovery_cache_invalidation_preserves_other_servers_and_pending_fetches() -> None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache - cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock()) + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock(), TypeAdapter(tuple[Prompt, ...])) entered: Final = asyncio.Event() release: Final = asyncio.Event() @@ -13432,3 +13432,16 @@ async def test_discovery_cache_invalidation_preserves_other_servers_and_pending_ assert (await cache.get(("other", None), refetched))[0].name == "pending" assert (await cache.get(("server-extra", None), refetched))[0].name == "original" assert (await cache.get(("server", None), refetched))[0].name == "refetched" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("description", ("x" * 96_000, "é" * 40_000), ids=("ascii", "unicode")) +async def test_discovery_cache_returns_oversized_results_without_retaining_them(description: str) -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache + + cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock(), TypeAdapter(tuple[Prompt, ...])) + fetch: Final = AsyncMock(return_value=[Prompt(name="large", description=description)]) + for _ in range(2): + result: Final = await cache.get(("server", None), fetch) + assert result[0].description == description + assert fetch.await_count == 2 From 105dc7710959e63264c99f0f68c581781e254494 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 11 Sep 2026 18:14:21 -0700 Subject: [PATCH 103/119] fix(realtime): probe Azure's GA realtime upstream in health checks when no protocol is pinned --- litellm/realtime_api/main.py | 8 ++--- tests/test_litellm/realtime_api/test_main.py | 36 +++++++++++++++----- 2 files changed, 31 insertions(+), 13 deletions(-) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index d42e7e18b75..44c47af57f4 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -586,9 +586,7 @@ def _azure_realtime_health_protocol( configured: Final = configured_raw if isinstance(configured_raw, str) else None if configured is not None: return configured, query_params - if query_params is not None: - return "GA", query_params - return "beta", None + return "GA", query_params def _realtime_health_check_auth_headers( @@ -621,8 +619,8 @@ async def _realtime_health_check( api_key: str - api key custom_llm_provider: str - custom llm provider realtime_protocol: Optional[str] - protocol version ("GA"/"v1" for GA path, "beta" for beta path); - None resolves it for Azure from model_params/env, with transcription-only models probing GA - plus intent=transcription the way real calls do + None resolves it for Azure from model_params/env and otherwise probes GA, the upstream a client + without the OpenAI-Beta header is bridged to, with transcription-only models adding intent=transcription Returns: bool - True if connection is successful, False otherwise diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py index 8a9abe819e5..0827bbcdc38 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/test_litellm/realtime_api/test_main.py @@ -266,7 +266,8 @@ def test_transcription_only_detection_rejects_speech_model(local_model_cost_map) @pytest.mark.asyncio -async def test_azure_health_check_keeps_beta_path_for_speech_model(): +async def test_azure_health_check_probes_the_ga_upstream_for_an_unconfigured_speech_model(monkeypatch): + monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) connect = _CapturingConnect() with patch("websockets.connect", connect): assert await realtime_main._realtime_health_check( @@ -276,14 +277,18 @@ async def test_azure_health_check_keeps_beta_path_for_speech_model(): api_base="https://my-endpoint.openai.azure.com", api_version="2024-10-01-preview", ) - assert connect.url == ( - "wss://my-endpoint.openai.azure.com/openai/realtime" - "?api-version=2024-10-01-preview&deployment=gpt-4o-realtime-preview" - ) + assert connect.url == "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=gpt-4o-realtime-preview" + + +_AZURE_BETA_HEALTH_URL: Final = ( + "wss://my-endpoint.openai.azure.com/openai/realtime" + "?api-version=2024-10-01-preview&deployment=gpt-4o-realtime-preview" +) @pytest.mark.asyncio -async def test_azure_health_check_honors_deployment_realtime_protocol(): +async def test_azure_health_check_honors_deployment_realtime_protocol(monkeypatch): + monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False) connect = _CapturingConnect() with patch("websockets.connect", connect): assert await realtime_main._realtime_health_check( @@ -292,9 +297,24 @@ async def test_azure_health_check_honors_deployment_realtime_protocol(): api_key="fake-key", api_base="https://my-endpoint.openai.azure.com", api_version="2024-10-01-preview", - model_params={"realtime_protocol": "GA"}, + model_params={"realtime_protocol": "beta"}, ) - assert connect.url == "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=gpt-4o-realtime-preview" + assert connect.url == _AZURE_BETA_HEALTH_URL + + +@pytest.mark.asyncio +async def test_azure_health_check_honors_env_realtime_protocol(monkeypatch): + monkeypatch.setenv("LITELLM_AZURE_REALTIME_PROTOCOL", "beta") + connect = _CapturingConnect() + with patch("websockets.connect", connect): + assert await realtime_main._realtime_health_check( + model="gpt-4o-realtime-preview", + custom_llm_provider="azure", + api_key="fake-key", + api_base="https://my-endpoint.openai.azure.com", + api_version="2024-10-01-preview", + ) + assert connect.url == _AZURE_BETA_HEALTH_URL class _ConnectThatStopsAfterCapturingTheUrl: From f84f986b4ee571c8b3db18a560e007801379a933 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 11 Sep 2026 18:16:33 -0700 Subject: [PATCH 104/119] fix(guardrails): keep post_call guardrail info on streamed chat completions (#40806) * fix(guardrails): sync logging_obj guardrail info on every record so post_call entries survive streamed chat completions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(guardrails): hoist regression test imports to module scope Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 1 + .../integrations/test_custom_guardrail.py | 61 ++++++++++++++++++- 2 files changed, 61 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index bb54767edef..8a976a966a6 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1217,6 +1217,7 @@ class CustomGuardrail(CustomLogger): _, metadata_bucket = get_or_create_metadata_bucket(request_data) _append_guardrail_info(metadata_bucket) + _sync_guardrail_info_to_logging_obj(request_data, request_data.get("litellm_logging_obj")) _guardrail_self_recorded.set(True) diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 1644d78ae37..bb4822eae57 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1,4 +1,5 @@ import asyncio +import datetime as dt from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional from unittest.mock import AsyncMock @@ -9,9 +10,16 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy._types import CallTypes, UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks, Mode -from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailTracingDetail +from litellm.types.utils import ( + Choices, + GenericGuardrailAPIInputs, + GuardrailTracingDetail, + Message, + ModelResponse, +) if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -2345,6 +2353,57 @@ class TestUndecoratedApplyGuardrailIsLogged: assert _Labelled.seen_label == "docs-style" + @pytest.mark.asyncio + async def test_post_call_recorded_outside_decorator_reaches_standard_logging_object(self): + """LIT-7608 regression: the auto-wrapped pre_call apply_guardrail copies the request bucket + into logging_obj.litellm_params["metadata"]. A post_call entry recorded later without the + decorator (the Bedrock streaming hook) must not be shadowed by that stale copy.""" + messages: Final = [{"role": "user", "content": "hello there"}] + litellm_metadata: Final[dict] = {"user_api_key_user_id": "u1"} + logging_obj: Final = Logging( + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=messages, + stream=True, + call_type=CallTypes.acompletion.value, + start_time=dt.datetime.now(), + litellm_call_id="call-1", + function_id="fn-1", + ) + logging_obj.update_environment_variables( + litellm_params={"litellm_metadata": litellm_metadata}, + optional_params={}, + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + custom_llm_provider="bedrock", + ) + request_data: Final = { + "model": "bedrock-haiku", + "messages": messages, + "litellm_metadata": litellm_metadata, + "litellm_logging_obj": logging_obj, + } + guardrail: Final = _UndecoratedGuardrail(guardrail_name="bedrock-pre", event_hook=GuardrailEventHooks.pre_call) + + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["hello there"]), + request_data=request_data, + input_type="request", + logging_obj=logging_obj, + ) + guardrail.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"action": "NONE"}, + request_data=request_data, + guardrail_status="success", + event_type=GuardrailEventHooks.post_call, + ) + await logging_obj.async_success_handler( + result=ModelResponse(choices=[Choices(message=Message(role="assistant", content="general kenobi"))]), + start_time=dt.datetime.now(), + end_time=dt.datetime.now(), + ) + + entries: Final = logging_obj.model_call_details["standard_logging_object"]["guardrail_information"] + assert [e["guardrail_mode"] for e in entries] == ["pre_call", "post_call"] + class _ApplyOnlyObserver(CustomGuardrail): """Overrides only apply_guardrail, like panw_prisma_airs; inherits async_logging_hook.""" From cf97b757a4356b02baef6b88b5b216ed01dcb515 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 11 Sep 2026 18:37:08 -0700 Subject: [PATCH 105/119] fix(ui): preserve cleared shared select values (#40795) Preserve explicit null when shared selectors clear and adapt affected forms, validation, and request payloads. Clear stale dependent relationships and retain required-selection checks. Document project detachment, user model-budget clearing, and routing-compression clearing as deferred follow-ups. --- .../add_agent_form.integration.test.tsx | 33 ++++++++++++++- .../agents/_components/add_agent_form.tsx | 16 ++++---- .../_components/ShadowEvalStartForm.tsx | 17 ++++---- .../_components/pricing_calculator/index.tsx | 4 +- .../_components/pricing_calculator/types.ts | 4 +- .../_components/TeamGuardrailsTab.tsx | 5 ++- .../components/AllModelsTable.tsx | 4 +- .../panels/AddModelPanel.tsx | 6 ++- ...I.test.tsx => ChatUI.integration.test.tsx} | 16 ++++++++ .../playground/components/chat_ui/ChatUI.tsx | 30 +++++++++----- .../components/chat_ui/EndpointSelector.tsx | 4 +- .../components/chat_ui/SessionManagement.tsx | 2 +- .../compareUI/components/ModelSelector.tsx | 3 +- ...x => add_policy_form.integration.test.tsx} | 0 .../policies/_components/add_policy_form.tsx | 14 +++---- .../_components/ai_suggestion_modal.tsx | 6 +-- .../_components/pipeline_flow_builder.tsx | 2 +- .../_components/template_parameter_modal.tsx | 6 +-- .../ProjectModals/EditProjectModal.tsx | 6 +-- .../ProjectModals/ProjectBaseForm.tsx | 6 +-- .../ProjectModals/projectFormSchema.ts | 10 +++-- .../prompt_editor_view/ModelConfigCard.tsx | 4 +- .../prompt_editor_view/PromptEditorHeader.tsx | 4 +- .../_components/prompt_editor_view/types.ts | 2 +- .../prompt_editor_view/utils.test.ts | 16 ++++++++ .../_components/prompt_editor_view/utils.ts | 4 +- .../DefaultUserSettingsForm.tsx | 11 +++-- .../default-user-settings/mapper.test.ts | 8 ++-- .../default-user-settings/mapper.ts | 11 +++-- .../default-user-settings/schema.ts | 12 ++++-- .../_components/view_users/UsersTable.tsx | 4 +- .../src/components/CreateUserButton.tsx | 2 +- .../MCPSemanticFilterSettings.tsx | 2 +- .../MCPSemanticFilterTestPanel.tsx | 4 +- .../semanticFilterTestUtils.ts | 6 +-- .../toolSearchForm.test.ts | 1 + .../MCPToolSearchSettings/toolSearchForm.ts | 6 +-- .../Fallbacks/FallbackGroupConfig.tsx | 10 ++--- ui/litellm-dashboard/src/components/Teams.tsx | 4 +- .../src/components/TeamsPage/TeamsTable.tsx | 2 +- .../VirtualKeysPage/VirtualKeysTable.tsx | 4 +- ....tsx => AddModelForm.integration.test.tsx} | 2 +- .../src/components/add_model/AddModelForm.tsx | 15 ++++--- .../add_model/ClassificationMethodConfig.tsx | 3 +- .../add_model/ComplexityRouterConfig.tsx | 2 +- .../add_model/CompressionControls.tsx | 8 ++-- ... RouterConfigBuilder.integration.test.tsx} | 2 +- .../add_model/RouterConfigBuilder.test.ts | 10 +++++ .../add_model/RouterConfigBuilder.tsx | 23 +++++++---- .../add_model/SemanticKeywordMatching.tsx | 4 +- .../add_model/add_auto_router_tab.tsx | 19 +++++---- .../add_model/litellm_model_name.tsx | 10 ++--- .../add_model/provider_specific_fields.tsx | 3 +- .../common_components/ModelAliasManager.tsx | 14 +++++-- .../common_components/ModelSelector.tsx | 16 ++++---- .../OrganizationDropdown.tsx | 4 +- .../common_components/ProjectDropdown.tsx | 4 +- .../common_components/UserDropdown.tsx | 4 +- .../common_components/team_dropdown.tsx | 8 ++-- ...=> user_search_modal.integration.test.tsx} | 18 +++++--- .../common_components/user_search_modal.tsx | 28 +++++++++---- .../edit_auto_router_modal.tsx | 8 ++-- .../BudgetFallbacksEditor.tsx | 4 +- .../key_team_helpers/ModelMaxBudgetEditor.tsx | 4 +- .../components/model_add/CredentialModal.tsx | 8 ++-- .../model_add/credential_form_helpers.ts | 6 +-- .../components/organisms/createKeyPayload.ts | 2 + .../create_key_button.integration.test.tsx | 32 ++++++++++++++- .../organisms/create_key_button.tsx | 24 +++++------ .../src/components/provider_info_helpers.tsx | 2 +- ...aginatedSearchSelect.integration.test.tsx} | 41 ++++++++++++++----- .../shared/PaginatedSearchSelect.tsx | 10 ++--- ....tsx => SearchSelect.integration.test.tsx} | 32 ++++++++++++--- .../src/components/shared/SearchSelect.tsx | 10 ++--- .../src/components/team/TeamInfo.tsx | 7 +++- ...tsx => key_edit_view.integration.test.tsx} | 32 +++++++++++++-- .../components/templates/key_edit_view.tsx | 15 ++++--- .../view_logs/RequestLogsFilters.tsx | 10 ++--- 78 files changed, 492 insertions(+), 263 deletions(-) rename ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/{ChatUI.test.tsx => ChatUI.integration.test.tsx} (95%) rename ui/litellm-dashboard/src/app/(dashboard)/policies/_components/{add_policy_form.test.tsx => add_policy_form.integration.test.tsx} (100%) rename ui/litellm-dashboard/src/components/add_model/{AddModelForm.test.tsx => AddModelForm.integration.test.tsx} (99%) rename ui/litellm-dashboard/src/components/add_model/{RouterConfigBuilder.test.tsx => RouterConfigBuilder.integration.test.tsx} (99%) create mode 100644 ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.test.ts rename ui/litellm-dashboard/src/components/common_components/{user_search_modal.test.tsx => user_search_modal.integration.test.tsx} (94%) rename ui/litellm-dashboard/src/components/shared/{PaginatedSearchSelect.test.tsx => PaginatedSearchSelect.integration.test.tsx} (89%) rename ui/litellm-dashboard/src/components/shared/{SearchSelect.test.tsx => SearchSelect.integration.test.tsx} (76%) rename ui/litellm-dashboard/src/components/templates/{key_edit_view.test.tsx => key_edit_view.integration.test.tsx} (98%) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx index dad8599e967..ebc97891744 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx @@ -1,11 +1,11 @@ import React from "react"; -import { render, screen, waitFor, within } from "@testing-library/react"; +import { screen, waitFor, within } from "@testing-library/react"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; import AddAgentForm from "./add_agent_form"; import * as networking from "@/components/networking"; import type { AgentCreateInfo } from "@/components/networking"; -import { chooseSelectOption } from "../../../../../tests/test-utils"; +import { chooseSelectOption, renderWithProviders as render } from "../../../../../tests/test-utils"; vi.mock("@/components/networking", () => ({ createAgentCall: vi.fn(), @@ -351,4 +351,33 @@ describe("AddAgentForm submit payload", () => { expect(await screen.findByText("Agent Created!")).toBeInTheDocument(); expect(within(screen.getByText("Agent Created!").parentElement!).getByText("created-agent")).toBeInTheDocument(); }); + it("blocks creation after clearing the existing key and assigns the reselected key", async () => { + vi.mocked(networking.keyListCall).mockResolvedValue({ + keys: [{ token: "key-maple", key_alias: "Maple key" }], + }); + const user = userEvent.setup(); + renderForm(); + await user.type(await screen.findByLabelText("Agent Name"), "key-selection-agent"); + await user.type(screen.getByLabelText("Display Name"), "Key selection"); + await user.type(screen.getByPlaceholderText("Describe what this agent does..."), "d"); + for (let step = 0; step < 3; step++) { + await user.click(screen.getByRole("button", { name: /^Next/ })); + } + await user.click(screen.getByRole("radio", { name: "Assign an existing key" })); + const keySelector = await screen.findByPlaceholderText("Search by key name…"); + await chooseSelectOption(user, keySelector, "Maple key"); + await user.click(screen.getByRole("button", { name: "Clear" })); + await user.click(screen.getByRole("button", { name: /Create Agent/ })); + expect(networking.createAgentCall).not.toHaveBeenCalled(); + expect(networking.keyUpdateCall).not.toHaveBeenCalled(); + await chooseSelectOption(user, keySelector, "Maple key"); + await user.click(screen.getByRole("button", { name: /Create Agent/ })); + await waitFor(() => + expect(networking.keyUpdateCall).toHaveBeenCalledWith("tok", { + key: "key-maple", + agent_id: "agent-1", + }), + ); + expect(networking.createAgentCall).toHaveBeenCalledTimes(1); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx index 108bae977e1..e71fed40209 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx @@ -338,6 +338,11 @@ const AddAgentForm: React.FC = ({ visible, onClose, accessTok return; } + if (keyAssignOption === "existing_key" && !selectedExistingKey) { + toast.error("Please select an existing key to assign"); + return; + } + setIsSubmitting(true); try { const isValid = await form.trigger(); @@ -406,12 +411,7 @@ const AddAgentForm: React.FC = ({ visible, onClose, accessTok selectedTeamId, ); setCreatedKeyValue(keyResponse.key || null); - } else if (keyAssignOption === "existing_key") { - if (!selectedExistingKey) { - toast.error("Please select an existing key to assign"); - setIsSubmitting(false); - return; - } + } else if (keyAssignOption === "existing_key" && selectedExistingKey) { await keyUpdateCall(accessToken, { key: selectedExistingKey, agent_id: agentId, @@ -963,8 +963,8 @@ const AddAgentForm: React.FC = ({ visible, onClose, accessTok setSelectedExistingKey(value || null)} + value={selectedExistingKey} + onValueChange={setSelectedExistingKey} options={existingKeys.map((k) => ({ label: k.key_alias || k.token?.slice(0, 12) + "…", value: k.token, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx index d85d26a21a8..06e332d2cfc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx @@ -170,8 +170,8 @@ interface StartFormValidityInputs { models: string[]; routerNames: string[]; direction: ShadowEvalDirection; - baselineModel: string; - judgeModel: string; + baselineModel: string | null; + judgeModel: string | null; percentage: string; maxBudget: string; } @@ -181,13 +181,13 @@ const startFormValidity = (inputs: StartFormValidityInputs) => { const percentageValid = parsedPct >= 0.1 && parsedPct <= 100; const parsedMaxBudget = Number.parseFloat(inputs.maxBudget); const maxBudgetValid = parsedMaxBudget >= 0.01 && parsedMaxBudget <= 10000; - const baselinePicked = inputs.direction === "forward" || inputs.baselineModel !== ""; + const baselinePicked = inputs.direction === "forward" || Boolean(inputs.baselineModel); const targetsPicked = inputs.apiKeyIds.length + inputs.teamIds.length + inputs.userIds.length > 0; const routerCountValid = inputs.routerNames.length >= 1 && inputs.routerNames.length <= MAX_ROUTERS; const routersMatchDirection = inputs.direction === "forward" || inputs.routerNames.length === 1; const routersValid = routerCountValid && routersMatchDirection; const scopeValid = routersValid && (inputs.direction === "reverse" || inputs.models.length <= MAX_MODELS); - const modelsPicked = scopeValid && inputs.judgeModel !== "" && baselinePicked; + const modelsPicked = scopeValid && Boolean(inputs.judgeModel) && baselinePicked; const filled = targetsPicked && modelsPicked; const boundsValid = percentageValid && maxBudgetValid; const valid = Boolean(inputs.accessToken) && filled && boundsValid; @@ -201,7 +201,7 @@ interface StartBodyInputs { models: string[]; routerNames: string[]; direction: ShadowEvalDirection; - baselineModel: string; + baselineModel: string | null; shadowPercentage: number; durationDays: number; maxBudget: number; @@ -215,7 +215,7 @@ const buildStartBody = (inputs: StartBodyInputs) => ({ models: inputs.direction === "forward" ? inputs.models : [], router_names: inputs.routerNames, direction: inputs.direction, - ...(inputs.direction === "reverse" ? { baseline_model: inputs.baselineModel } : {}), + ...(inputs.direction === "reverse" ? { baseline_model: inputs.baselineModel ?? undefined } : {}), shadow_percentage: inputs.shadowPercentage, duration_days: inputs.durationDays, max_budget: inputs.maxBudget, @@ -230,10 +230,10 @@ export const StartForm: React.FC = () => { const [models, setModels] = useState([]); const [routerNames, setRouterNames] = useState([]); const [direction, setDirection] = useState("forward"); - const [baselineModel, setBaselineModel] = useState(""); + const [baselineModel, setBaselineModel] = useState(null); const [percentage, setPercentage] = useState("10"); const [durationDays, setDurationDays] = useState("7"); - const [judgeModel, setJudgeModel] = useState(""); + const [judgeModel, setJudgeModel] = useState(null); const [maxBudget, setMaxBudget] = useState("10"); const { data: autoRouters } = useAutoRouters(); const configuredGroups = usePlainModelGroups(); @@ -286,6 +286,7 @@ export const StartForm: React.FC = () => { }; const { parsedPct, parsedMaxBudget, percentageValid, maxBudgetValid, valid } = startFormValidity(validityInputs); const handleStart = () => { + if (!valid || !judgeModel) return; const bodyInputs: StartBodyInputs = { apiKeyIds, teamIds, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.tsx index f3bd74260ad..b0612412059 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/index.tsx @@ -15,7 +15,7 @@ const generateId = () => `entry-${Date.now()}-${Math.random().toString(36).subst const createDefaultEntry = (): ModelEntry => ({ id: generateId(), - model: "", + model: null, input_tokens: 1000, output_tokens: 500, num_requests_per_day: undefined, @@ -28,7 +28,7 @@ const PricingCalculator: React.FC = ({ accessToken, mode const { debouncedFetchForEntry, removeEntry, getMultiModelResult } = useMultiCostEstimate(accessToken); const handleEntryChange = useCallback( - (id: string, field: keyof ModelEntry, value: string | number | undefined) => { + (id: string, field: keyof ModelEntry, value: string | number | null | undefined) => { setEntries((prev) => { const updated = prev.map((entry) => (entry.id === id ? { ...entry, [field]: value } : entry)); const changedEntry = updated.find((e) => e.id === id); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/types.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/types.ts index 250857a74f5..859d2f3e8d9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/types.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/pricing_calculator/types.ts @@ -4,7 +4,7 @@ export interface PricingCalculatorProps { } export interface PricingFormValues { - model: string; + model: string | null; input_tokens: number; output_tokens: number; num_requests_per_day?: number; @@ -13,7 +13,7 @@ export interface PricingFormValues { export interface ModelEntry { id: string; - model: string; + model: string | null; input_tokens: number; output_tokens: number; num_requests_per_day?: number; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx index 1de4e697f64..d45cfc3fe7d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx @@ -48,7 +48,10 @@ const GUARDRAIL_MODES = [ ] as const; const submitGuardrailSchema = z.object({ - team_id: z.string().min(1, "Select a team"), + team_id: z + .string() + .nullable() + .pipe(z.string({ error: "Select a team" }).min(1, "Select a team")), guardrail_name: z.string().min(1, "Enter a guardrail name"), mode: z.string().min(1, "Select a mode"), api_base: z.string().min(1, "Enter the API base URL").refine(isValidUrl, "Must be a valid URL"), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTable.tsx index 5c7dbb18428..8482d0832c3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTable.tsx @@ -278,7 +278,7 @@ export function AllModelsTable({ options={modelGroupOptions} value={(get(MODEL_NAME_COLUMN_ID) as string) ?? ALL_MODEL_GROUPS_VALUE} onValueChange={(value) => - set(MODEL_NAME_COLUMN_ID, value === ALL_MODEL_GROUPS_VALUE ? undefined : value) + set(MODEL_NAME_COLUMN_ID, value === ALL_MODEL_GROUPS_VALUE ? undefined : value ?? undefined) } placeholder="Filter by Public Model Name" emptyText="No models found" @@ -289,7 +289,7 @@ export function AllModelsTable({ options={accessGroupOptions} value={(get(ACCESS_GROUPS_COLUMN_ID) as string) ?? ALL_MODEL_GROUPS_VALUE} onValueChange={(value) => - set(ACCESS_GROUPS_COLUMN_ID, value === ALL_MODEL_GROUPS_VALUE ? undefined : value) + set(ACCESS_GROUPS_COLUMN_ID, value === ALL_MODEL_GROUPS_VALUE ? undefined : value ?? undefined) } placeholder="Filter by Model Access Group" emptyText="No model access groups found" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx index 1443c065d9b..1cdce04d07c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx @@ -26,7 +26,7 @@ export default function AddModelPanel() { const { data: modelCostMapData } = useModelCostMap(); const { data: credentialsResponse } = useCredentials(); const { data: teams } = useTeams(); - const [selectedProvider, setSelectedProvider] = useState(Providers.Anthropic); + const [selectedProvider, setSelectedProvider] = useState(Providers.Anthropic); const [providerModels, setProviderModels] = useState([]); const [showAdvancedSettings, setShowAdvancedSettings] = useState(false); @@ -57,7 +57,9 @@ export default function AddModelPanel() { selectedProvider={selectedProvider} setSelectedProvider={setSelectedProvider} providerModels={providerModels} - setProviderModelsFn={(provider) => setProviderModels(getProviderModels(provider, modelCostMapData))} + setProviderModelsFn={(provider) => + setProviderModels(provider === null ? [] : getProviderModels(provider, modelCostMapData)) + } getPlaceholder={getPlaceholder} showAdvancedSettings={showAdvancedSettings} setShowAdvancedSettings={setShowAdvancedSettings} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.integration.test.tsx similarity index 95% rename from ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.integration.test.tsx index bf6b092a2b0..984996351df 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.integration.test.tsx @@ -153,6 +153,7 @@ describe("ChatUI", () => { }); it("should allow the user to select a model", async () => { + const user = userEvent.setup(); render( { await waitFor(() => { expect(screen.getAllByText("Model 1").length).toBeGreaterThan(0); }); + await user.click(screen.getByRole("option", { name: "Model 1Mode: chat" })); + expect(screen.getByPlaceholderText("Select a Model")).toHaveValue("Model 1"); + + await user.click(screen.getAllByRole("button", { name: "Clear" })[0]); + const input = screen.getByPlaceholderText("Describe the image you want to generate..."); + fireEvent.change(input, { target: { value: "Contract endpoint check" } }); + expect(screen.getByRole("button", { name: "Send message" })).toBeDisabled(); + fireEvent.keyDown(input, { key: "Enter", code: "Enter" }); + expect(input).toHaveValue("Contract endpoint check"); + expect(makeOpenAIChatCompletionRequest).not.toHaveBeenCalled(); + expect(sessionStorage.getItem("endpointType")).toBeNull(); + + await selectComboboxOption("Select an endpoint", "/v1/chat/completions"); + await selectComboboxOption("Select a Model", "Model 1"); + expect(screen.getByRole("button", { name: "Send message" })).toBeEnabled(); }); it("shows only endpoint-compatible models when chat endpoint is selected", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx index eca9e2323ae..ed8679cfdc1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx @@ -196,17 +196,17 @@ const ChatUI: React.FC = ({ () => sessionStorage.getItem("customProxyBaseUrl") || "", ); const [inputMessage, setInputMessage] = useState(""); - const [selectedModel, setSelectedModel] = useState(simplified ? fixedModel : undefined); + const [selectedModel, setSelectedModel] = useState(simplified ? fixedModel : null); const [showCustomModelInput, setShowCustomModelInput] = useState(false); const [modelInfo, setModelInfo] = useState([]); const [isLoadingModels, setIsLoadingModels] = useState(false); const [modelLoadError, setModelLoadError] = useState(false); const [agentInfo, setAgentInfo] = useState([]); - const [selectedAgent, setSelectedAgent] = useState(undefined); + const [selectedAgent, setSelectedAgent] = useState(null); const debouncedSetSelectedModel = useDebouncedCallback((value: string) => setSelectedModel(value), { wait: CUSTOM_MODEL_DEBOUNCE_WAIT_MS, }); - const [endpointType, setEndpointType] = useState( + const [endpointType, setEndpointType] = useState( () => sessionStorage.getItem("endpointType") || EndpointType.CHAT, ); const [isLoading, setIsLoading] = useState(false); @@ -327,7 +327,7 @@ const ChatUI: React.FC = ({ }; useEffect(() => { - if (isGetCodeModalVisible) { + if (isGetCodeModalVisible && endpointType !== null) { const code = generateCodeSnippet({ apiKeySource, accessToken, @@ -342,7 +342,7 @@ const ChatUI: React.FC = ({ mcpServers, mcpServerToolRestrictions, endpointType, - selectedModel, + selectedModel: selectedModel ?? undefined, selectedSdk, selectedVoice, proxySettings, @@ -376,7 +376,8 @@ const ChatUI: React.FC = ({ } catch { // Storage full or unavailable — non-critical, skip persisting. } - sessionStorage.setItem("endpointType", endpointType); + if (endpointType === null) sessionStorage.removeItem("endpointType"); + else sessionStorage.setItem("endpointType", endpointType); sessionStorage.setItem("selectedTags", JSON.stringify(selectedTags)); sessionStorage.setItem("selectedVectorStores", JSON.stringify(selectedVectorStores)); sessionStorage.setItem("selectedGuardrails", JSON.stringify(selectedGuardrails)); @@ -493,7 +494,7 @@ const ChatUI: React.FC = ({ setAgentInfo(agents); // Clear selection if current agent not in list if (selectedAgent && !agents.some((a) => a.agent_name === selectedAgent)) { - setSelectedAgent(undefined); + setSelectedAgent(null); } } catch (error) { console.error("Error fetching agents:", error); @@ -616,10 +617,11 @@ const ChatUI: React.FC = ({ setUploadedAudio(file); }; - const handleEndpointChange = (value: string) => { + const handleEndpointChange = (value: string | null) => { setEndpointType(value); - setSelectedModel(undefined); - setSelectedAgent(undefined); + setGeneratedCode(""); + setSelectedModel(null); + setSelectedAgent(null); setShowCustomModelInput(false); setSelectedMCPDirectTool(undefined); if (value === EndpointType.MCP) { @@ -710,6 +712,11 @@ const ChatUI: React.FC = ({ }; const handleSendMessage = async () => { + if (endpointType === null) { + toast.fromError("Please select an endpoint before sending a request"); + return; + } + if (inputMessage.trim() === "" && endpointType !== EndpointType.TRANSCRIPTION && endpointType !== EndpointType.MCP) return; @@ -1152,7 +1159,7 @@ const ChatUI: React.FC = ({ toast.success("Chat history cleared."); }; - const onModelChange = (value: string) => { + const onModelChange = (value: string | null) => { setSelectedModel(value); setShowCustomModelInput(value === "custom"); @@ -1210,6 +1217,7 @@ const ChatUI: React.FC = ({ : "Describe the image you want to generate..."; const sendDisabled = + endpointType === null || isLoading || (endpointType === EndpointType.MCP ? !(selectedMCPServers.length === 1 && selectedMCPServers[0] !== "__all__" && selectedMCPDirectTool) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.tsx index 644bab5b5a4..6b286c403a8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.tsx @@ -3,8 +3,8 @@ import React from "react"; import { ENDPOINT_OPTIONS } from "./chatConstants"; interface EndpointSelectorProps { - endpointType: string; // Accept string to avoid type conflicts - onEndpointChange: (value: string) => void; + endpointType: string | null; + onEndpointChange: (value: string | null) => void; className?: string; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/SessionManagement.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/SessionManagement.tsx index 35c8e839159..4dbf9168897 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/SessionManagement.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/SessionManagement.tsx @@ -7,7 +7,7 @@ import { Switch } from "@/components/ui/switch"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; interface SessionManagementProps { - endpointType: string; + endpointType: string | null; responsesSessionId: string | null; useApiSessionManagement: boolean; onToggleSessionManagement: (useApi: boolean) => void; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ModelSelector.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ModelSelector.tsx index 5659875cf82..8ecc0e0c8bb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ModelSelector.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ModelSelector.tsx @@ -22,7 +22,8 @@ export function ModelSelector({ value, onChange, models, loading, disabled }: Mo const selectValue = isAddingCustom ? "__custom__" : value || undefined; - const handleSelectChange = (selected: string) => { + const handleSelectChange = (selected: string | null) => { + if (selected === null) return; if (selected === "__custom__") { setIsAddingCustom(true); if (value && !options.includes(value)) { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.integration.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.integration.test.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.tsx index 2830b0fda12..d14165dd849 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.tsx @@ -44,10 +44,10 @@ const policyShape = { .min(1, "Please enter a policy name") .regex(/^[a-zA-Z0-9_-]+$/, "Policy name can only contain letters, numbers, hyphens, and underscores"), description: z.string(), - inherit: z.string(), + inherit: z.string().nullable(), guardrails_add: z.array(z.string()), guardrails_remove: z.array(z.string()), - model_condition: z.string(), + model_condition: z.string().nullable(), }; const policySchema = z.object(policyShape); @@ -57,19 +57,19 @@ type PolicyFormValues = z.infer; const EMPTY_VALUES: PolicyFormValues = { policy_name: "", description: "", - inherit: "", + inherit: null, guardrails_add: [], guardrails_remove: [], - model_condition: "", + model_condition: null, }; const toFormValues = (policy: Policy): PolicyFormValues => ({ policy_name: policy.policy_name, description: policy.description ?? "", - inherit: policy.inherit ?? "", + inherit: policy.inherit ?? null, guardrails_add: policy.guardrails_add || [], guardrails_remove: policy.guardrails_remove || [], - model_condition: policy.condition?.model ?? "", + model_condition: policy.condition?.model ?? null, }); const buildPolicyRequest = (values: PolicyFormValues): PolicyCreateRequest | PolicyUpdateRequest => ({ @@ -529,7 +529,7 @@ const AddPolicyForm: React.FC = ({ {...control} id={id} ref={ref} - value={value} + value={value ?? ""} onChange={onChange} placeholder="Leave empty to apply to all models (e.g., gpt-4.* or bedrock/claude-.*)" /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx index 32e777949d0..c4f76b07a4e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx @@ -68,7 +68,7 @@ const AiSuggestionModal: React.FC = ({ const [suggestions, setSuggestions] = useState(null); const [explanation, setExplanation] = useState(null); const [selectedIds, setSelectedIds] = useState>(new Set()); - const [selectedModel, setSelectedModel] = useState(undefined); + const [selectedModel, setSelectedModel] = useState(null); const [availableModels, setAvailableModels] = useState([]); const [isLoadingModels, setIsLoadingModels] = useState(false); // Test panel state @@ -114,7 +114,7 @@ const AiSuggestionModal: React.FC = ({ setSuggestions(null); setExplanation(null); setSelectedIds(new Set()); - setSelectedModel(undefined); + setSelectedModel(null); setShowTestPanel(false); setTestInputText(""); setIsTestLoading(false); @@ -837,7 +837,7 @@ const AiSuggestionModal: React.FC = ({ ({ label: m, value: m }))} value={selectedModel} - onValueChange={(value) => setSelectedModel(value || undefined)} + onValueChange={setSelectedModel} placeholder={isLoadingModels ? "Loading models..." : "Select a model to analyze your requirements"} emptyText="No models found" disabled={isLoadingModels} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/pipeline_flow_builder.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/pipeline_flow_builder.tsx index 8651be39f0d..da77b348028 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/pipeline_flow_builder.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/pipeline_flow_builder.tsx @@ -343,7 +343,7 @@ const StepCard: React.FC = ({ onChange({ guardrail: value })} + onValueChange={(value) => onChange({ guardrail: value ?? undefined })} placeholder="Select a guardrail" emptyText="No guardrails found" /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/template_parameter_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/template_parameter_modal.tsx index 410d4ea8467..b90379b149d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/template_parameter_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/template_parameter_modal.tsx @@ -46,7 +46,7 @@ const TemplateParameterModal: React.FC = ({ }) => { const [parameterValues, setParameterValues] = useState>({}); const [competitorMode, setCompetitorMode] = useState<"ai" | "manual">("ai"); - const [selectedModel, setSelectedModel] = useState(undefined); + const [selectedModel, setSelectedModel] = useState(null); const [availableModels, setAvailableModels] = useState([]); const [isLoadingModels, setIsLoadingModels] = useState(false); const [competitorTags, setCompetitorTags] = useState([]); @@ -72,7 +72,7 @@ const TemplateParameterModal: React.FC = ({ }); setParameterValues(initial); setCompetitorMode("ai"); - setSelectedModel(undefined); + setSelectedModel(null); setCompetitorTags([]); setVariationsMap({}); setIsGenerating(false); @@ -297,7 +297,7 @@ const TemplateParameterModal: React.FC = ({ ({ label: m, value: m }))} value={selectedModel} - onValueChange={(value) => setSelectedModel(value || undefined)} + onValueChange={setSelectedModel} placeholder={isLoadingModels ? "Loading models..." : "Select a model to generate names"} emptyText="No models found" disabled={isLoadingModels} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx index 5c1a443d6c5..77f28b05ea5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx @@ -10,7 +10,7 @@ import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { ProjectResponse } from "@/app/(dashboard)/hooks/projects/useProjects"; import { useUpdateProject, ProjectUpdateParams } from "@/app/(dashboard)/hooks/projects/useUpdateProject"; import { ProjectBaseForm } from "./ProjectBaseForm"; -import { projectFormSchema, type ProjectFormValues } from "./projectFormSchema"; +import { projectFormSchema, type ProjectFormValues, type ProjectSubmitValues } from "./projectFormSchema"; import { buildProjectUpdateParams } from "./projectFormUtils"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; @@ -58,7 +58,7 @@ export const toFormValues = (project: ProjectResponse): ProjectFormValues => { return { project_alias: project.project_alias ?? "", - team_id: project.team_id ?? "", + team_id: project.team_id ?? null, description: project.description ?? "", models: project.models ?? [], max_budget: project.litellm_budget_table?.max_budget ?? undefined, @@ -81,7 +81,7 @@ function EditProjectForm({ project, onClose, onSuccess }: Omit { - const submitted: ProjectFormValues = advancedEverOpened + const submitted: ProjectSubmitValues = advancedEverOpened ? values : { ...values, guardrails: undefined, modelLimits: undefined, metadata: undefined }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/ProjectBaseForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/ProjectBaseForm.tsx index fed1a88fe98..1ad8a1e4953 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/ProjectBaseForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/ProjectBaseForm.tsx @@ -5,7 +5,7 @@ import { useFieldArray, useWatch, type UseFormReturn } from "react-hook-form"; import { ChevronDown, CircleAlert, Minus, Plus } from "lucide-react"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { ALL_TEAM_MODELS, type ProjectFormValues } from "./projectFormSchema"; +import { ALL_TEAM_MODELS, type ProjectFormValues, type ProjectSubmitValues } from "./projectFormSchema"; import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import { Team } from "@/components/key_team_helpers/key_list"; import { fetchTeamModels } from "@/components/organisms/create_key_button"; @@ -32,7 +32,7 @@ const toOptionalNumber = (raw: string): number | undefined => { }; interface ProjectBaseFormProps { - form: UseFormReturn; + form: UseFormReturn; advancedOpen: boolean; onAdvancedOpenChange: (open: boolean) => void; } @@ -94,7 +94,7 @@ export function ProjectBaseForm({ form, advancedOpen, onAdvancedOpenChange }: Pr } }, [selectedTeam, accessToken, userId, userRole]); - const handleTeamChange = (teamId: string) => { + const handleTeamChange = (teamId: string | null) => { const team = teams?.find((t) => t.team_id === teamId) ?? null; setSelectedTeam(team); form.setValue("models", []); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts index 6bfaa831bba..d4c85d89616 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts @@ -16,7 +16,10 @@ const modelLimitSchema = z.object({ export const projectFormSchema = z .object({ project_alias: z.string().min(1, "Please enter a project name"), - team_id: z.string().min(1, "Please select a team"), + team_id: z + .string() + .nullable() + .pipe(z.string({ error: "Please select a team" }).min(1, "Please select a team")), description: z.string().optional(), models: z.array(z.string()), max_budget: z.number().optional(), @@ -48,11 +51,12 @@ export const projectFormSchema = z }); }); -export type ProjectFormValues = z.output; +export type ProjectFormValues = z.input; +export type ProjectSubmitValues = z.output; export const emptyProjectFormValues: ProjectFormValues = { project_alias: "", - team_id: "", + team_id: null, description: undefined, models: [], max_budget: undefined, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/ModelConfigCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/ModelConfigCard.tsx index a14e45de99a..e9b526c9703 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/ModelConfigCard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/ModelConfigCard.tsx @@ -6,11 +6,11 @@ import { SettingsIcon } from "lucide-react"; import ModelSelector from "@/components/common_components/ModelSelector"; interface ModelConfigCardProps { - model: string; + model: string | null; temperature?: number; maxTokens?: number; accessToken: string | null; - onModelChange: (model: string) => void; + onModelChange: (model: string | null) => void; onTemperatureChange: (temp: number) => void; onMaxTokensChange: (tokens: number) => void; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.tsx index 04cac01365a..0f001979c28 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.tsx @@ -21,7 +21,7 @@ interface PromptEditorHeaderProps { editMode?: boolean; onShowHistory?: () => void; version?: string | null; - promptModel?: string; + promptModel?: string | null; promptVariables?: Record; accessToken: string | null; proxySettings?: { @@ -85,7 +85,7 @@ const PromptEditorHeader: React.FC = ({
{ expect(result).toContain("output:"); expect(result).toContain("format: text"); expect(result).toContain("User: Hello world"); + const cleared = convertToDotPrompt({ ...prompt, model: null }); + expect(cleared).toBe(result.replace("model: gpt-4\n", "")); }); it("should include config parameters when set", () => { @@ -203,6 +205,20 @@ describe("convertToDotPrompt", () => { }); describe("parseExistingPrompt", () => { + it("should keep saved prompts with missing or blank models unassigned", () => { + for (const modelLine of ["", "model: \n"]) { + const prompt = parseExistingPrompt({ + prompt_spec: { + prompt_id: "unassigned-prompt", + litellm_params: { dotprompt_content: `---\n${modelLine}temperature: 0\n---\nUser: Keep this message` }, + }, + }); + + expect(prompt.model).toBeNull(); + expect(convertToDotPrompt(prompt)).not.toMatch(/^model:/m); + } + }); + it("should parse basic dotprompt content", () => { const apiResponse = { prompt_spec: { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/utils.ts b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/utils.ts index ef327d93495..3d768e930a9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/utils.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/utils.ts @@ -23,7 +23,7 @@ export const extractVariables = (prompt: PromptType): string[] => { export const convertToDotPrompt = (prompt: PromptType): string => { const variables = extractVariables(prompt); - let result = `---\nmodel: ${prompt.model}\n`; + let result = prompt.model ? `---\nmodel: ${prompt.model}\n` : "---\n"; // Add temperature if set if (prompt.config.temperature !== undefined) { @@ -237,7 +237,7 @@ export const parseExistingPrompt = (apiResponse: any): PromptType => { return { name: baseName, - model: parsedFrontmatter.model || "gpt-4o", + model: parsedFrontmatter.model || null, config: parsedFrontmatter.config, tools: parsedFrontmatter.tools, developerMessage: parsedBody.developerMessage, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/DefaultUserSettingsForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/DefaultUserSettingsForm.tsx index 269f3b0af39..0239844851f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/DefaultUserSettingsForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/DefaultUserSettingsForm.tsx @@ -20,7 +20,12 @@ import { useZodForm } from "@/lib/forms/useZodForm"; import { fetchClient } from "@/lib/http/api"; import { buildBody, settingsToForm, type DefaultInternalUserParams, type InternalUserSettings } from "./mapper"; -import { defaultUserSettingsSchema, EMPTY_TEAM_ROW, type DefaultUserSettingsFormValues } from "./schema"; +import { + defaultUserSettingsSchema, + EMPTY_TEAM_ROW, + type DefaultUserSettingsFormValues, + type DefaultUserSettingsSubmitValues, +} from "./schema"; const NO_RESET = "never"; @@ -63,7 +68,7 @@ interface RoleOption { description: string; } -type SettingsControl = Control; +type SettingsControl = Control; const TeamPickerField = ({ control, index }: { control: SettingsControl; index: number }) => { const [search, setSearch] = React.useState(""); @@ -233,7 +238,7 @@ const SettingsForm = ({ initialValues, roleOptions, updateSettings, onCancel, on const { isDirty } = form.formState; const mutation = useMutation({ - mutationFn: (values: DefaultUserSettingsFormValues) => updateSettings(buildBody(values)), + mutationFn: (values: DefaultUserSettingsSubmitValues) => updateSettings(buildBody(values)), onSuccess: (_result, values) => { toast.success("Default user settings updated successfully"); queryClient.invalidateQueries({ queryKey: SETTINGS_QUERY_KEY }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.test.ts index e8b350332c8..f4a7f001480 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.test.ts @@ -1,7 +1,7 @@ import { describe, expect, it } from "vitest"; import { buildBody, settingsToForm } from "./mapper"; -import type { DefaultUserSettingsFormValues } from "./schema"; +import type { DefaultUserSettingsSubmitValues } from "./schema"; const CONFIGURED_SETTINGS = { user_role: "internal_user", @@ -59,13 +59,13 @@ describe("settingsToForm", () => { it("degrades an unrecognisable team entry to a blank row instead of throwing", () => { expect(settingsToForm({ teams: [{ max_budget_in_team: 5 }, 7] }).teams).toStrictEqual([ - { team_id: "", max_budget_in_team: "", user_role: "user" }, - { team_id: "", max_budget_in_team: "", user_role: "user" }, + { team_id: null, max_budget_in_team: "", user_role: "user" }, + { team_id: null, max_budget_in_team: "", user_role: "user" }, ]); }); }); -const formValues = (overrides: Partial = {}): DefaultUserSettingsFormValues => ({ +const formValues = (overrides: Partial = {}): DefaultUserSettingsSubmitValues => ({ user_role: "internal_user", max_budget: "100", budget_duration: "30d", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.ts b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.ts index 50365081afb..ac528542ae2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/mapper.ts @@ -2,7 +2,12 @@ import { z } from "zod/v4"; import type { components } from "@/lib/http/schema"; -import { EMPTY_TEAM_ROW, type DefaultTeamRowValues, type DefaultUserSettingsFormValues } from "./schema"; +import { + EMPTY_TEAM_ROW, + type DefaultTeamRowValues, + type DefaultUserSettingsFormValues, + type DefaultUserSettingsSubmitValues, +} from "./schema"; export type InternalUserSettings = components["schemas"]["InternalUserSettingsResponse"]; export type DefaultInternalUserParams = components["schemas"]["DefaultInternalUserParams"]; @@ -60,13 +65,13 @@ const textOrNull = (raw: string): string | null => (raw.trim() === "" ? null : r const listOrNull = (items: readonly T[]): T[] | null => (items.length === 0 ? null : [...items]); -const toTeamBody = (team: DefaultTeamRowValues): DefaultTeamBody => ({ +const toTeamBody = (team: DefaultUserSettingsSubmitValues["teams"][number]): DefaultTeamBody => ({ team_id: team.team_id, max_budget_in_team: numberOrNull(team.max_budget_in_team), user_role: team.user_role, }); -export const buildBody = (values: DefaultUserSettingsFormValues): DefaultInternalUserParams => ({ +export const buildBody = (values: DefaultUserSettingsSubmitValues): DefaultInternalUserParams => ({ user_role: asDefaultUserRole(values.user_role), max_budget: numberOrNull(values.max_budget), budget_duration: textOrNull(values.budget_duration), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/schema.ts b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/schema.ts index 7309e3745da..cd66c3addc8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/schema.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/default-user-settings/schema.ts @@ -10,14 +10,17 @@ const amountOrEmpty = z ); const defaultTeamRowSchema = z.object({ - team_id: z.string().min(1, "Select a team"), + team_id: z + .string() + .nullable() + .pipe(z.string({ error: "Select a team" }).min(1, "Select a team")), max_budget_in_team: amountOrEmpty, user_role: z.enum(["user", "admin"]), }); -export type DefaultTeamRowValues = z.output; +export type DefaultTeamRowValues = z.input; -export const EMPTY_TEAM_ROW: DefaultTeamRowValues = { team_id: "", max_budget_in_team: "", user_role: "user" }; +export const EMPTY_TEAM_ROW: DefaultTeamRowValues = { team_id: null, max_budget_in_team: "", user_role: "user" }; const defaultUserSettingsShape = { user_role: z.string(), @@ -41,4 +44,5 @@ export const defaultUserSettingsSchema = z.object(defaultUserSettingsShape).supe ); }); -export type DefaultUserSettingsFormValues = z.output; +export type DefaultUserSettingsFormValues = z.input; +export type DefaultUserSettingsSubmitValues = z.output; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.tsx index 875df74652c..8fa76a64af6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.tsx @@ -192,7 +192,7 @@ export function UsersTable({ set("user_role", value)} + onValueChange={(value) => set("user_role", value ?? undefined)} placeholder="Select a role…" emptyText="No roles found" /> @@ -201,7 +201,7 @@ export function UsersTable({ set("team", value)} + onValueChange={(value) => set("team", value ?? undefined)} placeholder="Select a team…" emptyText="No teams found" /> diff --git a/ui/litellm-dashboard/src/components/CreateUserButton.tsx b/ui/litellm-dashboard/src/components/CreateUserButton.tsx index 0f7c356b8cc..83a72e50f5c 100644 --- a/ui/litellm-dashboard/src/components/CreateUserButton.tsx +++ b/ui/litellm-dashboard/src/components/CreateUserButton.tsx @@ -59,7 +59,7 @@ interface UISettings { interface CreateUserFormValues { user_email?: string; user_role: string; - team_id?: string; + team_id?: string | null; organization_ids?: string[]; metadata?: string; send_invite_email: boolean; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx index 390684d3739..f78790dea88 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx @@ -115,7 +115,7 @@ export default function MCPSemanticFilterSettings({ accessToken }: MCPSemanticFi // Test section state const [testQuery, setTestQuery] = useState(""); - const [testModel, setTestModel] = useState("gpt-4o"); + const [testModel, setTestModel] = useState("gpt-4o"); const [testResult, setTestResult] = useState(null); const [testError, setTestError] = useState(null); const [isTesting, setIsTesting] = useState(false); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterTestPanel.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterTestPanel.tsx index 5dc661391ac..82c699d9983 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterTestPanel.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterTestPanel.tsx @@ -11,8 +11,8 @@ interface MCPSemanticFilterTestPanelProps { accessToken: string | null; testQuery: string; setTestQuery: (value: string) => void; - testModel: string; - setTestModel: (value: string) => void; + testModel: string | null; + setTestModel: (value: string | null) => void; isTesting: boolean; onTest: () => void; filterEnabled: boolean; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/semanticFilterTestUtils.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/semanticFilterTestUtils.ts index d3e585e8052..c41b081da88 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/semanticFilterTestUtils.ts +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/semanticFilterTestUtils.ts @@ -32,7 +32,7 @@ export const runSemanticFilterTest = async ({ setTestError, }: { accessToken: string; - testModel: string; + testModel: string | null; testQuery: string; setIsTesting: (value: boolean) => void; setTestResult: (result: TestResult | null) => void; @@ -68,12 +68,12 @@ export const runSemanticFilterTest = async ({ } }; -export const getCurlCommand = (testModel: string, testQuery: string) => +export const getCurlCommand = (testModel: string | null, testQuery: string) => `curl --location 'http://localhost:4000/v1/responses' \\ --header 'Content-Type: application/json' \\ --header 'Authorization: Bearer sk-1234' \\ --data '{ - "model": "${testModel}", + "model": "${testModel ?? "YOUR_MODEL"}", "input": [ { "role": "user", diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/toolSearchForm.test.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/toolSearchForm.test.ts index 2f9032537c1..6d19c17c577 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/toolSearchForm.test.ts +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/toolSearchForm.test.ts @@ -51,6 +51,7 @@ describe("parseCoreTools", () => { describe("formToPayload", () => { it("sends null for a cleared embedding model so the proxy returns to keyword matching", () => { expect(formToPayload({ ...DEFAULT_FORM_VALUES, embedding_model: " " })).toEqual(KEYWORD_PAYLOAD); + expect(formToPayload({ ...DEFAULT_FORM_VALUES, embedding_model: null })).toEqual(KEYWORD_PAYLOAD); }); it("clamps top_k into the range the proxy accepts and lists core tools", () => { diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/toolSearchForm.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/toolSearchForm.ts index f10bfd249d7..d3535fe91a3 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/toolSearchForm.ts +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/toolSearchForm.ts @@ -1,7 +1,7 @@ import type { MCPToolSearchSettings } from "@/app/(dashboard)/hooks/mcpToolSearchSettings/useMCPToolSearchSettings"; export interface ToolSearchFormValues { - embedding_model: string; + embedding_model: string | null; top_k: number; similarity_threshold: number; core_tools_text: string; @@ -11,7 +11,7 @@ export const TOP_K_MIN = 1; export const TOP_K_MAX = 100; export const DEFAULT_FORM_VALUES: ToolSearchFormValues = { - embedding_model: "", + embedding_model: null, top_k: 5, similarity_threshold: 0, core_tools_text: "", @@ -42,7 +42,7 @@ export const storedValuesToForm = (values: Record): ToolSearchF }); export const formToPayload = (form: ToolSearchFormValues): MCPToolSearchSettings => ({ - embedding_model: form.embedding_model.trim() === "" ? null : form.embedding_model.trim(), + embedding_model: form.embedding_model?.trim() || null, top_k: clampTopK(form.top_k), similarity_threshold: form.similarity_threshold, core_tools: parseCoreTools(form.core_tools_text), diff --git a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackGroupConfig.tsx b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackGroupConfig.tsx index a82e0fdfa88..3897a9d9fa2 100644 --- a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackGroupConfig.tsx +++ b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackGroupConfig.tsx @@ -32,12 +32,8 @@ export function FallbackGroupConfig({ // Filter available options for fallbacks (exclude primary only, allow already selected to be shown for deselection) const availableFallbackOptions = availableModels.filter((m) => m !== group.primaryModel); - const handlePrimaryChange = (value: string) => { - let newFallbacks = [...group.fallbackModels]; - // Remove from fallbacks if it was there - if (newFallbacks.includes(value)) { - newFallbacks = newFallbacks.filter((m) => m !== value); - } + const handlePrimaryChange = (value: string | null) => { + const newFallbacks = group.fallbackModels.filter((model) => model !== value); onChange({ ...group, primaryModel: value, @@ -76,7 +72,7 @@ export function FallbackGroupConfig({ ({ label: m, value: m }))} - value={group.primaryModel ?? ""} + value={group.primaryModel} onValueChange={handlePrimaryChange} placeholder="Select primary model" emptyText="No models found" diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index d2d18139893..f0d21cc5350 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -328,11 +328,11 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser }; const selectCreateTeamOrganization = ( - next: string, + next: string | null, currentOrganizationId: string | null, onChange: (organizationId: string | null) => void, ) => { - const nextOrganizationId = next === "" ? null : next; + const nextOrganizationId = next; if (nextOrganizationId === currentOrganizationId) return; onChange(nextOrganizationId); form.setValue("models", []); diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx index 3fb19e522f3..bde1c7724c3 100644 --- a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx @@ -203,7 +203,7 @@ export function TeamsTable({ userRole, userID, onSelectTeam, onEditTeam, onDelet set("org_id", value)} + onValueChange={(value) => set("org_id", value ?? undefined)} placeholder="Select an organization…" emptyText="No organizations found" /> diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index f28ae27b6bc..d0367cb4d5b 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -311,7 +311,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { set("team_id", value)} + onValueChange={(value) => set("team_id", value ?? undefined)} placeholder="Select a team…" emptyText="No teams found" /> @@ -320,7 +320,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { set("org_id", value)} + onValueChange={(value) => set("org_id", value ?? undefined)} placeholder="Select an organization…" emptyText="No organizations found" /> diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx rename to ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx index 3b879da2d68..536bc231008 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx @@ -164,7 +164,7 @@ const createTestProps = (userRole = "proxy_admin", userId = "user-1", isTeamAdmi handleOk: vi.fn().mockResolvedValue(true), setSelectedProvider: vi.fn(), setProviderModelsFn: vi.fn(), - getPlaceholder: vi.fn((provider: Providers) => `Enter ${provider} model name`), + getPlaceholder: vi.fn((provider: string) => `Enter ${provider} model name`), setShowAdvancedSettings: vi.fn(), selectedProvider: Providers.OpenAI, providerModels: ["gpt-4", "gpt-3.5-turbo"], diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx index 92951b68fd5..30b4911da3e 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx @@ -25,7 +25,6 @@ import { } from "../common_components/MountedFormField"; import type { Team } from "../key_team_helpers/key_list"; import { type CredentialItem, type ProviderCreateInfo, modelAvailableCall } from "../networking"; -import { Providers } from "../provider_info_helpers"; import { ProviderLogo } from "../molecules/models/ProviderLogo"; import AccessGroupTagsCombobox from "./AccessGroupTagsCombobox"; import AdvancedSettings from "./advanced_settings"; @@ -42,11 +41,11 @@ interface AddModelFormProps { registry: MountRegistry; mountedValues: () => MountedFormValues; handleOk: () => Promise; - selectedProvider: Providers; - setSelectedProvider: (provider: Providers) => void; + selectedProvider: string | null; + setSelectedProvider: (provider: string | null) => void; providerModels: string[]; - setProviderModelsFn: (provider: Providers) => void; - getPlaceholder: (provider: Providers) => string; + setProviderModelsFn: (provider: string | null) => void; + getPlaceholder: (provider: string) => string; showAdvancedSettings: boolean; setShowAdvancedSettings: (show: boolean) => void; teams: Team[] | null; @@ -140,7 +139,7 @@ const AddModelForm: React.FC = ({ [credentials], ); - const applyProviderSelection = (provider: Providers) => { + const applyProviderSelection = (provider: string | null) => { setSelectedProvider(provider); setProviderModelsFn(provider); form.setValue("model", []); @@ -227,10 +226,10 @@ const AddModelForm: React.FC = ({ options={providerOptions} emptyText={providerMetadataErrorText ?? "No providers found"} placeholder={isProviderMetadataLoading ? "Loading providers..." : "Select a provider"} - value={(control.value as string | undefined) ?? ""} + value={typeof control.value === "string" ? control.value : null} onValueChange={(value) => { control.onChange(value); - applyProviderSelection(value as Providers); + applyProviderSelection(value); }} /> )} diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index 594c6d022ce..93ccf561387 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -328,7 +328,8 @@ const ClassificationMethodConfig: React.FC = ({ onChange(nextValue); }; - const handleClassifierModelChange = (model: string) => { + const handleClassifierModelChange = (model: string | null) => { + if (model === null) return; if (model === value.classifier_llm_config?.model) return; const { reasoning_effort: _reasoningEffort, ...classifierLlmConfig } = value.classifier_llm_config ?? { model: "", diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 97d6739aa22..febcde269f7 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -632,7 +632,7 @@ const ComplexityRouterConfig: React.FC = ({ // Clearing the select drops the key entirely rather than storing "", so an emptied pin reads as // "track the tiers" everywhere downstream instead of as a blank model name. - const handleDefaultModelChange = (model: string | undefined) => { + const handleDefaultModelChange = (model: string | null | undefined) => { onChange({ ...value, default_model: model || undefined }); }; diff --git a/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx index c1817918f60..e67583febf3 100644 --- a/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx @@ -42,8 +42,8 @@ const CompressionControls: React.FC = ({ value, onChan
onRoutingChange(value === "" ? undefined : value)} + value={routing} + onValueChange={(value) => onRoutingChange(value ?? undefined)} placeholder="Inherit from the request's own compression guardrails" emptyText="No compression guardrails found" aria-label="Routing decision compression" @@ -74,8 +74,8 @@ const CompressionControls: React.FC = ({ value, onChan
onModelChange(value === "" ? undefined : value)} + value={model} + onValueChange={(value) => onModelChange(value ?? undefined)} placeholder="None (no compression)" emptyText="No compression guardrails found" aria-label="Model call compression" diff --git a/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.test.tsx b/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.integration.test.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.test.tsx rename to ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.integration.test.tsx index ecd5abfa451..495fd207266 100644 --- a/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.integration.test.tsx @@ -47,7 +47,7 @@ describe("RouterConfigBuilder", () => { expect(onChange).toHaveBeenCalledWith({ routes: [ expect.objectContaining({ - name: "", + name: null, utterances: [], description: "", score_threshold: 0.5, diff --git a/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.test.ts b/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.test.ts new file mode 100644 index 00000000000..d681c911c46 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.test.ts @@ -0,0 +1,10 @@ +import { describe, expect, it } from "vitest"; +import { serializeRouterConfig } from "./RouterConfigBuilder"; + +describe("serializeRouterConfig", () => { + it("rejects a cleared route model and preserves selected route settings", () => { + expect(() => serializeRouterConfig({ routes: [{ name: null }] })).toThrow("Please select a model for every route"); + const config = { routes: [{ name: "model-silver", utterances: [], description: "", score_threshold: 0 }] }; + expect(JSON.parse(serializeRouterConfig(config))).toEqual(config); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.tsx b/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.tsx index 6dfb9f32091..9713d18872b 100644 --- a/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.tsx +++ b/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.tsx @@ -15,7 +15,7 @@ import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/comp interface Route { id: string; - model: string; + model: string | null; utterances: string[]; description: string; score_threshold: number; @@ -23,21 +23,28 @@ interface Route { interface SavedRoute { id?: string; - name?: string; - model?: string; + name?: string | null; + model?: string | null; utterances?: string[]; description?: string; score_threshold?: number; } -interface RouterConfig { +export interface RouterConfig { routes?: SavedRoute[]; } +export function serializeRouterConfig(config: RouterConfig | null): string { + if (config?.routes?.some((route) => !(route.name ?? route.model))) { + throw new Error("Please select a model for every route"); + } + return JSON.stringify(config); +} + interface RouterConfigBuilderProps { modelInfo: ModelGroup[]; - value?: RouterConfig; - onChange?: (config: any) => void; + value?: RouterConfig | null; + onChange?: (config: RouterConfig) => void; } interface UtteranceInputProps { @@ -136,7 +143,7 @@ const RouterConfigBuilder: React.FC = ({ modelInfo, va routeIds.push(id); return { id, - model: route.name || route.model || "", + model: route.name || route.model || null, utterances: route.utterances || [], description: route.description || "", score_threshold: route.score_threshold ?? 0.5, @@ -165,7 +172,7 @@ const RouterConfigBuilder: React.FC = ({ modelInfo, va const newRouteId = `route-${Date.now()}`; const updatedRoutes = [ ...routes, - { id: newRouteId, model: "", utterances: [], description: "", score_threshold: 0.5 }, + { id: newRouteId, model: null, utterances: [], description: "", score_threshold: 0.5 }, ]; setRoutes(updatedRoutes); updateConfig(updatedRoutes); diff --git a/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.tsx b/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.tsx index 5684d32a319..0a5a9a5bfb6 100644 --- a/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.tsx +++ b/ui/litellm-dashboard/src/components/add_model/SemanticKeywordMatching.tsx @@ -61,7 +61,9 @@ const SemanticKeywordMatching: React.FC = ({ { + if (model !== null) onEmbeddingModelChange(model); + }} placeholder="Select an embedding model" emptyText="No embedding models found" aria-label="Embedding model" diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index d7c4a671317..632f4427a82 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -150,7 +150,10 @@ export const getSubmitBlockedReason = ( const autoRouterSchema = (requiresTeamScope: boolean) => z.object({ auto_router_name: z.string().min(1, "Auto router name is required"), - team_id: requiresTeamScope ? z.string().min(1, "Please select a team to continue") : z.string(), + team_id: z + .string() + .nullable() + .refine((teamId) => !requiresTeamScope || Boolean(teamId), "Please select a team to continue"), model_access_group: z.array(z.string()).optional(), }); @@ -158,12 +161,12 @@ type AddAutoRouterFormValues = z.infer>; const EMPTY_FORM_VALUES: AddAutoRouterFormValues = { auto_router_name: "", - team_id: "", + team_id: null, model_access_group: undefined, }; -const teamScopePayload = (requiresTeamScope: boolean, teamId: string): { team_id?: string } => - requiresTeamScope ? { team_id: teamId } : {}; +const teamScopePayload = (requiresTeamScope: boolean, teamId: string | null): { team_id?: string } => + requiresTeamScope && teamId ? { team_id: teamId } : {}; const BlockedReasonTooltip: React.FC<{ reason: string | null; children: React.ReactElement }> = ({ reason, @@ -458,7 +461,7 @@ const AddAutoRouterTab: React.FC = ({ const serverVerdict = await validateAutoRouterConfig( accessToken, complexityRouterConfigPayload as unknown as Record, - requiresTeamScope ? form.getValues("team_id") : undefined, + requiresTeamScope ? form.getValues("team_id") ?? undefined : undefined, ); const dryRunError = dryRunRejection(serverVerdict); if (dryRunError) { @@ -630,9 +633,7 @@ const AddAutoRouterTab: React.FC = ({ "Select the team this auto router belongs to. Only keys for this team will be able to call it.", )} > - {({ id, value, onChange }) => ( - onChange(next ?? "")} /> - )} + {({ id, value, onChange }) => } )} @@ -776,7 +777,7 @@ const AddAutoRouterTab: React.FC = ({ config={buildComplexityRouterConfig(complexityRouterConfigParams)} defaultModel={resolveComplexityDefaultModel(complexityRouterConfig, complexityRouterConfig.default_model)} routerName={watchedName} - teamId={requiresTeamScope ? watchedTeamId : undefined} + teamId={requiresTeamScope ? watchedTeamId ?? undefined : undefined} /> )} diff --git a/ui/litellm-dashboard/src/components/add_model/litellm_model_name.tsx b/ui/litellm-dashboard/src/components/add_model/litellm_model_name.tsx index bfa2df8f54d..8cf4b077f39 100644 --- a/ui/litellm-dashboard/src/components/add_model/litellm_model_name.tsx +++ b/ui/litellm-dashboard/src/components/add_model/litellm_model_name.tsx @@ -8,9 +8,9 @@ import { MountedFormField, type MountedFormValues } from "../common_components/M import { Providers } from "../provider_info_helpers"; interface LiteLLMModelNameFieldProps { - selectedProvider: Providers; + selectedProvider: string | null; providerModels: string[]; - getPlaceholder: (provider: Providers) => string; + getPlaceholder: (provider: string) => string; } const LiteLLMModelNameField: React.FC = ({ @@ -123,7 +123,7 @@ const LiteLLMModelNameField: React.FC = ({ id={control.id} value={(control.value as string | undefined) ?? ""} onBlur={control.onBlur} - placeholder={getPlaceholder(selectedProvider)} + placeholder={selectedProvider === null ? "Select a provider first" : getPlaceholder(selectedProvider)} onChange={(event) => { control.onChange(event); if (selectedProvider === Providers.Azure) { @@ -147,7 +147,7 @@ const LiteLLMModelNameField: React.FC = ({ value: "custom", }, { - label: `All ${selectedProvider} Models (Wildcard)`, + label: `All ${selectedProvider ?? "provider"} Models (Wildcard)`, value: "all-wildcard", }, ...providerModels.map((model) => ({ @@ -163,7 +163,7 @@ const LiteLLMModelNameField: React.FC = ({ value={(control.value as string | undefined) ?? ""} onChange={control.onChange} onBlur={control.onBlur} - placeholder={getPlaceholder(selectedProvider)} + placeholder={selectedProvider === null ? "Select a provider first" : getPlaceholder(selectedProvider)} /> ) } diff --git a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx index e1a47e9792d..4add7b918f7 100644 --- a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx +++ b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx @@ -18,7 +18,7 @@ import { provider_map, Providers } from "../provider_info_helpers"; import { labelWithHint } from "@/components/shared/form/LabelWithHint"; interface ProviderSpecificFieldsProps { - selectedProvider: Providers; + selectedProvider: string | null; } const readTextFile = (file: File, onLoaded: (contents: string) => void) => { @@ -168,6 +168,7 @@ const ProviderSpecificFields: React.FC = ({ selecte }, [cacheEntries]); const allFields = React.useMemo(() => { + if (selectedProvider === null) return []; // First try to resolve from the in-memory cache. We support both the // enum/display-name form and the raw provider slug (e.g. "petals"). const cachedFields = diff --git a/ui/litellm-dashboard/src/components/common_components/ModelAliasManager.tsx b/ui/litellm-dashboard/src/components/common_components/ModelAliasManager.tsx index c6e24d3b470..384fa44cfde 100644 --- a/ui/litellm-dashboard/src/components/common_components/ModelAliasManager.tsx +++ b/ui/litellm-dashboard/src/components/common_components/ModelAliasManager.tsx @@ -27,8 +27,13 @@ const ModelAliasManager: React.FC = ({ showExampleConfig = true, }) => { const [aliases, setAliases] = useState([]); - const [newAlias, setNewAlias] = useState({ aliasName: "", targetModel: "" }); - const [editingAlias, setEditingAlias] = useState(null); + const [newAlias, setNewAlias] = useState<{ aliasName: string; targetModel: string | null }>({ + aliasName: "", + targetModel: null, + }); + const [editingAlias, setEditingAlias] = useState< + (Omit & { targetModel: string | null }) | null + >(null); const aliasNameId = useId(); useEffect(() => { @@ -61,7 +66,7 @@ const ModelAliasManager: React.FC = ({ const updatedAliases = [...aliases, newAliasObj]; setAliases(updatedAliases); - setNewAlias({ aliasName: "", targetModel: "" }); + setNewAlias({ aliasName: "", targetModel: null }); // Convert array back to object format and notify parent const aliasObject: { [key: string]: string } = {}; @@ -94,7 +99,8 @@ const ModelAliasManager: React.FC = ({ return; } - const updatedAliases = aliases.map((alias) => (alias.id === editingAlias.id ? editingAlias : alias)); + const savedAlias: AliasItem = { ...editingAlias, targetModel: editingAlias.targetModel }; + const updatedAliases = aliases.map((alias) => (alias.id === savedAlias.id ? savedAlias : alias)); setAliases(updatedAliases); setEditingAlias(null); diff --git a/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx b/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx index 6a6807de5c0..15b6b1644b6 100644 --- a/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx +++ b/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx @@ -9,9 +9,9 @@ const MODEL_SELECT_DEBOUNCE_MS = 500; interface ModelSelectorProps { accessToken: string; - value?: string; + value?: string | null; placeholder?: string; - onChange?: (value: string) => void; + onChange?: (value: string | null) => void; disabled?: boolean; style?: React.CSSProperties; className?: string; @@ -30,12 +30,12 @@ const ModelSelector: React.FC = ({ showLabel = true, labelText = "Select Model", }) => { - const [selectedModel, setSelectedModel] = useState(value); + const [selectedModel, setSelectedModel] = useState(value ?? null); const [showCustomModelInput, setShowCustomModelInput] = useState(false); const [modelInfo, setModelInfo] = useState([]); useEffect(() => { - setSelectedModel(value); + setSelectedModel(value ?? null); }, [value]); useEffect(() => { @@ -56,13 +56,13 @@ const ModelSelector: React.FC = ({ loadModels(); }, [accessToken]); - const onModelChange = (value: string) => { + const onModelChange = (value: string | null) => { if (value === "custom") { setShowCustomModelInput(true); - setSelectedModel(undefined); + setSelectedModel(null); } else { setShowCustomModelInput(false); - setSelectedModel(value); + setSelectedModel(value ?? null); if (onChange) { onChange(value); } @@ -71,7 +71,7 @@ const ModelSelector: React.FC = ({ const debouncedSelect = useDebouncedCallback( (value: string) => { - setSelectedModel(value); + setSelectedModel(value ?? null); onChange?.(value); }, { wait: MODEL_SELECT_DEBOUNCE_MS }, diff --git a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx index 663028b2d92..f2f1504fed5 100644 --- a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx @@ -4,7 +4,7 @@ import { Organization } from "../networking"; interface OrganizationDropdownProps { organizations?: Organization[] | null; - value?: string; + value?: string | null; onChange?: (value: string | null) => void; disabled?: boolean; loading?: boolean; @@ -32,7 +32,7 @@ const OrganizationDropdown: React.FC = ({ sublabel: org.organization_id, }))} value={value} - onValueChange={(organizationId) => onChange?.(organizationId || null)} + onValueChange={(organizationId) => onChange?.(organizationId)} placeholder={placeholder} emptyText={loading ? "Loading organizations…" : "No organizations found"} disabled={disabled} diff --git a/ui/litellm-dashboard/src/components/common_components/ProjectDropdown.tsx b/ui/litellm-dashboard/src/components/common_components/ProjectDropdown.tsx index 0b5a84f9589..f27e1a683dd 100644 --- a/ui/litellm-dashboard/src/components/common_components/ProjectDropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/ProjectDropdown.tsx @@ -4,8 +4,8 @@ import { ProjectResponse } from "@/app/(dashboard)/hooks/projects/useProjects"; interface ProjectDropdownProps { projects?: ProjectResponse[] | null; - value?: string; - onChange?: (value: string) => void; + value?: string | null; + onChange?: (value: string | null) => void; disabled?: boolean; loading?: boolean; /** When set, only show projects belonging to this team */ diff --git a/ui/litellm-dashboard/src/components/common_components/UserDropdown.tsx b/ui/litellm-dashboard/src/components/common_components/UserDropdown.tsx index eef1a8233bc..c80f654bbd1 100644 --- a/ui/litellm-dashboard/src/components/common_components/UserDropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/UserDropdown.tsx @@ -47,8 +47,8 @@ const UserDropdown: React.FC = ({ value, onChange, disabled,
onChange(next === "" ? null : next)} + value={value} + onValueChange={onChange} onSearchChange={setSearch} onLoadMore={fetchNextPage} hasNextPage={hasNextPage} diff --git a/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx b/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx index 7d385c2a3f7..e10a8d1e562 100644 --- a/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx @@ -4,7 +4,7 @@ import { useInfiniteTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import { Team } from "../key_team_helpers/key_list"; interface TeamDropdownProps { - value?: string; + value?: string | null; onChange?: (value: string | null) => void; /** Callback with the full Team object (or null on clear). */ onTeamSelect?: (team: Team | null) => void; @@ -46,8 +46,8 @@ const TeamDropdown: React.FC = ({ return result; }, [data]); - const handleChange = (teamId: string) => { - onChange?.(teamId || null); + const handleChange = (teamId: string | null) => { + onChange?.(teamId); if (onTeamSelect) { onTeamSelect(teamId ? teams.find((t) => t.team_id === teamId) ?? null : null); } @@ -61,7 +61,7 @@ const TeamDropdown: React.FC = ({ value: team.team_id, sublabel: team.team_id, }))} - value={value || undefined} + value={value} onValueChange={handleChange} onSearchChange={setSearch} onLoadMore={fetchNextPage} diff --git a/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx b/ui/litellm-dashboard/src/components/common_components/user_search_modal.integration.test.tsx similarity index 94% rename from ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx rename to ui/litellm-dashboard/src/components/common_components/user_search_modal.integration.test.tsx index 7ca2b530235..5769d8c2f88 100644 --- a/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/user_search_modal.integration.test.tsx @@ -1,4 +1,5 @@ -import { act, fireEvent, render, screen, waitFor, within } from "@testing-library/react"; +import { act, fireEvent, screen, waitFor, within } from "@testing-library/react"; +import { renderWithProviders as render } from "../../../tests/test-utils"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import UserSearchModal from "./user_search_modal"; @@ -101,15 +102,18 @@ describe("UserSearchModal submit payload", () => { await user.click(await screen.findByRole("option", { name: "picked@example.com" })); }; - it("submits every registered field, with the untouched identity fields undefined", async () => { + it("should block submission without a selected user and after clearing the paired identity", async () => { const { user, onSubmit } = setup(); + expect(save()).toBeDisabled(); + await searchByEmail(user, "pick"); + await user.click(screen.getAllByRole("button", { name: "Clear" })[0]); + expect(screen.getByLabelText("Email")).toHaveValue(""); + expect(screen.getByLabelText("User ID")).toHaveValue(""); + expect(save()).toBeDisabled(); await user.click(save()); - await waitFor(() => expect(onSubmit).toHaveBeenCalledTimes(1)); - const values = onSubmit.mock.calls[0][0]; - expect(Object.keys(values).sort()).toEqual(["role", "user_email", "user_id"]); - expect(values).toStrictEqual({ user_email: undefined, user_id: undefined, role: "user" }); + expect(onSubmit).not.toHaveBeenCalled(); }); it("carries the picked user's email and id into the payload", async () => { @@ -130,6 +134,7 @@ describe("UserSearchModal submit payload", () => { const { onSubmit } = setup(); const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); + await searchByEmail(user, "pick"); await user.click(screen.getByLabelText("Member Role")); await user.click(await screen.findByRole("option", { name: /^admin/ })); await user.click(save()); @@ -184,6 +189,7 @@ describe("UserSearchModal submit payload", () => { it("does not submit on Enter in any field, while the button still does", async () => { const { user, onSubmit } = setup(); + await searchByEmail(user, "pick"); await user.click(getEmailSearchInput()); await user.keyboard("{Enter}"); await user.click(screen.getByLabelText("User ID")); diff --git a/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx b/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx index 20c45cd0d2e..04ac412397f 100644 --- a/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx +++ b/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx @@ -31,8 +31,8 @@ interface Role { } interface FormValues { - user_email: string | undefined; - user_id: string | undefined; + user_email: string | null | undefined; + user_id: string | null | undefined; role: string; } @@ -66,6 +66,8 @@ const UserSearchModal: React.FC = ({ }) => { const emptyValues: FormValues = { user_email: undefined, user_id: undefined, role: defaultRole }; const form = useForm({ defaultValues: emptyValues }); + const selectedUserId = form.watch("user_id"); + const selectedUserEmail = form.watch("user_email"); const [userOptions, setUserOptions] = useState([]); const [loading, setLoading] = useState(false); const [selectedField, setSelectedField] = useState<"user_email" | "user_id">("user_email"); @@ -143,19 +145,29 @@ const UserSearchModal: React.FC = ({ const renderUserSearch = ( fieldName: "user_email" | "user_id", placeholder: string, - controlProps: { id: string; value: string | undefined; onChange: (value: string | undefined) => void }, + controlProps: { + id: string; + value: string | null | undefined; + onChange: (value: string | null | undefined) => void; + }, testId?: string, ) => { const items = selectedField === fieldName ? userOptions : []; + const handleValueChange = (value: string | null) => { + if (value === null) { + form.setValue("user_email", null); + form.setValue("user_id", null); + return; + } + controlProps.onChange(value); + handleSelect(items.find((option) => option.value === value) ?? null); + }; return (
{ - controlProps.onChange(value === "" ? undefined : value); - handleSelect(items.find((option) => option.value === value) ?? null); - }} + onValueChange={handleValueChange} onSearchChange={(query: string) => handleSearch(query, fieldName)} autoHighlight="always" isLoading={loading} @@ -226,7 +238,7 @@ const UserSearchModal: React.FC = ({
- diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index 8a62e86e842..86f18ee9b12 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -13,7 +13,7 @@ import AccessGroupTagsCombobox from "../add_model/AccessGroupTagsCombobox"; import ModelChoiceCombobox, { type ModelChoice } from "../add_model/ModelChoiceCombobox"; import { modelAvailableCall, modelPatchUpdateCall, validateAutoRouterConfig } from "../networking"; import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; -import RouterConfigBuilder from "../add_model/RouterConfigBuilder"; +import RouterConfigBuilder, { type RouterConfig, serializeRouterConfig } from "../add_model/RouterConfigBuilder"; import { hydrateTierModelParams } from "../add_model/complexity_router_tiers"; import { type ActiveTierSet, @@ -452,7 +452,7 @@ const EditAutoRouterModal: React.FC = ({ const [modelInfo, setModelInfo] = useState([]); const [showValidationErrors, setShowValidationErrors] = useState(false); const [editingTiers, setEditingTiers] = useState(false); - const [routerConfig, setRouterConfig] = useState(null); + const [routerConfig, setRouterConfig] = useState(null); const [customTechnicalKeywords, setCustomTechnicalKeywords] = useState([]); const [keywordTierRules, setKeywordTierRules] = useState([]); const [escalationKeywords, setEscalationKeywords] = useState([]); @@ -706,7 +706,7 @@ const EditAutoRouterModal: React.FC = ({ // Prepare the updated litellm_params const updatedLitellmParams = { ...modelData.litellm_params, - auto_router_config: JSON.stringify(routerConfig), + auto_router_config: serializeRouterConfig(routerConfig), auto_router_default_model: values.auto_router_default_model, auto_router_embedding_model: values.auto_router_embedding_model || undefined, }; @@ -745,7 +745,7 @@ const EditAutoRouterModal: React.FC = ({ })(); } catch (error) { console.error("Error updating auto router:", error); - toast.fromError("Failed to update auto router configuration"); + toast.fromError(error); } finally { setLoading(false); } diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.tsx index c0c12356c6f..dbc270fa62e 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.tsx @@ -96,10 +96,10 @@ export function BudgetFallbacksEditor({ value, onChange, availableModels }: Budg ({ label: m, value: m }))} - value={entry.primaryModel ?? ""} + value={entry.primaryModel} onValueChange={(v) => { const newFallbacks = entry.fallbackModels.filter((m) => m !== v); - updateEntry(entry.id, { primaryModel: v === "" ? null : v, fallbackModels: newFallbacks }); + updateEntry(entry.id, { primaryModel: v, fallbackModels: newFallbacks }); }} placeholder="Select model" emptyText="No models found" diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx index ef5a3e15b52..0fab5555343 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx @@ -153,8 +153,8 @@ export function ModelMaxBudgetEditor({ ({ label: model, value: model }))} - value={entry.model ?? ""} - onValueChange={(model) => updateEntry(entry.id, { model: model === "" ? null : model })} + value={entry.model} + onValueChange={(model) => updateEntry(entry.id, { model })} placeholder="Select model" emptyText="No models found" disabled={!premiumUser} diff --git a/ui/litellm-dashboard/src/components/model_add/CredentialModal.tsx b/ui/litellm-dashboard/src/components/model_add/CredentialModal.tsx index 8108cc1c20e..fdc02b4fc41 100644 --- a/ui/litellm-dashboard/src/components/model_add/CredentialModal.tsx +++ b/ui/litellm-dashboard/src/components/model_add/CredentialModal.tsx @@ -42,7 +42,7 @@ export default function CredentialModal({ existingCredential = null, }: CredentialModalProps) { const isEdit = mode === "edit"; - const [selectedProvider, setSelectedProvider] = useState( + const [selectedProvider, setSelectedProvider] = useState( (existingCredential?.credential_info.custom_llm_provider as Providers) ?? Providers.OpenAI, ); @@ -110,7 +110,7 @@ export default function CredentialModal({ {(control) => ( { control.onChange(value); - resetCredentialFormOnProviderChange(formAdapter, value as Providers, setSelectedProvider); + resetCredentialFormOnProviderChange(formAdapter, value, setSelectedProvider); }} /> )} diff --git a/ui/litellm-dashboard/src/components/model_add/credential_form_helpers.ts b/ui/litellm-dashboard/src/components/model_add/credential_form_helpers.ts index 7190c539c81..e7b1e861811 100644 --- a/ui/litellm-dashboard/src/components/model_add/credential_form_helpers.ts +++ b/ui/litellm-dashboard/src/components/model_add/credential_form_helpers.ts @@ -1,5 +1,3 @@ -import { Providers } from "../provider_info_helpers"; - interface CredentialFormAdapter { getFieldValue: (field: string) => unknown; resetFields: () => void; @@ -25,8 +23,8 @@ interface CredentialFormAdapter { */ export function resetCredentialFormOnProviderChange( form: CredentialFormAdapter, - newProvider: Providers, - setSelectedProvider: (p: Providers) => void, + newProvider: string | null, + setSelectedProvider: (p: string | null) => void, ): void { const preservedName = form.getFieldValue("credential_name"); form.resetFields(); diff --git a/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts b/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts index 37a973d5def..311e9357825 100644 --- a/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts +++ b/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts @@ -202,6 +202,8 @@ export const buildKeyCreatePayload = (input: KeyCreateInput): KeyPayloadResult = endpoint: input.keyOwner === "service_account" ? "service_account" : "standard", payload: { ...withoutKeys(values, dropped), + ...(values.organization_id === null && { organization_id: undefined }), + ...(values.project_id === null && { project_id: undefined }), ...(input.keyOwner === "you" && { user_id: input.userID }), ...(input.keyOwner === "agent" && { agent_id: input.selectedAgentId }), ...(input.autoRotationEnabled && { auto_rotate: true, rotation_interval: input.rotationInterval }), diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx index 2bf39bf4cd7..0d5d9f5ec8d 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx @@ -16,7 +16,7 @@ const state = vi.hoisted(() => ({ can: {} as Record, uiSettings: {} as Record, tags: {} as Record, - teams: [] as { team_id: string; team_alias: string; models: string[] }[], + teams: [] as { team_id: string; team_alias: string; models: string[]; organization_id?: string }[], organizations: [] as { organization_id: string; organization_alias: string }[], accessGroups: [] as { access_group_id: string; access_group_name: string }[], projects: [] as { project_id: string; project_alias: string; team_id?: string; models?: string[] }[], @@ -806,6 +806,36 @@ describe("CreateKey", () => { expect((await createdPayload()).organization_id).toBe("org-1"); }); + it("discards the old project and team when the organization changes", async () => { + state.uiSettings = { enable_projects_ui: true }; + state.organizations = [ + { organization_id: "scope-silver", organization_alias: "Silver" }, + { organization_id: "scope-copper", organization_alias: "Copper" }, + ]; + state.teams = [{ team_id: "group-maple", team_alias: "Maple", organization_id: "scope-silver", models: [] }]; + state.projects = [{ project_id: "project-orbit", project_alias: "Orbit", team_id: "group-maple", models: [] }]; + await openModal({ teams: state.teams as Team[] }); + await nameTheKey(); + await userEvent.click(await screen.findByLabelText("Organization")); + await userEvent.click(await screen.findByRole("option", { name: /Silver/ })); + await userEvent.click(await screen.findByLabelText("Project")); + await userEvent.click(await screen.findByRole("option", { name: /Orbit/ })); + await waitFor(() => expect(screen.getByLabelText("Team")).toHaveValue("Maple")); + expect(screen.getByLabelText("Team")).toBeDisabled(); + + await userEvent.click(screen.getByLabelText("Organization")); + await userEvent.click(await screen.findByRole("option", { name: /Copper/ })); + expect(screen.getByLabelText("Project")).toHaveValue(""); + expect(screen.getByLabelText("Team")).toHaveValue(""); + expect(screen.getByLabelText("Team")).toBeEnabled(); + await submit(); + + const payload = JSON.parse(JSON.stringify(await createdPayload())); + expect(payload.organization_id).toBe("scope-copper"); + expect(payload.team_id).toBeNull(); + expect(payload).not.toHaveProperty("project_id"); + }); + it("drops organization_id when the chosen organization is cleared again", async () => { state.organizations = [{ organization_id: "org-1", organization_alias: "Engineering" }]; await openModal(); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 831cb0cf6a6..82580fd4667 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -592,35 +592,35 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp }; const changeOrganization = (write: FieldWrite) => (orgId: string | null) => { - write(orgId ?? undefined); + write(orgId); setSelectedOrganizationId(orgId); // Clear team and project when org changes setSelectedCreateKeyTeam(null); setSelectedProjectId(null); - form.setValue("team_id", undefined); - form.setValue("project_id", undefined); + form.setValue("team_id", null); + form.setValue("project_id", null); }; const selectTeam = (team: Team | null) => { setSelectedCreateKeyTeam(team); setSelectedProjectId(null); - form.setValue("project_id", undefined); + form.setValue("project_id", null); // Auto-populate org from team for non-admin users if (team?.organization_id) { setSelectedOrganizationId(team.organization_id); form.setValue("organization_id", team.organization_id); } else if (!team) { setSelectedOrganizationId(null); - form.setValue("organization_id", undefined); + form.setValue("organization_id", null); } }; - const changeProject = (write: FieldWrite) => (projectId: string) => { + const changeProject = (write: FieldWrite) => (projectId: string | null) => { write(projectId); if (!projectId) { setSelectedProjectId(null); setSelectedCreateKeyTeam(null); - form.setValue("team_id", undefined); + form.setValue("team_id", null); return; } setSelectedProjectId(projectId); @@ -756,8 +756,8 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp inputId="create-key-agent" placeholder="Select an agent" emptyText="No agents found" - value={selectedAgentId ?? undefined} - onValueChange={(value) => setSelectedAgentId(value === "" ? null : value)} + value={selectedAgentId} + onValueChange={setSelectedAgentId} options={agentsList.map((a) => ({ label: a.agent_name || a.agent_id, value: a.agent_id, @@ -783,7 +783,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp {(control) => ( = ({ team, teams, data, addKey, autoOp {(control) => ( = ({ team, teams, data, addKey, autoOp {(control) => ( { return providerPlaceholderMap[resolvedProvider] ?? "gpt-3.5-turbo"; }; -export const getProviderModels = (provider: Providers, modelMap: any): Array => { +export const getProviderModels = (provider: string, modelMap: any): Array => { let providerKey = provider; let custom_llm_provider = provider_map[providerKey]; diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx similarity index 89% rename from ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx rename to ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx index 04c8d3e9012..2b948ca8420 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx @@ -1,4 +1,4 @@ -import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { fireEvent, renderWithProviders as render, screen, waitFor } from "../../../tests/test-utils"; import userEvent from "@testing-library/user-event"; import { useState } from "react"; import { describe, expect, it, vi } from "vitest"; @@ -51,7 +51,7 @@ describe("PaginatedSearchSelect", () => { const onSearchChange = vi.fn(); function Controlled() { - const [value, setValue] = useState(""); + const [value, setValue] = useState(null); return ( { expect(onSearchChange).not.toHaveBeenCalled(); }); - it("still reports a cleared input so the unfiltered page comes back", async () => { + it("should keep a cleared selection empty after a late page arrives and reset the query", async () => { const user = userEvent.setup(); const onSearchChange = vi.fn(); - renderSelect({ onSearchChange, value: "alias-alpha" }); - - await user.click(document.querySelector('[data-slot="combobox-clear"]') as HTMLElement); - - await waitFor(() => expect(onSearchChange).toHaveBeenCalledWith("")); + const onValueChange = vi.fn(); + function Controlled({ options }: { options: SearchSelectOption[] }) { + const [value, setValue] = useState("alias-alpha"); + return ( + { + setValue(next); + onValueChange(next); + }} + /> + ); + } + const { rerender } = render(); + await user.click(screen.getByRole("button", { name: "Clear" })); + expect(onValueChange).toHaveBeenLastCalledWith(null); + rerender( ({ ...option }))} />); + await waitFor(() => expect(onSearchChange).toHaveBeenLastCalledWith("")); + expect(screen.getByRole("combobox")).toHaveValue(""); + expect(screen.queryByRole("button", { name: "Clear" })).not.toBeInTheDocument(); + await chooseSelectOption(user, screen.getByRole("combobox"), "alias-beta"); + expect(onValueChange).toHaveBeenLastCalledWith("alias-beta"); }); it("requests the next page once the list is scrolled near the bottom", async () => { @@ -147,7 +166,7 @@ describe("PaginatedSearchSelect", () => { function ServerBacked() { const [search, setSearch] = useState(""); - const [value, setValue] = useState("alias-alpha"); + const [value, setValue] = useState("alias-alpha"); const freshlyBuiltOptions = OPTIONS.filter((option) => option.label.includes(search)).map((option) => ({ ...option, })); @@ -236,7 +255,7 @@ describe("PaginatedSearchSelect", () => { function Refetching() { const [options, setOptions] = useState([{ label: "Beta Team", value: "team-2" }]); - const [value, setValue] = useState(""); + const [value, setValue] = useState(null); return ( <> { function ServerBacked() { const [search, setSearch] = useState(""); - const [value, setValue] = useState(""); + const [value, setValue] = useState(null); return ( option.label.includes(search))} diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx index 4b7ef7401c7..1314afdf2bf 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx @@ -17,8 +17,8 @@ import { usePaginatedCombobox } from "./usePaginatedCombobox"; interface PaginatedSearchSelectProps { options: SearchSelectOption[]; - value?: string; - onValueChange: (value: string) => void; + value?: string | null; + onValueChange: (value: string | null) => void; onSearchChange: (query: string) => void; onLoadMore?: () => void; hasNextPage?: boolean; @@ -86,7 +86,7 @@ export function PaginatedSearchSelect({ }; const selected = useMemo(() => { - if (value === undefined || value === "") return null; + if (value == null || value === "") return null; return ( options.find((option) => option.value === value) ?? (pickedOption?.value === value ? pickedOption : { label: value, value }) @@ -118,7 +118,7 @@ export function PaginatedSearchSelect({ inputValue={typedQuery ?? selected?.label ?? ""} onValueChange={(item: SearchSelectOption | null) => { setPickedOption(item); - onValueChange(item?.value ?? ""); + onValueChange(item?.value ?? null); }} onInputValueChange={(next, eventDetails) => handleTypedInput(next, eventDetails.reason)} onOpenChange={(nextOpen, eventDetails) => handleOpenChange(nextOpen, eventDetails.reason)} @@ -139,7 +139,7 @@ export function PaginatedSearchSelect({ onKeyDown={snapshotWholeSelection} onPaste={snapshotWholeSelection} placeholder={placeholder} - showClear={value !== undefined && value !== ""} + showClear={value != null && value !== ""} className={`w-full ${className ?? ""}`} /> diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx similarity index 76% rename from ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx rename to ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx index 62a510268dd..c981010dff9 100644 --- a/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx @@ -1,4 +1,5 @@ -import { fireEvent, render, screen } from "@testing-library/react"; +import { fireEvent, renderWithProviders as render, screen } from "../../../tests/test-utils"; +import { useState } from "react"; import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi } from "vitest"; @@ -36,11 +37,30 @@ describe("SearchSelect", () => { expect(screen.getByRole("combobox")).toHaveValue("Growth"); }); - it("shows a clear control only when a value is selected", () => { - const { rerender } = render(); - expect(document.querySelector('[data-slot="combobox-clear"]')).toBeNull(); - rerender(); - expect(document.querySelector('[data-slot="combobox-clear"]')).not.toBeNull(); + it("should clear to null and allow selecting again through the real control", async () => { + const onValueChange = vi.fn(); + const user = userEvent.setup(); + function Controlled() { + const [value, setValue] = useState(null); + return ( + { + setValue(next); + onValueChange(next); + }} + /> + ); + } + render(); + expect(screen.queryByRole("button", { name: "Clear" })).not.toBeInTheDocument(); + await chooseSelectOption(user, screen.getByRole("combobox"), "Growth"); + await user.click(screen.getByRole("button", { name: "Clear" })); + expect(onValueChange).toHaveBeenLastCalledWith(null); + expect(screen.getByRole("combobox")).toHaveValue(""); + await chooseSelectOption(user, screen.getByRole("combobox"), "Data Team"); + expect(onValueChange).toHaveBeenLastCalledWith("team-3"); }); it("filters the options client-side as you type", async () => { diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx index 53a5371d72d..21d38e9458c 100644 --- a/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx @@ -21,8 +21,8 @@ export interface SearchSelectOption { interface SearchSelectProps { options: SearchSelectOption[]; - value?: string; - onValueChange: (value: string) => void; + value?: string | null; + onValueChange: (value: string | null) => void; placeholder?: string; emptyText?: string; disabled?: boolean; @@ -51,9 +51,7 @@ export function SearchSelect({ "aria-label": ariaLabel, }: SearchSelectProps) { const selected = - value === undefined || value === "" - ? null - : options.find((option) => option.value === value) ?? { label: value, value }; + value == null || value === "" ? null : options.find((option) => option.value === value) ?? { label: value, value }; const items = selected !== null && !options.some((option) => option.value === selected.value) ? [selected, ...options] : options; @@ -61,7 +59,7 @@ export function SearchSelect({ onValueChange(item?.value ?? "")} + onValueChange={(item: SearchSelectOption | null) => onValueChange(item?.value ?? null)} isItemEqualToValue={(a: SearchSelectOption, b: SearchSelectOption) => a.value === b.value} itemToStringLabel={(item: SearchSelectOption) => item.label} filter={matchesQuery} diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 988e2e082aa..8e3a17c2622 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -333,7 +333,10 @@ const teamUpdateFieldsSchema = z.object({ modelLimits: z .array( z.object({ - model: z.string().min(1, "Missing model"), + model: z + .string() + .nullable() + .refine((model) => Boolean(model), "Missing model"), tpm: z.number().nullish(), rpm: z.number().nullish(), }), @@ -1879,7 +1882,7 @@ const TeamInfoView: React.FC = ({ onChange(next === "" ? null : next)} + onValueChange={onChange} options={userOrganizations.map((org) => ({ value: org.organization_id ?? "", label: org.organization_alias || org.organization_id || "", diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx rename to ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx index b2c2a381b42..7c2226f0369 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx @@ -1483,11 +1483,11 @@ describe("KeyEditView", () => { }); }); - it("submits organization_id as null after the organization is cleared", async () => { + it("clears the organization and its dependent team in the update payload", async () => { const onSubmit = vi.fn().mockResolvedValue(undefined); renderWithProviders( {}} onSubmit={onSubmit} accessToken="" @@ -1504,9 +1504,33 @@ describe("KeyEditView", () => { await userEvent.click(screen.getByRole("button", { name: /save changes/i })); await waitFor(() => { - expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining({ organization_id: null })); + expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining({ organization_id: null, team_id: null })); }); - expect(JSON.parse(JSON.stringify(onSubmit.mock.calls[0][0]))).toHaveProperty("organization_id", null); + expect(JSON.parse(JSON.stringify(onSubmit.mock.calls[0][0]))).toMatchObject({ + organization_id: null, + team_id: null, + }); + }); + + it("keeps project key relationships locked and omits unsupported project updates", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + renderWithProviders( + {}} + onSubmit={onSubmit} + accessToken="" + userID="" + userRole="Admin" + premiumUser={false} + />, + ); + expect(await screen.findByRole("combobox", { name: "Organization" })).toBeDisabled(); + expect(screen.getByRole("combobox", { name: "Team ID" })).toBeDisabled(); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + await waitFor(() => expect(onSubmit).toHaveBeenCalledTimes(1)); + expect(onSubmit.mock.calls[0][0]).toMatchObject({ organization_id: "org-1", team_id: "group-maple" }); + expect(onSubmit.mock.calls[0][0]).not.toHaveProperty("project_id"); }); }); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index e464dfe5008..f0d63f9671d 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -305,7 +305,7 @@ export function KeyEditView({ const handleOrganizationChange = (setField: (value: string | null) => void, orgId: string | null) => { setField(orgId); setSelectedOrganizationId(orgId); - form.setValue("team_id", undefined); + form.setValue("team_id", null); }; const handleTeamChange = (setField: (value: string | null) => void, teamId: string | null) => { @@ -316,7 +316,7 @@ export function KeyEditView({ form.setValue("organization_id", selectedTeam.organization_id); } else if (!teamId) { setSelectedOrganizationId(null); - form.setValue("organization_id", undefined); + form.setValue("organization_id", null); } }; @@ -769,14 +769,15 @@ export function KeyEditView({ "Organization", "The organization this key belongs to. Selecting an organization filters the available teams.", )} + description={hasProject ? "Organization is locked because this key belongs to a project" : undefined} > {({ value, onChange, id }) => ( handleOrganizationChange(onChange, orgId)} /> )} @@ -786,15 +787,13 @@ export function KeyEditView({ control={form.control} name="team_id" label="Team ID" - description={ - enableProjectsUI && hasProject ? "Team is locked because this key belongs to a project" : undefined - } + description={hasProject ? "Team is locked because this key belongs to a project" : undefined} > {({ value, onChange, id }) => (