From e2932d22c0b6d6e0a01c4346fe304525eed9bf8e Mon Sep 17 00:00:00 2001 From: Nitish Agarwal <1592163+nitishagar@users.noreply.github.com> Date: Thu, 18 Jun 2026 14:06:36 +0530 Subject: [PATCH 001/281] fix(spend-tracking): hash raw api keys before persisting to spend logs (#30670) The Request Logs "Key Hash" column binds to the SpendLogs metadata user_api_key field, which _get_spend_logs_metadata copied verbatim from request metadata without hashing. The top-level api_key column only hashed values starting with sk-. So a non-sk- passthrough provider key, or a Bearer-prefixed key on paths that do not strip it, could be persisted and shown in plaintext. A single helper now redacts both fields at the SpendLogs builder: it strips a Bearer prefix, passes through values already a sha256 hash or hashed-jwt- identifier, and otherwise hashes with hash_token. sk- keys still map to the same canonical hash, so log and spend correlation is unchanged --- .../spend_tracking/spend_tracking_utils.py | 42 ++- .../test_spend_tracking_utils.py | 270 ++++++++++++++++-- 2 files changed, 290 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index aef06a3c668..e8aa2325a18 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -24,6 +24,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_litellm_metadata_from_kwargs, reconstruct_model_name, ) +from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash from litellm.litellm_core_utils.safe_json_dumps import safe_dumps, strip_null_bytes from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error @@ -65,6 +66,24 @@ def _is_master_key(api_key: Optional[str], _master_key: Optional[str]) -> bool: return secrets.compare_digest(api_key, _master_key) +_HASHED_JWT_RE = re.compile(r"^hashed-jwt-[a-fA-F0-9]{64}$") + + +def _redact_logged_api_key( + value: str | None, *, already_hashed: bool = False +) -> str | None: + if not isinstance(value, str) or not value: + return None + stripped = re.sub(r"(?i)^bearer ", "", value) + if not stripped: + return None + if already_hashed and is_valid_sha256_hash(stripped): + return stripped + if _HASHED_JWT_RE.match(stripped): + return stripped + return hash_token(stripped) + + def _get_spend_logs_metadata( metadata: Optional[dict], applied_guardrails: Optional[List[str]] = None, @@ -121,6 +140,16 @@ def _get_spend_logs_metadata( key: metadata.get(key) for key in SpendLogsMetadata.__annotations__.keys() } ) + _raw_key = clean_metadata.get("user_api_key") + _trusted_hash = metadata.get("user_api_key_hash") if metadata else None + _already_hashed = ( + isinstance(_trusted_hash, str) + and is_valid_sha256_hash(_trusted_hash) + and _trusted_hash == _raw_key + ) + clean_metadata["user_api_key"] = _redact_logged_api_key( + _raw_key, already_hashed=_already_hashed + ) clean_metadata["applied_guardrails"] = applied_guardrails clean_metadata["batch_models"] = batch_models clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata @@ -283,10 +312,13 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs "completion_tokens", 0 ) standard_logging_total_tokens = standard_logging_payload.get("total_tokens", 0) - if api_key is not None and isinstance(api_key, str): - if api_key.startswith("sk-"): - # hash the api_key - api_key = hash_token(api_key) + _trusted_hash = metadata.get("user_api_key_hash") if metadata else None + _key_already_hashed = ( + isinstance(_trusted_hash, str) + and is_valid_sha256_hash(_trusted_hash) + and _trusted_hash == api_key + ) + api_key = _redact_logged_api_key(api_key, already_hashed=_key_already_hashed) or "" if ( standard_logging_payload is not None @@ -299,8 +331,6 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs end_user_id = end_user_id or standard_logging_payload["metadata"].get( "user_api_key_end_user_id" ) - # BUG FIX: Don't overwrite api_key when standard_logging_payload is None - # The api_key was already extracted from metadata (line 243) and hashed (lines 256-259) request_tags = ( safe_dumps(metadata.get("tags", [])) if isinstance(metadata.get("tags", []), list) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 0c7511589de..67d57d2d6bb 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -7,7 +7,6 @@ from datetime import timezone from typing import Any, cast import pytest -from fastapi.testclient import TestClient sys.path.insert( 0, os.path.abspath("../../../..") @@ -15,7 +14,6 @@ sys.path.insert( from unittest.mock import AsyncMock, MagicMock, patch -import litellm from litellm.constants import ( LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, @@ -30,12 +28,14 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( _get_spend_logs_metadata, _get_vector_store_request_for_spend_logs_payload, _is_master_key, + _redact_logged_api_key, _redact_prompt_leaks_in_error_string, _sanitize_error_information_for_spend_logs, _sanitize_request_body_for_spend_logs_payload, _should_store_prompts_and_responses_in_spend_logs, get_logging_payload, ) +from litellm.proxy.utils import hash_token from litellm.types.utils import ( StandardLoggingHiddenParams, StandardLoggingMetadata, @@ -661,8 +661,6 @@ def test_get_logging_payload_api_key_preserved_when_standard_logging_payload_is_ assert payload["model"] == "openai/gpt-4.1" assert payload["user"] == "test_user" - print(f"✅ Test passed! api_key preserved: {payload['api_key']}") - @pytest.mark.asyncio @patch("litellm.proxy.proxy_server.master_key", "sk-master-key") @@ -810,18 +808,6 @@ async def test_api_key_preserved_through_failure_hook_to_database(): assert payload.get("model") == "gpt-3.5-turbo" assert payload.get("user") == "test_user" - print("\n" + "=" * 80) - print("✅ CRITICAL E2E TEST PASSED") - print("=" * 80) - print(f"Token: {data['token']}") - print(f"Payload api_key: {payload_api_key}") - print(f"Match: {data['token'] == payload_api_key}") - print("=" * 80) - print("Production incident bug is FIXED and protected:") - print("- Failed requests preserve api_key through entire flow") - print("- Both SpendLogs AND DailyUserSpend will have correct api_key") - print("=" * 80 + "\n") - @patch("litellm.proxy.proxy_server.master_key", None) @patch("litellm.proxy.proxy_server.general_settings", {}) @@ -2073,3 +2059,255 @@ def test_sanitize_error_information_redacts_pydantic_assignment_form( assert sanitized is not None assert "leaked-via-pydantic-msg" not in sanitized["error_message"] assert REDACTED_BY_LITELM_STRING in sanitized["error_message"] + + +# ── _redact_logged_api_key unit tests ────────────────────────────────────── + + +def test_redact_logged_api_key_none_returns_none(): + assert _redact_logged_api_key(None) is None + + +def test_redact_logged_api_key_empty_string_returns_none(): + assert _redact_logged_api_key("") is None + + +def test_redact_logged_api_key_sk_key_is_hashed(): + raw = "sk-1234secret" + result = _redact_logged_api_key(raw) + assert result == hash_token(raw) + assert result is not None + assert not result.startswith("sk-") + assert len(result) == 64 + + +def test_redact_logged_api_key_bearer_sk_equals_sk_hash(): + raw = "sk-1234secret" + result_plain = _redact_logged_api_key(raw) + result_bearer = _redact_logged_api_key(f"Bearer {raw}") + assert result_bearer == result_plain + + +def test_redact_logged_api_key_bearer_case_insensitive(): + raw = "sk-1234secret" + result_lower = _redact_logged_api_key(f"bearer {raw}") + result_upper = _redact_logged_api_key(f"BEARER {raw}") + expected = hash_token(raw) + assert result_lower == expected + assert result_upper == expected + + +def test_redact_logged_api_key_non_sk_raw_key_is_hashed(): + raw = "anthropic-raw-key-xyz" + result = _redact_logged_api_key(raw) + assert result is not None + assert result != raw + assert len(result) == 64 + assert result == hash_token(raw) + + +def test_redact_logged_api_key_already_valid_sha256_passes_through_with_flag(): + already_hashed = hash_token("sk-some-key") + assert len(already_hashed) == 64 + result = _redact_logged_api_key(already_hashed, already_hashed=True) + assert result == already_hashed + assert hash_token(already_hashed) != result # no double-hash + + +def test_redact_logged_api_key_sha256_without_flag_is_hashed(): + already_hashed = hash_token("sk-some-key") + assert len(already_hashed) == 64 + result = _redact_logged_api_key(already_hashed) + assert result is not None + assert result != already_hashed + assert len(result) == 64 + assert result == hash_token(already_hashed) + + +def test_redact_logged_api_key_hashed_jwt_passes_through(): + jwt_hash = "hashed-jwt-" + "a" * 64 + result = _redact_logged_api_key(jwt_hash) + assert result == jwt_hash + + +def test_redact_logged_api_key_hashed_jwt_short_suffix_is_hashed(): + short_jwt = "hashed-jwt-tooshort" + result = _redact_logged_api_key(short_jwt) + assert result is not None + assert result != short_jwt + assert len(result) == 64 + assert result == hash_token(short_jwt) + + +def test_redact_logged_api_key_bearer_only_returns_none(): + # "bearer " with nothing after stripping is equivalent to no key + assert _redact_logged_api_key("bearer ") is None + assert _redact_logged_api_key("Bearer ") is None + assert _redact_logged_api_key("BEARER ") is None + + +# ── _get_spend_logs_metadata key-hash invariant tests ───────────────────── + + +def test_get_spend_logs_metadata_sk_key_hashed(): + raw = "sk-1234secret" + meta = _get_spend_logs_metadata({"user_api_key": raw}) + assert meta["user_api_key"] == hash_token(raw) + assert meta["user_api_key"] is not None + result = meta["user_api_key"] + assert result is not None + assert not result.startswith("sk-") + assert len(result) == 64 + + +def test_get_spend_logs_metadata_bearer_sk_key_hashed_same_as_plain(): + raw = "sk-1234secret" + meta_plain = _get_spend_logs_metadata({"user_api_key": raw}) + meta_bearer = _get_spend_logs_metadata({"user_api_key": f"Bearer {raw}"}) + assert meta_bearer["user_api_key"] == meta_plain["user_api_key"] + + +def test_get_spend_logs_metadata_non_sk_raw_key_hashed(): + raw = "anthropic-raw-key-xyz" + meta = _get_spend_logs_metadata({"user_api_key": raw}) + result = meta["user_api_key"] + assert result is not None + assert result != raw + assert len(result) == 64 + + +def test_get_spend_logs_metadata_already_hashed_unchanged_with_provenance(): + already_hashed = hash_token("sk-some-key") + meta = _get_spend_logs_metadata( + {"user_api_key": already_hashed, "user_api_key_hash": already_hashed} + ) + assert meta["user_api_key"] == already_hashed + assert hash_token(already_hashed) != meta["user_api_key"] # no double-hash + + +def test_get_spend_logs_metadata_already_hashed_no_provenance_is_rehashed(): + already_hashed = hash_token("sk-some-key") + meta = _get_spend_logs_metadata({"user_api_key": already_hashed}) + assert meta["user_api_key"] != already_hashed + assert meta["user_api_key"] == hash_token(already_hashed) + + +def test_get_spend_logs_metadata_provenance_bypass_requires_hash_match(): + already_hashed = hash_token("sk-some-key") + different_hash = hash_token("sk-other-key") + meta = _get_spend_logs_metadata( + {"user_api_key": already_hashed, "user_api_key_hash": different_hash} + ) + assert meta["user_api_key"] == hash_token(already_hashed) + + +def test_get_spend_logs_metadata_hashed_jwt_unchanged(): + jwt_hash = "hashed-jwt-" + "b" * 64 + meta = _get_spend_logs_metadata({"user_api_key": jwt_hash}) + assert meta["user_api_key"] == jwt_hash + + +def test_get_spend_logs_metadata_none_key_is_none(): + meta = _get_spend_logs_metadata({"user_api_key": None}) + assert meta["user_api_key"] is None + + +# ── get_logging_payload key-hash invariant tests ─────────────────────────── + + +@patch("litellm.proxy.proxy_server.master_key", None) +@patch("litellm.proxy.proxy_server.general_settings", {}) +def test_get_logging_payload_non_sk_raw_key_both_fields_hashed(): + raw = "anthropic-raw-key-xyz" + kwargs = { + "model": "openai/gpt-4.1", + "messages": [{"role": "user", "content": "Hello"}], + "call_type": "acompletion", + "litellm_params": { + "metadata": { + "user_api_key": raw, + "user_api_key_user_id": "test_user", + "user_api_key_team_id": "test_team", + } + }, + } + payload = get_logging_payload( + kwargs=kwargs, + response_obj=Exception("error"), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert payload["api_key"] != raw + assert len(payload["api_key"]) == 64 + + parsed_meta = json.loads(payload["metadata"]) + assert parsed_meta["user_api_key"] != raw + assert parsed_meta["user_api_key"] is not None + assert len(parsed_meta["user_api_key"]) == 64 + + +@patch("litellm.proxy.proxy_server.master_key", None) +@patch("litellm.proxy.proxy_server.general_settings", {}) +def test_get_logging_payload_empty_key_slp_none_is_empty_string_not_none_literal(): + kwargs = { + "model": "openai/gpt-4.1", + "messages": [{"role": "user", "content": "Hello"}], + "call_type": "acompletion", + "litellm_params": { + "metadata": { + "user_api_key_user_id": "test_user", + } + }, + } + payload = get_logging_payload( + kwargs=kwargs, + response_obj=Exception("error"), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert payload["api_key"] == "", ( + f"Expected empty string but got {payload['api_key']!r}; " + "dropping _redact_logged_api_key's 'or \"\"' guard would yield 'None' here" + ) + + +def test_get_spend_logs_metadata_sibling_fields_preserved(): + raw = "anthropic-raw-key-xyz" + meta = _get_spend_logs_metadata( + { + "user_api_key": raw, + "user_api_key_alias": "my-alias", + "user_api_key_team_id": "team-123", + } + ) + assert meta["user_api_key"] == hash_token(raw) + assert meta["user_api_key_alias"] == "my-alias" + assert meta["user_api_key_team_id"] == "team-123" + + +def test_redact_logged_api_key_partial_sha256_is_hashed(): + partial_hex = "a" * 63 + result = _redact_logged_api_key(partial_hex) + assert result is not None + assert result != partial_hex + assert len(result) == 64 + assert result == hash_token(partial_hex) + + +def test_redact_logged_api_key_bearer_already_hashed_passes_through_with_flag(): + already_hashed = hash_token("sk-some-key") + assert len(already_hashed) == 64 + result = _redact_logged_api_key(f"Bearer {already_hashed}", already_hashed=True) + assert result == already_hashed + assert hash_token(already_hashed) != result + + +def test_redact_logged_api_key_bearer_sha256_without_flag_is_hashed(): + already_hashed = hash_token("sk-some-key") + assert len(already_hashed) == 64 + result = _redact_logged_api_key(f"Bearer {already_hashed}") + assert result is not None + assert result != already_hashed + assert result == hash_token(already_hashed) From 7e1f44f0cc15e639d8fe95fc8831c5a0b25cc76b Mon Sep 17 00:00:00 2001 From: mateo Date: Thu, 30 Jul 2026 02:38:28 +0000 Subject: [PATCH 002/281] feat(proxy): add admin toggle to block requests for models without pricing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/__init__.py | 1 + litellm/proxy/_types.py | 2 + litellm/proxy/auth/auth_checks.py | 33 +++++ .../cost_tracking_settings.py | 65 ++++++++ .../proxy/auth/test_auth_checks.py | 139 +++++++++++++++++- .../test_cost_tracking_settings.py | 56 +++++++ .../_components/cost_tracking_settings.tsx | 47 +++++- .../_components/use_block_unpriced_config.ts | 63 ++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 81 ++++++++++ 9 files changed, 482 insertions(+), 5 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/use_block_unpriced_config.ts diff --git a/litellm/__init__.py b/litellm/__init__.py index 3f8c742c5a2..d14f41ad49c 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -443,6 +443,7 @@ max_end_user_budget_id: Optional[str] = None # backwards compatibility — arbitrary client-supplied identifiers still # pass through unchanged. validate_end_user_id_in_db: bool = False +block_requests_for_models_without_pricing: bool = False disable_end_user_cost_tracking: Optional[bool] = None disable_end_user_cost_tracking_prometheus_only: Optional[bool] = None enable_end_user_cost_tracking_prometheus_only: Optional[bool] = None diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c6d4ee1120a..450868edf06 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3540,6 +3540,8 @@ class ProxyErrorTypes(str, enum.Enum): Project does not have access to the model """ + model_cost_map_missing = "model_cost_map_missing" + expired_key = "expired_key" """ Key has expired diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index c46bc110ca8..11a45e78410 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -274,6 +274,22 @@ def _is_cost_explicitly_configured(model: str, llm_router: "Router") -> bool: return False +def model_has_no_cost_mapping(model: Optional[str], llm_router: Optional[Router]) -> bool: + if not model or llm_router is None: + return False + + model_group_info = llm_router.get_model_group_info(model_group=model) + if model_group_info is None: + return False + + input_cost = model_group_info.input_cost_per_token or 0 + output_cost = model_group_info.output_cost_per_token or 0 + if input_cost > 0 or output_cost > 0: + return False + + return not _is_cost_explicitly_configured(model, llm_router) + + async def _run_project_checks( project_object: Optional[LiteLLM_ProjectTableCachedObj], _model: Optional[Union[str, List[str]]], @@ -534,6 +550,23 @@ async def common_checks( if route in MODEL_DISCOVERY_ROUTES: skip_budget_checks = True + if ( + litellm.block_requests_for_models_without_pricing + and isinstance(_model, str) + and RouteChecks.is_llm_api_route(route=route) + and model_has_no_cost_mapping(model=_model, llm_router=llm_router) + ): + raise ProxyException( + message=( + f"Model '{_model}' has no pricing in the cost map, so its spend would be tracked as $0. " + "Requests for unpriced models are blocked because 'block_requests_for_models_without_pricing' " + "is enabled. Add pricing for this model (input_cost_per_token/output_cost_per_token) to allow it." + ), + type=ProxyErrorTypes.model_cost_map_missing, + param="model", + code=status.HTTP_403_FORBIDDEN, + ) + # 1. If team is blocked if team_object is not None and team_object.blocked is True: raise Exception(f"Team={team_object.team_id} is blocked. Update via `/team/unblock` if you're an admin.") diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index cd2c5704778..2c18de6b903 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -13,6 +13,7 @@ POST /cost/estimate - Estimate cost for a given model and token counts from typing import Dict, Optional, Tuple, Union from fastapi import APIRouter, Depends, HTTPException +from pydantic import BaseModel import litellm from litellm._logging import verbose_proxy_logger @@ -407,6 +408,70 @@ async def update_cost_margin_config( ) +class BlockUnpricedModelsRequest(BaseModel): + enabled: bool + + +class BlockUnpricedModelsResponse(BaseModel): + enabled: bool + + +@router.get( + "/config/block_requests_for_models_without_pricing", + tags=["Cost Tracking"], + dependencies=[Depends(user_api_key_auth)], + response_model=BlockUnpricedModelsResponse, +) +async def get_block_requests_for_models_without_pricing() -> BlockUnpricedModelsResponse: + return BlockUnpricedModelsResponse(enabled=bool(litellm.block_requests_for_models_without_pricing)) + + +@router.patch( + "/config/block_requests_for_models_without_pricing", + tags=["Cost Tracking"], + dependencies=[Depends(user_api_key_auth)], + response_model=BlockUnpricedModelsResponse, +) +async def update_block_requests_for_models_without_pricing( + request: BlockUnpricedModelsRequest, +) -> BlockUnpricedModelsResponse: + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_config, + store_model_in_db, + ) + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + if store_model_in_db is not True: + raise HTTPException( + status_code=500, + detail={"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."}, + ) + + try: + config = await proxy_config.get_config() + if "litellm_settings" not in config: + config["litellm_settings"] = {} + config["litellm_settings"]["block_requests_for_models_without_pricing"] = request.enabled + await proxy_config.save_config(new_config=config) + + litellm.block_requests_for_models_without_pricing = request.enabled + verbose_proxy_logger.info(f"Updated block_requests_for_models_without_pricing: {request.enabled}") + + return BlockUnpricedModelsResponse(enabled=request.enabled) + except Exception as e: + verbose_proxy_logger.error(f"Error updating block_requests_for_models_without_pricing: {str(e)}") + raise HTTPException( + status_code=500, + detail={"error": f"Failed to update setting: {str(e)}"}, + ) + + @router.post( "/cost/estimate", tags=["Cost Tracking"], diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 34a353966bf..aec0ddc55f8 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -2,6 +2,7 @@ import asyncio import json import os import sys +from typing import Optional from unittest.mock import AsyncMock, MagicMock, patch sys.path.insert( @@ -5164,4 +5165,140 @@ async def test_get_project_object_db_fetch_returns_cached_obj(): assert isinstance(result, LiteLLM_ProjectTableCachedObj) assert result.project_id == "p-1" - assert result.project_alias == "proj" + + +UNPRICED_UNDERLYING_MODEL = "openai/unpriced-model-lit4984-xyz" + + +def _router_with_priced_and_unpriced_models() -> "Router": + from litellm.router import Router + + return Router( + model_list=[ + { + "model_name": "priced-group", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-test"}, + }, + { + "model_name": "unpriced-group", + "litellm_params": {"model": UNPRICED_UNDERLYING_MODEL, "api_key": "sk-test"}, + }, + ] + ) + + +def test_model_has_no_cost_mapping_priced_model_is_false(): + from litellm.proxy.auth.auth_checks import model_has_no_cost_mapping + + router = _router_with_priced_and_unpriced_models() + + assert model_has_no_cost_mapping(model="priced-group", llm_router=router) is False + + +def test_model_has_no_cost_mapping_unpriced_model_is_true(): + from litellm.proxy.auth.auth_checks import model_has_no_cost_mapping + + router = _router_with_priced_and_unpriced_models() + + assert model_has_no_cost_mapping(model="unpriced-group", llm_router=router) is True + + +def test_model_has_no_cost_mapping_no_model_or_router_is_false(): + from litellm.proxy.auth.auth_checks import model_has_no_cost_mapping + + router = _router_with_priced_and_unpriced_models() + + assert model_has_no_cost_mapping(model=None, llm_router=router) is False + assert model_has_no_cost_mapping(model="unpriced-group", llm_router=None) is False + + +async def _run_common_checks( + model: Optional[str], llm_router: Optional["Router"], route: str = "/chat/completions" +) -> bool: + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + return await common_checks( + request_body={"model": model, "messages": [{"role": "user", "content": "hi"}]}, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route=route, + llm_router=llm_router, + proxy_logging_obj=MagicMock(), + valid_token=UserAPIKeyAuth(token="test-token"), + request=MagicMock(spec=Request), + ) + + +@pytest.mark.asyncio +async def test_common_checks_blocks_unpriced_model_when_enabled(monkeypatch): + monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", True) + router = _router_with_priced_and_unpriced_models() + + with pytest.raises(ProxyException) as exc_info: + await _run_common_checks(model="unpriced-group", llm_router=router) + + assert exc_info.value.code == "403" + assert exc_info.value.type == ProxyErrorTypes.model_cost_map_missing + assert exc_info.value.param == "model" + assert "unpriced-group" in exc_info.value.message + assert "pricing" in exc_info.value.message.lower() + + +@pytest.mark.asyncio +async def test_common_checks_allows_unpriced_model_when_disabled(monkeypatch): + monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", False) + router = _router_with_priced_and_unpriced_models() + + result = await _run_common_checks(model="unpriced-group", llm_router=router) + + assert result is True + + +@pytest.mark.asyncio +async def test_common_checks_allows_priced_model_when_enabled(monkeypatch): + monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", True) + router = _router_with_priced_and_unpriced_models() + + result = await _run_common_checks(model="priced-group", llm_router=router) + + assert result is True + + +@pytest.mark.asyncio +async def test_common_checks_ignores_non_llm_route_when_enabled(monkeypatch): + monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", True) + router = _router_with_priced_and_unpriced_models() + + result = await _run_common_checks( + model="unpriced-group", llm_router=router, route="/model/new" + ) + + assert result is True + + +@pytest.mark.asyncio +async def test_common_checks_blocks_alias_resolving_to_unpriced_model(monkeypatch): + from litellm.router import Router + + monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", True) + router = Router( + model_list=[ + { + "model_name": "billed-underlying-group", + "litellm_params": {"model": UNPRICED_UNDERLYING_MODEL, "api_key": "sk-test"}, + } + ], + model_group_alias={"public-alias": "billed-underlying-group"}, + ) + + with pytest.raises(ProxyException) as exc_info: + await _run_common_checks(model="public-alias", llm_router=router) + + assert exc_info.value.code == "403" + assert exc_info.value.type == ProxyErrorTypes.model_cost_map_missing + assert "public-alias" in exc_info.value.message diff --git a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py index bc463d5e75d..4fb90e9fb2d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py @@ -500,3 +500,59 @@ class TestResolveModelForCostLookup: assert resolved_model == "openai/gpt-4" assert provider is None + + +class TestBlockRequestsForModelsWithoutPricing: + """Test suite for the block_requests_for_models_without_pricing toggle endpoints""" + + @pytest.mark.asyncio + async def test_get_reflects_in_memory_flag(self): + with patch.object(litellm, "block_requests_for_models_without_pricing", True): + response = client.get( + "/config/block_requests_for_models_without_pricing", + headers={"Authorization": "Bearer sk-1234"}, + ) + + assert response.status_code == 200 + assert response.json() == {"enabled": True} + + @pytest.mark.asyncio + async def test_patch_persists_and_updates_flag(self): + mock_proxy_config = AsyncMock() + mock_proxy_config.get_config = AsyncMock(return_value={"litellm_settings": {}}) + mock_proxy_config.save_config = AsyncMock() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch.object(litellm, "block_requests_for_models_without_pricing", False), + ): + response = client.patch( + "/config/block_requests_for_models_without_pricing", + headers={"Authorization": "Bearer sk-1234"}, + json={"enabled": True}, + ) + + assert response.status_code == 200 + assert response.json() == {"enabled": True} + assert litellm.block_requests_for_models_without_pricing is True + + saved_config = mock_proxy_config.save_config.call_args.kwargs["new_config"] + assert saved_config["litellm_settings"]["block_requests_for_models_without_pricing"] is True + + @pytest.mark.asyncio + async def test_patch_requires_store_model_in_db(self): + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_config", AsyncMock()), + patch("litellm.proxy.proxy_server.store_model_in_db", False), + ): + response = client.patch( + "/config/block_requests_for_models_without_pricing", + headers={"Authorization": "Bearer sk-1234"}, + json={"enabled": True}, + ) + + assert response.status_code == 500 + assert "error" in response.json()["detail"] diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx index b32e7afd756..f6a3d487ade 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx @@ -12,7 +12,7 @@ import { TabPanels, TabPanel, } from "@tremor/react"; -import { Modal, Form } from "antd"; +import { Modal, Form, Switch } from "antd"; import { CostTrackingSettingsProps } from "./types"; import ProviderDiscountTable from "./provider_discount_table"; import AddProviderForm from "./add_provider_form"; @@ -24,6 +24,7 @@ import { DocsMenu } from "@/components/HelpLink"; import HowItWorks from "./how_it_works"; import { useDiscountConfig } from "./use_discount_config"; import { useMarginConfig } from "./use_margin_config"; +import { useBlockUnpricedConfig } from "./use_block_unpriced_config"; import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; const DOCS_LINKS = [ @@ -65,9 +66,16 @@ const CostTrackingSettings: React.FC = ({ userID, use handleMarginChange, } = useMarginConfig({ accessToken }); + const { + blockUnpriced, + isUpdating: isUpdatingBlockUnpriced, + fetchBlockUnpriced, + setBlockUnpriced, + } = useBlockUnpricedConfig({ accessToken }); + useEffect(() => { if (accessToken) { - Promise.all([fetchDiscountConfig(), fetchMarginConfig()]).finally(() => { + Promise.all([fetchDiscountConfig(), fetchMarginConfig(), fetchBlockUnpriced()]).finally(() => { setIsFetching(false); }); @@ -82,7 +90,7 @@ const CostTrackingSettings: React.FC = ({ userID, use }; loadModels(); } - }, [accessToken, fetchDiscountConfig, fetchMarginConfig]); + }, [accessToken, fetchDiscountConfig, fetchMarginConfig, fetchBlockUnpriced]); const handleAddProvider = async () => { const success = await addProvider(selectedProvider, newDiscount); @@ -293,7 +301,38 @@ const CostTrackingSettings: React.FC = ({ userID, use )} - {/* Accordion 3: Pricing Calculator - Available to all roles */} + {isProxyAdmin && ( + + +
+ Block Unpriced Models + + Reject requests for models that have no pricing in the cost map instead of logging them as $0 spend + +
+ + +
+
+
+ Block requests for models without pricing + + When enabled, a request whose resolved model has no cost mapping is rejected with a 403 so an + admin can add pricing for it. Off by default + +
+ setBlockUnpriced(checked)} + /> +
+
+ + + )} + + {/* Accordion 4: Pricing Calculator - Available to all roles */}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/use_block_unpriced_config.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/use_block_unpriced_config.ts new file mode 100644 index 00000000000..f4a110c7878 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/use_block_unpriced_config.ts @@ -0,0 +1,63 @@ +import { useState, useCallback } from "react"; +import { apiClient } from "@/components/networking"; +import NotificationsManager from "@/components/molecules/notifications_manager"; + +export interface UseBlockUnpricedConfigProps { + accessToken: string | null; +} + +export interface UseBlockUnpricedConfigReturn { + blockUnpriced: boolean; + isUpdating: boolean; + fetchBlockUnpriced: () => Promise; + setBlockUnpriced: (enabled: boolean) => Promise; +} + +interface BlockUnpricedResponse { + enabled: boolean; +} + +const ENDPOINT = "/config/block_requests_for_models_without_pricing"; + +export function useBlockUnpricedConfig({ accessToken }: UseBlockUnpricedConfigProps): UseBlockUnpricedConfigReturn { + const [blockUnpriced, setBlockUnpricedState] = useState(false); + const [isUpdating, setIsUpdating] = useState(false); + + const fetchBlockUnpriced = useCallback(async () => { + if (!accessToken) return; + try { + const data = await apiClient.get(ENDPOINT, { accessToken }); + setBlockUnpricedState(Boolean(data?.enabled)); + } catch (error) { + console.error("Error fetching block-unpriced-models setting:", error); + } + }, [accessToken]); + + const setBlockUnpriced = useCallback( + async (enabled: boolean) => { + if (!accessToken) return; + setIsUpdating(true); + try { + const data = await apiClient.patch(ENDPOINT, { accessToken, body: { enabled } }); + setBlockUnpricedState(Boolean(data?.enabled)); + NotificationsManager.success( + enabled + ? "Requests for models without pricing will now be blocked" + : "Requests for models without pricing are now allowed", + ); + } catch (error) { + console.error("Error updating block-unpriced-models setting:", error); + } finally { + setIsUpdating(false); + } + }, + [accessToken], + ); + + return { + blockUnpriced, + isUpdating, + fetchBlockUnpriced, + setBlockUnpriced, + }; +} diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index ed975c6be0a..5abe3b4d587 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -1844,6 +1844,24 @@ export interface paths { patch?: never; trace?: never; }; + "/config/block_requests_for_models_without_pricing": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** Get Block Requests For Models Without Pricing */ + get: operations["get_block_requests_for_models_without_pricing_config_block_requests_for_models_without_pricing_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + /** Update Block Requests For Models Without Pricing */ + patch: operations["update_block_requests_for_models_without_pricing_config_block_requests_for_models_without_pricing_patch"]; + trace?: never; + }; "/config/callback/delete": { parameters: { query?: never; @@ -21247,6 +21265,16 @@ export interface components { /** Team Id */ team_id: string; }; + /** BlockUnpricedModelsRequest */ + BlockUnpricedModelsRequest: { + /** Enabled */ + enabled: boolean; + }; + /** BlockUnpricedModelsResponse */ + BlockUnpricedModelsResponse: { + /** Enabled */ + enabled: boolean; + }; /** BlockUsers */ BlockUsers: { /** User Ids */ @@ -37097,6 +37125,59 @@ export interface operations { }; }; }; + get_block_requests_for_models_without_pricing_config_block_requests_for_models_without_pricing_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["BlockUnpricedModelsResponse"]; + }; + }; + }; + }; + update_block_requests_for_models_without_pricing_config_block_requests_for_models_without_pricing_patch: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["BlockUnpricedModelsRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["BlockUnpricedModelsResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; delete_callback_config_callback_delete_post: { parameters: { query?: never; From 84c41dcdc3e55c180255433538f6df8969a56297 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 31 Jul 2026 03:02:10 +0000 Subject: [PATCH 003/281] fix(proxy): treat non-token pricing as priced and propagate the unpriced-model toggle across workers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/proxy/auth/auth_checks.py | 58 +++++++++++++++++-- .../proxy/auth/test_auth_checks.py | 45 ++++++++++++++ .../test_cost_tracking_settings.py | 14 +++++ 4 files changed, 112 insertions(+), 6 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 1014b472c61..b0ff992931f 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1525,6 +1525,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES = [ "public_model_groups_links", "cost_discount_config", "cost_margin_config", + "block_requests_for_models_without_pricing", "budget_exceeded_throttle_percentage", # Every field editable from the Admin UI (proxy_server._GENERAL_SETTINGS_UI_LITELLM_FIELDS) # must be listed here so a DB write from one worker overrides the live litellm attribute on diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 11a45e78410..136b8d19c04 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -13,7 +13,18 @@ import asyncio import math import re import time -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Literal, + Mapping, + Optional, + Type, + Union, + cast, +) from fastapi import HTTPException, Request, status from pydantic import BaseModel @@ -274,17 +285,52 @@ def _is_cost_explicitly_configured(model: str, llm_router: "Router") -> bool: return False +def _has_positive_cost(value: object) -> bool: + if isinstance(value, bool): + return False + if isinstance(value, (int, float)): + return value > 0 + if isinstance(value, dict): + return any(_has_positive_cost(nested) for nested in value.values()) + return False + + +def _entry_has_priced_metric(entry: Mapping[str, object]) -> bool: + return any("cost_per" in key and _has_positive_cost(value) for key, value in entry.items()) + + +def _model_group_has_pricing(model: str, llm_router: "Router") -> bool: + """ + Check every deployment behind a model group for a positive price on any billed + metric (tokens, characters, seconds, pages, images, queries, ...), so models that + are billed by a non-token metric are not treated as unpriced. + """ + for deployment in llm_router.get_model_list(model_name=model) or []: + litellm_params = deployment.get("litellm_params") or {} + if _entry_has_priced_metric(litellm_params): + return True + + model_id = (deployment.get("model_info") or {}).get("id") + if model_id is None: + continue + + model_info = llm_router.get_deployment_model_info( + model_id=model_id, model_name=litellm_params.get("model") or "" + ) + if model_info is not None and _entry_has_priced_metric(model_info): + return True + + return False + + def model_has_no_cost_mapping(model: Optional[str], llm_router: Optional[Router]) -> bool: if not model or llm_router is None: return False - model_group_info = llm_router.get_model_group_info(model_group=model) - if model_group_info is None: + if llm_router.get_model_group_info(model_group=model) is None: return False - input_cost = model_group_info.input_cost_per_token or 0 - output_cost = model_group_info.output_cost_per_token or 0 - if input_cost > 0 or output_cost > 0: + if _model_group_has_pricing(model=model, llm_router=llm_router): return False return not _is_cost_explicitly_configured(model, llm_router) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index aec0ddc55f8..a078e041aa7 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5165,6 +5165,7 @@ async def test_get_project_object_db_fetch_returns_cached_obj(): assert isinstance(result, LiteLLM_ProjectTableCachedObj) assert result.project_id == "p-1" + assert result.project_alias == "proj" UNPRICED_UNDERLYING_MODEL = "openai/unpriced-model-lit4984-xyz" @@ -5212,6 +5213,50 @@ def test_model_has_no_cost_mapping_no_model_or_router_is_false(): assert model_has_no_cost_mapping(model="unpriced-group", llm_router=None) is False +@pytest.mark.parametrize( + "underlying_model", + [ + "azure/speech/azure-tts", + "mistral/mistral-ocr-latest", + "vertex_ai/imagen-3.0-generate-001", + ], +) +def test_model_has_no_cost_mapping_non_token_priced_model_is_false(underlying_model): + from litellm.proxy.auth.auth_checks import model_has_no_cost_mapping + from litellm.router import Router + + router = Router( + model_list=[ + { + "model_name": "non-token-priced-group", + "litellm_params": {"model": underlying_model, "api_key": "sk-test"}, + } + ] + ) + + assert model_has_no_cost_mapping(model="non-token-priced-group", llm_router=router) is False + + +def test_model_has_no_cost_mapping_non_token_price_from_litellm_params_is_false(): + from litellm.proxy.auth.auth_checks import model_has_no_cost_mapping + from litellm.router import Router + + router = Router( + model_list=[ + { + "model_name": "custom-tts", + "litellm_params": { + "model": UNPRICED_UNDERLYING_MODEL, + "api_key": "sk-test", + "input_cost_per_second": 0.0001, + }, + } + ] + ) + + assert model_has_no_cost_mapping(model="custom-tts", llm_router=router) is False + + async def _run_common_checks( model: Optional[str], llm_router: Optional["Router"], route: str = "/chat/completions" ) -> bool: diff --git a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py index 4fb90e9fb2d..1124a4a31d4 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py @@ -541,6 +541,20 @@ class TestBlockRequestsForModelsWithoutPricing: saved_config = mock_proxy_config.save_config.call_args.kwargs["new_config"] assert saved_config["litellm_settings"]["block_requests_for_models_without_pricing"] is True + def test_peer_workers_pick_up_persisted_flag_on_config_reload(self): + """A PATCH only mutates the flag on the worker that served it; peer workers must pick the + persisted value up when they reload litellm_settings from the DB.""" + from litellm.proxy.proxy_server import ProxyConfig + + with patch.object(litellm, "block_requests_for_models_without_pricing", False): + ProxyConfig()._update_config_fields( + current_config={}, + param_name="litellm_settings", + db_param_value={"block_requests_for_models_without_pricing": True}, + ) + + assert litellm.block_requests_for_models_without_pricing is True + @pytest.mark.asyncio async def test_patch_requires_store_model_in_db(self): with ( From 074b37b4f987f239725a165565be16ff5e1f9686 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 31 Jul 2026 03:08:29 +0000 Subject: [PATCH 004/281] refactor(proxy): flatten the pricing-metric check to avoid recursion Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 136b8d19c04..3639ef245cf 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -285,18 +285,19 @@ def _is_cost_explicitly_configured(model: str, llm_router: "Router") -> bool: return False -def _has_positive_cost(value: object) -> bool: - if isinstance(value, bool): - return False - if isinstance(value, (int, float)): - return value > 0 - if isinstance(value, dict): - return any(_has_positive_cost(nested) for nested in value.values()) - return False +def _is_positive_cost(value: object) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0 def _entry_has_priced_metric(entry: Mapping[str, object]) -> bool: - return any("cost_per" in key and _has_positive_cost(value) for key, value in entry.items()) + for key, value in entry.items(): + if "cost_per" not in key: + continue + if _is_positive_cost(value): + return True + if isinstance(value, dict) and any(_is_positive_cost(nested) for nested in value.values()): + return True + return False def _model_group_has_pricing(model: str, llm_router: "Router") -> bool: From 658c67c1523a54d8d2e51ad848af4df3edd40b40 Mon Sep 17 00:00:00 2001 From: Shifat Islam Santo Date: Fri, 14 Aug 2026 14:19:48 -0500 Subject: [PATCH 005/281] fix: preserve prompt cache for mid-conversation system on unflagged Claude models --- .../messages/transformation.py | 56 +++++++++++++------ ...st_messages_mid_conversation_system_e2e.py | 6 +- ...onversation_system_native_providers_e2e.py | 16 +++--- ...azure_anthropic_messages_transformation.py | 20 +++++-- .../test_anthropic_claude3_transformation.py | 55 +++++++++++++----- ...artner_models_anthropic_messages_config.py | 29 +++++++++- 6 files changed, 130 insertions(+), 52 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 4d3354c58b7..1e016cffb0e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -159,8 +159,22 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): def _is_system_role_message(message: Any) -> bool: return isinstance(message, dict) and message.get("role") == "system" + _CONVERTED_SYSTEM_NOTE: Final = ( + "Operator note (not from the user): the following was originally a mid-conversation system-role reminder." + ) + + def _system_role_message_as_user(self, message: dict) -> dict: + return { + **message, + "role": "user", + "content": [ + {"type": "text", "text": self._CONVERTED_SYSTEM_NOTE}, + *self._as_system_content_blocks(message.get("content")), + ], + } + def _normalize_system_role_messages(self, anthropic_messages_request: dict, model: str) -> None: - """Move ``role: "system"`` entries out of ``messages`` per the Anthropic + """Normalize ``role: "system"`` entries in ``messages`` per the Anthropic ``/v1/messages`` contract, which the first-party API, Bedrock Invoke, Vertex, and Azure Foundry all enforce identically. @@ -173,9 +187,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): stay: hoisting one mutates the ``system`` prefix and invalidates the prompt cache for the whole message history. Older Claude models reject the role in every position ("role 'system' is not supported on this model"), - so without the flag every system entry is hoisted to keep the request from - 400-ing. Billing-header system blocks are stripped from the top-level - ``system`` field regardless of whether anything was hoisted. + so without the flag a mid-conversation entry is converted to a user turn + in place (prefixed with an operator note) rather than hoisted: hoisting + would mutate the ``system`` prefix and likewise collapse the cache, while + the in-place conversion keeps everything before it byte-identical. + Billing-header system blocks are stripped from the top-level ``system`` + field regardless of whether anything was hoisted. Subclasses whose upstream rejects the role opt in by calling this from their ``transform_anthropic_messages_request``; the first-party Anthropic @@ -185,21 +202,24 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): messages: Final = anthropic_messages_request.get("messages") if not isinstance(messages, list): return - if _supports_factory( - model=model, - custom_llm_provider=self.custom_llm_provider, - key="supports_mid_conversation_system", - ): - leading_count: Final = next( - (i for i, m in enumerate(messages) if not self._is_system_role_message(m)), - len(messages), + leading_count: Final = next( + (i for i, m in enumerate(messages) if not self._is_system_role_message(m)), + len(messages), + ) + hoisted: Final = messages[:leading_count] + remaining: Final = ( + messages[leading_count:] + if _supports_factory( + model=model, + custom_llm_provider=self.custom_llm_provider, + key="supports_mid_conversation_system", ) - hoisted = messages[:leading_count] - remaining = messages[leading_count:] - else: - hoisted = [m for m in messages if self._is_system_role_message(m)] - remaining = [m for m in messages if not self._is_system_role_message(m)] - if hoisted: + else [ + self._system_role_message_as_user(m) if self._is_system_role_message(m) else m + for m in messages[leading_count:] + ] + ) + if hoisted or remaining != messages: anthropic_messages_request["messages"] = remaining system_content: Final = [ block diff --git a/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py b/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py index fff2109b0cf..e1d836108e0 100644 --- a/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py +++ b/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py @@ -6,7 +6,7 @@ and the 5 family) must keep a mid-conversation system reminder in place inside ``messages`` so the top-level ``system`` prefix stays byte-identical and the prompt cache written on turn one is read back in full on turn two. Models without the flag (Claude 4.7 and older) reject the role inside ``messages`` -outright, so the proxy must hoist the reminder into the top-level ``system`` +outright, so the proxy must convert the reminder to a user turn in place field and the call must still return a completion instead of a provider 400. The conversation shape mirrors what Claude Code sends mid-session: a cached @@ -201,7 +201,7 @@ class TestBedrockInvokeMidConversationSystem: "llm.messages.bedrock_invoke.mid_conversation_system.nonstream.works", exercised_on=[], ) - def test_unflagged_model_hoists_system_reminder_and_succeeds( + def test_unflagged_model_converts_system_reminder_and_succeeds( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: model = _register_invoke_deployment( @@ -227,5 +227,5 @@ class TestBedrockInvokeMidConversationSystem: assert completion.text.strip(), ( f"{model}: conversation with a mid-conversation system reminder " f"returned no text; the reminder was forwarded in place to a model " - f"that rejects role 'system' inside messages instead of being hoisted" + f"that rejects role 'system' inside messages instead of being converted to a user turn" ) diff --git a/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py b/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py index 35ed3dc881a..6eae0d66c59 100644 --- a/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py +++ b/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py @@ -7,13 +7,13 @@ accepted in place on Claude 4.8+/5 (200) but rejected on Claude 4.7 and older ("role 'system' is not supported on this model", 400), and a *leading* system entry is rejected on every model ("messages.0: use the top-level 'system' parameter"). This mirrors Bedrock Invoke (PRs #32578/#32831/#32882); the same -model-gated hoist now runs for these two providers (customer RCA gap #3). +model-gated normalization now runs for these two providers (customer RCA gap #3). Flagged models (``supports_mid_conversation_system`` in the cost map: Claude 4.8+ and the 5 family) must keep the reminder in ``messages`` so the top-level ``system`` prefix stays byte-identical and the prompt cache written on turn one is read back in full on turn two. Unflagged models (Claude 4.7 and older) must -have the reminder hoisted into the top-level ``system`` field so the call +have the reminder converted to a user turn in place so the call returns a completion instead of a provider 400. The conversation shape mirrors what Claude Code sends mid-session: a cached @@ -206,7 +206,7 @@ def _assert_flagged_model_keeps_cache( ) -def _assert_unflagged_model_hoists_and_succeeds( +def _assert_unflagged_model_converts_and_succeeds( client: EndpointsClient, resources: ResourceManager, params: LiteLLMParamsBody ) -> None: model = _register_deployment(client, resources, params) @@ -228,7 +228,7 @@ def _assert_unflagged_model_hoists_and_succeeds( assert completion.text.strip(), ( f"{model}: conversation with a mid-conversation system reminder returned " f"no text; the reminder was forwarded in place to a model that rejects " - f"role 'system' inside messages instead of being hoisted" + f"role 'system' inside messages instead of being converted to a user turn" ) @@ -250,10 +250,10 @@ class TestAzureFoundryMidConversationSystem: "llm.messages.azure_foundry.mid_conversation_system.nonstream.works", exercised_on=[], ) - def test_unflagged_model_hoists_system_reminder_and_succeeds( + def test_unflagged_model_converts_system_reminder_and_succeeds( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: - _assert_unflagged_model_hoists_and_succeeds( + _assert_unflagged_model_converts_and_succeeds( endpoints_client, resources, _azure_params(self.UNFLAGGED_MODEL) ) @@ -276,9 +276,9 @@ class TestVertexMidConversationSystem: "llm.messages.vertex.mid_conversation_system.nonstream.works", exercised_on=[], ) - def test_unflagged_model_hoists_system_reminder_and_succeeds( + def test_unflagged_model_converts_system_reminder_and_succeeds( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: - _assert_unflagged_model_hoists_and_succeeds( + _assert_unflagged_model_converts_and_succeeds( endpoints_client, resources, _vertex_params(self.UNFLAGGED_MODEL) ) diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py index f6446b43fab..add1e9967db 100644 --- a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py @@ -425,7 +425,7 @@ class TestAzureAnthropicMidConversationSystem: {"type": "text", "text": "Cite sources."}, ] - def test_unsupported_model_hoists_mid_conversation_system(self, local_model_cost_map): + def test_unsupported_model_converts_mid_conversation_system_in_place(self, local_model_cost_map): messages = [ {"role": "user", "content": "read the file"}, {"role": "system", "content": "[Truncated: PARTIAL view of big1.txt]"}, @@ -437,13 +437,23 @@ class TestAzureAnthropicMidConversationSystem: ) assert result["messages"] == [ {"role": "user", "content": "read the file"}, + { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + "Operator note (not from the user): the following was " + "originally a mid-conversation system-role reminder." + ), + }, + {"type": "text", "text": "[Truncated: PARTIAL view of big1.txt]"}, + ], + }, {"role": "assistant", "content": "reading"}, {"role": "user", "content": "continue"}, ] - assert result["system"] == [ - {"type": "text", "text": "Base."}, - {"type": "text", "text": "[Truncated: PARTIAL view of big1.txt]"}, - ] + assert result["system"] == [{"type": "text", "text": "Base."}] def test_azure_claude_4_8_plus_cost_map_entries_carry_mid_conversation_system_flag(): diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index fd66667af64..84df706b722 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -2125,13 +2125,14 @@ def test_bedrock_invoke_transform_hoists_only_leading_system_run(local_model_cos ] -def test_bedrock_invoke_transform_hoists_mid_conversation_system_for_older_claude(local_model_cost_map): - """Regression test for Claude Code 400s on pre-Opus-4.8 Bedrock models: - Invoke rejects ``role: "system"`` in every position on Opus 4.7, Sonnet 4.6, - Haiku 4.5, etc. ("role 'system' is not supported on this model"), so on - models without ``supports_mid_conversation_system`` every system entry must - be hoisted into the top-level ``system`` field, mid-conversation ones - included.""" +def test_bedrock_invoke_transform_converts_mid_conversation_system_for_older_claude(local_model_cost_map): + """Invoke rejects ``role: "system"`` in every position on Opus 4.7, Sonnet + 4.6, Haiku 4.5, etc. ("role 'system' is not supported on this model"), but + hoisting a mid-conversation reminder into the top-level ``system`` field + mutates the cached prefix and reprocesses the whole history. On models + without ``supports_mid_conversation_system`` the reminder is converted to a + user turn in place instead: the request stays valid and a cache breakpoint + before the reminder still hits.""" from litellm.types.router import GenericLiteLLMParams cfg = AmazonAnthropicClaudeMessagesConfig() @@ -2156,19 +2157,30 @@ def test_bedrock_invoke_transform_hoists_mid_conversation_system_for_older_claud assert result["messages"] == [ {"role": "user", "content": "read the file"}, + { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + "Operator note (not from the user): the following was " + "originally a mid-conversation system-role reminder." + ), + }, + {"type": "text", "text": "[Truncated: PARTIAL view of big1.txt]"}, + ], + }, {"role": "assistant", "content": "reading"}, {"role": "user", "content": "continue"}, ] - assert result["system"] == [ - {"type": "text", "text": "Base."}, - {"type": "text", "text": "[Truncated: PARTIAL view of big1.txt]"}, - ] + assert result["system"] == [{"type": "text", "text": "Base."}] -def test_bedrock_invoke_transform_hoists_all_system_for_unmapped_model(local_model_cost_map): +def test_bedrock_invoke_transform_converts_system_for_unmapped_model(local_model_cost_map): """A model with no cost-map entry and no fallback-generalization rule gets - the hoist-everything behavior: the safe default is a mutated cache prefix, - never a provider 400 from forwarding a role the model may not accept.""" + the unsupported-model treatment: the safe default converts the reminder to + a user turn in place, never a provider 400 from forwarding a role the model + may not accept, and never a mutated cache prefix.""" from litellm.types.router import GenericLiteLLMParams cfg = AmazonAnthropicClaudeMessagesConfig() @@ -2189,10 +2201,23 @@ def test_bedrock_invoke_transform_hoists_all_system_for_unmapped_model(local_mod assert result["messages"] == [ {"role": "user", "content": "hi"}, + { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + "Operator note (not from the user): the following was " + "originally a mid-conversation system-role reminder." + ), + }, + {"type": "text", "text": "mid-conversation reminder"}, + ], + }, {"role": "assistant", "content": "hello"}, {"role": "user", "content": "continue"}, ] - assert result["system"] == [{"type": "text", "text": "mid-conversation reminder"}] + assert "system" not in result def test_bedrock_invoke_transform_keeps_system_in_place_for_unmapped_future_claude(local_model_cost_map): diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index 292bddf1274..ef7db337a74 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -622,7 +622,7 @@ class TestVertexAnthropicMidConversationSystem: {"type": "text", "text": "Cite sources."}, ] - def test_unsupported_model_hoists_mid_conversation_system(self, local_model_cost_map): + def test_unsupported_model_converts_mid_conversation_system_in_place(self, local_model_cost_map): messages = [ {"role": "user", "content": "read the file"}, {"role": "system", "content": "[Truncated: PARTIAL view of big1.txt]"}, @@ -634,12 +634,35 @@ class TestVertexAnthropicMidConversationSystem: ) assert result["messages"] == [ {"role": "user", "content": "read the file"}, + { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + "Operator note (not from the user): the following was " + "originally a mid-conversation system-role reminder." + ), + }, + {"type": "text", "text": "[Truncated: PARTIAL view of big1.txt]"}, + ], + }, {"role": "assistant", "content": "reading"}, {"role": "user", "content": "continue"}, ] + assert result["system"] == [{"type": "text", "text": "Base."}] + + def test_unsupported_model_still_hoists_leading_system_run(self, local_model_cost_map): + messages = [ + {"role": "system", "content": "You are terse."}, + {"role": "system", "content": "Cite sources."}, + {"role": "user", "content": "hi"}, + ] + result = _vertex_transform("claude-sonnet-4-6", messages) + assert result["messages"] == [{"role": "user", "content": "hi"}] assert result["system"] == [ - {"type": "text", "text": "Base."}, - {"type": "text", "text": "[Truncated: PARTIAL view of big1.txt]"}, + {"type": "text", "text": "You are terse."}, + {"type": "text", "text": "Cite sources."}, ] From 1b5e50727c62cdec0c06ab7c1f54006ad8c8f8a9 Mon Sep 17 00:00:00 2001 From: Shifat Islam Santo Date: Fri, 14 Aug 2026 14:40:10 -0500 Subject: [PATCH 006/281] fix: reuse block builder for lint budget, assert unflagged cache e2e --- .../messages/transformation.py | 6 ++-- ...st_messages_mid_conversation_system_e2e.py | 30 +++++++++++++------ ...onversation_system_native_providers_e2e.py | 28 ++++++++++++----- 3 files changed, 43 insertions(+), 21 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 1e016cffb0e..84fb7a1f45e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -167,10 +167,8 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): return { **message, "role": "user", - "content": [ - {"type": "text", "text": self._CONVERTED_SYSTEM_NOTE}, - *self._as_system_content_blocks(message.get("content")), - ], + "content": self._as_system_content_blocks(self._CONVERTED_SYSTEM_NOTE) + + self._as_system_content_blocks(message.get("content")), } def _normalize_system_role_messages(self, anthropic_messages_request: dict, model: str) -> None: diff --git a/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py b/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py index e1d836108e0..bd59df959c4 100644 --- a/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py +++ b/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py @@ -208,24 +208,36 @@ class TestBedrockInvokeMidConversationSystem: endpoints_client, resources, UNFLAGGED_INVOKE_MODEL ) key = resources.key(models=[model]) + system_block = _cacheable_system_block(unique_marker()) - body = RichMessagesRequest( + primed = _prime_prompt_cache(endpoints_client, key, model, system_block) + + reminder_turn_body = RichMessagesRequest( model=model, - system=[TextBlock(text="You are terse.")], + system=[system_block], messages=[ - _user_turn(f"Say hi. Run {unique_marker()}."), + _user_turn(primed.first_user_text, cached=True), _system_reminder_turn(), - RichMessage(role="assistant", content=[TextBlock(text="Hi.")]), - _user_turn("Say bye."), + RichMessage(role="assistant", content=[TextBlock(text="OK.")]), + _user_turn("Reply with one word again.", cached=True), ], ) - completion = unwrap(_post_messages(endpoints_client, key, body)) + second = unwrap(_post_messages(endpoints_client, key, reminder_turn_body)) - assert completion.role == "assistant", ( - f"{model}: unexpected role {completion.role!r}" + assert second.role == "assistant", ( + f"{model}: unexpected role {second.role!r}" ) - assert completion.text.strip(), ( + assert second.text.strip(), ( f"{model}: conversation with a mid-conversation system reminder " f"returned no text; the reminder was forwarded in place to a model " f"that rejects role 'system' inside messages instead of being converted to a user turn" ) + assert second.usage.cache_read_input_tokens >= primed.full_prefix_tokens, ( + f"{model}: reminder turn read {second.usage.cache_read_input_tokens} " + f"cached tokens, expected at least the {primed.full_prefix_tokens} " + f"cached on turn one ({primed.prefix_read_tokens} system prefix + " + f"{primed.first_turn_creation_tokens} first user turn); the reminder " + f"was hoisted into the top-level system field instead of being " + f"converted to a user turn in place, mutating the cached prefix and " + f"re-billing the conversation at cache-write pricing" + ) diff --git a/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py b/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py index 6eae0d66c59..bf4e662e02e 100644 --- a/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py +++ b/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py @@ -211,25 +211,37 @@ def _assert_unflagged_model_converts_and_succeeds( ) -> None: model = _register_deployment(client, resources, params) key = resources.key(models=[model]) + system_block = _cacheable_system_block(unique_marker()) - body = RichMessagesRequest( + primed = _prime_prompt_cache(client, key, model, system_block) + + reminder_turn_body = RichMessagesRequest( model=model, - system=[TextBlock(text="You are terse.")], + system=[system_block], messages=[ - _user_turn(f"Say hi. Run {unique_marker()}."), + _user_turn(primed.first_user_text, cached=True), _system_reminder_turn(), - RichMessage(role="assistant", content=[TextBlock(text="Hi.")]), - _user_turn("Say bye."), + RichMessage(role="assistant", content=[TextBlock(text="OK.")]), + _user_turn("Reply with one word again.", cached=True), ], ) - completion = unwrap(_post_messages(client, key, body)) + second = unwrap(_post_messages(client, key, reminder_turn_body)) - assert completion.role == "assistant", f"{model}: unexpected role {completion.role!r}" - assert completion.text.strip(), ( + assert second.role == "assistant", f"{model}: unexpected role {second.role!r}" + assert second.text.strip(), ( f"{model}: conversation with a mid-conversation system reminder returned " f"no text; the reminder was forwarded in place to a model that rejects " f"role 'system' inside messages instead of being converted to a user turn" ) + assert second.usage.cache_read_input_tokens >= primed.full_prefix_tokens, ( + f"{model}: reminder turn read {second.usage.cache_read_input_tokens} cached " + f"tokens, expected at least the {primed.full_prefix_tokens} cached on turn " + f"one ({primed.prefix_read_tokens} system prefix + " + f"{primed.first_turn_creation_tokens} first user turn); the reminder was " + f"hoisted into the top-level system field instead of being converted to a " + f"user turn in place, mutating the cached prefix and re-billing the " + f"conversation at cache-write pricing" + ) class TestAzureFoundryMidConversationSystem: From 4cd1c81a2dcd905112a7bc388afb5ed1412b8aa7 Mon Sep 17 00:00:00 2001 From: Shifat Islam Santo Date: Fri, 14 Aug 2026 14:53:11 -0500 Subject: [PATCH 007/281] fix: add supports_mid_conversation_system to bare first-party Claude cost-map keys --- ...odel_prices_and_context_window_backup.json | 5 +++ model_prices_and_context_window.json | 5 +++ ...erimental_pass_through_messages_handler.py | 37 +++++++++++++++++++ 3 files changed, 47 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b288269b0a2..4c1b46f85ab 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -12106,6 +12106,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -12493,6 +12494,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -12528,6 +12530,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -12566,6 +12569,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -47684,6 +47688,7 @@ }, "source": "https://docs.claude.com/en/docs/about-claude/models/overview", "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b288269b0a2..4c1b46f85ab 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -12106,6 +12106,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -12493,6 +12494,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -12528,6 +12530,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -12566,6 +12569,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -47684,6 +47688,7 @@ }, "source": "https://docs.claude.com/en/docs/about-claude/models/overview", "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index f11324ca376..91f5023496a 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -960,3 +960,40 @@ def test_gate_passthrough_skipped_when_only_chat_completions_supported(monkeypat assert result == "translated" assert translation_calls["count"] == 1 assert "config" not in captured + + +def test_first_party_claude_4_8_plus_cost_map_entries_carry_mid_conversation_system_flag(): + """Regional and provider-prefixed Claude 4.8+/5 entries carry + ``supports_mid_conversation_system``, but the bare first-party keys + (``claude-opus-4-8``) that a plain ``custom_llm_provider="anthropic"`` + lookup resolves were missed, so that lookup reports the capability as + unset. Every mapped first-party entry the fallback rule matches must + carry the flag.""" + import json + import os + import re + + import litellm + + cost_map_path = os.path.join( + os.path.dirname(litellm.__file__), "model_prices_and_context_window_backup.json" + ) + with open(cost_map_path) as f: + cost_map = json.load(f) + rules = cost_map["fallback_generalizations"]["rules"] + rule_pattern = next( + (r["pattern"] for r in rules if r["name"] == "claude-mid-conversation-system"), + None, + ) + assert rule_pattern is not None, "claude-mid-conversation-system rule not found in fallback_generalizations" + pattern = re.compile(rule_pattern, re.IGNORECASE) + missing = [ + key + for key, info in cost_map.items() + if isinstance(info, dict) + and info.get("litellm_provider") == "anthropic" + and "claude" in key + and pattern.search(key) + and info.get("supports_mid_conversation_system") is not True + ] + assert missing == [] From 9c922f4aa4f6ab3a261326332d01b74e3717e63c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 15 Aug 2026 02:49:10 +0000 Subject: [PATCH 008/281] fix(bedrock): forward provider response headers on chat completions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/streaming_handler.py | 2 +- litellm/llms/bedrock/chat/converse_handler.py | 18 ++++-- litellm/llms/bedrock/chat/invoke_handler.py | 8 +-- .../base_invoke_transformation.py | 64 +++++++++---------- litellm/llms/custom_httpx/llm_http_handler.py | 2 + litellm/types/utils.py | 7 +- .../llm_translation/test_bedrock_moonshot.py | 19 ++---- .../llms/bedrock/chat/test_invoke_handler.py | 25 ++++++++ .../llms/chat/test_converse_handler.py | 55 ++++++++++++++++ 9 files changed, 143 insertions(+), 57 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 99b1c1a2ab7..fe107d349c6 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -164,7 +164,7 @@ class CustomStreamWrapper: custom_llm_provider: str | None = None, stream_options=None, make_call: Callable | None = None, - _response_headers: dict | None = None, + _response_headers: dict | httpx.Headers | None = None, ): self.model = model self.make_call = make_call diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 25e544f4521..c9ea06d37dc 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -33,7 +33,7 @@ def make_sync_call( json_mode: bool | None = False, fake_stream: bool = False, stream_chunk_size: int | None = None, -): +) -> tuple[Any, httpx.Headers]: if client is None: client = _get_httpx_client() # Create a new client if none provided @@ -74,7 +74,7 @@ def make_sync_call( additional_args={"complete_input_dict": data}, ) - return completion_stream + return completion_stream, response.headers class BedrockConverseLLM(BaseAWSLLM): @@ -132,7 +132,7 @@ class BedrockConverseLLM(BaseAWSLLM): }, ) - completion_stream: Final = await make_call( + completion_stream, response_headers = await make_call( client=client, api_base=api_base, headers=dict(prepped.headers), @@ -149,6 +149,7 @@ class BedrockConverseLLM(BaseAWSLLM): model=model, custom_llm_provider="bedrock", logging_obj=logging_obj, + _response_headers=response_headers, ) return streaming_response @@ -225,7 +226,7 @@ class BedrockConverseLLM(BaseAWSLLM): except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") - return litellm.AmazonConverseConfig()._transform_response( + transformed_response: Final = litellm.AmazonConverseConfig()._transform_response( model=model, response=response, model_response=model_response, @@ -237,6 +238,8 @@ class BedrockConverseLLM(BaseAWSLLM): optional_params=optional_params, encoding=encoding, ) + transformed_response.set_provider_response_headers(response.headers) + return transformed_response def completion( self, @@ -440,7 +443,7 @@ class BedrockConverseLLM(BaseAWSLLM): client = client if stream is not None and stream is True: - completion_stream: Final = make_sync_call( + completion_stream, response_headers = make_sync_call( client=(client if client is not None and isinstance(client, HTTPHandler) else None), api_base=proxy_endpoint_url, headers=prepped.headers, @@ -457,6 +460,7 @@ class BedrockConverseLLM(BaseAWSLLM): model=model, custom_llm_provider="bedrock", logging_obj=logging_obj, + _response_headers=response_headers, ) return streaming_response @@ -477,7 +481,7 @@ class BedrockConverseLLM(BaseAWSLLM): except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") - return litellm.AmazonConverseConfig()._transform_response( + sync_transformed_response: Final = litellm.AmazonConverseConfig()._transform_response( model=model, response=response, model_response=model_response, @@ -489,3 +493,5 @@ class BedrockConverseLLM(BaseAWSLLM): optional_params=optional_params, encoding=encoding, ) + sync_transformed_response.set_provider_response_headers(response.headers) + return sync_transformed_response diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 8d2b3dae71b..21892af3aae 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -163,7 +163,7 @@ async def make_call( json_mode: bool | None = False, bedrock_invoke_provider: litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None = None, stream_chunk_size: int | None = None, -): +) -> tuple[Any, httpx.Headers]: try: if client is None: client = get_async_httpx_client( @@ -225,7 +225,7 @@ async def make_call( additional_args={"complete_input_dict": data}, ) - return completion_stream + return completion_stream, response.headers except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code raise BedrockError(status_code=error_code, message=err.response.text) @@ -248,7 +248,7 @@ def make_sync_call( json_mode: bool | None = False, bedrock_invoke_provider: litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None = None, stream_chunk_size: int | None = None, -): +) -> tuple[Any, httpx.Headers]: try: if client is None: client = _get_httpx_client( @@ -309,7 +309,7 @@ def make_sync_call( additional_args={"complete_input_dict": data}, ) - return completion_stream + return completion_stream, response.headers except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code raise BedrockError(status_code=error_code, message=err.response.text) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 430d0a92b51..0b5855ccbd4 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -1,7 +1,6 @@ import copy import json import time -from functools import partial from typing import TYPE_CHECKING, Any, Final, cast, get_args import httpx @@ -444,24 +443,24 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): json_mode: bool | None = None, signed_json_body: bytes | None = None, ) -> CustomStreamWrapper: + completion_stream, response_headers = await make_call( + client=client, + api_base=api_base, + headers=headers, + data=json.dumps(data), + model=model, + messages=messages, + logging_obj=logging_obj, + fake_stream=True if "ai21" in api_base else False, + bedrock_invoke_provider=self.get_bedrock_invoke_provider(model), + json_mode=json_mode, + ) streaming_response: Final = CustomStreamWrapper( - completion_stream=None, - make_call=partial( - make_call, - client=client, - api_base=api_base, - headers=headers, - data=json.dumps(data), - model=model, - messages=messages, - logging_obj=logging_obj, - fake_stream=True if "ai21" in api_base else False, - bedrock_invoke_provider=self.get_bedrock_invoke_provider(model), - json_mode=json_mode, - ), + completion_stream=completion_stream, model=model, custom_llm_provider="bedrock", logging_obj=logging_obj, + _response_headers=response_headers, ) return streaming_response @@ -479,27 +478,28 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): json_mode: bool | None = None, signed_json_body: bytes | None = None, ) -> CustomStreamWrapper: - if client is None or isinstance(client, AsyncHTTPHandler): - client = _get_httpx_client(params={}) + sync_client: Final = ( + _get_httpx_client(params={}) if client is None or isinstance(client, AsyncHTTPHandler) else client + ) + completion_stream, response_headers = make_sync_call( + client=sync_client, + api_base=api_base, + headers=headers, + data=json.dumps(data), + signed_json_body=signed_json_body, + model=model, + messages=messages, + logging_obj=logging_obj, + fake_stream=True if "ai21" in api_base else False, + bedrock_invoke_provider=self.get_bedrock_invoke_provider(model), + json_mode=json_mode, + ) streaming_response: Final = CustomStreamWrapper( - completion_stream=None, - make_call=partial( - make_sync_call, - client=client, - api_base=api_base, - headers=headers, - data=json.dumps(data), - signed_json_body=signed_json_body, - model=model, - messages=messages, - logging_obj=logging_obj, - fake_stream=True if "ai21" in api_base else False, - bedrock_invoke_provider=self.get_bedrock_invoke_provider(model), - json_mode=json_mode, - ), + completion_stream=completion_stream, model=model, custom_llm_provider="bedrock", logging_obj=logging_obj, + _response_headers=response_headers, ) return streaming_response diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 721b9545ac1..36bcf58a243 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -608,6 +608,7 @@ class BaseLLMHTTPHandler: model=model, custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, + _response_headers=headers, ) if client is None or not isinstance(client, HTTPHandler): @@ -771,6 +772,7 @@ class BaseLLMHTTPHandler: model=model, custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, + _response_headers=_response_headers, ) return streamwrapper diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 220826ccbca..f49b5735f04 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -11,6 +11,7 @@ from typing import ( get_args, ) +import httpx from openai._models import BaseModel as OpenAIObject from openai.types.audio.transcription_create_params import ( FileTypes as FileTypes, @@ -49,7 +50,7 @@ from litellm.types.llms.base import ( ) from litellm.types.mcp import MCPServerCostInfo -from ..litellm_core_utils.core_helpers import map_finish_reason +from ..litellm_core_utils.core_helpers import map_finish_reason, process_response_headers from .agents import LiteLLMSendMessageResponse from .guardrails import GuardrailEventHooks from .llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse @@ -1896,6 +1897,10 @@ class ModelResponseBase(OpenAIObject): _response_headers: dict | None = None + def set_provider_response_headers(self, headers: httpx.Headers) -> None: + """Surface a provider's raw response headers to the caller as `llm_provider-*` headers.""" + self._hidden_params["additional_headers"] = process_response_headers(headers) + def model_dump(self, **kwargs): """Default to exclude_unset to avoid Pydantic serializer warnings for OpenAIObject-derived types.""" if "exclude_unset" not in kwargs and "exclude_none" not in kwargs: diff --git a/tests/llm_translation/test_bedrock_moonshot.py b/tests/llm_translation/test_bedrock_moonshot.py index a9f4a86b3b6..c777305d562 100644 --- a/tests/llm_translation/test_bedrock_moonshot.py +++ b/tests/llm_translation/test_bedrock_moonshot.py @@ -209,14 +209,12 @@ class TestBedrockMoonshotInvoke(BaseLLMChatTest): not exercised here — moonshot streaming delegates to the OpenAI parser and is covered by the OpenAI test suite. - Note: bedrock invoke streaming cannot be intercepted by patching - the caller-supplied client, because ``CustomStreamWrapper.fetch_sync_stream`` - at streaming_handler.py invokes the stored ``make_call`` partial with - ``client=litellm.module_level_client``, which overrides any client the - caller passed. Patch ``make_sync_call`` at its import site in - ``base_invoke_transformation`` so we observe the exact kwargs the - partial was built with at stream-wrapper construction time. + Patch ``make_sync_call`` at its import site in + ``base_invoke_transformation`` so we observe the exact kwargs it is + called with at stream-wrapper construction time. """ + import httpx + from litellm.utils import CustomStreamWrapper captured: dict = {} @@ -225,7 +223,7 @@ class TestBedrockMoonshotInvoke(BaseLLMChatTest): captured.update(kwargs) # Return an empty iterator so the stream wrapper's iteration # doesn't try to parse real bytes. - return iter([]) + return iter([]), httpx.Headers() with patch( "litellm.llms.bedrock.chat.invoke_transformations." @@ -246,11 +244,6 @@ class TestBedrockMoonshotInvoke(BaseLLMChatTest): aws_region_name="us-west-2", ) assert isinstance(response, CustomStreamWrapper) - # Trigger fetch_sync_stream → make_call(...) → fake_make_sync_call. - try: - next(iter(response)) - except StopIteration: - pass assert captured, "make_sync_call was never invoked" assert captured["api_base"].endswith("/invoke-with-response-stream") diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py index ee50b9db015..f211cdac475 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -2,17 +2,20 @@ import os import sys from unittest.mock import AsyncMock, MagicMock +import httpx import pytest sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path +import litellm from litellm.llms.bedrock.chat.invoke_handler import ( AWSEventStreamDecoder, make_call, make_sync_call, ) +from litellm.llms.custom_httpx.http_handler import HTTPHandler def test_transform_thinking_blocks_with_redacted_content(): @@ -293,3 +296,25 @@ def test_make_sync_call_honors_explicit_stream_chunk_size(): response.iter_bytes.assert_called_once_with(chunk_size=2048) + +def test_invoke_streaming_forwards_bedrock_response_headers(): + """Streaming callers need `x-amzn-requestid` to correlate a LiteLLM request with AWS support.""" + response = MagicMock() + response.status_code = 200 + response.iter_bytes = MagicMock(return_value=iter([])) + response.headers = httpx.Headers({"x-amzn-requestid": "req-789"}) + client = HTTPHandler() + client.post = MagicMock(return_value=response) + + stream = litellm.completion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert stream._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-789" + diff --git a/tests/test_litellm/llms/chat/test_converse_handler.py b/tests/test_litellm/llms/chat/test_converse_handler.py index 2a3db5982ef..89ab29122d7 100644 --- a/tests/test_litellm/llms/chat/test_converse_handler.py +++ b/tests/test_litellm/llms/chat/test_converse_handler.py @@ -1,7 +1,9 @@ +import json import os import sys from unittest.mock import MagicMock +import httpx import pytest import litellm @@ -202,6 +204,59 @@ def test_make_sync_call_honors_explicit_stream_chunk_size(): response.iter_bytes.assert_called_once_with(chunk_size=2048) +def _converse_response_body() -> dict: + return { + "output": {"message": {"role": "assistant", "content": [{"text": "hi"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + } + + +def test_converse_completion_forwards_bedrock_response_headers(): + """Bedrock returns x-amzn-requestid on every converse call, which customers need to + correlate proxy requests with AWS support cases, so it must reach the caller as + llm_provider-x-amzn-requestid.""" + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json = MagicMock(return_value=_converse_response_body()) + mock_response.text = json.dumps(_converse_response_body()) + mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-123"}) + client = HTTPHandler() + client.post = MagicMock(return_value=mock_response) + + response = litellm.completion( + model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-123" + + +def test_converse_streaming_forwards_bedrock_response_headers(): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.iter_bytes = MagicMock(return_value=iter([])) + mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-456"}) + client = HTTPHandler() + client.post = MagicMock(return_value=mock_response) + + response = litellm.completion( + model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-456" + + def test_completion_plumbs_stream_chunk_size_through_converse(): iter_bytes_spy = _stream_completion_with_spied_iter_bytes( model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" From 726db1a4c118a2ab92214ba741efacdce7fe25b5 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 15 Aug 2026 03:07:14 +0000 Subject: [PATCH 009/281] test(bedrock): cover async header forwarding for converse and invoke Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_translation/test_bedrock_moonshot.py | 3 +- .../llms/bedrock/chat/test_invoke_handler.py | 29 +++++++++- .../llms/chat/test_converse_handler.py | 55 +++++++++++++++++-- 3 files changed, 78 insertions(+), 9 deletions(-) diff --git a/tests/llm_translation/test_bedrock_moonshot.py b/tests/llm_translation/test_bedrock_moonshot.py index c777305d562..61364cf2caa 100644 --- a/tests/llm_translation/test_bedrock_moonshot.py +++ b/tests/llm_translation/test_bedrock_moonshot.py @@ -12,6 +12,7 @@ This test suite verifies: """ from base_llm_unit_tests import BaseLLMChatTest +import httpx import pytest import sys import os @@ -213,8 +214,6 @@ class TestBedrockMoonshotInvoke(BaseLLMChatTest): ``base_invoke_transformation`` so we observe the exact kwargs it is called with at stream-wrapper construction time. """ - import httpx - from litellm.utils import CustomStreamWrapper captured: dict = {} diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py index f211cdac475..e8964910c69 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -15,7 +15,7 @@ from litellm.llms.bedrock.chat.invoke_handler import ( make_call, make_sync_call, ) -from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler def test_transform_thinking_blocks_with_redacted_content(): @@ -298,7 +298,6 @@ def test_make_sync_call_honors_explicit_stream_chunk_size(): def test_invoke_streaming_forwards_bedrock_response_headers(): - """Streaming callers need `x-amzn-requestid` to correlate a LiteLLM request with AWS support.""" response = MagicMock() response.status_code = 200 response.iter_bytes = MagicMock(return_value=iter([])) @@ -318,3 +317,29 @@ def test_invoke_streaming_forwards_bedrock_response_headers(): assert stream._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-789" + +@pytest.mark.asyncio +async def test_async_invoke_streaming_forwards_bedrock_response_headers(): + async def _no_bytes(chunk_size=None): + return + yield b"" + + response = MagicMock() + response.status_code = 200 + response.aiter_bytes = _no_bytes + response.headers = httpx.Headers({"x-amzn-requestid": "req-987"}) + client = AsyncHTTPHandler() + client.post = AsyncMock(return_value=response) + + stream = await litellm.acompletion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert stream._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-987" + diff --git a/tests/test_litellm/llms/chat/test_converse_handler.py b/tests/test_litellm/llms/chat/test_converse_handler.py index 89ab29122d7..6f8a2788c38 100644 --- a/tests/test_litellm/llms/chat/test_converse_handler.py +++ b/tests/test_litellm/llms/chat/test_converse_handler.py @@ -1,7 +1,7 @@ import json import os import sys -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import httpx import pytest @@ -10,7 +10,7 @@ import litellm from litellm.llms.bedrock.chat import BedrockConverseLLM from litellm.llms.bedrock.chat.converse_handler import make_sync_call from litellm.llms.bedrock.common_utils import _get_all_bedrock_regions -from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler sys.path.insert( 0, os.path.abspath("../../../../..") @@ -213,9 +213,6 @@ def _converse_response_body() -> dict: def test_converse_completion_forwards_bedrock_response_headers(): - """Bedrock returns x-amzn-requestid on every converse call, which customers need to - correlate proxy requests with AWS support cases, so it must reach the caller as - llm_provider-x-amzn-requestid.""" mock_response = MagicMock() mock_response.status_code = 200 mock_response.json = MagicMock(return_value=_converse_response_body()) @@ -257,6 +254,54 @@ def test_converse_streaming_forwards_bedrock_response_headers(): assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-456" +@pytest.mark.asyncio +async def test_async_converse_completion_forwards_bedrock_response_headers(): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json = MagicMock(return_value=_converse_response_body()) + mock_response.text = json.dumps(_converse_response_body()) + mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-abc"}) + client = AsyncHTTPHandler() + client.post = AsyncMock(return_value=mock_response) + + response = await litellm.acompletion( + model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-abc" + + +@pytest.mark.asyncio +async def test_async_converse_streaming_forwards_bedrock_response_headers(): + async def _no_bytes(chunk_size=None): + return + yield b"" + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.aiter_bytes = _no_bytes + mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-def"}) + client = AsyncHTTPHandler() + client.post = AsyncMock(return_value=mock_response) + + response = await litellm.acompletion( + model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-def" + + def test_completion_plumbs_stream_chunk_size_through_converse(): iter_bytes_spy = _stream_completion_with_spied_iter_bytes( model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" From ffa37d05b7cf16c9874101f3c737da19b4154aca Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Sun, 16 Aug 2026 14:28:50 -0400 Subject: [PATCH 010/281] feat(mistral): add zai-glm-5-2 model pricing and metadata --- .../model_prices_and_context_window_backup.json | 14 ++++++++++++++ model_prices_and_context_window.json | 14 ++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e6c6cab0631..b73feae90d3 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -29045,6 +29045,20 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "mistral/zai-glm-5-2": { + "input_cost_per_token": 1.4e-06, + "litellm_provider": "mistral", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.mistral.ai/models/zai-glm-5-2", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "mistral/magistral-medium-2506": { "deprecation_date": "2025-11-30", "input_cost_per_token": 2e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e6c6cab0631..b73feae90d3 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -29045,6 +29045,20 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "mistral/zai-glm-5-2": { + "input_cost_per_token": 1.4e-06, + "litellm_provider": "mistral", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.mistral.ai/models/zai-glm-5-2", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "mistral/magistral-medium-2506": { "deprecation_date": "2025-11-30", "input_cost_per_token": 2e-06, From 539a61be080b9ed8100ef00fca77f5b6175d8853 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Sun, 16 Aug 2026 14:44:09 -0400 Subject: [PATCH 011/281] feat(perplexity): add Agent API third-party models (DeepSeek V4 Flash, GLM 5.2, Kimi K3, Kimi K2.7 Code) --- ...odel_prices_and_context_window_backup.json | 44 +++++++++++++++++++ model_prices_and_context_window.json | 44 +++++++++++++++++++ 2 files changed, 88 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e6c6cab0631..b786a84ada7 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -34286,6 +34286,50 @@ "supports_reasoning": false, "supports_function_calling": true }, + "perplexity/perplexity/deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 2.8e-08, + "input_cost_per_token": 1.3e-07, + "litellm_provider": "perplexity", + "mode": "responses", + "output_cost_per_token": 2.6e-07, + "source": "https://docs.perplexity.ai/docs/agent-api/models", + "supports_web_search": true, + "supports_reasoning": true, + "supports_function_calling": true + }, + "perplexity/perplexity/glm-5.2": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "perplexity", + "mode": "responses", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.perplexity.ai/docs/agent-api/models", + "supports_web_search": true, + "supports_reasoning": true, + "supports_function_calling": true + }, + "perplexity/perplexity/kimi-k3": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "perplexity", + "mode": "responses", + "output_cost_per_token": 1.5e-05, + "source": "https://docs.perplexity.ai/docs/agent-api/models", + "supports_web_search": true, + "supports_reasoning": true, + "supports_function_calling": true + }, + "perplexity/perplexity/kimi-k2.7-code": { + "cache_read_input_token_cost": 1.9e-07, + "input_cost_per_token": 9.5e-07, + "litellm_provider": "perplexity", + "mode": "responses", + "output_cost_per_token": 4e-06, + "source": "https://docs.perplexity.ai/docs/agent-api/models", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true + }, "perplexity/pplx-embed-v1-0.6b": { "input_cost_per_token": 4e-09, "litellm_provider": "perplexity", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e6c6cab0631..b786a84ada7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -34286,6 +34286,50 @@ "supports_reasoning": false, "supports_function_calling": true }, + "perplexity/perplexity/deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 2.8e-08, + "input_cost_per_token": 1.3e-07, + "litellm_provider": "perplexity", + "mode": "responses", + "output_cost_per_token": 2.6e-07, + "source": "https://docs.perplexity.ai/docs/agent-api/models", + "supports_web_search": true, + "supports_reasoning": true, + "supports_function_calling": true + }, + "perplexity/perplexity/glm-5.2": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "perplexity", + "mode": "responses", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.perplexity.ai/docs/agent-api/models", + "supports_web_search": true, + "supports_reasoning": true, + "supports_function_calling": true + }, + "perplexity/perplexity/kimi-k3": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "perplexity", + "mode": "responses", + "output_cost_per_token": 1.5e-05, + "source": "https://docs.perplexity.ai/docs/agent-api/models", + "supports_web_search": true, + "supports_reasoning": true, + "supports_function_calling": true + }, + "perplexity/perplexity/kimi-k2.7-code": { + "cache_read_input_token_cost": 1.9e-07, + "input_cost_per_token": 9.5e-07, + "litellm_provider": "perplexity", + "mode": "responses", + "output_cost_per_token": 4e-06, + "source": "https://docs.perplexity.ai/docs/agent-api/models", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true + }, "perplexity/pplx-embed-v1-0.6b": { "input_cost_per_token": 4e-09, "litellm_provider": "perplexity", From 782746553c109bdec6aa5ecddf8e674f5167f575 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Sun, 16 Aug 2026 14:57:11 -0400 Subject: [PATCH 012/281] fix(perplexity): accept float usage.cost in cost_per_token, not just dict ResponseAPIUsage.parse_cost already flattens Perplexity's usage.cost.total_cost dict down to a float before it reaches the perplexity cost calculator, so the isinstance(cost_info, dict) check was always False on that path. Every Responses-mode Perplexity model was silently falling back to manual token-rate calculation and recording $0 spend whenever static per-token rates were missing. --- litellm/llms/perplexity/cost_calculator.py | 19 ++++++++------ .../test_perplexity_cost_calculator.py | 25 +++++++++++++++++++ 2 files changed, 37 insertions(+), 7 deletions(-) diff --git a/litellm/llms/perplexity/cost_calculator.py b/litellm/llms/perplexity/cost_calculator.py index 337fa8e630d..27835ecbfe8 100644 --- a/litellm/llms/perplexity/cost_calculator.py +++ b/litellm/llms/perplexity/cost_calculator.py @@ -21,14 +21,19 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd """ ## USE PRE-CALCULATED COST FROM PERPLEXITY IF AVAILABLE - ## Perplexity returns accurate cost in usage.cost.total_cost including request fees + ## Perplexity returns accurate cost in usage.cost.total_cost including request fees. + ## By the time it reaches here, ResponseAPIUsage.parse_cost has already flattened + ## that dict down to a float, so both shapes must be accepted. cost_info: Final = getattr(usage, "cost", None) - if cost_info is not None and isinstance(cost_info, dict): - total_cost: Final = cost_info.get("total_cost") - if total_cost is not None: - # Return total cost as completion_cost (prompt_cost=0) since Perplexity - # doesn't break down by input/output in their cost object - return (0.0, float(total_cost)) + total_cost: float | None = None + if isinstance(cost_info, dict): + total_cost = cost_info.get("total_cost") + elif isinstance(cost_info, (int, float)) and not isinstance(cost_info, bool): + total_cost = float(cost_info) + if total_cost is not None: + # Return total cost as completion_cost (prompt_cost=0) since Perplexity + # doesn't break down by input/output in their cost object + return (0.0, float(total_cost)) ## FALLBACK: Calculate cost manually if Perplexity doesn't provide it ## GET MODEL INFO diff --git a/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py b/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py index 46c1e457d7c..71ccb494cc7 100644 --- a/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py +++ b/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py @@ -400,6 +400,31 @@ class TestPerplexityCostCalculator: assert completion_cost == 0.008 assert prompt_cost + completion_cost == 0.008 + def test_uses_perplexity_provided_cost_when_normalized_to_float(self): + """ + Regression: for Responses API / Agent API models, `ResponseAPIUsage.parse_cost` + (litellm/types/llms/openai.py) already flattens Perplexity's + `usage.cost.total_cost` dict down to a plain float before + `_transform_response_api_usage_to_chat_usage` (litellm/responses/utils.py) copies + it onto the chat `Usage` object. So `usage.cost` arrives here as a float, not a + dict, on that path. + + Pre-fix, the `isinstance(cost_info, dict)` check was always False for a float, + so the pre-calculated cost branch was dead code for every Responses-mode + Perplexity model and it silently fell back to manual token-rate calculation, + recording $0 for any model missing static per-token rates (e.g. + perplexity/openai/gpt-5.2 before rates existed). + """ + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + usage.cost = 0.008 + + prompt_cost, completion_cost = perplexity_cost_per_token( + model="sonar-pro", usage=usage + ) + + assert prompt_cost == 0.0 + assert completion_cost == 0.008 + def test_falls_back_to_manual_calculation_when_no_cost_provided(self): """ Test that manual cost calculation is used when Perplexity doesn't From 5d5dc4523fb950e131a235bea9a4f767ba7e0e17 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 19 Aug 2026 03:56:32 +0000 Subject: [PATCH 013/281] fix(cost): price streamed Messages usage via calculate_usage and the logging obj Streamed `/v1/messages` `usage.cost` disagreed with the cost the logging callback recorded in three ways: `input_tokens` was read as the whole prompt total, but Anthropic reports it excluding cache tokens, so the non-cached input went unbilled on cache hits; the `cache_creation` 5m/1h split was dropped, billing 1h writes at the 5m rate; and costing by model name alone ignored the deployment's custom pricing, so a negotiated discount still streamed sticker price. Anthropic usage now goes through `AnthropicConfig.calculate_usage`, the same transformation the non-streaming path uses, and the chunk is priced through the call's logging object when there is one so it inherits `custom_pricing`, `custom_llm_provider`, `base_model` and `router_model_id`, falling back to `completion_cost` by model name. `calculate_usage` only reads its `usage_object`, so it now takes a `Mapping`. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/anthropic/chat/transformation.py | 2 +- litellm/proxy/common_request_processing.py | 130 ++++++++++------ .../streaming_handler.py | 2 +- .../proxy/test_common_request_processing.py | 140 ++++++++++++++++++ 4 files changed, 226 insertions(+), 48 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index ef4ad7011c5..b4f040b0a8c 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -2152,7 +2152,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def calculate_usage( self, - usage_object: dict, + usage_object: Mapping[str, Any], reasoning_content: str | None, completion_response: dict | None = None, speed: str | None = None, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index a0b69ecb0bf..f5299b273b4 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -158,7 +158,7 @@ ProxyRouteType: TypeAlias = Literal[ "acancel_run", "adelete_run", ] -from litellm.types.utils import ServerToolUse +from litellm.llms.anthropic.chat.transformation import AnthropicConfig # Type alias for streaming chunk serializer (chunk after hooks + cost injection -> wire format) StreamChunkSerializer = Callable[[Any], str] @@ -3321,7 +3321,9 @@ class ProxyBaseLLMRequestProcessing: str_so_far += str(chunk.get("content", "")) model_name = request_data.get("model", "") - chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(chunk, model_name) + chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( + chunk, model_name, request_data.get("litellm_logging_obj") + ) # Set before the yield: an async generator suspends at the yield, # so a GeneratorExit on client disconnect is raised there and any @@ -3418,20 +3420,28 @@ class ProxyBaseLLMRequestProcessing: @overload @staticmethod - def _process_chunk_with_cost_injection(chunk: bytes, model_name: str) -> bytes: ... + def _process_chunk_with_cost_injection( + chunk: bytes, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + ) -> bytes: ... @overload @staticmethod - def _process_chunk_with_cost_injection(chunk: object, model_name: str) -> object: ... + def _process_chunk_with_cost_injection( + chunk: object, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + ) -> object: ... @staticmethod - def _process_chunk_with_cost_injection(chunk: object, model_name: str) -> object: + def _process_chunk_with_cost_injection( + chunk: object, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + ) -> object: """ Process a streaming chunk and inject cost information if enabled. Args: chunk: The streaming chunk (dict, str, bytes, or bytearray) model_name: Model name for cost calculation + litellm_logging_obj: The call's logging object, used to price the chunk with + the same custom/deployment pricing as the logging callback Returns: The processed chunk with cost information injected if applicable @@ -3441,21 +3451,27 @@ class ProxyBaseLLMRequestProcessing: try: if isinstance(chunk, dict): - maybe_modified: Final = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(chunk, model_name) + maybe_modified: Final = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict( + chunk, model_name, litellm_logging_obj + ) if maybe_modified is not None: return maybe_modified elif isinstance(chunk, (bytes, bytearray)): try: s: Final = chunk.decode("utf-8") if s.endswith(("\n\n", "\r\n\r\n")): - maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(s, model_name) + maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str( + s, model_name, litellm_logging_obj + ) if maybe_mod is not None: return maybe_mod.encode("utf-8") except Exception: pass elif isinstance(chunk, str): # Try to parse SSE frame and inject cost into the data line - maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(chunk, model_name) + maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str( + chunk, model_name, litellm_logging_obj + ) if maybe_mod is not None: # Ensure trailing frame separator return maybe_mod if maybe_mod.endswith("\n\n") else (maybe_mod + "\n\n") @@ -3466,13 +3482,16 @@ class ProxyBaseLLMRequestProcessing: return chunk @staticmethod - def _inject_cost_into_sse_frame_str(frame_str: str, model_name: str) -> str | None: + def _inject_cost_into_sse_frame_str( + frame_str: str, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + ) -> str | None: """ Inject cost information into an SSE frame string by modifying the JSON in the 'data:' line. Args: frame_str: SSE frame string that may contain multiple lines model_name: Model name for cost calculation + litellm_logging_obj: The call's logging object, forwarded for pricing Returns: Modified SSE frame string with cost injected, or None if no modification needed @@ -3486,7 +3505,9 @@ class ProxyBaseLLMRequestProcessing: json_part = stripped_ln.split("data:", 1)[1].strip() if json_part and json_part != "[DONE]": obj = json.loads(json_part) - maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(obj, model_name) + maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict( + obj, model_name, litellm_logging_obj + ) if maybe_modified is not None: lines[idx] = "data: " + safe_dumps(maybe_modified) + ("\r" if ln.endswith("\r") else "") return "\n".join(lines) @@ -3494,34 +3515,6 @@ class ProxyBaseLLMRequestProcessing: except Exception: return None - @staticmethod - def _anthropic_stream_usage_kwargs(usage: Mapping[str, Any]) -> Mapping[str, Any]: - prompt_tokens: Final = int(usage.get("input_tokens", 0) or 0) - completion_tokens: Final = int(usage.get("output_tokens", 0) or 0) - total_tokens: Final = int( - usage.get("total_tokens", prompt_tokens + completion_tokens) or (prompt_tokens + completion_tokens) - ) - web_search_requests: Final = usage.get("web_search_requests") - server_tool_use: Final = ( - ServerToolUse(web_search_requests=web_search_requests) if web_search_requests is not None else None - ) - return MappingProxyType( - { - key: value - for key, value in ( - ("prompt_tokens", prompt_tokens), - ("completion_tokens", completion_tokens), - ("total_tokens", total_tokens), - ("completion_tokens_details", usage.get("completion_tokens_details")), - ("prompt_tokens_details", usage.get("prompt_tokens_details")), - ("cache_creation_input_tokens", usage.get("cache_creation_input_tokens")), - ("cache_read_input_tokens", usage.get("cache_read_input_tokens")), - ("server_tool_use", server_tool_use), - ) - if value is not None - } - ) - @staticmethod def _openai_stream_usage_kwargs(usage: Mapping[str, Any]) -> Mapping[str, Any]: prompt_tokens: Final = int(usage.get("prompt_tokens", 0) or 0) @@ -3544,11 +3537,19 @@ class ProxyBaseLLMRequestProcessing: ) @staticmethod - def _stream_usage_kwargs_for_event(obj: Mapping[str, object], usage: Mapping[str, Any]) -> Mapping[str, Any] | None: + def _stream_usage_for_event(obj: Mapping[str, object], usage: Mapping[str, Any]) -> Usage | None: + """ + Build the ``Usage`` to price a streamed usage event. + + Anthropic goes through ``AnthropicConfig.calculate_usage``, the same transformation + the non-streaming path uses, so ``prompt_tokens`` is the full input total and the + cache read plus 5m/1h cache creation split land in ``prompt_tokens_details`` where + the pricer looks for them. + """ if obj.get("type") == "message_delta": - return ProxyBaseLLMRequestProcessing._anthropic_stream_usage_kwargs(usage) + return AnthropicConfig().calculate_usage(usage_object=usage, reasoning_content=None) if obj.get("object") == "chat.completion.chunk": - return ProxyBaseLLMRequestProcessing._openai_stream_usage_kwargs(usage) + return Usage(**ProxyBaseLLMRequestProcessing._openai_stream_usage_kwargs(usage)) return None @staticmethod @@ -3563,7 +3564,41 @@ class ProxyBaseLLMRequestProcessing: return None @staticmethod - def _inject_cost_into_usage_dict(obj: dict, model_name: str) -> dict | None: + def _logging_obj_cost_or_none( + model_response: ModelResponse, litellm_logging_obj: LiteLLMLoggingObj + ) -> float | None: + try: + cost: Final = litellm_logging_obj._response_cost_calculator(result=model_response) # pyright: ignore[reportPrivateUsage] # reuse the call's own cost calc for pricing parity with the logging callback + except Exception: + return None + return float(cost) if isinstance(cost, (int, float)) and not isinstance(cost, bool) else None + + @staticmethod + def _streamed_usage_cost( + model_response: ModelResponse, + model_name: str, + service_tier: str | None, + litellm_logging_obj: LiteLLMLoggingObj | None, + ) -> float | None: + """ + Price a streamed usage response through the call's logging object when there is one, + so the streamed cost picks up the same custom/deployment pricing (``custom_pricing``, + ``custom_llm_provider``, ``base_model``, ``router_model_id``) as the logging callback + rather than the model's sticker price; fall back to ``completion_cost`` by model name. + """ + cost_from_logging_obj: Final = ( + ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, litellm_logging_obj) + if litellm_logging_obj is not None + else None + ) + if cost_from_logging_obj is not None: + return cost_from_logging_obj + return ProxyBaseLLMRequestProcessing._completion_cost_or_none(model_response, model_name, service_tier) + + @staticmethod + def _inject_cost_into_usage_dict( + obj: dict, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + ) -> dict | None: """ Inject cost information into the usage object of a streamed usage event (Anthropic ``message_delta`` or OpenAI ``chat.completion.chunk``). @@ -3571,6 +3606,8 @@ class ProxyBaseLLMRequestProcessing: Args: obj: Dictionary containing the SSE event data model_name: Model name for cost calculation + litellm_logging_obj: The call's logging object, used so the injected cost + matches the cost the logging callback records Returns: Modified dictionary with cost injected, or None if no modification needed @@ -3578,14 +3615,15 @@ class ProxyBaseLLMRequestProcessing: usage: Final = obj.get("usage") if not isinstance(usage, dict): return None - usage_kwargs: Final = ProxyBaseLLMRequestProcessing._stream_usage_kwargs_for_event(obj, usage) - if usage_kwargs is None: + stream_usage: Final = ProxyBaseLLMRequestProcessing._stream_usage_for_event(obj, usage) + if stream_usage is None: return None service_tier: Final = obj.get("service_tier") - cost_val: Final = ProxyBaseLLMRequestProcessing._completion_cost_or_none( - ModelResponse(usage=Usage(**usage_kwargs)), + cost_val: Final = ProxyBaseLLMRequestProcessing._streamed_usage_cost( + ModelResponse(usage=stream_usage), model_name, service_tier if isinstance(service_tier, str) else None, + litellm_logging_obj, ) if cost_val is None: return None diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 697eb7b96eb..b71622fc33d 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -106,7 +106,7 @@ class PassThroughStreamingHandler: ) # rebind-ok: SSE frame reassembly buffer across transport chunks if complete_frames: yield ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( - complete_frames, resolved_model_name + complete_frames, resolved_model_name, litellm_logging_obj ) if pending: yield pending diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 355c6d27eb2..083982778b9 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -6082,6 +6082,121 @@ class TestInjectCostIntoUsageDict: injected = json.loads(result.split("\n")[0].split("data:", 1)[1].strip()) assert injected["usage"]["cost"] == pytest.approx(self._expected_cost("gpt-4o-mini", 11, 4)) + def test_message_delta_cost_charges_the_non_cached_input_tokens(self): + """Anthropic reports ``input_tokens`` excluding cache tokens, so reading it as the whole + prompt total drops the non-cached input from the bill on every cache hit.""" + model = "claude-haiku-4-5" + pricing = litellm.model_cost[model] + event = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": { + "input_tokens": 14, + "output_tokens": 8, + "cache_read_input_tokens": 3202, + "cache_creation_input_tokens": 0, + }, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, model) + + assert result is not None + expected = ( + 14 * pricing["input_cost_per_token"] + + 3202 * pricing["cache_read_input_token_cost"] + + 8 * pricing["output_cost_per_token"] + ) + dropped_input = expected - 14 * pricing["input_cost_per_token"] + assert result["usage"]["cost"] == pytest.approx(expected) + assert result["usage"]["cost"] > dropped_input + + def test_message_delta_prices_1h_cache_creation_above_the_5m_rate(self): + """The ``cache_creation`` 5m/1h split has to survive into ``prompt_tokens_details``, + otherwise a 1h write is billed at the cheaper 5m rate.""" + model = "claude-haiku-4-5" + pricing = litellm.model_cost[model] + event = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": { + "input_tokens": 14, + "output_tokens": 8, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 2000, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 2000}, + }, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, model) + + assert result is not None + base = 14 * pricing["input_cost_per_token"] + 8 * pricing["output_cost_per_token"] + expected_1h = base + 2000 * pricing["cache_creation_input_token_cost_above_1hr"] + flat_5m = base + 2000 * pricing["cache_creation_input_token_cost"] + assert expected_1h != pytest.approx(flat_5m) + assert result["usage"]["cost"] == pytest.approx(expected_1h) + + def test_message_delta_prices_through_the_logging_obj_so_custom_pricing_applies(self): + """Costing by model name alone yields sticker price, so a deployment with a negotiated + discount streamed a ``usage.cost`` that disagreed with the callback's ``response_cost``.""" + + class _StubLoggingObj: + def __init__(self, cost): + self._cost = cost + self.captured_result = None + + def _response_cost_calculator(self, result): + self.captured_result = result + return self._cost + + model = "claude-haiku-4-5" + discounted_cost = 0.00099 + stub = _StubLoggingObj(discounted_cost) + event = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": { + "input_tokens": 14, + "output_tokens": 8, + "cache_read_input_tokens": 3202, + "cache_creation_input_tokens": 500, + "cache_creation": {"ephemeral_5m_input_tokens": 100, "ephemeral_1h_input_tokens": 400}, + }, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, model, stub) + + assert result is not None + assert result["usage"]["cost"] == discounted_cost + assert result["usage"]["cost"] != pytest.approx(self._expected_cost(model, 14 + 500 + 3202, 8)) + usage = stub.captured_result.usage + assert usage.prompt_tokens == 14 + 500 + 3202 + details = usage.prompt_tokens_details.cache_creation_token_details + assert details.ephemeral_5m_input_tokens == 100 + assert details.ephemeral_1h_input_tokens == 400 + + def test_message_delta_falls_back_to_model_pricing_when_the_logging_obj_returns_no_cost(self): + class _StubLoggingObj: + def _response_cost_calculator(self, result): + return None + + model = "claude-haiku-4-5" + pricing = litellm.model_cost[model] + event = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"input_tokens": 14, "output_tokens": 8, "cache_read_input_tokens": 3202}, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, model, _StubLoggingObj()) + + assert result is not None + assert result["usage"]["cost"] == pytest.approx( + 14 * pricing["input_cost_per_token"] + + 3202 * pricing["cache_read_input_token_cost"] + + 8 * pricing["output_cost_per_token"] + ) + class TestProcessChunkWithCostInjection: def test_complete_usage_frame_chunk_is_injected(self, monkeypatch): @@ -6116,6 +6231,31 @@ class TestProcessChunkWithCostInjection: assert ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(chunk, "gpt-4o-mini") == chunk + def test_message_delta_frame_is_priced_with_the_logging_obj(self, monkeypatch): + """Pins that the logging object reaches the pricer through the byte-frame entry point, + which is how the proxy actually calls this on a streamed Messages API request.""" + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + + class _StubLoggingObj: + def _response_cost_calculator(self, result): + return 0.00042 + + chunk = ( + b"event: message_delta\n" + b'data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},' + b'"usage":{"input_tokens":14,"output_tokens":8,"cache_read_input_tokens":3202}}\n\n' + ) + + result = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( + chunk, "claude-haiku-4-5", _StubLoggingObj() + ) + + assert result != chunk + data_line = next(ln for ln in result.decode("utf-8").splitlines() if ln.startswith("data:")) + payload = json.loads(data_line.split("data:", 1)[1].strip()) + assert payload["usage"]["cost"] == 0.00042 + assert payload["usage"]["cache_read_input_tokens"] == 3202 + # --------------------------------------------------------------------------- # SSE keepalive during the time-to-first-token (issue #34819) From 6f4844bfd9dd50901e3e4542f296602c9c8c9e56 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 19 Aug 2026 03:58:24 +0000 Subject: [PATCH 014/281] refactor(cost): trim streamed cost helper docstrings to the non-obvious bits Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 24 ++++++---------------- 1 file changed, 6 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f5299b273b4..d533df8616d 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3440,8 +3440,7 @@ class ProxyBaseLLMRequestProcessing: Args: chunk: The streaming chunk (dict, str, bytes, or bytearray) model_name: Model name for cost calculation - litellm_logging_obj: The call's logging object, used to price the chunk with - the same custom/deployment pricing as the logging callback + litellm_logging_obj: The call's logging object, used for pricing Returns: The processed chunk with cost information injected if applicable @@ -3538,14 +3537,8 @@ class ProxyBaseLLMRequestProcessing: @staticmethod def _stream_usage_for_event(obj: Mapping[str, object], usage: Mapping[str, Any]) -> Usage | None: - """ - Build the ``Usage`` to price a streamed usage event. - - Anthropic goes through ``AnthropicConfig.calculate_usage``, the same transformation - the non-streaming path uses, so ``prompt_tokens`` is the full input total and the - cache read plus 5m/1h cache creation split land in ``prompt_tokens_details`` where - the pricer looks for them. - """ + # Anthropic reports input_tokens excluding cache tokens, so reuse the non-streaming + # transformation to total the prompt and keep the 5m/1h cache creation split if obj.get("type") == "message_delta": return AnthropicConfig().calculate_usage(usage_object=usage, reasoning_content=None) if obj.get("object") == "chat.completion.chunk": @@ -3580,12 +3573,8 @@ class ProxyBaseLLMRequestProcessing: service_tier: str | None, litellm_logging_obj: LiteLLMLoggingObj | None, ) -> float | None: - """ - Price a streamed usage response through the call's logging object when there is one, - so the streamed cost picks up the same custom/deployment pricing (``custom_pricing``, - ``custom_llm_provider``, ``base_model``, ``router_model_id``) as the logging callback - rather than the model's sticker price; fall back to ``completion_cost`` by model name. - """ + # Pricing via the logging object inherits the deployment's custom pricing, so the + # streamed cost matches what the logging callback records instead of sticker price cost_from_logging_obj: Final = ( ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, litellm_logging_obj) if litellm_logging_obj is not None @@ -3606,8 +3595,7 @@ class ProxyBaseLLMRequestProcessing: Args: obj: Dictionary containing the SSE event data model_name: Model name for cost calculation - litellm_logging_obj: The call's logging object, used so the injected cost - matches the cost the logging callback records + litellm_logging_obj: The call's logging object, used for pricing Returns: Modified dictionary with cost injected, or None if no modification needed From 087cdcff07ae34cd2d5ae9b2b4cd2282610d3953 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 19 Aug 2026 21:18:12 +0000 Subject: [PATCH 015/281] feat: add bedrock grok 4.6 to model cost map Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 54 +++++++++++++++++++ model_prices_and_context_window.json | 54 +++++++++++++++++++ 2 files changed, 108 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 07f9027313b..b00ab87bf43 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -47379,6 +47379,60 @@ "supports_vision": true, "source": "https://aws.amazon.com/bedrock/pricing/" }, + "bedrock_mantle/xai.grok-4.6": { + "use_openai_responses_path": true, + "input_cost_per_token": 2.2e-06, + "output_cost_per_token": 6.6e-06, + "cache_read_input_token_cost": 5.5e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 500000, + "max_output_tokens": 500000, + "max_tokens": 500000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-6.html" + }, + "us.xai.grok-4.6": { + "input_cost_per_token": 2.2e-06, + "output_cost_per_token": 6.6e-06, + "cache_read_input_token_cost": 5.5e-07, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 500000, + "max_output_tokens": 500000, + "max_tokens": 500000, + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-6.html" + }, + "global.xai.grok-4.6": { + "input_cost_per_token": 2e-06, + "output_cost_per_token": 6e-06, + "cache_read_input_token_cost": 5e-07, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 500000, + "max_output_tokens": 500000, + "max_tokens": 500000, + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-6.html" + }, "volcengine/doubao-seed-2-0-pro-260215": { "litellm_provider": "volcengine", "max_input_tokens": 256000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 07f9027313b..b00ab87bf43 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -47379,6 +47379,60 @@ "supports_vision": true, "source": "https://aws.amazon.com/bedrock/pricing/" }, + "bedrock_mantle/xai.grok-4.6": { + "use_openai_responses_path": true, + "input_cost_per_token": 2.2e-06, + "output_cost_per_token": 6.6e-06, + "cache_read_input_token_cost": 5.5e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 500000, + "max_output_tokens": 500000, + "max_tokens": 500000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-6.html" + }, + "us.xai.grok-4.6": { + "input_cost_per_token": 2.2e-06, + "output_cost_per_token": 6.6e-06, + "cache_read_input_token_cost": 5.5e-07, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 500000, + "max_output_tokens": 500000, + "max_tokens": 500000, + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-6.html" + }, + "global.xai.grok-4.6": { + "input_cost_per_token": 2e-06, + "output_cost_per_token": 6e-06, + "cache_read_input_token_cost": 5e-07, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 500000, + "max_output_tokens": 500000, + "max_tokens": 500000, + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-6.html" + }, "volcengine/doubao-seed-2-0-pro-260215": { "litellm_provider": "volcengine", "max_input_tokens": 256000, From 1f6bef79cac5d33a39678e4e4df3bf11cba0b5b1 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 19 Aug 2026 21:24:29 +0000 Subject: [PATCH 016/281] fix: drop source url from grok 4.6 cost map entries Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 9 +++------ model_prices_and_context_window.json | 9 +++------ 2 files changed, 6 insertions(+), 12 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b00ab87bf43..872ef142123 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -47398,8 +47398,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-6.html" + "supports_vision": true }, "us.xai.grok-4.6": { "input_cost_per_token": 2.2e-06, @@ -47414,8 +47413,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-6.html" + "supports_vision": true }, "global.xai.grok-4.6": { "input_cost_per_token": 2e-06, @@ -47430,8 +47428,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-6.html" + "supports_vision": true }, "volcengine/doubao-seed-2-0-pro-260215": { "litellm_provider": "volcengine", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b00ab87bf43..872ef142123 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -47398,8 +47398,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-6.html" + "supports_vision": true }, "us.xai.grok-4.6": { "input_cost_per_token": 2.2e-06, @@ -47414,8 +47413,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-6.html" + "supports_vision": true }, "global.xai.grok-4.6": { "input_cost_per_token": 2e-06, @@ -47430,8 +47428,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-6.html" + "supports_vision": true }, "volcengine/doubao-seed-2-0-pro-260215": { "litellm_provider": "volcengine", From aa832d81e91bb17e0f5ab081431cb03d2ef4083f Mon Sep 17 00:00:00 2001 From: ljogeiger <39981740+ljogeiger@users.noreply.github.com> Date: Wed, 19 Aug 2026 22:08:42 +0000 Subject: [PATCH 017/281] fix(vertex_ai): only fall back to a placeholder thought signature on the first parallel function call Gemini returns a thoughtSignature on the first function call of a parallel batch and leaves the siblings bare. When replaying that assistant turn, litellm gave every unsigned call the skip_thought_signature_validator placeholder, so a three-call turn went back with three signatures where Gemini had produced one. Keep the placeholder for the first call only and forward the siblings with whatever signature they actually carry, which is usually none. --- .../prompt_templates/factory.py | 23 ++-- .../test_vertex_ai_gemini_transformation.py | 115 ++++++++++++++++++ 2 files changed, 126 insertions(+), 12 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 0ed15c43ccf..0a9d7c427b4 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1200,13 +1200,14 @@ def _encode_tool_call_id_with_signature(tool_call_id: str, thought_signature: st return tool_call_id -def _get_thought_signature_from_tool(tool: dict, model: str | None = None) -> str | None: +def _get_thought_signature_from_tool(tool: dict) -> str | None: """Extract thought signature from tool call's provider_specific_fields. If not provided try to extract thought signature from tool call id Checks both tool.provider_specific_fields and tool.function.provider_specific_fields. - If no signature is found and model is gemini-3, returns a dummy signature. + Returns None when the tool call carries no signature; callers decide whether a + placeholder signature is needed. """ # First check tool's provider_specific_fields provider_fields: Final = tool.get("provider_specific_fields") or {} @@ -1236,13 +1237,6 @@ def _get_thought_signature_from_tool(tool: dict, model: str | None = None) -> st if len(parts) == 2: _, signature = parts return signature - # If no signature found and model is gemini-3, return dummy signature - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, - ) - - if model and VertexGeminiConfig._is_gemini_3_or_newer(model): - return _get_dummy_thought_signature() return None @@ -1312,8 +1306,10 @@ def convert_to_gemini_tool_call_invoke( VertexGeminiConfig, ) + needs_dummy_signature: Final = model is not None and VertexGeminiConfig._is_gemini_3_or_newer(model) + if tool_calls is not None: - for idx, tool in enumerate(tool_calls): + for tool in tool_calls: if "function" in tool: gemini_function_call: VertexFunctionCall | None = _gemini_tool_call_invoke_helper( function_call_params=tool["function"], @@ -1321,7 +1317,10 @@ def convert_to_gemini_tool_call_invoke( ) if gemini_function_call is not None: part_dict: VertexPartType = {"function_call": gemini_function_call} - thought_signature = _get_thought_signature_from_tool(dict(tool), model=model) + thought_signature = _get_thought_signature_from_tool(dict(tool)) + is_first_function_call = len(_parts_list) == 0 + if not thought_signature and is_first_function_call and needs_dummy_signature: + thought_signature = _get_dummy_thought_signature() if thought_signature: part_dict["thoughtSignature"] = thought_signature @@ -1344,7 +1343,7 @@ def convert_to_gemini_tool_call_invoke( thought_signature = provider_fields.get("thought_signature") # If no signature found and model is gemini-3, use dummy signature - if not thought_signature and model and VertexGeminiConfig._is_gemini_3_or_newer(model): + if not thought_signature and needs_dummy_signature: thought_signature = _get_dummy_thought_signature() if thought_signature: diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index 8ee8186f6bb..28c551ab824 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -1,3 +1,5 @@ +import base64 + from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_gemini_tool_call_result, ) @@ -784,6 +786,119 @@ def test_dummy_signature_with_function_call_mode(): assert gemini_parts[0]["thoughtSignature"] == expected_dummy +def _parallel_tool_calls(*signatures): + return [ + { + "id": f"call_{idx}", + "type": "function", + "function": { + "name": f"tool_{idx}", + "arguments": '{"location": "Paris"}', + **( + {"provider_specific_fields": {"thought_signature": signature}} + if signature is not None + else {} + ), + }, + "index": idx, + } + for idx, signature in enumerate(signatures) + ] + + +REAL_THOUGHT_SIGNATURE = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n" + + +def test_dummy_signature_only_on_first_parallel_tool_call(): + """Gemini only returns a thought signature on the first of N parallel function calls. + + The sibling calls carry no signature, so replaying them must not fabricate one. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(None, None, None), + }, + model="gemini-3-pro-preview", + ) + + expected_dummy = base64.b64encode(b"skip_thought_signature_validator").decode( + "utf-8" + ) + assert len(gemini_parts) == 3 + assert gemini_parts[0]["thoughtSignature"] == expected_dummy + assert "thoughtSignature" not in gemini_parts[1] + assert "thoughtSignature" not in gemini_parts[2] + + +def test_real_signature_on_first_parallel_tool_call_leaves_siblings_empty(): + """The real signature from Gemini rides on the first call; siblings stay signature-free.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(REAL_THOUGHT_SIGNATURE, None, None), + }, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 3 + assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + assert "thoughtSignature" not in gemini_parts[2] + + +def test_real_signature_on_later_parallel_tool_call_is_preserved(): + """A signature attached to a non-first call is still forwarded as-is.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(None, REAL_THOUGHT_SIGNATURE), + }, + model="gemini-3-pro-preview", + ) + + expected_dummy = base64.b64encode(b"skip_thought_signature_validator").decode( + "utf-8" + ) + assert len(gemini_parts) == 2 + assert gemini_parts[0]["thoughtSignature"] == expected_dummy + assert gemini_parts[1]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + + +def test_no_signatures_on_parallel_tool_calls_for_gemini_2_5(): + """Non-gemini-3 models never get a placeholder signature, on any call.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(None, None), + }, + model="gemini-2.5-flash", + ) + + assert len(gemini_parts) == 2 + assert all("thoughtSignature" not in part for part in gemini_parts) + + # Tests for media_resolution (detail parameter) handling - Issue #17084 class TestMediaResolution: """Tests for media_resolution handling in Gemini 2.x models""" From 7a6a677b72cd3822e3fe0ae29329fa8b0e17d577 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:09:05 -0700 Subject: [PATCH 018/281] feat(proxy): enqueued-token rate limiting for batches with refund on completion and cancellation --- litellm/constants.py | 6 + litellm/proxy/hooks/batch_enqueued_tokens.py | 357 ++++++++++++++++++ litellm/proxy/hooks/batch_rate_limiter.py | 80 +++- .../hooks/parallel_request_limiter_v3.py | 42 +++ tests/e2e/batches/test_batches_e2e.py | 146 ++++++- .../coverage_registry/quota_management.yaml | 3 + tests/e2e/models.py | 1 + .../proxy/hooks/test_batch_enqueued_tokens.py | 190 ++++++++++ .../proxy/hooks/test_batch_file_validation.py | 192 ++++++++++ .../hooks/test_parallel_request_limiter_v3.py | 120 ++++++ 10 files changed, 1133 insertions(+), 4 deletions(-) create mode 100644 litellm/proxy/hooks/batch_enqueued_tokens.py create mode 100644 tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py diff --git a/litellm/constants.py b/litellm/constants.py index 39a49e55f0d..d7ddacf5fac 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1766,6 +1766,12 @@ PTU_LAPSED_ALERT_LIMIT: Final[int] = 10 # one is seconds old, so a few minutes separates them. PTU_PRUNE_SKEW_GRACE_SECONDS: Final[int] = 300 +# How long enqueued-token reservations for batches live without a refund. Providers +# complete or expire batches within their completion window (24h for OpenAI), so a +# reservation still unrefunded after 8 days belongs to a batch whose terminal state +# was never observed (e.g. proxy restart); expiry returns the tokens to the caller. +BATCH_ENQUEUED_TOKEN_TTL_SECONDS: Final[int] = 8 * 24 * 60 * 60 + # Shared read-only empty mapping, for defaulting optional Mapping parameters without # constructing a fresh mutable dict at each call site. EMPTY_MAPPING: Final = MappingProxyType({}) diff --git a/litellm/proxy/hooks/batch_enqueued_tokens.py b/litellm/proxy/hooks/batch_enqueued_tokens.py new file mode 100644 index 00000000000..040c8b687a7 --- /dev/null +++ b/litellm/proxy/hooks/batch_enqueued_tokens.py @@ -0,0 +1,357 @@ +""" +Enqueued-token accounting for batch submissions. + +Opt-in via ``batch_enqueued_token_limit`` in key or team metadata: batch +submissions reserve their estimated token count against a long-lived +enqueued-token allowance instead of the per-minute rate-limit windows, and +the reservation is refunded when the batch reaches a terminal state +(completed, failed, expired, or cancelled). +""" + +import asyncio +from collections.abc import Awaitable, Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.constants import BATCH_ENQUEUED_TOKEN_TTL_SECONDS +from litellm.proxy._types import UserAPIKeyAuth + +if TYPE_CHECKING: + from opentelemetry.trace import Span as _Span + + from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache + + Span = _Span + InternalUsageCache = _InternalUsageCache + +BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY: Final = "batch_enqueued_token_limit" + +BATCH_ENQUEUED_REFUND_STATUSES: Final[frozenset[str]] = frozenset( + {"completed", "complete", "failed", "expired", "cancelled", "cancelling"} +) + +ScopeKey: TypeAlias = Literal["api_key", "team"] + +RESERVE_ENQUEUED_TOKENS_SCRIPT: Final = """ +local amount = tonumber(ARGV[1]) +local ttl = tonumber(ARGV[2]) +for i = 1, #KEYS do + local limit = tonumber(ARGV[2 + i]) + local current = tonumber(redis.call('GET', KEYS[i]) or '0') + if current + amount > limit then + return {0, i - 1, current} + end +end +for i = 1, #KEYS do + redis.call('INCRBY', KEYS[i], amount) + redis.call('EXPIRE', KEYS[i], ttl) +end +return {1, -1, 0} +""" + +REFUND_ENQUEUED_TOKENS_SCRIPT: Final = """ +local amount = tonumber(ARGV[1]) +for i = 1, #KEYS do + local updated = redis.call('DECRBY', KEYS[i], amount) + if updated <= 0 then + redis.call('DEL', KEYS[i]) + end +end +return 1 +""" + +SAVE_RESERVATION_SCRIPT: Final = """ +redis.call('SET', KEYS[1], ARGV[1], 'EX', tonumber(ARGV[2])) +return 1 +""" + +POP_RESERVATION_SCRIPT: Final = """ +local value = redis.call('GET', KEYS[1]) +if value then + redis.call('DEL', KEYS[1]) +end +return value +""" + + +@dataclass(frozen=True, slots=True) +class BatchEnqueuedTokenScope: + key: ScopeKey + value: str + limit: int + + +@dataclass(frozen=True, slots=True) +class BatchEnqueuedTokenReservation: + tokens: int + scopes: tuple[BatchEnqueuedTokenScope, ...] + + +@dataclass(frozen=True, slots=True) +class BatchEnqueuedTokenOverLimit: + scope: BatchEnqueuedTokenScope + enqueued: int + + +BatchEnqueuedTokenOutcome: TypeAlias = BatchEnqueuedTokenReservation | BatchEnqueuedTokenOverLimit + +_LIMIT_ADAPTER: Final = TypeAdapter(Annotated[int, Field(gt=0)]) +_RESERVE_RESULT_ADAPTER: Final = TypeAdapter(tuple[int, int, int]) +_POPPED_VALUE_ADAPTER: Final = TypeAdapter(str | bytes | None) +_STORED_COUNTER_ADAPTER: Final = TypeAdapter(int | None) +_RESERVATION_ADAPTER: Final = TypeAdapter(BatchEnqueuedTokenReservation) + + +class _ScriptRunner(Protocol): + def __call__(self, keys: Sequence[str], args: Sequence[str | bytes | int | float]) -> Awaitable[object]: ... + + +def _read_metadata_limit(metadata: Mapping[str, object] | None) -> int | None: + if not metadata: + return None + raw: Final = metadata.get(BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY) + if raw is None: + return None + try: + return _LIMIT_ADAPTER.validate_python(raw) + except ValidationError: + verbose_proxy_logger.warning( + "Ignoring invalid %s value %r; expected a positive integer", + BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY, + raw, + ) + return None + + +def resolve_batch_enqueued_token_scopes( + user_api_key_dict: UserAPIKeyAuth, +) -> tuple[BatchEnqueuedTokenScope, ...]: + key_limit: Final = _read_metadata_limit(user_api_key_dict.metadata) + team_limit: Final = _read_metadata_limit(user_api_key_dict.team_metadata) + candidates: Final = ( + BatchEnqueuedTokenScope(key="api_key", value=user_api_key_dict.api_key, limit=key_limit) + if key_limit is not None and user_api_key_dict.api_key + else None, + BatchEnqueuedTokenScope(key="team", value=user_api_key_dict.team_id, limit=team_limit) + if team_limit is not None and user_api_key_dict.team_id + else None, + ) + return tuple(scope for scope in candidates if scope is not None) + + +def canonical_provider_batch_id(batch_id: str) -> str: + from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, # pyright: ignore[reportPrivateUsage] # canonical unified-id decoder has no public wrapper + get_batch_id_from_unified_batch_id, + get_original_file_id, + ) + + decoded: Final = _is_base64_encoded_unified_file_id(batch_id) + if isinstance(decoded, str): + if "llm_batch_id" in decoded or "generic_response_id" in decoded: + return get_batch_id_from_unified_batch_id(decoded) + return decoded + return get_original_file_id(batch_id) + + +class _BatchResponseView(BaseModel): + model_config = ConfigDict(extra="ignore") + + id: str + status: str + object: Literal["batch"] + + +def batch_response_view(response: object) -> _BatchResponseView | None: + try: + return _BatchResponseView.model_validate(response, from_attributes=True) + except ValidationError: + return None + + +class BatchEnqueuedTokenStore: + """Tracks enqueued batch tokens per scope, plus per-batch reservation records for refunds. + + Counters and records live in Redis (via atomic Lua scripts) when Redis is + configured; otherwise a single-process in-memory fallback guarded by one + asyncio lock is used. Everything expires after + ``BATCH_ENQUEUED_TOKEN_TTL_SECONDS`` so a crash between submission and the + terminal-state refund can never leak tokens forever. + """ + + def __init__(self, internal_usage_cache: "InternalUsageCache") -> None: + self.internal_usage_cache = internal_usage_cache + self._lock = asyncio.Lock() + redis_cache = internal_usage_cache.dual_cache.redis_cache + self._reserve_script: _ScriptRunner | None = ( + redis_cache.async_register_script(RESERVE_ENQUEUED_TOKENS_SCRIPT) if redis_cache is not None else None + ) + self._refund_script: _ScriptRunner | None = ( + redis_cache.async_register_script(REFUND_ENQUEUED_TOKENS_SCRIPT) if redis_cache is not None else None + ) + self._save_script: _ScriptRunner | None = ( + redis_cache.async_register_script(SAVE_RESERVATION_SCRIPT) if redis_cache is not None else None + ) + self._pop_script: _ScriptRunner | None = ( + redis_cache.async_register_script(POP_RESERVATION_SCRIPT) if redis_cache is not None else None + ) + + @staticmethod + def _counter_key(scope: BatchEnqueuedTokenScope) -> str: + return f"batch_enqueued_tokens:{scope.key}:{scope.value}" + + @staticmethod + def _record_key(batch_id: str) -> str: + return f"batch_enqueued_token_reservation:{batch_id}" + + async def reserve( + self, + tokens: int, + scopes: tuple[BatchEnqueuedTokenScope, ...], + litellm_parent_otel_span: "Span | None" = None, + ) -> BatchEnqueuedTokenOutcome: + if tokens <= 0 or not scopes: + return BatchEnqueuedTokenReservation(tokens=max(tokens, 0), scopes=scopes) + if self._reserve_script is not None: + try: + raw_result = await self._reserve_script( + tuple(self._counter_key(scope) for scope in scopes), + (tokens, BATCH_ENQUEUED_TOKEN_TTL_SECONDS, *(scope.limit for scope in scopes)), + ) + result = _RESERVE_RESULT_ADAPTER.validate_python(raw_result) + if result[0] == 1: + return BatchEnqueuedTokenReservation(tokens=tokens, scopes=scopes) + return BatchEnqueuedTokenOverLimit(scope=scopes[result[1]], enqueued=result[2]) + except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory counters + verbose_proxy_logger.warning( + "Redis enqueued-token reserve failed, falling back to in-memory: %s", str(e) + ) + return await self._reserve_in_memory(tokens=tokens, scopes=scopes, span=litellm_parent_otel_span) + + async def _reserve_in_memory( + self, + tokens: int, + scopes: tuple[BatchEnqueuedTokenScope, ...], + span: "Span | None", + ) -> BatchEnqueuedTokenOutcome: + async with self._lock: + currents: Final = tuple([await self._get_local_counter(scope, span) for scope in scopes]) + for scope, current in zip(scopes, currents): + if current + tokens > scope.limit: + return BatchEnqueuedTokenOverLimit(scope=scope, enqueued=current) + for scope, current in zip(scopes, currents): + await self._set_local_counter(scope, current + tokens, span) + return BatchEnqueuedTokenReservation(tokens=tokens, scopes=scopes) + + async def refund( + self, + reservation: BatchEnqueuedTokenReservation, + litellm_parent_otel_span: "Span | None" = None, + ) -> None: + if reservation.tokens <= 0 or not reservation.scopes: + return + if self._refund_script is not None: + try: + await self._refund_script( + tuple(self._counter_key(scope) for scope in reservation.scopes), + (reservation.tokens,), + ) + except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory counters + verbose_proxy_logger.warning( + "Redis enqueued-token refund failed, falling back to in-memory: %s", str(e) + ) + else: + return + async with self._lock: + for scope in reservation.scopes: + current = await self._get_local_counter(scope, litellm_parent_otel_span) + remaining = current - reservation.tokens + if remaining <= 0: + self.internal_usage_cache.dual_cache.in_memory_cache.delete_cache(key=self._counter_key(scope)) + else: + await self._set_local_counter(scope, remaining, litellm_parent_otel_span) + + async def save_reservation( + self, + batch_id: str, + reservation: BatchEnqueuedTokenReservation, + litellm_parent_otel_span: "Span | None" = None, + ) -> None: + serialized: Final = _RESERVATION_ADAPTER.dump_json(reservation).decode("utf-8") + if self._save_script is not None: + try: + await self._save_script( + (self._record_key(batch_id),), + (serialized, BATCH_ENQUEUED_TOKEN_TTL_SECONDS), + ) + except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory record + verbose_proxy_logger.warning( + "Redis enqueued-token reservation save failed, falling back to in-memory: %s", str(e) + ) + else: + return + await self.internal_usage_cache.async_set_cache( + key=self._record_key(batch_id), + value=serialized, + ttl=BATCH_ENQUEUED_TOKEN_TTL_SECONDS, + litellm_parent_otel_span=litellm_parent_otel_span, + local_only=True, + ) + + async def pop_reservation( + self, + batch_id: str, + litellm_parent_otel_span: "Span | None" = None, + ) -> BatchEnqueuedTokenReservation | None: + raw: object = None + if self._pop_script is not None: + try: + raw = _POPPED_VALUE_ADAPTER.validate_python(await self._pop_script((self._record_key(batch_id),), ())) + except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory record + verbose_proxy_logger.warning( + "Redis enqueued-token reservation pop failed, falling back to in-memory: %s", str(e) + ) + raw = await self._pop_local_record(batch_id, litellm_parent_otel_span) + else: + raw = await self._pop_local_record(batch_id, litellm_parent_otel_span) + if raw is None: + return None + try: + if isinstance(raw, (str, bytes)): + return _RESERVATION_ADAPTER.validate_json(raw) + return _RESERVATION_ADAPTER.validate_python(raw) + except ValidationError: + verbose_proxy_logger.warning("Discarding malformed enqueued-token reservation record for %s", batch_id) + return None + + async def _pop_local_record(self, batch_id: str, span: "Span | None") -> object: + async with self._lock: + stored = await self.internal_usage_cache.async_get_cache( + key=self._record_key(batch_id), + litellm_parent_otel_span=span, + local_only=True, + ) + if stored is None: + return None + self.internal_usage_cache.dual_cache.in_memory_cache.delete_cache(key=self._record_key(batch_id)) + return stored + + async def _get_local_counter(self, scope: BatchEnqueuedTokenScope, span: "Span | None") -> int: + stored = await self.internal_usage_cache.async_get_cache( + key=self._counter_key(scope), + litellm_parent_otel_span=span, + local_only=True, + ) + return _STORED_COUNTER_ADAPTER.validate_python(stored) or 0 + + async def _set_local_counter(self, scope: BatchEnqueuedTokenScope, value: int, span: "Span | None") -> None: + await self.internal_usage_cache.async_set_cache( + key=self._counter_key(scope), + value=value, + ttl=BATCH_ENQUEUED_TOKEN_TTL_SECONDS, + litellm_parent_otel_span=span, + local_only=True, + ) diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index d6229fb80a6..5b814ad28fd 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -46,9 +46,16 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ( ProxyRateLimitError, map_v3_rate_limit_type, ) +from litellm.proxy.hooks.batch_enqueued_tokens import ( + BatchEnqueuedTokenOverLimit, + BatchEnqueuedTokenReservation, + BatchEnqueuedTokenScope, + resolve_batch_enqueued_token_scopes, +) from litellm.proxy.hooks.parallel_request_limiter_v3 import ( PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY, + get_or_create_request_stash, ) from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit @@ -291,6 +298,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): self, data: dict, user_api_key_dict: UserAPIKeyAuth, + has_enqueued_scopes: bool = False, ) -> tuple[bool, list["RateLimitDescriptor"] | None]: """ Skip downloading batch input files when the operator disabled batch @@ -343,8 +351,10 @@ class _PROXY_BatchRateLimiter(CustomLogger): user_api_key_dict=user_api_key_dict, data=data, ) - if not self._has_applicable_batch_rate_limits(descriptors) and not self._project_has_any_io_token_limits( - user_api_key_dict + if ( + not has_enqueued_scopes + and not self._has_applicable_batch_rate_limits(descriptors) + and not self._project_has_any_io_token_limits(user_api_key_dict) ): verbose_proxy_logger.debug("Skipping batch input file processing: no rate limits configured") return True, None @@ -511,6 +521,59 @@ class _PROXY_BatchRateLimiter(CustomLogger): return file_id, fetch_kwargs + async def _reserve_batch_enqueued_tokens( + self, + user_api_key_dict: UserAPIKeyAuth, + data: Mapping[str, object], + batch_usage: BatchFileUsage, + scopes: tuple[BatchEnqueuedTokenScope, ...], + ) -> None: + """Reserve the batch's estimated tokens against the caller's enqueued-token allowance. + + Runs instead of the per-minute counter charge when the key or team + opted in via ``batch_enqueued_token_limit`` metadata. The reservation + is stashed on the request so the v3 limiter's post-call hooks can + persist it (keyed by the provider batch id) and refund it when the + batch reaches a terminal state. + """ + outcome: Final = await self.parallel_request_limiter.batch_enqueued_token_store.reserve( + tokens=batch_usage.total_tokens, + scopes=scopes, + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + ) + match outcome: + case BatchEnqueuedTokenOverLimit(): + self._raise_enqueued_limit_error(over_limit=outcome, data=data, batch_usage=batch_usage) + case BatchEnqueuedTokenReservation(): + get_or_create_request_stash().batch_enqueued_reservation = outcome + + def _raise_enqueued_limit_error( + self, + over_limit: BatchEnqueuedTokenOverLimit, + data: Mapping[str, object], + batch_usage: BatchFileUsage, + ) -> NoReturn: + scope: Final = over_limit.scope + remaining: Final = max(0, scope.limit - over_limit.enqueued) + detail: Final = ( + f"Batch enqueued token limit exceeded for {scope.key}: {scope.value}. " + f"Batch requires {batch_usage.total_tokens} tokens but only {remaining} enqueued tokens remaining " + f"out of {scope.limit} enqueued token limit. " + f"Tokens free up as running batches complete or are cancelled." + ) + raw_model: Final = data.get("model") + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( + raw_model if isinstance(raw_model, str) else None + ) + raise ProxyRateLimitError( + detail=detail, + headers=MappingProxyType({"rate_limit_type": "tokens"}), + category=RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT, + rate_limit_type=map_v3_rate_limit_type("tokens"), + model=resolved_model, + llm_provider=llm_provider, + ) + def _raise_rate_limit_error( self, status: "RateLimitStatus", @@ -1039,8 +1102,9 @@ class _PROXY_BatchRateLimiter(CustomLogger): verbose_proxy_logger.debug("No input_file_id in batch request, skipping rate limiting") return data + enqueued_scopes: Final = resolve_batch_enqueued_token_scopes(user_api_key_dict) should_skip, batch_rate_limit_descriptors = self._should_skip_batch_input_file_processing( - data=data, user_api_key_dict=user_api_key_dict + data=data, user_api_key_dict=user_api_key_dict, has_enqueued_scopes=bool(enqueued_scopes) ) if should_skip: return data @@ -1066,6 +1130,16 @@ class _PROXY_BatchRateLimiter(CustomLogger): data["_batch_token_count"] = batch_usage.total_tokens data["_batch_request_count"] = batch_usage.request_count + if enqueued_scopes: + await self._reserve_batch_enqueued_tokens( + user_api_key_dict=user_api_key_dict, + data=data, + batch_usage=batch_usage, + scopes=enqueued_scopes, + ) + verbose_proxy_logger.debug("Batch enqueued-token reservation succeeded") + return data + # Directly increment counters by batch amounts (check happens atomically) # This will raise HTTPException if limits are exceeded await self._check_and_increment_batch_counters( diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 2d799ded752..b295da69255 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -44,6 +44,13 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ( ProxyRateLimitError, map_v3_rate_limit_type, ) +from litellm.proxy.hooks.batch_enqueued_tokens import ( + BATCH_ENQUEUED_REFUND_STATUSES, + BatchEnqueuedTokenReservation, + BatchEnqueuedTokenStore, + batch_response_view, + canonical_provider_batch_id, +) from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage @@ -515,6 +522,7 @@ class RequestRateLimiterStash: otpm_reserved_window_identities: frozenset[tuple[str, str, Literal["redis", "local"]]] = field( default_factory=frozenset ) + batch_enqueued_reservation: BatchEnqueuedTokenReservation | None = None reservation_released: bool = False @@ -619,6 +627,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Batch rate limiter (lazy loaded) self._batch_rate_limiter: CallTypeRateLimiter | None = None + self.batch_enqueued_token_store = BatchEnqueuedTokenStore(internal_usage_cache=internal_usage_cache) # Serializes multi-phase check+increment sequences (batch + dynamic # limiters) within this process to close the TOCTOU window between @@ -4673,6 +4682,32 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): except Exception as e: verbose_proxy_logger.exception("Error in rate limit post-call hook: %s", e) + try: + await self._handle_batch_enqueued_post_call(user_api_key_dict=user_api_key_dict, response=response) + except Exception as e: # noqa: BLE001 # post-call batch accounting must never fail the response + verbose_proxy_logger.exception("Error in batch enqueued-token post-call hook: %s", e) + + async def _handle_batch_enqueued_post_call(self, user_api_key_dict: UserAPIKeyAuth, response: object) -> None: + view: Final = batch_response_view(response) + if view is None: + return + span: Final = user_api_key_dict.parent_otel_span + stash: Final = get_request_stash() + if stash is not None and stash.batch_enqueued_reservation is not None: + await self.batch_enqueued_token_store.save_reservation( + batch_id=canonical_provider_batch_id(view.id), + reservation=stash.batch_enqueued_reservation, + litellm_parent_otel_span=span, + ) + stash.batch_enqueued_reservation = None + if view.status in BATCH_ENQUEUED_REFUND_STATUSES: + popped: Final = await self.batch_enqueued_token_store.pop_reservation( + batch_id=canonical_provider_batch_id(view.id), + litellm_parent_otel_span=span, + ) + if popped is not None: + await self.batch_enqueued_token_store.refund(reservation=popped, litellm_parent_otel_span=span) + async def async_post_call_failure_hook( self, request_data: dict, @@ -4706,6 +4741,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) stash.parallel_slot = None + if stash.batch_enqueued_reservation is not None: + await self.batch_enqueued_token_store.refund( + reservation=stash.batch_enqueued_reservation, + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + ) + stash.batch_enqueued_reservation = None + if stash.reservation_released: return reserved_tokens: Final = stash.reserved_tokens diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 1376bdbed38..53bf9739983 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -17,6 +17,7 @@ from __future__ import annotations import json import os +import re import time from datetime import datetime, timedelta, timezone from typing import Callable @@ -57,7 +58,7 @@ from e2e_http import ( unwrap, ) from lifecycle import ResourceManager -from models import KeyGenerateBody, LiteLLMParamsBody, SpendLogRow +from models import KeyGenerateBody, KeyMetadata, LiteLLMParamsBody, SpendLogRow pytestmark = pytest.mark.e2e @@ -685,6 +686,149 @@ class TestBatchRateLimitErrorMapping: ) +BATCH_ENQUEUED_HEADROOM_TOKENS = 100_000 +_BATCH_REQUIRES_TOKENS = re.compile(r"Batch requires (\d+) tokens") + + +class TestBatchEnqueuedTokenLimit: + """Opt-in enqueued-token allowance governs batch submission instead of RPM/TPM. + + A key whose metadata carries batch_enqueued_token_limit reserves the batch's + token estimate against that allowance at create time: per-minute limits no + longer gate batch submission, exhausting the allowance rejects the create + before it reaches the provider, and cancelling a running batch refunds its + reservation so blocked submissions go through again (LIT-5273). + """ + + def _upload_batch_file( + self, client: BatchClient, resources: ResourceManager, key: str + ) -> FileObject: + file = unwrap( + client.upload_file( + content=_multi_request_jsonl("gpt-4o-mini", BATCH_RL_REQUEST_LINES), + form=FileUploadForm(purpose="batch"), + model=OPENAI_BATCH_MODEL, + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + return file + + def _generate_enqueued_key( + self, + client: BatchClient, + resources: ResourceManager, + *, + limit: int, + marker: str, + rpm_limit: int | None = None, + ) -> str: + key = client.proxy.generate_key( + KeyGenerateBody( + models=[], + rpm_limit=rpm_limit, + user_id=f"e2e-batch-enq-{marker}-{unique_marker()}", + metadata=KeyMetadata(batch_enqueued_token_limit=limit), + ) + ) + resources.defer(lambda: client.proxy.delete_key(key)) + return key + + @pytest.mark.covers( + "quota_management.ratelimit.batch_enqueued_tokens.accepts_over_rpm", + exercised_on=["batches"], + ) + def test_enqueued_allowance_accepts_batch_over_key_rpm( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = self._generate_enqueued_key( + client, + resources, + limit=BATCH_ENQUEUED_HEADROOM_TOKENS, + marker="rpm", + rpm_limit=BATCH_RL_RPM_LIMIT, + ) + file = self._upload_batch_file(client, resources, key) + + created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) + + assert created.status_code != 429, ( + f"enqueued-token allowance must govern batch submission instead of the " + f"key RPM ({BATCH_RL_RPM_LIMIT} < {BATCH_RL_REQUEST_LINES} rows); " + f"got 429: {created.body[:400]}" + ) + require_successful_call(created) + batch = BatchObject.model_validate_json(created.body) + resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + + @pytest.mark.covers( + "quota_management.ratelimit.batch_enqueued_tokens.blocks_when_exhausted", + exercised_on=["batches"], + ) + @pytest.mark.covers( + "quota_management.ratelimit.batch_enqueued_tokens.refunds_on_cancel", + exercised_on=["batches"], + ) + def test_exhausted_allowance_blocks_until_cancel_refunds( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + sizing_key = self._generate_enqueued_key( + client, resources, limit=1, marker="size" + ) + sizing_file = self._upload_batch_file(client, resources, sizing_key) + sized = client.create_batch( + body=BatchCreateBody(input_file_id=sizing_file.id), key=sizing_key + ) + assert sized.status_code == 429, ( + f"a 1-token allowance must reject any batch before it reaches the " + f"provider, got {sized.status_code}: {sized.body[:400]}" + ) + assert "batch enqueued token limit exceeded" in sized.body.lower(), ( + f"429 body must name the enqueued token limit, got: {sized.body[:400]}" + ) + requires = _BATCH_REQUIRES_TOKENS.search(sized.body) + assert requires is not None, ( + f"429 body must report the batch token requirement so callers can size " + f"allowances, got: {sized.body[:400]}" + ) + batch_tokens = int(requires.group(1)) + assert batch_tokens > 1 + + key = self._generate_enqueued_key( + client, resources, limit=batch_tokens + batch_tokens // 2, marker="refund" + ) + file = self._upload_batch_file(client, resources, key) + + first = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) + require_successful_call(first) + first_batch = BatchObject.model_validate_json(first.body) + resources.defer(quietly(lambda: client.cancel_batch(first_batch.id, key=key))) + + blocked = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) + assert blocked.status_code == 429, ( + f"second batch must not fit the remaining allowance while the first is " + f"enqueued, got {blocked.status_code}: {blocked.body[:400]}" + ) + assert "batch enqueued token limit exceeded" in blocked.body.lower(), ( + f"429 body must name the enqueued token limit, got: {blocked.body[:400]}" + ) + + cancelled = cancel_batch(client, first_batch.id, key=key, provider=None) + assert cancelled.status in {"cancelling", "cancelled"}, ( + f"cancel must reach a cancel state for the refund to fire, " + f"got {cancelled.status}" + ) + + retried = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) + assert retried.status_code != 429, ( + f"cancelling the first batch must refund its reservation so the retry " + f"fits the allowance, got 429: {retried.body[:400]}" + ) + require_successful_call(retried) + retry_batch = BatchObject.model_validate_json(retried.body) + resources.defer(quietly(lambda: client.cancel_batch(retry_batch.id, key=key))) + + ASSUME_ROLE_RAW_MODEL = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index 2dfa7adddea..4b8aa1da002 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -2,6 +2,9 @@ # litellm/proxy/hooks/ + litellm/proxy/auth/auth_checks.py + litellm/proxy/spend_tracking/. - {id: quota_management.ratelimit.rpm.blocks_over_limit, module: quota_management, tier: P0, behavior: ratelimit, variant: rpm, assertions: [blocks_over_limit], exercised_on: [chat_completions, messages], source: "parallel_request_limiter_v3.py", rationale: "v3 limiter enforces RPM per key/team/model; 429 on breach"} - {id: quota_management.ratelimit.batch_rpm.blocks_over_limit, module: quota_management, tier: P0, behavior: ratelimit, variant: batch_rpm, assertions: [blocks_over_limit], exercised_on: [batches], source: "batch_rate_limiter.py", rationale: "Batch create that exceeds key RPM returns mapped 429 with retry-after"} +- {id: quota_management.ratelimit.batch_enqueued_tokens.accepts_over_rpm, module: quota_management, tier: P0, behavior: ratelimit, variant: batch_enqueued_tokens, assertions: [accepts_over_rpm], exercised_on: [batches], source: "batch_rate_limiter.py + batch_enqueued_tokens.py", rationale: "Key with an enqueued-token allowance submits a batch whose row count exceeds its RPM and the create is accepted"} +- {id: quota_management.ratelimit.batch_enqueued_tokens.blocks_when_exhausted, module: quota_management, tier: P0, behavior: ratelimit, variant: batch_enqueued_tokens, assertions: [blocks_when_exhausted], exercised_on: [batches], source: "batch_rate_limiter.py + batch_enqueued_tokens.py", rationale: "Batch create is rejected with a 429 naming the enqueued token limit once the allowance cannot fit the file"} +- {id: quota_management.ratelimit.batch_enqueued_tokens.refunds_on_cancel, module: quota_management, tier: P0, behavior: ratelimit, variant: batch_enqueued_tokens, assertions: [refunds_on_cancel], exercised_on: [batches], source: "batch_rate_limiter.py + batch_enqueued_tokens.py", rationale: "Cancelling a running batch returns its reserved tokens so a previously blocked submission succeeds"} - {id: quota_management.ratelimit.tpm.blocks_over_limit, module: quota_management, tier: P0, behavior: ratelimit, variant: tpm, assertions: [blocks_over_limit], exercised_on: [chat_completions, messages], source: "parallel_request_limiter_v3.py", rationale: "v3 limiter enforces TPM per key/team/model; 429 on breach"} - {id: quota_management.ratelimit.tpm.excludes_cached_tokens, module: quota_management, tier: P0, behavior: ratelimit, variant: tpm, assertions: [excludes_cached_tokens], exercised_on: [chat_completions], source: "parallel_request_limiter_v3.py:_get_total_tokens_from_usage", rationale: "Cached prompt tokens must not count toward TPM (LIT-1930)"} - {id: quota_management.ratelimit.redis_backed.blocks_over_limit, module: quota_management, tier: P0, behavior: ratelimit, variant: redis_backed, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "parallel_request_limiter_v3.py", rationale: "With Redis configured, RPM still enforces 429 across the shared limiter path customers run multi-replica"} diff --git a/tests/e2e/models.py b/tests/e2e/models.py index df5cb841fad..150be8966ee 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -46,6 +46,7 @@ class KeyLoggingCallback(BaseModel): class KeyMetadata(BaseModel): logging: list[KeyLoggingCallback] | None = None priority: str | None = None + batch_enqueued_token_limit: int | None = None class ObjectPermission(BaseModel): diff --git a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py new file mode 100644 index 00000000000..a917206f33a --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py @@ -0,0 +1,190 @@ +""" +LIT-5273: enqueued-token accounting for batch submissions. + +Covers the ``BatchEnqueuedTokenStore`` (reserve / refund / reservation +records), the metadata-driven scope resolution, and the batch-id and +response-shape helpers the v3 limiter's post-call hooks rely on. +""" + +import base64 +import socket +import uuid +from types import SimpleNamespace + +import pytest + +from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.batch_enqueued_tokens import ( + BatchEnqueuedTokenOverLimit, + BatchEnqueuedTokenReservation, + BatchEnqueuedTokenScope, + BatchEnqueuedTokenStore, + batch_response_view, + canonical_provider_batch_id, + resolve_batch_enqueued_token_scopes, +) +from litellm.proxy.utils import InternalUsageCache + + +def _in_memory_store() -> BatchEnqueuedTokenStore: + return BatchEnqueuedTokenStore(internal_usage_cache=InternalUsageCache(DualCache(default_in_memory_ttl=60))) + + +def _scope(limit: int, key: str = "api_key") -> BatchEnqueuedTokenScope: + return BatchEnqueuedTokenScope(key=key, value=f"{key}-{uuid.uuid4().hex}", limit=limit) + + +def test_scope_resolution_reads_key_and_team_metadata(): + user = UserAPIKeyAuth( + api_key="hashed-key", + metadata={"batch_enqueued_token_limit": 100}, + team_id="team-1", + team_metadata={"batch_enqueued_token_limit": "150"}, + ) + scopes = resolve_batch_enqueued_token_scopes(user) + assert scopes == ( + BatchEnqueuedTokenScope(key="api_key", value="hashed-key", limit=100), + BatchEnqueuedTokenScope(key="team", value="team-1", limit=150), + ) + + +def test_scope_resolution_returns_empty_without_opt_in(): + assert resolve_batch_enqueued_token_scopes(UserAPIKeyAuth(api_key="k")) == () + assert resolve_batch_enqueued_token_scopes(UserAPIKeyAuth(api_key="k", metadata={}, team_metadata=None)) == () + + +@pytest.mark.parametrize("bad_value", ["not-a-number", 0, -5, None, [1000]]) +def test_scope_resolution_ignores_invalid_limits(bad_value): + user = UserAPIKeyAuth(api_key="k", metadata={"batch_enqueued_token_limit": bad_value}) + assert resolve_batch_enqueued_token_scopes(user) == () + + +def test_scope_resolution_skips_team_scope_without_team_id(): + user = UserAPIKeyAuth(api_key="k", team_metadata={"batch_enqueued_token_limit": 100}) + assert resolve_batch_enqueued_token_scopes(user) == () + + +@pytest.mark.asyncio +async def test_reserve_rejects_once_allowance_is_exhausted(): + store = _in_memory_store() + scope = _scope(limit=100) + first = await store.reserve(tokens=80, scopes=(scope,)) + assert isinstance(first, BatchEnqueuedTokenReservation) + second = await store.reserve(tokens=30, scopes=(scope,)) + assert second == BatchEnqueuedTokenOverLimit(scope=scope, enqueued=80) + third = await store.reserve(tokens=20, scopes=(scope,)) + assert isinstance(third, BatchEnqueuedTokenReservation) + + +@pytest.mark.asyncio +async def test_reserve_is_all_or_nothing_across_scopes(): + store = _in_memory_store() + key_scope = _scope(limit=100, key="api_key") + team_scope = _scope(limit=50, key="team") + over = await store.reserve(tokens=60, scopes=(key_scope, team_scope)) + assert over == BatchEnqueuedTokenOverLimit(scope=team_scope, enqueued=0) + exact_fit = await store.reserve(tokens=50, scopes=(key_scope, team_scope)) + assert isinstance(exact_fit, BatchEnqueuedTokenReservation) + + +@pytest.mark.asyncio +async def test_refund_restores_allowance_and_never_goes_negative(): + store = _in_memory_store() + scope = _scope(limit=100) + reservation = await store.reserve(tokens=30, scopes=(scope,)) + assert isinstance(reservation, BatchEnqueuedTokenReservation) + await store.refund(reservation) + await store.refund(reservation) + refill = await store.reserve(tokens=100, scopes=(scope,)) + assert isinstance(refill, BatchEnqueuedTokenReservation) + assert isinstance(await store.reserve(tokens=1, scopes=(scope,)), BatchEnqueuedTokenOverLimit) + + +@pytest.mark.asyncio +async def test_reservation_record_roundtrip_pops_exactly_once(): + store = _in_memory_store() + scope = _scope(limit=100) + reservation = await store.reserve(tokens=40, scopes=(scope,)) + assert isinstance(reservation, BatchEnqueuedTokenReservation) + await store.save_reservation("batch_abc", reservation) + assert await store.pop_reservation("batch_abc") == reservation + assert await store.pop_reservation("batch_abc") is None + assert await store.pop_reservation("batch_never_saved") is None + + +@pytest.mark.asyncio +async def test_zero_token_reserve_charges_nothing(): + store = _in_memory_store() + scope = _scope(limit=100) + empty = await store.reserve(tokens=0, scopes=(scope,)) + assert empty == BatchEnqueuedTokenReservation(tokens=0, scopes=(scope,)) + full = await store.reserve(tokens=100, scopes=(scope,)) + assert isinstance(full, BatchEnqueuedTokenReservation) + + +def test_canonical_provider_batch_id_passes_raw_ids_through(): + assert canonical_provider_batch_id("batch_abc123") == "batch_abc123" + + +def test_canonical_provider_batch_id_decodes_unified_batch_ids(): + unified = "litellm_proxy;model_id:m-1;llm_batch_id:batch_prov_9;llm_output_file_id:file-9" + encoded = base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + assert canonical_provider_batch_id(encoded) == "batch_prov_9" + + +def test_canonical_provider_batch_id_decodes_model_embedded_ids(): + from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model + + encoded = encode_file_id_with_model(file_id="batch_prov_7", model="my-alias", id_type="batch") + assert canonical_provider_batch_id(encoded) == "batch_prov_7" + + +def test_batch_response_view_accepts_batch_objects_only(): + batch = SimpleNamespace(id="batch_1", status="completed", object="batch") + view = batch_response_view(batch) + assert view is not None and view.id == "batch_1" and view.status == "completed" + assert batch_response_view({"id": "chatcmpl-1", "object": "chat.completion"}) is None + assert batch_response_view(None) is None + assert batch_response_view("batch_1") is None + + +def _local_redis_port() -> int | None: + for port in (6379,): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(0.2) + if sock.connect_ex(("127.0.0.1", port)) == 0: + return port + return None + + +@pytest.mark.asyncio +@pytest.mark.skipif(_local_redis_port() is None, reason="requires a local Redis on 6379 for the Lua script path") +async def test_redis_lua_path_full_lifecycle(): + from litellm.caching.redis_cache import RedisCache + + port = _local_redis_port() + redis_cache = RedisCache(host="127.0.0.1", port=port) + store = BatchEnqueuedTokenStore( + internal_usage_cache=InternalUsageCache(DualCache(redis_cache=redis_cache, default_in_memory_ttl=60)) + ) + key_scope = _scope(limit=100, key="api_key") + team_scope = _scope(limit=50, key="team") + + over = await store.reserve(tokens=60, scopes=(key_scope, team_scope)) + assert over == BatchEnqueuedTokenOverLimit(scope=team_scope, enqueued=0) + + reservation = await store.reserve(tokens=50, scopes=(key_scope, team_scope)) + assert isinstance(reservation, BatchEnqueuedTokenReservation) + assert isinstance(await store.reserve(tokens=1, scopes=(key_scope, team_scope)), BatchEnqueuedTokenOverLimit) + + batch_id = f"batch_{uuid.uuid4().hex}" + await store.save_reservation(batch_id, reservation) + popped = await store.pop_reservation(batch_id) + assert popped == reservation + assert await store.pop_reservation(batch_id) is None + + await store.refund(popped) + refill = await store.reserve(tokens=50, scopes=(key_scope, team_scope)) + assert isinstance(refill, BatchEnqueuedTokenReservation) + await store.refund(refill) diff --git a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py index 1ce1a2f3e51..8698ee8ba12 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -2094,3 +2094,195 @@ def test_estimate_entry_output_tokens_multiplies_candidate_count(body_extra, exp } assert rate_limiter._estimate_entry_output_tokens(entry, None) == expected + + +# --------------------------------------------------------------------------- +# LIT-5273: enqueued-token limits govern batch submission when opted in +# --------------------------------------------------------------------------- + + +def _enqueued_rate_limiter(): + from litellm import DualCache + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, + ) + from litellm.proxy.utils import InternalUsageCache + + local_cache = DualCache(default_in_memory_ttl=60) + internal_usage_cache = InternalUsageCache(local_cache) + parallel_request_limiter = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=internal_usage_cache) + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=internal_usage_cache, + parallel_request_limiter=parallel_request_limiter, + ) + return rate_limiter, local_cache + + +_ENQUEUED_BATCH_FILE_CONTENT = ( + b'{"body": {"model": "gpt-4o-mini", "max_tokens": 40, "messages": [{"role": "user", "content": "hi"}]}}\n' + b'{"body": {"model": "gpt-4o-mini", "max_tokens": 40, "messages": [{"role": "user", "content": "hi"}]}}\n' + b'{"body": {"model": "gpt-4o-mini", "max_tokens": 40, "messages": [{"role": "user", "content": "hi"}]}}\n' +) + + +def _enqueued_batch_patches(): + mock_content = MagicMock() + mock_content.content = _ENQUEUED_BATCH_FILE_CONTENT + afile_content_mock = AsyncMock(return_value=mock_content) + return afile_content_mock, ( + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_credentials_for_model", + return_value={"custom_llm_provider": "openai"}, + ), + ) + + +@pytest.mark.asyncio +async def test_enqueued_limit_accepts_batch_over_per_minute_limits(): + """The headline LIT-5273 behavior: a key that opted into an enqueued-token + allowance submits a batch whose row count and token count both exceed its + per-minute RPM/TPM limits, and the batch is accepted (repeatedly) because + only the enqueued allowance governs. Without the opt-in the same key is + rejected on RPM before the batch reaches the provider.""" + from litellm.proxy.hooks.parallel_request_limiter_v3 import get_request_stash + + rate_limiter, local_cache = _enqueued_rate_limiter() + afile_content_mock, patches = _enqueued_batch_patches() + + legacy_user = UserAPIKeyAuth(api_key="sk-legacy-rpm", models=["*"], rpm_limit=1, tpm_limit=10) + opted_in_user = UserAPIKeyAuth( + api_key="sk-enqueued-rpm", + models=["*"], + rpm_limit=1, + tpm_limit=10, + metadata={"batch_enqueued_token_limit": 100000}, + ) + + with patches[0], patches[1], patches[2], patch("litellm.afile_content", new=afile_content_mock): + with pytest.raises(HTTPException) as legacy_exc: + await rate_limiter.async_pre_call_hook( + user_api_key_dict=legacy_user, + cache=local_cache, + data={"input_file_id": "file-abc123", "model": "gpt-4o-mini"}, + call_type="acreate_batch", + ) + assert legacy_exc.value.status_code == 429 + + first_data = {"input_file_id": "file-abc123", "model": "gpt-4o-mini"} + result = await rate_limiter.async_pre_call_hook( + user_api_key_dict=opted_in_user, + cache=local_cache, + data=first_data, + call_type="acreate_batch", + ) + assert result is first_data + stash = get_request_stash() + assert stash is not None and stash.batch_enqueued_reservation is not None + assert stash.batch_enqueued_reservation.tokens == first_data["_batch_token_count"] > 0 + + second = await rate_limiter.async_pre_call_hook( + user_api_key_dict=opted_in_user, + cache=local_cache, + data={"input_file_id": "file-abc123", "model": "gpt-4o-mini"}, + call_type="acreate_batch", + ) + assert second is not None + + +@pytest.mark.asyncio +async def test_enqueued_limit_rejects_when_allowance_is_exhausted(): + """Submissions are rejected pre-provider once the enqueued allowance can't + fit the batch, even for a key with no per-minute limits at all (which + previously skipped batch rate limiting entirely).""" + rate_limiter, local_cache = _enqueued_rate_limiter() + afile_content_mock, patches = _enqueued_batch_patches() + + sizing_user = UserAPIKeyAuth( + api_key="sk-enqueued-sizing", models=["*"], metadata={"batch_enqueued_token_limit": 1000000} + ) + with patches[0], patches[1], patches[2], patch("litellm.afile_content", new=afile_content_mock): + sizing_data = {"input_file_id": "file-abc123", "model": "gpt-4o-mini"} + await rate_limiter.async_pre_call_hook( + user_api_key_dict=sizing_user, + cache=local_cache, + data=sizing_data, + call_type="acreate_batch", + ) + batch_tokens = sizing_data["_batch_token_count"] + assert batch_tokens > 0 + + capped_user = UserAPIKeyAuth( + api_key="sk-enqueued-capped", + models=["*"], + metadata={"batch_enqueued_token_limit": batch_tokens + batch_tokens // 2}, + ) + await rate_limiter.async_pre_call_hook( + user_api_key_dict=capped_user, + cache=local_cache, + data={"input_file_id": "file-abc123", "model": "gpt-4o-mini"}, + call_type="acreate_batch", + ) + with pytest.raises(HTTPException) as exc: + await rate_limiter.async_pre_call_hook( + user_api_key_dict=capped_user, + cache=local_cache, + data={"input_file_id": "file-abc123", "model": "gpt-4o-mini"}, + call_type="acreate_batch", + ) + + assert exc.value.status_code == 429 + assert "Batch enqueued token limit exceeded for api_key" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_enqueued_team_limit_applies_to_batch_submission(): + rate_limiter, local_cache = _enqueued_rate_limiter() + afile_content_mock, patches = _enqueued_batch_patches() + + team_user = UserAPIKeyAuth( + api_key="sk-enqueued-team-key", + models=["*"], + team_id="team-enqueued-batch", + team_metadata={"batch_enqueued_token_limit": 10}, + ) + with patches[0], patches[1], patches[2], patch("litellm.afile_content", new=afile_content_mock): + with pytest.raises(HTTPException) as exc: + await rate_limiter.async_pre_call_hook( + user_api_key_dict=team_user, + cache=local_cache, + data={"input_file_id": "file-abc123", "model": "gpt-4o-mini"}, + call_type="acreate_batch", + ) + + assert exc.value.status_code == 429 + assert "Batch enqueued token limit exceeded for team: team-enqueued-batch" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_disable_flag_still_skips_batch_processing_with_enqueued_limits(): + rate_limiter, local_cache = _enqueued_rate_limiter() + afile_content_mock, _ = _enqueued_batch_patches() + + opted_in_user = UserAPIKeyAuth( + api_key="sk-enqueued-disabled", + models=["*"], + metadata={"batch_enqueued_token_limit": 10}, + ) + with ( + patch("litellm.proxy.proxy_server.general_settings", {"disable_batch_input_file_rate_limiting": True}), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch("litellm.afile_content", new=afile_content_mock), + ): + data = {"input_file_id": "file-abc123", "model": "gpt-4o-mini"} + result = await rate_limiter.async_pre_call_hook( + user_api_key_dict=opted_in_user, + cache=local_cache, + data=data, + call_type="acreate_batch", + ) + + assert result is data + afile_content_mock.assert_not_awaited() diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 54366226dfb..68c219fd72d 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -5858,3 +5858,123 @@ async def test_conflicting_token_limits_cannot_bypass_tpm_reservation(): ) assert exc_info.value.status_code == 429 + + +# --------------------------------------------------------------------------- +# LIT-5273: batch enqueued-token reservations in the post-call hooks +# --------------------------------------------------------------------------- + + +def _enqueued_test_handler() -> _PROXY_MaxParallelRequestsHandler: + return _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache(default_in_memory_ttl=60))) + + +def _batch_response(batch_id: str, status: str): + from types import SimpleNamespace + + return SimpleNamespace(id=batch_id, status=status, object="batch") + + +@pytest.mark.asyncio +async def test_success_hook_persists_batch_enqueued_reservation_and_refunds_on_completion(): + from litellm.proxy.hooks.batch_enqueued_tokens import ( + BatchEnqueuedTokenOverLimit, + BatchEnqueuedTokenReservation, + BatchEnqueuedTokenScope, + ) + + handler = _enqueued_test_handler() + store = handler.batch_enqueued_token_store + scope = BatchEnqueuedTokenScope(key="api_key", value="hashed-enqueued-key", limit=100) + user = UserAPIKeyAuth(api_key="hashed-enqueued-key") + + reservation = await store.reserve(tokens=60, scopes=(scope,)) + assert isinstance(reservation, BatchEnqueuedTokenReservation) + get_or_create_request_stash().batch_enqueued_reservation = reservation + + await handler.async_post_call_success_hook( + data={}, user_api_key_dict=user, response=_batch_response("batch_enq_1", "validating") + ) + assert get_request_stash().batch_enqueued_reservation is None + assert isinstance(await store.reserve(tokens=50, scopes=(scope,)), BatchEnqueuedTokenOverLimit) + + await handler.async_post_call_success_hook( + data={}, user_api_key_dict=user, response=_batch_response("batch_enq_1", "completed") + ) + refill = await store.reserve(tokens=40, scopes=(scope,)) + assert isinstance(refill, BatchEnqueuedTokenReservation) + + await handler.async_post_call_success_hook( + data={}, user_api_key_dict=user, response=_batch_response("batch_enq_1", "completed") + ) + assert isinstance(await store.reserve(tokens=70, scopes=(scope,)), BatchEnqueuedTokenOverLimit) + + +@pytest.mark.asyncio +async def test_success_hook_refunds_batch_enqueued_reservation_on_cancellation(): + from litellm.proxy.hooks.batch_enqueued_tokens import ( + BatchEnqueuedTokenReservation, + BatchEnqueuedTokenScope, + ) + + handler = _enqueued_test_handler() + store = handler.batch_enqueued_token_store + scope = BatchEnqueuedTokenScope(key="team", value="team-enqueued", limit=100) + user = UserAPIKeyAuth(api_key="hashed-enqueued-key", team_id="team-enqueued") + + reservation = await store.reserve(tokens=90, scopes=(scope,)) + assert isinstance(reservation, BatchEnqueuedTokenReservation) + get_or_create_request_stash().batch_enqueued_reservation = reservation + await handler.async_post_call_success_hook( + data={}, user_api_key_dict=user, response=_batch_response("batch_enq_2", "validating") + ) + + await handler.async_post_call_success_hook( + data={}, user_api_key_dict=user, response=_batch_response("batch_enq_2", "cancelling") + ) + assert isinstance(await store.reserve(tokens=100, scopes=(scope,)), BatchEnqueuedTokenReservation) + + +@pytest.mark.asyncio +async def test_failure_hook_refunds_stashed_batch_enqueued_reservation(): + from litellm.proxy.hooks.batch_enqueued_tokens import ( + BatchEnqueuedTokenReservation, + BatchEnqueuedTokenScope, + ) + + handler = _enqueued_test_handler() + store = handler.batch_enqueued_token_store + scope = BatchEnqueuedTokenScope(key="api_key", value="hashed-failing-key", limit=100) + user = UserAPIKeyAuth(api_key="hashed-failing-key") + + reservation = await store.reserve(tokens=80, scopes=(scope,)) + assert isinstance(reservation, BatchEnqueuedTokenReservation) + get_or_create_request_stash().batch_enqueued_reservation = reservation + + await handler.async_post_call_failure_hook( + request_data={}, original_exception=Exception("guardrail rejected"), user_api_key_dict=user + ) + assert get_request_stash().batch_enqueued_reservation is None + assert isinstance(await store.reserve(tokens=100, scopes=(scope,)), BatchEnqueuedTokenReservation) + + +@pytest.mark.asyncio +async def test_success_hook_leaves_stash_untouched_for_non_batch_responses(): + from litellm.proxy.hooks.batch_enqueued_tokens import ( + BatchEnqueuedTokenReservation, + BatchEnqueuedTokenScope, + ) + + handler = _enqueued_test_handler() + store = handler.batch_enqueued_token_store + scope = BatchEnqueuedTokenScope(key="api_key", value="hashed-chat-key", limit=100) + user = UserAPIKeyAuth(api_key="hashed-chat-key") + + reservation = await store.reserve(tokens=10, scopes=(scope,)) + assert isinstance(reservation, BatchEnqueuedTokenReservation) + get_or_create_request_stash().batch_enqueued_reservation = reservation + + await handler.async_post_call_success_hook( + data={}, user_api_key_dict=user, response=ModelResponse(usage=Usage(total_tokens=5)) + ) + assert get_request_stash().batch_enqueued_reservation == reservation From 677ef1e317080c20aec895a7ed25c058a1e18582 Mon Sep 17 00:00:00 2001 From: ljogeiger <39981740+ljogeiger@users.noreply.github.com> Date: Wed, 19 Aug 2026 22:21:31 +0000 Subject: [PATCH 019/281] docs(vertex_ai): drop stale note about the removed model argument --- litellm/llms/vertex_ai/gemini/transformation.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index f2d318a9ffd..11c026010ee 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -645,10 +645,9 @@ def _collect_tool_call_thought_signatures( the text part as well would send two copies and double-bill the previous turn's reasoning tokens on gemini-3 and newer models. - Detection deliberately calls _get_thought_signature_from_tool without the - model argument: with a gemini-3 model that helper synthesizes a dummy - signature for unsigned tool calls, which must not suppress a real - text-part signature (e.g. replaying gemini-2.5 history to a newer model). + Only real signatures count here; a synthesized placeholder must not + suppress a genuine text-part signature (e.g. replaying gemini-2.5 history + to a newer model). """ signatures: tuple[str, ...] = () From d5af42717e9713771b8a014455b79d22b0756754 Mon Sep 17 00:00:00 2001 From: ljogeiger <39981740+ljogeiger@users.noreply.github.com> Date: Wed, 19 Aug 2026 22:32:24 +0000 Subject: [PATCH 020/281] test(vertex_ai): cover id-embedded, tool-level, and end-to-end parallel signature replay --- .../test_vertex_ai_gemini_transformation.py | 181 ++++++++++++++++++ 1 file changed, 181 insertions(+) diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index 28c551ab824..562f27c11cf 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -806,6 +806,27 @@ def _parallel_tool_calls(*signatures): ] +def _parallel_tool_calls_signed_via_id(*signatures): + """Parallel tool calls in the shape LiteLLM actually hands back to clients. + + The signature rides in the tool call id behind __thought__, which is what an + OpenAI-format client echoes back on the next turn. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + _encode_tool_call_id_with_signature, + ) + + return [ + { + "id": _encode_tool_call_id_with_signature(f"call_{idx}", signature), + "type": "function", + "function": {"name": f"tool_{idx}", "arguments": '{"location": "Paris"}'}, + "index": idx, + } + for idx, signature in enumerate(signatures) + ] + + REAL_THOUGHT_SIGNATURE = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n" @@ -899,6 +920,166 @@ def test_no_signatures_on_parallel_tool_calls_for_gemini_2_5(): assert all("thoughtSignature" not in part for part in gemini_parts) +def test_signature_embedded_in_tool_call_id_only_on_first_parallel_call(): + """The production shape: the signature arrives inside the first call's id, siblings have bare ids.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls_signed_via_id( + REAL_THOUGHT_SIGNATURE, None, None + ), + }, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 3 + assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + assert "thoughtSignature" not in gemini_parts[2] + + +def test_tool_level_provider_specific_fields_signature_leaves_siblings_empty(): + """A signature on the tool call itself, rather than on its function, behaves the same way.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + tool_calls = _parallel_tool_calls(None, None) + tool_calls[0]["provider_specific_fields"] = { + "thought_signature": REAL_THOUGHT_SIGNATURE + } + + gemini_parts = convert_to_gemini_tool_call_invoke( + {"role": "assistant", "content": None, "tool_calls": tool_calls}, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 2 + assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + + +def test_placeholder_lands_on_first_emitted_part_not_first_tool_call_entry(): + """A non-function entry (e.g. an OpenAI custom tool call) emits no part, so it must not + consume the one placeholder slot and leave the real first function call bare.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + _get_dummy_thought_signature, + convert_to_gemini_tool_call_invoke, + ) + + tool_calls = [ + {"id": "call_custom", "type": "custom", "custom": {"name": "noop", "input": ""}} + ] + _parallel_tool_calls(None, None) + + gemini_parts = convert_to_gemini_tool_call_invoke( + {"role": "assistant", "content": None, "tool_calls": tool_calls}, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 2 + assert gemini_parts[0]["thoughtSignature"] == _get_dummy_thought_signature() + assert "thoughtSignature" not in gemini_parts[1] + + +def test_no_placeholder_when_model_is_unknown(): + """Without a model there is nothing to prove the target needs a placeholder, so none is added.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(None, None), + }, + ) + + assert len(gemini_parts) == 2 + assert all("thoughtSignature" not in part for part in gemini_parts) + + +def test_real_signature_forwarded_to_gemini_2_5_without_placeholder_siblings(): + """Older models still receive a real signature that a client replays, and still get no placeholder.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(REAL_THOUGHT_SIGNATURE, None), + }, + model="gemini-2.5-flash", + ) + + assert len(gemini_parts) == 2 + assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + + +def test_parallel_tool_call_history_replayed_through_full_message_conversion(): + """End to end through the message-history converter, the path a real /chat/completions replay takes.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + messages = [ + {"role": "user", "content": "Weather in Paris, London and Tokyo?"}, + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls_signed_via_id( + REAL_THOUGHT_SIGNATURE, None, None + ), + }, + ] + + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-3-pro-preview" + ) + + model_parts = contents[1]["parts"] + assert len(model_parts) == 3 + assert model_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in model_parts[1] + assert "thoughtSignature" not in model_parts[2] + + +def test_signed_text_part_survives_alongside_unsigned_parallel_tool_calls(): + """gemini-2.5 history with a signed text part and parallel unsigned calls: the text keeps its real + signature, only the first call gets the placeholder, and the siblings stay bare.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + _get_dummy_thought_signature, + ) + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": "Checking all three cities.", + "provider_specific_fields": {"thought_signatures": ["real_25_signature"]}, + "tool_calls": _parallel_tool_calls(None, None, None), + } + + parts = _gemini_convert_messages_with_history( + messages=[msg], model="gemini-3-pro-preview" + )[0]["parts"] + + assert parts[0]["text"] == "Checking all three cities." + assert parts[0]["thoughtSignature"] == "real_25_signature" + assert parts[1]["thoughtSignature"] == _get_dummy_thought_signature() + assert "thoughtSignature" not in parts[2] + assert "thoughtSignature" not in parts[3] + + # Tests for media_resolution (detail parameter) handling - Issue #17084 class TestMediaResolution: """Tests for media_resolution handling in Gemini 2.x models""" From 160d3dac42ddbcde36a538303bfc9612f6e02487 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:40:15 -0700 Subject: [PATCH 021/281] fix(proxy): issue enqueued-token Lua calls one key at a time for Redis Cluster compatibility --- litellm/proxy/hooks/batch_enqueued_tokens.py | 85 +++++++++++-------- .../proxy/hooks/test_batch_enqueued_tokens.py | 60 ++++++++++++- 2 files changed, 109 insertions(+), 36 deletions(-) diff --git a/litellm/proxy/hooks/batch_enqueued_tokens.py b/litellm/proxy/hooks/batch_enqueued_tokens.py index 040c8b687a7..e484b6a59fe 100644 --- a/litellm/proxy/hooks/batch_enqueued_tokens.py +++ b/litellm/proxy/hooks/batch_enqueued_tokens.py @@ -38,27 +38,20 @@ ScopeKey: TypeAlias = Literal["api_key", "team"] RESERVE_ENQUEUED_TOKENS_SCRIPT: Final = """ local amount = tonumber(ARGV[1]) local ttl = tonumber(ARGV[2]) -for i = 1, #KEYS do - local limit = tonumber(ARGV[2 + i]) - local current = tonumber(redis.call('GET', KEYS[i]) or '0') - if current + amount > limit then - return {0, i - 1, current} - end +local limit = tonumber(ARGV[3]) +local current = tonumber(redis.call('GET', KEYS[1]) or '0') +if current + amount > limit then + return {0, current} end -for i = 1, #KEYS do - redis.call('INCRBY', KEYS[i], amount) - redis.call('EXPIRE', KEYS[i], ttl) -end -return {1, -1, 0} +local updated = redis.call('INCRBY', KEYS[1], amount) +redis.call('EXPIRE', KEYS[1], ttl) +return {1, updated} """ REFUND_ENQUEUED_TOKENS_SCRIPT: Final = """ -local amount = tonumber(ARGV[1]) -for i = 1, #KEYS do - local updated = redis.call('DECRBY', KEYS[i], amount) - if updated <= 0 then - redis.call('DEL', KEYS[i]) - end +local updated = redis.call('DECRBY', KEYS[1], tonumber(ARGV[1])) +if updated <= 0 then + redis.call('DEL', KEYS[1]) end return 1 """ @@ -99,7 +92,7 @@ class BatchEnqueuedTokenOverLimit: BatchEnqueuedTokenOutcome: TypeAlias = BatchEnqueuedTokenReservation | BatchEnqueuedTokenOverLimit _LIMIT_ADAPTER: Final = TypeAdapter(Annotated[int, Field(gt=0)]) -_RESERVE_RESULT_ADAPTER: Final = TypeAdapter(tuple[int, int, int]) +_RESERVE_RESULT_ADAPTER: Final = TypeAdapter(tuple[int, int]) _POPPED_VALUE_ADAPTER: Final = TypeAdapter(str | bytes | None) _STORED_COUNTER_ADAPTER: Final = TypeAdapter(int | None) _RESERVATION_ADAPTER: Final = TypeAdapter(BatchEnqueuedTokenReservation) @@ -175,9 +168,11 @@ def batch_response_view(response: object) -> _BatchResponseView | None: class BatchEnqueuedTokenStore: """Tracks enqueued batch tokens per scope, plus per-batch reservation records for refunds. - Counters and records live in Redis (via atomic Lua scripts) when Redis is - configured; otherwise a single-process in-memory fallback guarded by one - asyncio lock is used. Everything expires after + Counters and records live in Redis when Redis is configured, through + single-key Lua scripts issued one scope at a time (Redis Cluster safe: no + cross-slot commands), with an over-limit scope rolling back the scopes + reserved before it; otherwise a single-process in-memory fallback guarded + by one asyncio lock is used. Everything expires after ``BATCH_ENQUEUED_TOKEN_TTL_SECONDS`` so a crash between submission and the terminal-state refund can never leak tokens forever. """ @@ -215,22 +210,44 @@ class BatchEnqueuedTokenStore: ) -> BatchEnqueuedTokenOutcome: if tokens <= 0 or not scopes: return BatchEnqueuedTokenReservation(tokens=max(tokens, 0), scopes=scopes) - if self._reserve_script is not None: + reserve_script: Final = self._reserve_script + refund_script: Final = self._refund_script + if reserve_script is not None and refund_script is not None: try: - raw_result = await self._reserve_script( - tuple(self._counter_key(scope) for scope in scopes), - (tokens, BATCH_ENQUEUED_TOKEN_TTL_SECONDS, *(scope.limit for scope in scopes)), - ) - result = _RESERVE_RESULT_ADAPTER.validate_python(raw_result) - if result[0] == 1: - return BatchEnqueuedTokenReservation(tokens=tokens, scopes=scopes) - return BatchEnqueuedTokenOverLimit(scope=scopes[result[1]], enqueued=result[2]) + return await self._reserve_via_redis(reserve_script, refund_script, tokens=tokens, scopes=scopes) except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory counters verbose_proxy_logger.warning( "Redis enqueued-token reserve failed, falling back to in-memory: %s", str(e) ) return await self._reserve_in_memory(tokens=tokens, scopes=scopes, span=litellm_parent_otel_span) + async def _reserve_via_redis( + self, + reserve_script: _ScriptRunner, + refund_script: _ScriptRunner, + tokens: int, + scopes: tuple[BatchEnqueuedTokenScope, ...], + ) -> BatchEnqueuedTokenOutcome: + for index, scope in enumerate(scopes): + raw_result = await reserve_script( + (self._counter_key(scope),), + (tokens, BATCH_ENQUEUED_TOKEN_TTL_SECONDS, scope.limit), + ) + result = _RESERVE_RESULT_ADAPTER.validate_python(raw_result) + if result[0] != 1: + await self._refund_via_redis(refund_script, tokens=tokens, scopes=scopes[:index]) + return BatchEnqueuedTokenOverLimit(scope=scope, enqueued=result[1]) + return BatchEnqueuedTokenReservation(tokens=tokens, scopes=scopes) + + async def _refund_via_redis( + self, + refund_script: _ScriptRunner, + tokens: int, + scopes: tuple[BatchEnqueuedTokenScope, ...], + ) -> None: + for scope in scopes: + await refund_script((self._counter_key(scope),), (tokens,)) + async def _reserve_in_memory( self, tokens: int, @@ -253,12 +270,10 @@ class BatchEnqueuedTokenStore: ) -> None: if reservation.tokens <= 0 or not reservation.scopes: return - if self._refund_script is not None: + refund_script: Final = self._refund_script + if refund_script is not None: try: - await self._refund_script( - tuple(self._counter_key(scope) for scope in reservation.scopes), - (reservation.tokens,), - ) + await self._refund_via_redis(refund_script, tokens=reservation.tokens, scopes=reservation.scopes) except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory counters verbose_proxy_logger.warning( "Redis enqueued-token refund failed, falling back to in-memory: %s", str(e) diff --git a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py index a917206f33a..1bb4798eaa0 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py +++ b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py @@ -9,7 +9,9 @@ response-shape helpers the v3 limiter's post-call hooks rely on. import base64 import socket import uuid -from types import SimpleNamespace +from collections.abc import Mapping, Sequence +from types import MappingProxyType, SimpleNamespace +from typing import Final import pytest @@ -123,6 +125,62 @@ async def test_zero_token_reserve_charges_nothing(): assert isinstance(full, BatchEnqueuedTokenReservation) +class _SingleKeyRedisFake: + """Emulates the Redis script path one single-key call at a time, recording every call.""" + + def __init__(self) -> None: + self.script_calls: tuple[tuple[str, tuple[str, ...]], ...] = () + self.counters: Mapping[str, int] = MappingProxyType({}) + + def async_register_script(self, script: str): + kind: Final = "reserve" if "INCRBY" in script else "refund" if "DECRBY" in script else "record" + + async def run(keys: Sequence[str], args: Sequence[str | bytes | int | float]) -> object: + self.script_calls = (*self.script_calls, (kind, tuple(keys))) + return self._run(kind, tuple(keys), tuple(args)) + + return run + + def _run(self, kind: str, keys: tuple[str, ...], args: tuple[str | bytes | int | float, ...]) -> object: + if kind == "reserve": + amount, limit = int(args[0]), int(args[2]) + current: Final = self.counters.get(keys[0], 0) + if current + amount > limit: + return (0, current) + self.counters = MappingProxyType({**self.counters, keys[0]: current + amount}) + return (1, current + amount) + if kind == "refund": + remaining: Final = self.counters.get(keys[0], 0) - int(args[0]) + self.counters = MappingProxyType( + {key: value for key, value in self.counters.items() if key != keys[0]} + if remaining <= 0 + else {**self.counters, keys[0]: remaining} + ) + return 1 + raise AssertionError(f"unexpected {kind} script call for keys {keys}") + + +@pytest.mark.asyncio +async def test_redis_reserve_issues_single_key_calls_and_rolls_back_on_over_limit(): + fake = _SingleKeyRedisFake() + store = BatchEnqueuedTokenStore( + internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60)) + ) + key_scope = _scope(limit=100, key="api_key") + team_scope = _scope(limit=50, key="team") + + over = await store.reserve(tokens=60, scopes=(key_scope, team_scope)) + assert over == BatchEnqueuedTokenOverLimit(scope=team_scope, enqueued=0) + assert tuple(kind for kind, _ in fake.script_calls) == ("reserve", "reserve", "refund") + assert not fake.counters + + fits = await store.reserve(tokens=50, scopes=(key_scope, team_scope)) + assert isinstance(fits, BatchEnqueuedTokenReservation) + await store.refund(fits) + assert not fake.counters + assert all(len(keys) == 1 for _, keys in fake.script_calls) + + def test_canonical_provider_batch_id_passes_raw_ids_through(): assert canonical_provider_batch_id("batch_abc123") == "batch_abc123" From db50e123d5f23d475dab0bc62e33364ea19df817 Mon Sep 17 00:00:00 2001 From: ljogeiger <39981740+ljogeiger@users.noreply.github.com> Date: Wed, 19 Aug 2026 22:42:24 +0000 Subject: [PATCH 022/281] test(vertex_ai): parametrize placeholder scoping across gemini-3 model variants --- .../test_vertex_ai_gemini_transformation.py | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index 562f27c11cf..9ebf6db11de 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -1,5 +1,7 @@ import base64 +import pytest + from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_gemini_tool_call_result, ) @@ -1052,6 +1054,40 @@ def test_parallel_tool_call_history_replayed_through_full_message_conversion(): assert "thoughtSignature" not in model_parts[2] +@pytest.mark.parametrize( + "model", + [ + "gemini-3-pro-preview", + "gemini-3-flash-preview", + "gemini-3.1-pro-preview", + "gemini-3.6-flash", + "gemini-3.7-flash", + "vertex_ai/gemini-3.7-flash", + "gemini/gemini-3.7-flash", + ], +) +def test_placeholder_scoped_to_first_call_across_gemini_3_variants(model): + """Every gemini-3 family member, bare or provider-prefixed, gets one placeholder at most.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + _get_dummy_thought_signature, + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(None, None, None), + }, + model=model, + ) + + assert len(gemini_parts) == 3 + assert gemini_parts[0]["thoughtSignature"] == _get_dummy_thought_signature() + assert "thoughtSignature" not in gemini_parts[1] + assert "thoughtSignature" not in gemini_parts[2] + + def test_signed_text_part_survives_alongside_unsigned_parallel_tool_calls(): """gemini-2.5 history with a signed text part and parallel unsigned calls: the text keeps its real signature, only the first call gets the placeholder, and the siblings stay bare.""" From 50896f21b346a08d99047a0f4e72f16e1e7dc54f Mon Sep 17 00:00:00 2001 From: mateo-berri Date: Wed, 19 Aug 2026 16:16:08 -0700 Subject: [PATCH 023/281] fix(batch_enqueued_tokens): roll back partial reserves, route refunds by backend, lowercase terminal statuses Reserve-script failures now roll back the scopes already incremented before re-raising into the in-memory fallback, so a partial redis outage no longer leaks counter increments that shrink the shared allowance. Reservations record which backend granted them, so a refund never debits redis counters an in-memory grant did not charge. Terminal-status matching is now case-insensitive because the Bedrock async-invoke retrieve path returns raw AWS-cased statuses like Completed. --- litellm/proxy/hooks/batch_enqueued_tokens.py | 59 +++++++++++++++---- .../hooks/parallel_request_limiter_v3.py | 2 +- .../proxy/hooks/test_batch_enqueued_tokens.py | 41 ++++++++++++- .../hooks/test_parallel_request_limiter_v3.py | 27 +++++++++ 4 files changed, 117 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/hooks/batch_enqueued_tokens.py b/litellm/proxy/hooks/batch_enqueued_tokens.py index e484b6a59fe..570291aa452 100644 --- a/litellm/proxy/hooks/batch_enqueued_tokens.py +++ b/litellm/proxy/hooks/batch_enqueued_tokens.py @@ -77,10 +77,14 @@ class BatchEnqueuedTokenScope: limit: int +ReservationBackend: TypeAlias = Literal["redis", "memory"] + + @dataclass(frozen=True, slots=True) class BatchEnqueuedTokenReservation: tokens: int scopes: tuple[BatchEnqueuedTokenScope, ...] + backend: ReservationBackend = "redis" @dataclass(frozen=True, slots=True) @@ -170,9 +174,10 @@ class BatchEnqueuedTokenStore: Counters and records live in Redis when Redis is configured, through single-key Lua scripts issued one scope at a time (Redis Cluster safe: no - cross-slot commands), with an over-limit scope rolling back the scopes - reserved before it; otherwise a single-process in-memory fallback guarded - by one asyncio lock is used. Everything expires after + cross-slot commands), with an over-limit or failing scope rolling back the + scopes reserved before it; otherwise a single-process in-memory fallback + guarded by one asyncio lock is used. Reservations remember which backend + granted them so a refund never debits counters the grant did not charge. Everything expires after ``BATCH_ENQUEUED_TOKEN_TTL_SECONDS`` so a crash between submission and the terminal-state refund can never leak tokens forever. """ @@ -229,15 +234,49 @@ class BatchEnqueuedTokenStore: scopes: tuple[BatchEnqueuedTokenScope, ...], ) -> BatchEnqueuedTokenOutcome: for index, scope in enumerate(scopes): - raw_result = await reserve_script( - (self._counter_key(scope),), - (tokens, BATCH_ENQUEUED_TOKEN_TTL_SECONDS, scope.limit), + result = await self._run_reserve_script( + reserve_script, + refund_script, + tokens=tokens, + scope=scope, + already_reserved=scopes[:index], ) - result = _RESERVE_RESULT_ADAPTER.validate_python(raw_result) if result[0] != 1: await self._refund_via_redis(refund_script, tokens=tokens, scopes=scopes[:index]) return BatchEnqueuedTokenOverLimit(scope=scope, enqueued=result[1]) - return BatchEnqueuedTokenReservation(tokens=tokens, scopes=scopes) + return BatchEnqueuedTokenReservation(tokens=tokens, scopes=scopes, backend="redis") + + async def _run_reserve_script( + self, + reserve_script: _ScriptRunner, + refund_script: _ScriptRunner, + tokens: int, + scope: BatchEnqueuedTokenScope, + already_reserved: tuple[BatchEnqueuedTokenScope, ...], + ) -> tuple[int, int]: + try: + raw_result: Final = await reserve_script( + (self._counter_key(scope),), + (tokens, BATCH_ENQUEUED_TOKEN_TTL_SECONDS, scope.limit), + ) + return _RESERVE_RESULT_ADAPTER.validate_python(raw_result) + except Exception: + await self._rollback_partial_reserve(refund_script, tokens=tokens, scopes=already_reserved) + raise + + async def _rollback_partial_reserve( + self, + refund_script: _ScriptRunner, + tokens: int, + scopes: tuple[BatchEnqueuedTokenScope, ...], + ) -> None: + try: + await self._refund_via_redis(refund_script, tokens=tokens, scopes=scopes) + except Exception as e: # noqa: BLE001 # best-effort rollback: the leak is TTL-bounded and only tightens the allowance + verbose_proxy_logger.warning( + "Rollback of partially reserved enqueued tokens failed; leaked increments expire with the TTL: %s", + str(e), + ) async def _refund_via_redis( self, @@ -261,7 +300,7 @@ class BatchEnqueuedTokenStore: return BatchEnqueuedTokenOverLimit(scope=scope, enqueued=current) for scope, current in zip(scopes, currents): await self._set_local_counter(scope, current + tokens, span) - return BatchEnqueuedTokenReservation(tokens=tokens, scopes=scopes) + return BatchEnqueuedTokenReservation(tokens=tokens, scopes=scopes, backend="memory") async def refund( self, @@ -271,7 +310,7 @@ class BatchEnqueuedTokenStore: if reservation.tokens <= 0 or not reservation.scopes: return refund_script: Final = self._refund_script - if refund_script is not None: + if reservation.backend == "redis" and refund_script is not None: try: await self._refund_via_redis(refund_script, tokens=reservation.tokens, scopes=reservation.scopes) except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory counters diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index b295da69255..1e65da5b867 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -4700,7 +4700,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): litellm_parent_otel_span=span, ) stash.batch_enqueued_reservation = None - if view.status in BATCH_ENQUEUED_REFUND_STATUSES: + if view.status.lower() in BATCH_ENQUEUED_REFUND_STATUSES: popped: Final = await self.batch_enqueued_token_store.pop_reservation( batch_id=canonical_provider_batch_id(view.id), litellm_parent_otel_span=span, diff --git a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py index 1bb4798eaa0..bc924c32ba3 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py +++ b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py @@ -128,9 +128,10 @@ async def test_zero_token_reserve_charges_nothing(): class _SingleKeyRedisFake: """Emulates the Redis script path one single-key call at a time, recording every call.""" - def __init__(self) -> None: + def __init__(self, fail_reserve_keys: frozenset[str] = frozenset()) -> None: self.script_calls: tuple[tuple[str, tuple[str, ...]], ...] = () self.counters: Mapping[str, int] = MappingProxyType({}) + self.fail_reserve_keys = fail_reserve_keys def async_register_script(self, script: str): kind: Final = "reserve" if "INCRBY" in script else "refund" if "DECRBY" in script else "record" @@ -143,6 +144,8 @@ class _SingleKeyRedisFake: def _run(self, kind: str, keys: tuple[str, ...], args: tuple[str | bytes | int | float, ...]) -> object: if kind == "reserve": + if keys[0] in self.fail_reserve_keys: + raise ConnectionError(f"simulated redis failure for {keys[0]}") amount, limit = int(args[0]), int(args[2]) current: Final = self.counters.get(keys[0], 0) if current + amount > limit: @@ -181,6 +184,42 @@ async def test_redis_reserve_issues_single_key_calls_and_rolls_back_on_over_limi assert all(len(keys) == 1 for _, keys in fake.script_calls) +@pytest.mark.asyncio +async def test_partial_redis_reserve_failure_rolls_back_and_grants_in_memory(): + key_scope = _scope(limit=100, key="api_key") + team_scope = _scope(limit=50, key="team") + fake = _SingleKeyRedisFake(fail_reserve_keys=frozenset({f"batch_enqueued_tokens:team:{team_scope.value}"})) + store = BatchEnqueuedTokenStore( + internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60)) + ) + + outcome = await store.reserve(tokens=10, scopes=(key_scope, team_scope)) + assert isinstance(outcome, BatchEnqueuedTokenReservation) + assert outcome.backend == "memory" + assert tuple(kind for kind, _ in fake.script_calls) == ("reserve", "reserve", "refund") + assert not fake.counters + + await store.refund(outcome) + assert tuple(kind for kind, _ in fake.script_calls) == ("reserve", "reserve", "refund") + + refilled = await store.reserve(tokens=50, scopes=(team_scope,)) + assert isinstance(refilled, BatchEnqueuedTokenReservation) + assert refilled.backend == "memory" + + +@pytest.mark.asyncio +async def test_pop_reservation_defaults_legacy_records_to_redis_backend(): + store = _in_memory_store() + legacy = '{"tokens": 5, "scopes": [{"key": "api_key", "value": "k", "limit": 10}]}' + store.internal_usage_cache.dual_cache.in_memory_cache.set_cache( + key="batch_enqueued_token_reservation:batch_legacy", value=legacy + ) + popped = await store.pop_reservation("batch_legacy") + assert popped == BatchEnqueuedTokenReservation( + tokens=5, scopes=(BatchEnqueuedTokenScope(key="api_key", value="k", limit=10),), backend="redis" + ) + + def test_canonical_provider_batch_id_passes_raw_ids_through(): assert canonical_provider_batch_id("batch_abc123") == "batch_abc123" diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 68c219fd72d..fee892ad342 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -5935,6 +5935,33 @@ async def test_success_hook_refunds_batch_enqueued_reservation_on_cancellation() assert isinstance(await store.reserve(tokens=100, scopes=(scope,)), BatchEnqueuedTokenReservation) +@pytest.mark.asyncio +async def test_success_hook_refunds_on_provider_cased_terminal_status(): + from litellm.proxy.hooks.batch_enqueued_tokens import ( + BatchEnqueuedTokenOverLimit, + BatchEnqueuedTokenReservation, + BatchEnqueuedTokenScope, + ) + + handler = _enqueued_test_handler() + store = handler.batch_enqueued_token_store + scope = BatchEnqueuedTokenScope(key="api_key", value="hashed-enqueued-key", limit=100) + user = UserAPIKeyAuth(api_key="hashed-enqueued-key") + + reservation = await store.reserve(tokens=90, scopes=(scope,)) + assert isinstance(reservation, BatchEnqueuedTokenReservation) + get_or_create_request_stash().batch_enqueued_reservation = reservation + await handler.async_post_call_success_hook( + data={}, user_api_key_dict=user, response=_batch_response("batch_enq_cased", "InProgress") + ) + assert isinstance(await store.reserve(tokens=100, scopes=(scope,)), BatchEnqueuedTokenOverLimit) + + await handler.async_post_call_success_hook( + data={}, user_api_key_dict=user, response=_batch_response("batch_enq_cased", "Completed") + ) + assert isinstance(await store.reserve(tokens=100, scopes=(scope,)), BatchEnqueuedTokenReservation) + + @pytest.mark.asyncio async def test_failure_hook_refunds_stashed_batch_enqueued_reservation(): from litellm.proxy.hooks.batch_enqueued_tokens import ( From 4333d528136571ae11a2f998ac0fafede111d520 Mon Sep 17 00:00:00 2001 From: mateo-berri Date: Wed, 19 Aug 2026 16:36:27 -0700 Subject: [PATCH 024/281] fix(batch_enqueued_tokens): scope in-memory refunds to the granting worker In-memory grants now record an owner token, and a refund only debits local counters when the popping worker is the one that granted them, so a terminal response handled elsewhere can no longer shrink another worker's unrelated fallback reservations. A Redis-granted refund that fails no longer falls back to decrementing local counters either: the leaked Redis increments expire with the TTL and only tighten the allowance. --- litellm/proxy/hooks/batch_enqueued_tokens.py | 40 +++++++++++----- .../proxy/hooks/test_batch_enqueued_tokens.py | 47 ++++++++++++++++++- 2 files changed, 74 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/hooks/batch_enqueued_tokens.py b/litellm/proxy/hooks/batch_enqueued_tokens.py index 570291aa452..bbc007bb00d 100644 --- a/litellm/proxy/hooks/batch_enqueued_tokens.py +++ b/litellm/proxy/hooks/batch_enqueued_tokens.py @@ -9,6 +9,7 @@ the reservation is refunded when the batch reaches a terminal state """ import asyncio +import uuid from collections.abc import Awaitable, Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias @@ -85,6 +86,7 @@ class BatchEnqueuedTokenReservation: tokens: int scopes: tuple[BatchEnqueuedTokenScope, ...] backend: ReservationBackend = "redis" + owner: str = "" @dataclass(frozen=True, slots=True) @@ -177,7 +179,8 @@ class BatchEnqueuedTokenStore: cross-slot commands), with an over-limit or failing scope rolling back the scopes reserved before it; otherwise a single-process in-memory fallback guarded by one asyncio lock is used. Reservations remember which backend - granted them so a refund never debits counters the grant did not charge. Everything expires after + granted them, and in-memory grants also remember the granting worker, so a + refund never debits counters the grant did not charge. Everything expires after ``BATCH_ENQUEUED_TOKEN_TTL_SECONDS`` so a crash between submission and the terminal-state refund can never leak tokens forever. """ @@ -185,6 +188,7 @@ class BatchEnqueuedTokenStore: def __init__(self, internal_usage_cache: "InternalUsageCache") -> None: self.internal_usage_cache = internal_usage_cache self._lock = asyncio.Lock() + self._owner_token = uuid.uuid4().hex redis_cache = internal_usage_cache.dual_cache.redis_cache self._reserve_script: _ScriptRunner | None = ( redis_cache.async_register_script(RESERVE_ENQUEUED_TOKENS_SCRIPT) if redis_cache is not None else None @@ -300,7 +304,7 @@ class BatchEnqueuedTokenStore: return BatchEnqueuedTokenOverLimit(scope=scope, enqueued=current) for scope, current in zip(scopes, currents): await self._set_local_counter(scope, current + tokens, span) - return BatchEnqueuedTokenReservation(tokens=tokens, scopes=scopes, backend="memory") + return BatchEnqueuedTokenReservation(tokens=tokens, scopes=scopes, backend="memory", owner=self._owner_token) async def refund( self, @@ -309,16 +313,14 @@ class BatchEnqueuedTokenStore: ) -> None: if reservation.tokens <= 0 or not reservation.scopes: return - refund_script: Final = self._refund_script - if reservation.backend == "redis" and refund_script is not None: - try: - await self._refund_via_redis(refund_script, tokens=reservation.tokens, scopes=reservation.scopes) - except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory counters - verbose_proxy_logger.warning( - "Redis enqueued-token refund failed, falling back to in-memory: %s", str(e) - ) - else: - return + if reservation.backend == "redis": + await self._refund_redis_reservation(reservation) + return + if reservation.owner != self._owner_token: + verbose_proxy_logger.warning( + "Skipping enqueued-token refund granted in another worker's memory; its counters expire with the TTL" + ) + return async with self._lock: for scope in reservation.scopes: current = await self._get_local_counter(scope, litellm_parent_otel_span) @@ -328,6 +330,20 @@ class BatchEnqueuedTokenStore: else: await self._set_local_counter(scope, remaining, litellm_parent_otel_span) + async def _refund_redis_reservation(self, reservation: BatchEnqueuedTokenReservation) -> None: + refund_script: Final = self._refund_script + if refund_script is None: + verbose_proxy_logger.warning( + "No Redis client for a Redis-granted enqueued-token refund; leaked increments expire with the TTL" + ) + return + try: + await self._refund_via_redis(refund_script, tokens=reservation.tokens, scopes=reservation.scopes) + except Exception as e: # noqa: BLE001 # best-effort refund: the leak is TTL-bounded and only tightens the allowance + verbose_proxy_logger.warning( + "Redis enqueued-token refund failed; leaked increments expire with the TTL: %s", str(e) + ) + async def save_reservation( self, batch_id: str, diff --git a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py index bc924c32ba3..d6ae200dfaf 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py +++ b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py @@ -128,10 +128,15 @@ async def test_zero_token_reserve_charges_nothing(): class _SingleKeyRedisFake: """Emulates the Redis script path one single-key call at a time, recording every call.""" - def __init__(self, fail_reserve_keys: frozenset[str] = frozenset()) -> None: + def __init__( + self, + fail_reserve_keys: frozenset[str] = frozenset(), + fail_refund_keys: frozenset[str] = frozenset(), + ) -> None: self.script_calls: tuple[tuple[str, tuple[str, ...]], ...] = () self.counters: Mapping[str, int] = MappingProxyType({}) self.fail_reserve_keys = fail_reserve_keys + self.fail_refund_keys = fail_refund_keys def async_register_script(self, script: str): kind: Final = "reserve" if "INCRBY" in script else "refund" if "DECRBY" in script else "record" @@ -153,6 +158,8 @@ class _SingleKeyRedisFake: self.counters = MappingProxyType({**self.counters, keys[0]: current + amount}) return (1, current + amount) if kind == "refund": + if keys[0] in self.fail_refund_keys: + raise ConnectionError(f"simulated redis failure for {keys[0]}") remaining: Final = self.counters.get(keys[0], 0) - int(args[0]) self.counters = MappingProxyType( {key: value for key, value in self.counters.items() if key != keys[0]} @@ -207,6 +214,44 @@ async def test_partial_redis_reserve_failure_rolls_back_and_grants_in_memory(): assert refilled.backend == "memory" +@pytest.mark.asyncio +async def test_memory_refund_skips_reservations_granted_by_another_worker(): + store = _in_memory_store() + scope = _scope(limit=100) + reservation = await store.reserve(tokens=60, scopes=(scope,)) + assert isinstance(reservation, BatchEnqueuedTokenReservation) + assert reservation.backend == "memory" + assert reservation.owner + + foreign: Final = BatchEnqueuedTokenReservation( + tokens=60, scopes=reservation.scopes, backend="memory", owner="another-worker" + ) + await store.refund(foreign) + assert await store.reserve(tokens=50, scopes=(scope,)) == BatchEnqueuedTokenOverLimit(scope=scope, enqueued=60) + + await store.refund(reservation) + assert isinstance(await store.reserve(tokens=100, scopes=(scope,)), BatchEnqueuedTokenReservation) + + +@pytest.mark.asyncio +async def test_failed_redis_refund_leaves_local_counters_untouched(): + scope = _scope(limit=100) + counter_key: Final = f"batch_enqueued_tokens:api_key:{scope.value}" + fake = _SingleKeyRedisFake(fail_refund_keys=frozenset({counter_key})) + store = BatchEnqueuedTokenStore( + internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60)) + ) + + reservation = await store.reserve(tokens=60, scopes=(scope,)) + assert isinstance(reservation, BatchEnqueuedTokenReservation) + assert reservation.backend == "redis" + store.internal_usage_cache.dual_cache.in_memory_cache.set_cache(key=counter_key, value=45) + + await store.refund(reservation) + assert store.internal_usage_cache.dual_cache.in_memory_cache.get_cache(key=counter_key) == 45 + assert fake.counters == {counter_key: 60} + + @pytest.mark.asyncio async def test_pop_reservation_defaults_legacy_records_to_redis_backend(): store = _in_memory_store() From 504112d5ca0f94dde42ec4a736bb18a6036c76d9 Mon Sep 17 00:00:00 2001 From: mateo-berri Date: Wed, 19 Aug 2026 16:52:46 -0700 Subject: [PATCH 025/281] fix(batch_enqueued_tokens): keep the over-limit verdict on rollback failure, find locally saved records on pop A Redis over-limit verdict now survives a failing rollback DECRBY instead of escaping into the in-memory fallback and granting tokens the counter already rejected; the unrolled increments expire with the TTL. pop_reservation now falls through to the local record when the Redis pop succeeds but finds nothing, so a reservation saved in memory after a transient Redis save failure still refunds on cancel or completion. --- litellm/proxy/hooks/batch_enqueued_tokens.py | 29 +++++----- .../proxy/hooks/test_batch_enqueued_tokens.py | 54 ++++++++++++++++++- 2 files changed, 70 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/hooks/batch_enqueued_tokens.py b/litellm/proxy/hooks/batch_enqueued_tokens.py index bbc007bb00d..9d6a599af4a 100644 --- a/litellm/proxy/hooks/batch_enqueued_tokens.py +++ b/litellm/proxy/hooks/batch_enqueued_tokens.py @@ -246,7 +246,7 @@ class BatchEnqueuedTokenStore: already_reserved=scopes[:index], ) if result[0] != 1: - await self._refund_via_redis(refund_script, tokens=tokens, scopes=scopes[:index]) + await self._rollback_partial_reserve(refund_script, tokens=tokens, scopes=scopes[:index]) return BatchEnqueuedTokenOverLimit(scope=scope, enqueued=result[1]) return BatchEnqueuedTokenReservation(tokens=tokens, scopes=scopes, backend="redis") @@ -376,17 +376,10 @@ class BatchEnqueuedTokenStore: batch_id: str, litellm_parent_otel_span: "Span | None" = None, ) -> BatchEnqueuedTokenReservation | None: - raw: object = None - if self._pop_script is not None: - try: - raw = _POPPED_VALUE_ADAPTER.validate_python(await self._pop_script((self._record_key(batch_id),), ())) - except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory record - verbose_proxy_logger.warning( - "Redis enqueued-token reservation pop failed, falling back to in-memory: %s", str(e) - ) - raw = await self._pop_local_record(batch_id, litellm_parent_otel_span) - else: - raw = await self._pop_local_record(batch_id, litellm_parent_otel_span) + redis_raw: Final = await self._pop_redis_record(batch_id) + raw: Final = ( + redis_raw if redis_raw is not None else await self._pop_local_record(batch_id, litellm_parent_otel_span) + ) if raw is None: return None try: @@ -397,6 +390,18 @@ class BatchEnqueuedTokenStore: verbose_proxy_logger.warning("Discarding malformed enqueued-token reservation record for %s", batch_id) return None + async def _pop_redis_record(self, batch_id: str) -> str | bytes | None: + pop_script: Final = self._pop_script + if pop_script is None: + return None + try: + return _POPPED_VALUE_ADAPTER.validate_python(await pop_script((self._record_key(batch_id),), ())) + except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory record + verbose_proxy_logger.warning( + "Redis enqueued-token reservation pop failed, falling back to in-memory: %s", str(e) + ) + return None + async def _pop_local_record(self, batch_id: str, span: "Span | None") -> object: async with self._lock: stored = await self.internal_usage_cache.async_get_cache( diff --git a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py index d6ae200dfaf..5167856359d 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py +++ b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py @@ -132,14 +132,21 @@ class _SingleKeyRedisFake: self, fail_reserve_keys: frozenset[str] = frozenset(), fail_refund_keys: frozenset[str] = frozenset(), + fail_save_keys: frozenset[str] = frozenset(), ) -> None: self.script_calls: tuple[tuple[str, tuple[str, ...]], ...] = () self.counters: Mapping[str, int] = MappingProxyType({}) + self.records: Mapping[str, str] = MappingProxyType({}) self.fail_reserve_keys = fail_reserve_keys self.fail_refund_keys = fail_refund_keys + self.fail_save_keys = fail_save_keys def async_register_script(self, script: str): - kind: Final = "reserve" if "INCRBY" in script else "refund" if "DECRBY" in script else "record" + kind: Final = ( + "reserve" + if "INCRBY" in script + else "refund" if "DECRBY" in script else "save" if "SET" in script else "pop" + ) async def run(keys: Sequence[str], args: Sequence[str | bytes | int | float]) -> object: self.script_calls = (*self.script_calls, (kind, tuple(keys))) @@ -167,6 +174,15 @@ class _SingleKeyRedisFake: else {**self.counters, keys[0]: remaining} ) return 1 + if kind == "save": + if keys[0] in self.fail_save_keys: + raise ConnectionError(f"simulated redis failure for {keys[0]}") + self.records = MappingProxyType({**self.records, keys[0]: str(args[0])}) + return 1 + if kind == "pop": + popped: Final = self.records.get(keys[0]) + self.records = MappingProxyType({key: value for key, value in self.records.items() if key != keys[0]}) + return popped raise AssertionError(f"unexpected {kind} script call for keys {keys}") @@ -214,6 +230,42 @@ async def test_partial_redis_reserve_failure_rolls_back_and_grants_in_memory(): assert refilled.backend == "memory" +@pytest.mark.asyncio +async def test_over_limit_verdict_survives_a_failing_rollback(): + key_scope = _scope(limit=100, key="api_key") + team_scope = _scope(limit=5, key="team") + key_counter: Final = f"batch_enqueued_tokens:api_key:{key_scope.value}" + fake = _SingleKeyRedisFake(fail_refund_keys=frozenset({key_counter})) + store = BatchEnqueuedTokenStore( + internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60)) + ) + + outcome = await store.reserve(tokens=10, scopes=(key_scope, team_scope)) + assert outcome == BatchEnqueuedTokenOverLimit(scope=team_scope, enqueued=0) + assert fake.counters == {key_counter: 10} + + +@pytest.mark.asyncio +async def test_pop_falls_back_to_local_record_when_redis_pop_finds_nothing(): + scope = _scope(limit=100) + record_key: Final = "batch_enqueued_token_reservation:batch_local_record" + fake = _SingleKeyRedisFake(fail_save_keys=frozenset({record_key})) + store = BatchEnqueuedTokenStore( + internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60)) + ) + + reservation = await store.reserve(tokens=60, scopes=(scope,)) + assert isinstance(reservation, BatchEnqueuedTokenReservation) + await store.save_reservation("batch_local_record", reservation) + assert not fake.records + + popped = await store.pop_reservation("batch_local_record") + assert popped == reservation + await store.refund(popped) + assert not fake.counters + assert await store.pop_reservation("batch_local_record") is None + + @pytest.mark.asyncio async def test_memory_refund_skips_reservations_granted_by_another_worker(): store = _in_memory_store() From 579291774b0a8b5e98c33140ae89fe60a30360b0 Mon Sep 17 00:00:00 2001 From: ljogeiger <39981740+ljogeiger@users.noreply.github.com> Date: Thu, 20 Aug 2026 00:16:36 +0000 Subject: [PATCH 026/281] docs(vertex_ai): cite Google's thought signature rules for parallel calls Link the Gemini Enterprise Agent Platform docs at both places the behavior is decided. The docs state that only the first functionCall part of a parallel batch carries a thought_signature, and that setting skip_thought_signature_validator "should be a last resort as it will negatively impact model performance". --- .../litellm_core_utils/prompt_templates/factory.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 0a9d7c427b4..b676077ab0e 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1245,10 +1245,14 @@ def _get_dummy_thought_signature() -> str: This is used when transferring conversation history from older models (like gemini-2.5-flash) to gemini-3, which requires thought_signature - for strict validation. + for strict validation. Google documents it as a last resort that "will + negatively impact model performance", so callers must only fall back to it + when no real signature is available. + + See: + https://ai.google.dev/gemini-api/docs/thought-signatures#faqs + https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/thinking/thought-signatures """ - # Return a base64-encoded dummy signature string - # Below dummy signature is recommended by google - https://ai.google.dev/gemini-api/docs/thought-signatures#faqs dummy_data: Final = b"skip_thought_signature_validator" return base64.b64encode(dummy_data).decode("utf-8") @@ -1318,6 +1322,9 @@ def convert_to_gemini_tool_call_invoke( if gemini_function_call is not None: part_dict: VertexPartType = {"function_call": gemini_function_call} thought_signature = _get_thought_signature_from_tool(dict(tool)) + # Gemini signs only the first functionCall part of a parallel batch, so scope the + # placeholder fallback to that part instead of fabricating one per sibling call: + # https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/thinking/thought-signatures#parallel_function_calling_example is_first_function_call = len(_parts_list) == 0 if not thought_signature and is_first_function_call and needs_dummy_signature: thought_signature = _get_dummy_thought_signature() From a5ad22b8a3f734b09ce5e05a68dab5c5205907a4 Mon Sep 17 00:00:00 2001 From: ljogeiger <39981740+ljogeiger@users.noreply.github.com> Date: Thu, 20 Aug 2026 00:16:36 +0000 Subject: [PATCH 027/281] test(vertex_ai): cover gemini-3.5-flash and drop assertion-echoing docstrings Add gemini-3.5-flash to the placeholder-scoping matrix and a regression test that a natively signed parallel turn replays with no skip_thought_signature_validator anywhere in the payload, the shape that was producing empty text responses on 3.5. Hoist the repeated placeholder expression into one constant and rewrite the docstrings that restated their own assertions to say why the case matters instead. --- .../test_vertex_ai_gemini_transformation.py | 83 +++++++++++++------ 1 file changed, 58 insertions(+), 25 deletions(-) diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index 9ebf6db11de..8c1de12e7d9 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -830,13 +830,14 @@ def _parallel_tool_calls_signed_via_id(*signatures): REAL_THOUGHT_SIGNATURE = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n" +PLACEHOLDER_SIGNATURE = base64.b64encode(b"skip_thought_signature_validator").decode( + "utf-8" +) def test_dummy_signature_only_on_first_parallel_tool_call(): - """Gemini only returns a thought signature on the first of N parallel function calls. - - The sibling calls carry no signature, so replaying them must not fabricate one. - """ + """Google documents the placeholder as a last resort that degrades quality, so an unsigned + parallel turn replayed to gemini-3 gets a budget of exactly one.""" from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_gemini_tool_call_invoke, ) @@ -850,17 +851,15 @@ def test_dummy_signature_only_on_first_parallel_tool_call(): model="gemini-3-pro-preview", ) - expected_dummy = base64.b64encode(b"skip_thought_signature_validator").decode( - "utf-8" - ) assert len(gemini_parts) == 3 - assert gemini_parts[0]["thoughtSignature"] == expected_dummy + assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE assert "thoughtSignature" not in gemini_parts[1] assert "thoughtSignature" not in gemini_parts[2] def test_real_signature_on_first_parallel_tool_call_leaves_siblings_empty(): - """The real signature from Gemini rides on the first call; siblings stay signature-free.""" + """Gemini signs only the first of N parallel function calls, so a faithful replay has + nothing to attach to the siblings.""" from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_gemini_tool_call_invoke, ) @@ -881,7 +880,8 @@ def test_real_signature_on_first_parallel_tool_call_leaves_siblings_empty(): def test_real_signature_on_later_parallel_tool_call_is_preserved(): - """A signature attached to a non-first call is still forwarded as-is.""" + """Clients may reorder or drop calls, so a signature that lands on a non-first call is + still the model's own and must survive the round trip.""" from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_gemini_tool_call_invoke, ) @@ -895,11 +895,8 @@ def test_real_signature_on_later_parallel_tool_call_is_preserved(): model="gemini-3-pro-preview", ) - expected_dummy = base64.b64encode(b"skip_thought_signature_validator").decode( - "utf-8" - ) assert len(gemini_parts) == 2 - assert gemini_parts[0]["thoughtSignature"] == expected_dummy + assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE assert gemini_parts[1]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE @@ -970,7 +967,6 @@ def test_placeholder_lands_on_first_emitted_part_not_first_tool_call_entry(): """A non-function entry (e.g. an OpenAI custom tool call) emits no part, so it must not consume the one placeholder slot and leave the real first function call bare.""" from litellm.litellm_core_utils.prompt_templates.factory import ( - _get_dummy_thought_signature, convert_to_gemini_tool_call_invoke, ) @@ -984,7 +980,7 @@ def test_placeholder_lands_on_first_emitted_part_not_first_tool_call_entry(): ) assert len(gemini_parts) == 2 - assert gemini_parts[0]["thoughtSignature"] == _get_dummy_thought_signature() + assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE assert "thoughtSignature" not in gemini_parts[1] @@ -1054,22 +1050,62 @@ def test_parallel_tool_call_history_replayed_through_full_message_conversion(): assert "thoughtSignature" not in model_parts[2] +@pytest.mark.parametrize( + "model", + ["gemini-3.5-flash", "vertex_ai/gemini-3.5-flash", "gemini/gemini-3.5-flash"], +) +def test_natively_signed_parallel_turn_never_carries_a_placeholder(model): + """A native gemini-3.5 parallel turn replays with zero skip_thought_signature_validator parts. + + Fabricating the placeholder alongside a real signature is what produced empty text responses + on gemini-3.5 parallel function calling, so the whole payload has to stay placeholder-free. + """ + import json + + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + messages = [ + {"role": "user", "content": "Weather in Paris, London and Tokyo?"}, + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls_signed_via_id( + REAL_THOUGHT_SIGNATURE, None, None + ), + }, + ] + + contents = _gemini_convert_messages_with_history(messages=messages, model=model) + + model_parts = contents[1]["parts"] + assert len(model_parts) == 3 + assert model_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in model_parts[1] + assert "thoughtSignature" not in model_parts[2] + assert PLACEHOLDER_SIGNATURE not in json.dumps(contents) + + @pytest.mark.parametrize( "model", [ "gemini-3-pro-preview", "gemini-3-flash-preview", "gemini-3.1-pro-preview", + "gemini-3.5-flash", "gemini-3.6-flash", "gemini-3.7-flash", + "vertex_ai/gemini-3.5-flash", "vertex_ai/gemini-3.7-flash", + "gemini/gemini-3.5-flash", "gemini/gemini-3.7-flash", ], ) def test_placeholder_scoped_to_first_call_across_gemini_3_variants(model): - """Every gemini-3 family member, bare or provider-prefixed, gets one placeholder at most.""" + """The gemini-3 gate is a substring match, so every family member and prefix form has to + land on the same one-placeholder budget rather than only the versions we happened to try.""" from litellm.litellm_core_utils.prompt_templates.factory import ( - _get_dummy_thought_signature, convert_to_gemini_tool_call_invoke, ) @@ -1083,17 +1119,14 @@ def test_placeholder_scoped_to_first_call_across_gemini_3_variants(model): ) assert len(gemini_parts) == 3 - assert gemini_parts[0]["thoughtSignature"] == _get_dummy_thought_signature() + assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE assert "thoughtSignature" not in gemini_parts[1] assert "thoughtSignature" not in gemini_parts[2] def test_signed_text_part_survives_alongside_unsigned_parallel_tool_calls(): - """gemini-2.5 history with a signed text part and parallel unsigned calls: the text keeps its real - signature, only the first call gets the placeholder, and the siblings stay bare.""" - from litellm.litellm_core_utils.prompt_templates.factory import ( - _get_dummy_thought_signature, - ) + """Text-part and function-call signatures are collected by separate code paths, so scoping the + placeholder must not disturb a real signature that arrived on the text part.""" from litellm.llms.vertex_ai.gemini.transformation import ( _gemini_convert_messages_with_history, ) @@ -1111,7 +1144,7 @@ def test_signed_text_part_survives_alongside_unsigned_parallel_tool_calls(): assert parts[0]["text"] == "Checking all three cities." assert parts[0]["thoughtSignature"] == "real_25_signature" - assert parts[1]["thoughtSignature"] == _get_dummy_thought_signature() + assert parts[1]["thoughtSignature"] == PLACEHOLDER_SIGNATURE assert "thoughtSignature" not in parts[2] assert "thoughtSignature" not in parts[3] From 6b17b8a8f2489085f8923b5a84cdb79412108c24 Mon Sep 17 00:00:00 2001 From: hiraku-miyoshi Date: Wed, 19 Aug 2026 17:22:15 -0700 Subject: [PATCH 028/281] fix(proxy): restrict batch_enqueued_token_limit metadata writes to proxy admins The field replaces the standard RPM/TPM checks for batch submissions, so a key holder or team admin writing it could pick their own batch quota. Mirrors the output-token-estimate admin gate: change-based, so resending the stored value stays allowed, and enforced on key generate, update, bulk team-key update, regenerate, and team new/update. --- litellm/constants.py | 5 + litellm/proxy/auth/auth_utils.py | 47 +++- litellm/proxy/hooks/batch_enqueued_tokens.py | 6 +- .../key_management_endpoints.py | 27 +++ .../management_endpoints/team_endpoints.py | 17 +- .../test_key_management_endpoints.py | 201 ++++++++++++++++++ .../test_team_endpoints.py | 57 +++++ 7 files changed, 354 insertions(+), 6 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index d7ddacf5fac..facfc6f7c19 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1772,6 +1772,11 @@ PTU_PRUNE_SKEW_GRACE_SECONDS: Final[int] = 300 # was never observed (e.g. proxy restart); expiry returns the tokens to the caller. BATCH_ENQUEUED_TOKEN_TTL_SECONDS: Final[int] = 8 * 24 * 60 * 60 +# Key/team metadata field that opts batches into enqueued-token limiting. Only proxy +# admins may write it: when present it replaces the standard RPM/TPM checks for +# batch submissions. +BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY: Final = "batch_enqueued_token_limit" + # Shared read-only empty mapping, for defaulting optional Mapping parameters without # constructing a fresh mutable dict at each call site. EMPTY_MAPPING: Final = MappingProxyType({}) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index a105bf19458..883f986f6fd 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -12,7 +12,12 @@ from pydantic import PositiveInt, TypeAdapter, ValidationError import litellm from litellm import Router, provider_list from litellm._logging import verbose_proxy_logger -from litellm.constants import MINIMUM_CUSTOM_KEY_LENGTH, STANDARD_CUSTOMER_ID_HEADERS +from litellm.constants import ( + BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY, + EMPTY_MAPPING, + MINIMUM_CUSTOM_KEY_LENGTH, + STANDARD_CUSTOMER_ID_HEADERS, +) from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.url_utils import ( SSRFError, @@ -1169,6 +1174,46 @@ def enforce_output_token_estimates_are_admin_only( ) +class BatchEnqueuedTokenLimitRequest(Protocol): + """The shape of any management request that can carry a batch enqueued-token limit.""" + + @property + def metadata(self) -> Mapping[str, object] | None: ... + + @property + def model_fields_set(self) -> Collection[str]: ... + + +def enforce_batch_enqueued_token_limit_is_admin_only( + data: BatchEnqueuedTokenLimitRequest, + existing_metadata: Mapping[str, object] | None, + user_api_key_dict: UserAPIKeyAuth, + entity: Literal["key", "team"], +) -> None: + """Only a proxy admin may change a key or team's batch enqueued-token limit. + + When set, ``batch_enqueued_token_limit`` replaces the standard RPM/TPM checks + for batch submissions, so a holder-writable copy would let a caller lift their + own batch quota. Gated on the resulting value rather than on presence, so a + form resending the stored value stays a no-op. + """ + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return + stored: Final[Mapping[str, object]] = existing_metadata or EMPTY_MAPPING + requested: Final[Mapping[str, object]] = ( + (data.metadata or EMPTY_MAPPING) if "metadata" in data.model_fields_set else stored + ) + if requested.get(BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY) == stored.get(BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY): + return + raise HTTPException( + status_code=403, + detail={ # mutable-ok: HTTPException.detail has no immutable form + "error": f"Only proxy admins can set {BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY} on a {entity}. " + "It replaces the standard rate limit checks for batch submissions." + }, + ) + + def get_model_rate_limit_from_metadata( user_api_key_dict: UserAPIKeyAuth, metadata_accessor_key: Literal["team_metadata", "organization_metadata", "project_metadata"], diff --git a/litellm/proxy/hooks/batch_enqueued_tokens.py b/litellm/proxy/hooks/batch_enqueued_tokens.py index 9d6a599af4a..54b0111682d 100644 --- a/litellm/proxy/hooks/batch_enqueued_tokens.py +++ b/litellm/proxy/hooks/batch_enqueued_tokens.py @@ -1,7 +1,7 @@ """ Enqueued-token accounting for batch submissions. -Opt-in via ``batch_enqueued_token_limit`` in key or team metadata: batch +Opt-in via admin-set ``batch_enqueued_token_limit`` in key or team metadata: batch submissions reserve their estimated token count against a long-lived enqueued-token allowance instead of the per-minute rate-limit windows, and the reservation is refunded when the batch reaches a terminal state @@ -17,7 +17,7 @@ from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger -from litellm.constants import BATCH_ENQUEUED_TOKEN_TTL_SECONDS +from litellm.constants import BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY, BATCH_ENQUEUED_TOKEN_TTL_SECONDS from litellm.proxy._types import UserAPIKeyAuth if TYPE_CHECKING: @@ -28,8 +28,6 @@ if TYPE_CHECKING: Span = _Span InternalUsageCache = _InternalUsageCache -BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY: Final = "batch_enqueued_token_limit" - BATCH_ENQUEUED_REFUND_STATUSES: Final[frozenset[str]] = frozenset( {"completed", "complete", "failed", "expired", "cancelled", "cancelling"} ) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 67a836b8c92..ade568ad07f 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -57,6 +57,7 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.proxy.auth.auth_utils import ( abbreviate_api_key, + enforce_batch_enqueued_token_limit_is_admin_only, enforce_output_token_estimates_are_admin_only, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -901,6 +902,12 @@ async def _common_key_generation_helper( user_api_key_dict=user_api_key_dict, entity="key", ) + enforce_batch_enqueued_token_limit_is_admin_only( + data=data, + existing_metadata=None, + user_api_key_dict=user_api_key_dict, + entity="key", + ) if data.metadata is not None and data.metadata.get("service_account_id") is not None and data.team_id is None: await validate_team_id_used_in_service_account_request( @@ -2296,6 +2303,14 @@ async def _process_single_key_update( prisma_client=prisma_client, ) + _existing_row_metadata: Final = getattr(existing_key_row, "metadata", None) + enforce_batch_enqueued_token_limit_is_admin_only( + data=update_key_request, + existing_metadata=_existing_row_metadata if isinstance(_existing_row_metadata, dict) else None, + user_api_key_dict=user_api_key_dict, + entity="key", + ) + # Check team member permissions if prisma_client is not None: await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( @@ -2558,6 +2573,12 @@ async def _validate_update_key_data( user_api_key_dict=user_api_key_dict, entity="key", ) + enforce_batch_enqueued_token_limit_is_admin_only( + data=data, + existing_metadata=_existing_metadata if isinstance(_existing_metadata, dict) else None, + user_api_key_dict=user_api_key_dict, + entity="key", + ) # Personal-key bypass: the caller both created the key AND still owns it # (user_id == caller). Checking only created_by would let a demoted admin @@ -4754,6 +4775,12 @@ async def _execute_virtual_key_regeneration( user_api_key_dict=user_api_key_dict, entity="key", ) + enforce_batch_enqueued_token_limit_is_admin_only( + data=data, + existing_metadata=_existing_key_metadata if isinstance(_existing_key_metadata, dict) else None, + user_api_key_dict=user_api_key_dict, + entity="key", + ) new_token: Final = await get_new_token(data=data) new_token_hash: Final = hash_token(new_token) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 95632d7cb35..31006958ba1 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -85,7 +85,10 @@ from litellm.proxy.auth.auth_checks import ( get_team_object, get_user_object, ) -from litellm.proxy.auth.auth_utils import enforce_output_token_estimates_are_admin_only +from litellm.proxy.auth.auth_utils import ( + enforce_batch_enqueued_token_limit_is_admin_only, + enforce_output_token_estimates_are_admin_only, +) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch @@ -1303,6 +1306,12 @@ async def new_team( user_api_key_dict=user_api_key_dict, entity="team", ) + enforce_batch_enqueued_token_limit_is_admin_only( + data=data, + existing_metadata=None, + user_api_key_dict=user_api_key_dict, + entity="team", + ) # Check if license is over limit total_teams: Final = await _team_db(prisma_client).count() @@ -2007,6 +2016,12 @@ async def update_team( user_api_key_dict=user_api_key_dict, entity="team", ) + enforce_batch_enqueued_token_limit_is_admin_only( + data=data, + existing_metadata=_existing_team_metadata if isinstance(_existing_team_metadata, dict) else None, + user_api_key_dict=user_api_key_dict, + entity="team", + ) _check_passthrough_routes_caller_permission(data, user_api_key_dict, entity="team") diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index bdf09a95e4b..93a4788caf8 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -15515,6 +15515,207 @@ async def test_regenerate_key_output_token_estimate_lowered_rejected_for_non_adm assert "Only proxy admins can set" in str(exc.value.detail) +_BATCH_LIMIT = "batch_enqueued_token_limit" + + +@pytest.mark.parametrize( + "label, request_body, existing_metadata, allowed", + [ + ("set on a key with none stored", {"metadata": {_BATCH_LIMIT: 50000}}, None, False), + ("raised above the stored limit", {"metadata": {_BATCH_LIMIT: 200000}}, {_BATCH_LIMIT: 100000}, False), + ("cleared by replacing the blob", {"metadata": {}}, {_BATCH_LIMIT: 100000}, False), + ("resent unchanged", {"metadata": {_BATCH_LIMIT: 100000}}, {_BATCH_LIMIT: 100000}, True), + ("left untouched", {}, {_BATCH_LIMIT: 100000}, True), + ], +) +def test_batch_enqueued_token_limit_admin_gate_matrix(label, request_body, existing_metadata, allowed): + """A non-admin may only leave a key's stored batch enqueued-token limit as it is. + + When set, the limit replaces the standard RPM/TPM checks for batch + submissions, so a key holder writing it would pick their own batch quota. + Resending the stored value is what the edit form produces on every save + and has to stay allowed. + """ + from litellm.proxy.auth.auth_utils import ( + enforce_batch_enqueued_token_limit_is_admin_only, + ) + + def _call(caller): + enforce_batch_enqueued_token_limit_is_admin_only( + data=UpdateKeyRequest(key="sk-1", **request_body), + existing_metadata=existing_metadata, + user_api_key_dict=caller, + entity="key", + ) + + non_admin = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-non-admin", + user_id="alice", + ) + if allowed: + _call(non_admin) + else: + with pytest.raises(HTTPException) as exc: + _call(non_admin) + assert exc.value.status_code == 403 + assert "Only proxy admins can set" in str(exc.value.detail) + + _call( + UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin", + ) + ) + + +@pytest.mark.asyncio +async def test_generate_key_batch_enqueued_token_limit_rejected_for_non_admin(): + """A non-admin self-minting a key with the limit would replace the standard + batch RPM/TPM checks with a cap of their own choosing.""" + with patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()): + with pytest.raises(HTTPException) as exc: + await _common_key_generation_helper( + data=GenerateKeyRequest(metadata={_BATCH_LIMIT: 100000}, rpm_limit=2), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="alice", + ), + litellm_changed_by=None, + team_table=None, + ) + assert int(getattr(exc.value, "status_code", 0)) == 403 + assert "Only proxy admins can set" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_update_key_batch_enqueued_token_limit_raised_rejected_for_non_admin(monkeypatch): + """/key/update is reachable by the key's own holder, so the gate has to + fire inside the update path itself rather than only at generation.""" + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + token = "d1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + _wire_update_key_fn(monkeypatch, _estimate_key_row(token, {_BATCH_LIMIT: 100000})) + + mock_request = MagicMock() + mock_request.query_params = {} + + with pytest.raises(ProxyException) as exc: + await update_key_fn( + request=mock_request, + data=UpdateKeyRequest(key=token, metadata={_BATCH_LIMIT: 10**12}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ), + litellm_changed_by=None, + ) + + assert str(exc.value.code) == "403" + assert "Only proxy admins can set" in str(exc.value.message) + + +@pytest.mark.asyncio +async def test_update_key_batch_enqueued_token_limit_unchanged_allows_non_admin_edit(monkeypatch): + """The edit form resends every field it renders, so gating on presence + would 403 a key owner renaming a key that carries an admin-set limit.""" + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + token = "e1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + _wire_update_key_fn(monkeypatch, _estimate_key_row(token, {_BATCH_LIMIT: 100000})) + + mock_request = MagicMock() + mock_request.query_params = {} + + result = await update_key_fn( + request=mock_request, + data=UpdateKeyRequest(key=token, key_alias="my-alias", metadata={_BATCH_LIMIT: 100000}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ), + litellm_changed_by=None, + ) + + assert result is not None + + +@pytest.mark.asyncio +async def test_regenerate_key_batch_enqueued_token_limit_rejected_for_non_admin(): + """/key/regenerate runs the request body through prepare_key_update_data + exactly as an update does, so it is a third write path into the field.""" + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + + token = "f1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + key_in_db = LiteLLM_VerificationToken( + token=token, + user_id="internal_user", + metadata={_BATCH_LIMIT: 100000}, + ) + + with pytest.raises(HTTPException) as exc: + await _execute_virtual_key_regeneration( + prisma_client=AsyncMock(), + key_in_db=key_in_db, + hashed_api_key=token, + key="sk-original", + data=RegenerateKeyRequest(key="sk-original", metadata={_BATCH_LIMIT: 10**12}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert exc.value.status_code == 403 + assert "Only proxy admins can set" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_bulk_key_update_batch_enqueued_token_limit_rejected_for_non_admin(): + """Bulk team-key updates run through _process_single_key_update, not + /key/update's validator, so the gate must also live on that path.""" + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _process_single_key_update, + ) + + token = "a2b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + existing = _estimate_key_row(token, {_BATCH_LIMIT: 100000}) + + with pytest.raises(HTTPException) as exc: + await _process_single_key_update( + update_key_request=UpdateKeyRequest(key=token, metadata={_BATCH_LIMIT: 10**12}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ), + litellm_changed_by=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + llm_router=None, + existing_key_row=existing, + ) + + assert exc.value.status_code == 403 + assert "Only proxy admins can set" in str(exc.value.detail) + + @pytest.mark.asyncio async def test_execute_virtual_key_regeneration_stamps_settings_updated_at(): """Regenerate rewrites the key's config, so it must move settings_updated_at.""" 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 db39fcd3799..05072a0a6d7 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -11649,6 +11649,63 @@ async def test_new_team_output_token_estimate_rejected_for_non_admin(): assert "on a team" in str(exc.value.message) +_TEAM_BATCH_LIMIT = "batch_enqueued_token_limit" + + +@pytest.mark.asyncio +async def test_update_team_batch_enqueued_token_limit_raised_rejected_for_team_admin(): + """_verify_team_access admits a team admin, so the gate has to fire inside + update_team itself to keep the team's batch quota admin-owned.""" + import contextlib + from unittest.mock import Mock + + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import update_team + + with contextlib.ExitStack() as stack: + _wire_update_team(stack, {_TEAM_BATCH_LIMIT: 100000}) + with pytest.raises(ProxyException) as exc: + await update_team( + data=UpdateTeamRequest(team_id="test_team_id", metadata={_TEAM_BATCH_LIMIT: 10**12}), + http_request=Mock(spec=Request), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-team-admin", + user_id="team-admin", + ), + ) + + assert str(exc.value.code) == "403" + assert "on a team" in str(exc.value.message) + + +@pytest.mark.asyncio +async def test_new_team_batch_enqueued_token_limit_rejected_for_non_admin(): + """/team/new is the other write path into the same stored metadata.""" + from unittest.mock import Mock + + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + with pytest.raises(ProxyException) as exc: + await new_team( + data=NewTeamRequest(team_alias="t", metadata={_TEAM_BATCH_LIMIT: 100000}), + http_request=Mock(spec=Request), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="alice", + ), + ) + + assert str(exc.value.code) == "403" + assert "on a team" in str(exc.value.message) + + @pytest.mark.asyncio async def test_get_team_daily_activity_aggregated_scopes_and_flags(mock_db_client): """The aggregated endpoint must apply the same non-admin key scoping as the From 5ab20c3678756bb5a04ff2131e1de11a2619583c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 17:47:06 -0700 Subject: [PATCH 029/281] fix(batch_enqueued_tokens): tombstone popped Redis reservation records so local ghosts cannot double-refund --- litellm/proxy/hooks/batch_enqueued_tokens.py | 14 +++++-- .../proxy/hooks/test_batch_enqueued_tokens.py | 41 ++++++++++++++++++- 2 files changed, 50 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/hooks/batch_enqueued_tokens.py b/litellm/proxy/hooks/batch_enqueued_tokens.py index 54b0111682d..82d4e7c66fc 100644 --- a/litellm/proxy/hooks/batch_enqueued_tokens.py +++ b/litellm/proxy/hooks/batch_enqueued_tokens.py @@ -62,8 +62,8 @@ return 1 POP_RESERVATION_SCRIPT: Final = """ local value = redis.call('GET', KEYS[1]) -if value then - redis.call('DEL', KEYS[1]) +if value and value ~= '' then + redis.call('SET', KEYS[1], '', 'EX', tonumber(ARGV[1])) end return value """ @@ -375,6 +375,12 @@ class BatchEnqueuedTokenStore: litellm_parent_otel_span: "Span | None" = None, ) -> BatchEnqueuedTokenReservation | None: redis_raw: Final = await self._pop_redis_record(batch_id) + if redis_raw is not None and not redis_raw: + # The Redis pop tombstones popped records in place, so a hit on the empty + # tombstone means the batch was already refunded elsewhere; a local copy + # left behind by a save that raised after landing must not refund again. + await self._pop_local_record(batch_id, litellm_parent_otel_span) + return None raw: Final = ( redis_raw if redis_raw is not None else await self._pop_local_record(batch_id, litellm_parent_otel_span) ) @@ -393,7 +399,9 @@ class BatchEnqueuedTokenStore: if pop_script is None: return None try: - return _POPPED_VALUE_ADAPTER.validate_python(await pop_script((self._record_key(batch_id),), ())) + return _POPPED_VALUE_ADAPTER.validate_python( + await pop_script((self._record_key(batch_id),), (BATCH_ENQUEUED_TOKEN_TTL_SECONDS,)) + ) except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory record verbose_proxy_logger.warning( "Redis enqueued-token reservation pop failed, falling back to in-memory: %s", str(e) diff --git a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py index 5167856359d..440b3050a39 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py +++ b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py @@ -133,6 +133,7 @@ class _SingleKeyRedisFake: fail_reserve_keys: frozenset[str] = frozenset(), fail_refund_keys: frozenset[str] = frozenset(), fail_save_keys: frozenset[str] = frozenset(), + raise_after_landing_save_keys: frozenset[str] = frozenset(), ) -> None: self.script_calls: tuple[tuple[str, tuple[str, ...]], ...] = () self.counters: Mapping[str, int] = MappingProxyType({}) @@ -140,12 +141,13 @@ class _SingleKeyRedisFake: self.fail_reserve_keys = fail_reserve_keys self.fail_refund_keys = fail_refund_keys self.fail_save_keys = fail_save_keys + self.raise_after_landing_save_keys = raise_after_landing_save_keys def async_register_script(self, script: str): kind: Final = ( "reserve" if "INCRBY" in script - else "refund" if "DECRBY" in script else "save" if "SET" in script else "pop" + else "refund" if "DECRBY" in script else "pop" if "GET" in script else "save" ) async def run(keys: Sequence[str], args: Sequence[str | bytes | int | float]) -> object: @@ -178,10 +180,13 @@ class _SingleKeyRedisFake: if keys[0] in self.fail_save_keys: raise ConnectionError(f"simulated redis failure for {keys[0]}") self.records = MappingProxyType({**self.records, keys[0]: str(args[0])}) + if keys[0] in self.raise_after_landing_save_keys: + raise TimeoutError(f"simulated redis timeout after landing for {keys[0]}") return 1 if kind == "pop": popped: Final = self.records.get(keys[0]) - self.records = MappingProxyType({key: value for key, value in self.records.items() if key != keys[0]}) + if popped: + self.records = MappingProxyType({**self.records, keys[0]: ""}) return popped raise AssertionError(f"unexpected {kind} script call for keys {keys}") @@ -266,6 +271,38 @@ async def test_pop_falls_back_to_local_record_when_redis_pop_finds_nothing(): assert await store.pop_reservation("batch_local_record") is None +@pytest.mark.asyncio +async def test_local_ghost_left_by_landed_save_never_refunds_twice(): + scope = _scope(limit=100) + record_key: Final = "batch_enqueued_token_reservation:batch_ghost" + fake = _SingleKeyRedisFake(raise_after_landing_save_keys=frozenset({record_key})) + store = BatchEnqueuedTokenStore( + internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60)) + ) + + reservation = await store.reserve(tokens=60, scopes=(scope,)) + assert isinstance(reservation, BatchEnqueuedTokenReservation) + await store.save_reservation("batch_ghost", reservation) + assert fake.records[record_key] + + first = await store.pop_reservation("batch_ghost") + assert first == reservation + await store.refund(first) + assert not fake.counters + assert fake.records[record_key] == "" + + assert await store.pop_reservation("batch_ghost") is None + assert ( + await store.internal_usage_cache.async_get_cache( + key=record_key, litellm_parent_otel_span=None, local_only=True + ) + is None + ) + assert await store.pop_reservation("batch_ghost") is None + assert isinstance(await store.reserve(tokens=100, scopes=(scope,)), BatchEnqueuedTokenReservation) + assert fake.counters[f"batch_enqueued_tokens:{scope.key}:{scope.value}"] == 100 + + @pytest.mark.asyncio async def test_memory_refund_skips_reservations_granted_by_another_worker(): store = _in_memory_store() From 367dd537b995d23421fd99bdb1093e443e98911e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 18:39:15 -0700 Subject: [PATCH 030/281] feat(e2e): move record/replay to the provider edge (LIT-5745) Replaces the test-side fixture transport with an in-process provider-edge HTTP server the proxy's deployments point their api_base at. Record forwards provider calls verbatim and writes them to the bundle; replay answers them from the bundle with zero provider calls while key auth, routing, cost calculation, and spend-log writes still execute against the live proxy and database. Drift comes back as HTTP 599 naming the computed and closest recorded keys. Request headers are never stored and responses are kept byte-identical between modes from the proxy's side of the socket. --- tests/e2e/CLAUDE.md | 12 +- tests/e2e/CONTRIBUTING.md | 10 +- tests/e2e/conftest.py | 14 +- tests/e2e/e2e_config.py | 32 +- tests/e2e/e2e_http.py | 34 + tests/e2e/fixture_bundle.py | 133 +--- tests/e2e/fixture_mode.py | 132 ++++ tests/e2e/fixture_transport.py | 724 ------------------ tests/e2e/provider_edge.py | 546 +++++++++++++ tests/e2e/proxy_client.py | 16 +- .../test_provider_edge_spend_e2e.py | 50 ++ tests/e2e/test_fixture_bundle.py | 46 +- tests/e2e/test_fixture_mode.py | 114 +++ tests/e2e/test_fixture_transport.py | 676 ---------------- tests/e2e/test_provider_edge.py | 492 ++++++++++++ 15 files changed, 1453 insertions(+), 1578 deletions(-) create mode 100644 tests/e2e/fixture_mode.py delete mode 100644 tests/e2e/fixture_transport.py create mode 100644 tests/e2e/provider_edge.py create mode 100644 tests/e2e/quota_management/spend_tracking/test_provider_edge_spend_e2e.py create mode 100644 tests/e2e/test_fixture_mode.py delete mode 100644 tests/e2e/test_fixture_transport.py create mode 100644 tests/e2e/test_provider_edge.py diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index d7334552d0c..840a40a54cd 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -73,13 +73,17 @@ Mark live tests with `@pytest.mark.e2e` (on the class or the module). Pure cover ## Record and replay fixtures -`E2E_FIXTURE_MODE` selects the transport every client is built on: `live` (the default, and what an unset variable means: nothing changes), `record` (run against the live proxy and write every interaction to a fixture bundle), or `replay` (serve every interaction back from the bundle with no HTTP at all, so a replay run needs no proxy and cannot bill a provider). The seam is `select_transport` in `fixture_transport.py`, applied inside `build_proxy_client`; both transports fulfil the same `Transport` protocol, so no test or client changes shape in any mode +`E2E_FIXTURE_MODE` scopes the proxy's provider-bound traffic: `live` (the default, and what an unset variable means: nothing changes), `record` (the proxy's provider calls are forwarded to the real provider through a local edge server and written to a fixture bundle), or `replay` (the edge answers those calls from the bundle, so the run makes zero provider calls and spends nothing). Test-to-proxy traffic always goes over the wire in every mode: record and replay both need the live proxy and database, because the point is that key auth, routing, cost calculation, and spend-log writes execute for real while only the provider is swapped out. Breaking any of those in the proxy turns a replay run red -A bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`) is a directory: `manifest.json` carries the record timestamp, harness git version, and format version, and each test gets a subdirectory holding one JSON file per transport call in call order (`0000-post-chat-completions.json`). Auth header values and credential request fields (`api_key`, `*_secret_key`, `static_headers`, and the like; the list is `fixture_canonical.py`'s) are redacted on write, and file uploads store a sha256 digest instead of the bytes; response bodies are stored verbatim (a /key/generate response keeps the ephemeral virtual key it minted), which is part of why bundles are gitignored. `fixture_bundle.py` owns the format +The seam is `provider_edge.py`: `start_provider_edge` boots an in-process HTTP server (one shared instance per pytest process, `e2e_config.provider_edge_base` is the accessor) that mounts each supported provider under a path prefix (`EDGE_MOUNTS`: `/openai` -> `https://api.openai.com`, `/anthropic` -> `https://api.anthropic.com`). A test participates by registering its deployment with `api_base=provider_edge_base("openai")` plus the provider's path suffix; `quota_management/spend_tracking/test_provider_edge_spend_e2e.py` is the reference. In live mode the accessor returns None and the deployment defaults to the real provider, so an edge-wired test runs in all three modes unchanged. Non-wired tests hit their providers live in every mode. The edge binds `E2E_PROVIDER_EDGE_BIND_HOST` (default 127.0.0.1) and advertises `E2E_PROVIDER_EDGE_ADVERTISE_HOST` in the api_base it hands out, for proxies running in containers -Replay matches calls per test by canonical key: `fixture_canonical.py` canonicalizes the recorded request (volatile headers and credential fields out, unique markers, generated ids, uuids, and timestamps replaced with fixed placeholders, object keys sorted) and the key is the method, path, and a content hash, so identity survives re-records and machine changes while any real content drift is a `ReplayMiss` that names the computed key, the closest recorded key with its file, and a content diff, and never falls through to a live call. Matching is order-independent across distinct keys (concurrent calls may interleave) and FIFO within one key (a poll loop replays its responses in recorded order); a passed test must also consume its whole recording, or teardown fails it naming a leftover key. Either way the fix is always to re-record with `E2E_FIXTURE_MODE=record`. Every rewrite rule lives in `fixture_canonical.py`, so a new volatile header, credential field name, or generated-id shape is one edit there. Record starts fresh every time: it wipes the previous bundle (refusing to wipe a directory that is not a bundle) and never reads it. A replay bundle whose manifest is older than seven days hard-fails at collection time naming the bundle's age, so replay can never certify against fixtures that have drifted more than a week from the live proxy +A bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`) is a directory: `manifest.json` carries the record timestamp, harness git version, and format version, and each test gets a subdirectory holding one JSON file per provider call in call order (`0000-post-openai-v1-chat-completions.json`). Request headers are never stored (provider credentials never touch disk), non-JSON request bodies store a canonicalized sha256 digest instead of the bytes, and responses store status, filtered headers, and the verbatim body base64-encoded, which is part of why bundles are gitignored. `fixture_bundle.py` owns the format. Record serves the proxy the same filtered stored response replay will serve later, so the two modes are byte-identical from the proxy's side of the socket -Deliberately not here yet: streaming chunk fidelity (LIT-5742) and scoping record/replay to provider-bound traffic (LIT-5745) +Replay matches calls per test by canonical key: `fixture_canonical.py` canonicalizes the recorded request (volatile headers and credential fields out, unique markers, generated ids, uuids, and timestamps replaced with fixed placeholders, object keys sorted) and the key is the method, edge path, and a content hash, so identity survives re-records and machine changes while any real content drift comes back as an HTTP 599 naming the computed key, the closest recorded key with its file, and a content diff, and never falls through to a live call. Matching is order-independent across distinct keys (concurrent calls may interleave) and FIFO within one key (a retry loop replays its responses in recorded order); a passed test must also consume its whole recording, or teardown fails it naming a leftover key. Either way the fix is always to re-record with `E2E_FIXTURE_MODE=record`. Every rewrite rule lives in `fixture_canonical.py`, so a new volatile header, credential field name, or generated-id shape is one edit there. Record starts fresh every time: it wipes the previous bundle (refusing to wipe a directory that is not a bundle) and never reads it. A replay bundle whose manifest is older than seven days hard-fails at collection time naming the bundle's age, so replay can never certify against fixtures that have drifted more than a week from the live providers + +A replayed response carries the recorded provider response id, and `LiteLLM_SpendLogs.request_id` (the table's primary key) is that id, so a replay against a database that still holds the record run's rows silently dedupes its spend inserts and any spend assertion goes red with zero matching rows and nothing in the proxy log. Run both modes with `E2E_RESET_SPEND_LOGS=1` (plus `DATABASE_URL` in the runner env) so each session truncates the table after itself, or replay against a fresh database, which is the CI shape + +Current limits: streaming chunk fidelity is LIT-5742 (a streamed response records as one buffered body), CI wiring is LIT-5748, Bedrock cannot be mounted (SigV4 signs the Host header, so a rewritten api_base fails signature verification), multipart uploads have per-run random boundaries (the digest changes every run, so they always miss), and deployments baked into the proxy's config file cannot be edge-wired (only `/model/new` registrations can carry the edge api_base) ## Typing diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 67da1be9562..9096050a45a 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -54,14 +54,16 @@ Some suites need extra services the bare proxy does not start. The `logging/` OT ### Record and replay -`E2E_FIXTURE_MODE=record` runs a suite against the live proxy as usual while writing every request/response pair to a fixture bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`); `E2E_FIXTURE_MODE=replay` then runs the same suite entirely from that bundle, with no proxy traffic and no provider spend; the proxy liveness gate is skipped, so replay runs with no proxy up at all. Unset (or `live`) behaves exactly as before the knob existed +Record/replay scopes to the proxy's provider-bound traffic only. In `E2E_FIXTURE_MODE=record` the harness boots a local provider-edge server, edge-wired tests register their deployments with an `api_base` pointing at it, and every provider call the proxy makes is forwarded verbatim and written to a fixture bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`). `E2E_FIXTURE_MODE=replay` runs the same tests against the same live proxy and database, but the edge answers the proxy's provider calls from the bundle instead of the provider, so the run makes zero provider calls and spends nothing while key auth, routing, cost calculation, and spend-log writes all still execute for real. Unset (or `live`) behaves exactly as before the knob existed. Both record and replay need the proxy up; only the provider is taken out of the loop ```bash -E2E_FIXTURE_MODE=record uv run pytest tests/e2e/llm_translation/ -v -E2E_FIXTURE_MODE=replay uv run pytest tests/e2e/llm_translation/ -v +E2E_FIXTURE_MODE=record uv run pytest tests/e2e/quota_management/spend_tracking/test_provider_edge_spend_e2e.py -v +E2E_FIXTURE_MODE=replay uv run pytest tests/e2e/quota_management/spend_tracking/test_provider_edge_spend_e2e.py -v ``` -Replay fails hard (`ReplayMiss`) when the tests drift from the recording, and a bundle older than seven days fails at collection time naming its age; either way the fix is to re-record. See `CLAUDE.md` in this directory for the bundle format and the transport seam +One sharp edge: a replayed response reuses the recorded provider response id, and that id is the primary key of `LiteLLM_SpendLogs`, so replaying against a database that still holds the record run's rows silently dedupes the spend writes and a spend assertion fails with zero rows. Run both commands above with `E2E_RESET_SPEND_LOGS=1` (and `DATABASE_URL` set in the pytest env) so each session truncates the spend log table after itself, or point replay at a fresh database + +Replay answers any provider call that drifted from the recording with an HTTP 599 whose body names the computed and closest recorded keys, so the test fails loudly instead of silently going live, and a bundle older than seven days fails at collection time naming its age; either way the fix is to re-record. Only tests that register edge-wired deployments participate: everything else hits its provider live in every mode, so record exactly the suite you replay. If the proxy runs in a container, set `E2E_PROVIDER_EDGE_ADVERTISE_HOST` (e.g. `host.docker.internal`) so the api_base the proxy stores can reach the edge on the pytest host, and `E2E_PROVIDER_EDGE_BIND_HOST=0.0.0.0` so the edge accepts it. See `CLAUDE.md` in this directory for the bundle format, the edge design, and the current limits (streaming, Bedrock, multipart) Tests marked `@pytest.mark.e2e` hard-fail when no proxy answers `/health/liveliness`, so a run that goes red with `No live proxy` at setup means the proxy isn't up; they never skip for a missing proxy, so an absent proxy can't be mistaken for a pass diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index da2a7da0bfa..dbe2d6e514e 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -23,12 +23,8 @@ 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_transport import ( - fixture_mode_collection_error, - fixture_report_lines, - parse_fixture_mode, - replay_leftover_error, -) +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 proxy_client import ProxyClient, build_proxy_client @@ -114,12 +110,10 @@ def _proxy_fail_reason() -> str | None: def pytest_runtest_setup(item: pytest.Item) -> None: """Hard-fail `e2e`-marked tests unless a proxy answers its liveness probe. Unmarked tests (unit coverage of the harness) don't touch the proxy, so they - run even when none is up. Never skip for a missing proxy. Replay mode serves - every call from the fixture bundle, so it needs no live proxy either.""" + run even when none is up. Never skip for a missing proxy. Replay mode needs + the proxy too: only provider-bound traffic replays from the bundle.""" if item.get_closest_marker("e2e") is None: return - if parse_fixture_mode(FIXTURE_MODE_RAW) == "replay": - return reason = _proxy_fail_reason() if reason is not None: pytest.fail(reason) diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index a5c3729f4be..8bf39f6021f 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -13,7 +13,8 @@ from pathlib import Path from dotenv import load_dotenv -from fixture_transport import deterministic_marker, parse_fixture_mode +from fixture_mode import deterministic_marker, parse_fixture_mode +from provider_edge import provider_edge_api_base # Local runs keep provider / DataDog keys in tests/e2e/.env (see CONTRIBUTING.md). # Compose injects them into the proxy container, but pytest on the host does not @@ -92,15 +93,24 @@ PROPAGATION_TIMEOUT = float(os.environ.get("E2E_PROPAGATION_TIMEOUT", "15")) EXPECT_RUST = os.environ.get("E2E_EXPECT_RUST", "").strip().lower() in ("1", "true", "yes") -# Record/replay fixture selection (see fixture_transport.py). The raw mode value -# is parsed and validated there; "live" (the default, also for empty values) -# means the harness behaves exactly as before this knob existed. +# Record/replay fixture selection (see fixture_mode.py and provider_edge.py). +# The raw mode value is parsed and validated there; "live" (the default, also +# for empty values) means the harness behaves exactly as before this knob +# existed. FIXTURE_MODE_RAW = os.environ.get("E2E_FIXTURE_MODE", "live") FIXTURE_DIR = Path( os.environ.get("E2E_FIXTURE_DIR", "").strip() or str(Path(__file__).resolve().parent / ".fixtures") ) +# Where the provider-edge server binds, and the host name edge api_base URLs +# advertise to the proxy. They differ when the proxy runs in a container and +# reaches the pytest host via a gateway name like host.docker.internal. +PROVIDER_EDGE_BIND_HOST = os.environ.get("E2E_PROVIDER_EDGE_BIND_HOST", "").strip() or "127.0.0.1" +PROVIDER_EDGE_ADVERTISE_HOST = ( + os.environ.get("E2E_PROVIDER_EDGE_ADVERTISE_HOST", "").strip() or PROVIDER_EDGE_BIND_HOST +) + # Deliberately modest concurrency. The suite shares its proxy with every other # suite in the run, and 750 users at spawn rate 50 saturated the request path hard # enough to distort latency-sensitive neighbours (and to spend real provider money @@ -157,6 +167,20 @@ def datadog_mcp_url(*, toolsets: str = "core") -> str: return f"{base}?toolsets={toolsets}" if toolsets else base +def provider_edge_base(mount: str) -> str | None: + """The api_base an edge-wired deployment should register with, using this + process's fixture-mode and edge-host configuration: None in live mode, the + shared edge server's mount URL in record and replay.""" + return provider_edge_api_base( + mount, + mode_raw=FIXTURE_MODE_RAW, + bundle_dir=FIXTURE_DIR, + bind_host=PROVIDER_EDGE_BIND_HOST, + advertise_host=PROVIDER_EDGE_ADVERTISE_HOST, + forward_timeout=REQUEST_TIMEOUT, + ) + + def unique_marker() -> str: """A short unique token per call/run, so concurrent runs and the shared response cache never collide on prompts, tags, or customer ids. In record diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index cb6fc7a01e5..03f201e946e 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -647,3 +647,37 @@ def download( content_type=_hdr(resp, "content-type"), body=resp.text, ) + + +class RawResponse(BaseModel): + """A verbatim upstream HTTP response for the provider edge (provider_edge.py): + status, lowercased headers, raw bytes. No Result classification because the + edge relays provider errors to the proxy untouched.""" + + status_code: int + headers: dict[str, str] + body: bytes + + +def forward( + method: str, + url: str, + *, + headers: dict[str, str], + body: bytes | None, + timeout: float = 60.0, +) -> RawResponse | NetworkError: + """Relay one provider-bound request verbatim for the provider edge's record + mode. No retries, no redirects, no schema: the proxy owns retry policy and + the recorded bundle must hold exactly what the provider returned.""" + try: + resp = requests.request( + method, url, headers=headers, data=body, timeout=timeout, allow_redirects=False + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return RawResponse( + status_code=resp.status_code, + headers={name.lower(): value for name, value in resp.headers.items()}, + body=resp.content, + ) diff --git a/tests/e2e/fixture_bundle.py b/tests/e2e/fixture_bundle.py index 615ae8df1a4..6feb40fc8bc 100644 --- a/tests/e2e/fixture_bundle.py +++ b/tests/e2e/fixture_bundle.py @@ -1,17 +1,18 @@ -"""On-disk fixture bundle format for record/replay e2e runs (LIT-5729). +"""On-disk fixture bundle format for record/replay e2e runs (LIT-5729/LIT-5745). A bundle is a directory: one ``manifest.json`` (record timestamp + harness version + format version) plus one subdirectory per test, holding one JSON file -per transport interaction in call order. Bundles older than +per provider-bound interaction in call order. Bundles older than ``MAX_BUNDLE_AGE`` hard-fail replay at collection time (see conftest), so a green replay run can never certify against fixtures that have drifted more than -a week from the live proxy. +a week from the live providers. -This module owns the format only. The transports that produce and consume it -live in fixture_transport.py and the canonical match keys they compute live in -fixture_canonical.py (LIT-5741); streaming chunk fidelity and provider-scoping -are follow-ups (LIT-5742/5745). Every interaction file stores the full redacted -request because replay matches on its canonicalized content. +This module owns the format only. The provider-edge server that produces and +consumes it lives in provider_edge.py (LIT-5745) and the canonical match keys +it computes live in fixture_canonical.py (LIT-5741); streaming chunk fidelity +is a follow-up (LIT-5742). Every interaction file stores the full redacted +request because replay matches on its canonicalized content, and the response +as the raw HTTP status, filtered headers, and base64 body the provider sent. """ from __future__ import annotations @@ -23,29 +24,14 @@ import subprocess from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from pathlib import Path -from typing import Annotated, Final, Literal +from typing import Final -from pydantic import BaseModel, Field, JsonValue, TypeAdapter +from pydantic import BaseModel, JsonValue -from e2e_http import ( - BinaryStream, - NetworkError, - ProbeResult, - RateLimitedError, - Result, - StreamingResponse, - Success, - UnauthorizedError, - UnknownApiError, - ValidationError, -) - -BUNDLE_FORMAT_VERSION: Final = 1 +BUNDLE_FORMAT_VERSION: Final = 2 MAX_BUNDLE_AGE: Final = timedelta(days=7) MANIFEST_FILENAME: Final = "manifest.json" -_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) - class Manifest(BaseModel): format_version: int @@ -54,13 +40,14 @@ class Manifest(BaseModel): class RecordedRequest(BaseModel): - """The request as the transport saw it, auth header values and credential - body/form fields redacted. + """The provider-bound request as the edge saw it, headers empty (SDK + telemetry headers vary run to run and auth material never touches disk). Replay matches on the canonical content key fixture_canonical.py computes - over ``method`` (the transport verb, not the HTTP verb), ``path``, and the - canonicalized headers, params, body, form, and file identity. File uploads - store a content digest instead of the bytes.""" + over ``method``, ``path`` (the edge path including the provider mount, + query string excluded), and the canonicalized headers, params, body, form, + and file identity. Non-JSON bodies store a canonicalized content digest + instead of the bytes.""" method: str path: str @@ -73,85 +60,19 @@ class RecordedRequest(BaseModel): file_bytes: int | None = None -class RecordedResult(BaseModel): - """A ``Result[R]`` flattened for disk. ``data`` holds the success payload as - raw JSON; replay re-validates it against the ``response_type`` the caller - passes, exactly like a live response body.""" +class RecordedHttpResponse(BaseModel): + """The provider's raw HTTP response: status, headers minus hop-by-hop and + volatile entries (see provider_edge.py), and the body as base64 so binary + payloads survive JSON.""" - shape: Literal["result"] = "result" - kind: Literal["success", "network", "unauthorized", "rate_limited", "validation", "unknown"] - status_code: int | None = None - data: JsonValue | None = None - message: str | None = None - body: str | None = None - retry_after_seconds: int | None = None - - -class RecordedStreaming(BaseModel): - shape: Literal["streaming"] = "streaming" - payload: StreamingResponse - - -class RecordedBinary(BaseModel): - shape: Literal["binary"] = "binary" - payload: BinaryStream - - -class RecordedProbe(BaseModel): - shape: Literal["probe"] = "probe" - payload: ProbeResult - - -type RecordedResponse = RecordedResult | RecordedStreaming | RecordedBinary | RecordedProbe + status_code: int + headers: dict[str, str] + body_b64: str class Interaction(BaseModel): request: RecordedRequest - response: Annotated[ - RecordedResult | RecordedStreaming | RecordedBinary | RecordedProbe, - Field(discriminator="shape"), - ] - - -def to_json_value(model: BaseModel) -> JsonValue: - return _JSON.validate_json(model.model_dump_json(by_alias=True)) - - -def from_result[R: BaseModel](result: Result[R]) -> RecordedResult: - match result: - case Success(status_code=status_code, data=data): - return RecordedResult(kind="success", status_code=status_code, data=to_json_value(data)) - case NetworkError(message=message): - return RecordedResult(kind="network", message=message) - case UnauthorizedError(): - return RecordedResult(kind="unauthorized") - case RateLimitedError(retry_after_seconds=retry_after_seconds, body=body): - return RecordedResult(kind="rate_limited", retry_after_seconds=retry_after_seconds, body=body) - case ValidationError(message=message): - return RecordedResult(kind="validation", message=message) - case UnknownApiError(status_code=status_code, body=body): - return RecordedResult(kind="unknown", status_code=status_code, body=body) - - -def to_result[R: BaseModel](recorded: RecordedResult, response_type: type[R]) -> Result[R]: - match recorded.kind: - case "success": - return Success( - status_code=recorded.status_code or 200, - data=response_type.model_validate(recorded.data), - ) - case "network": - return NetworkError(message=recorded.message or "") - case "unauthorized": - return UnauthorizedError() - case "rate_limited": - return RateLimitedError( - retry_after_seconds=recorded.retry_after_seconds, body=recorded.body or "" - ) - case "validation": - return ValidationError(message=recorded.message or "") - case "unknown": - return UnknownApiError(status_code=recorded.status_code or 0, body=recorded.body or "") + response: RecordedHttpResponse def slugify(raw: str, *, limit: int = 60) -> str: @@ -198,7 +119,7 @@ class BundleRecorder: root: Path _ordinals: dict[str, int] = field(default_factory=dict) - def record(self, *, test_key: str, request: RecordedRequest, response: RecordedResponse) -> None: + def record(self, *, test_key: str, request: RecordedRequest, response: RecordedHttpResponse) -> None: slug = slug_for_test(test_key) ordinal = self._ordinals.get(slug, 0) self._ordinals[slug] = ordinal + 1 diff --git a/tests/e2e/fixture_mode.py b/tests/e2e/fixture_mode.py new file mode 100644 index 00000000000..110f44380b4 --- /dev/null +++ b/tests/e2e/fixture_mode.py @@ -0,0 +1,132 @@ +"""Fixture-mode selection and per-test determinism for record/replay e2e runs. + +``E2E_FIXTURE_MODE`` is live (the default; nothing changes), record, or replay. +This module owns everything mode-shaped that is independent of the provider +edge itself: parsing the raw env value, the collection-time gate that aborts a +run whose mode can never work (unknown value, or replay against a missing or +stale bundle), the pytest report-header lines, the running test's node id, and +the deterministic per-test marker that lets a replay run regenerate exactly +the requests the record run sent. The provider-edge server that records and +serves provider traffic lives in provider_edge.py (LIT-5745). +""" + +from __future__ import annotations + +import hashlib +import os +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Final, Literal, assert_never + +from fixture_bundle import ( + FreshBundle, + StaleBundle, + UnreadableBundle, + check_freshness, + format_age, +) + +type FixtureMode = Literal["live", "record", "replay"] + +FIXTURE_MODES: Final[tuple[FixtureMode, ...]] = ("live", "record", "replay") + +SESSION_TEST_KEY: Final = "session" + + +@dataclass(frozen=True, slots=True) +class InvalidFixtureMode: + value: str + + +def parse_fixture_mode(raw: str) -> FixtureMode | InvalidFixtureMode: + normalized = raw.strip().lower() or "live" + match normalized: + case "live" | "record" | "replay": + return normalized + case _: + return InvalidFixtureMode(value=raw) + + +def current_test_key() -> str: + """The pytest node id of the running test, from the PYTEST_CURRENT_TEST env + var pytest maintains (`` (setup|call|teardown)``); ``session`` for + calls outside any test (e.g. session-finish cleanup).""" + raw = os.environ.get("PYTEST_CURRENT_TEST", "") + if not raw: + return SESSION_TEST_KEY + return raw.rsplit(" (", 1)[0] + + +class ReplayMiss(AssertionError): + """Replay had no recorded interaction for a provider call the proxy made. + The suite drifted from the bundle (or the bundle from the suite): re-record.""" + + +_marker_ordinals: Final[dict[str, int]] = {} + + +def deterministic_marker() -> str: + """Stable stand-in for uuid-based unique markers in record and replay modes: + the Nth marker of a test is a pure function of the test's node id and N, so a + replay run regenerates exactly the model names, prompts, and tags the record + run sent and every recorded provider interaction still matches its key.""" + test_key = current_test_key() + ordinal = _marker_ordinals.get(test_key, 0) + _marker_ordinals[test_key] = ordinal + 1 + return hashlib.sha1(f"{test_key}#{ordinal}".encode()).hexdigest()[:12] + + +def fixture_mode_collection_error(mode_raw: str, bundle_dir: Path, *, now: datetime) -> str | None: + """Session-abort reason for a fixture-mode setup that can never work, or None. + Called at collection time (conftest pytest_sessionstart) so a stale or missing + bundle fails the whole run up front, naming the bundle age, instead of failing + every test individually.""" + mode = parse_fixture_mode(mode_raw) + match mode: + case InvalidFixtureMode(value=value): + return f"E2E_FIXTURE_MODE={value!r} is not one of {', '.join(FIXTURE_MODES)}" + case "live" | "record": + return None + case "replay": + freshness = check_freshness(bundle_dir, now=now) + match freshness: + case FreshBundle(): + return None + case StaleBundle(recorded_at=recorded_at, age=age, limit=limit): + return ( + f"fixture bundle at {bundle_dir} is stale: recorded {recorded_at.isoformat()}, " + f"age {format_age(age)} exceeds the {limit.days}-day limit; " + "re-record with E2E_FIXTURE_MODE=record" + ) + case UnreadableBundle(reason=reason): + return f"E2E_FIXTURE_MODE=replay cannot use bundle at {bundle_dir}: {reason}" + case _: + assert_never(freshness) + case _: + assert_never(mode) + + +def fixture_report_lines(mode_raw: str, bundle_dir: Path, *, now: datetime) -> list[str]: + """pytest report-header lines; empty in live mode so an unset + E2E_FIXTURE_MODE keeps today's output byte-identical.""" + mode = parse_fixture_mode(mode_raw) + match mode: + case InvalidFixtureMode() | "live": + return [] + case "record": + return [f"e2e fixture mode: record -> {bundle_dir}"] + case "replay": + freshness = check_freshness(bundle_dir, now=now) + match freshness: + case FreshBundle(manifest=manifest): + return [ + f"e2e fixture mode: replay <- {bundle_dir} " + f"(recorded {manifest.recorded_at.isoformat()}, harness {manifest.harness_version})" + ] + case StaleBundle() | UnreadableBundle(): + return [f"e2e fixture mode: replay <- {bundle_dir}"] + case _: + assert_never(freshness) + case _: + assert_never(mode) diff --git a/tests/e2e/fixture_transport.py b/tests/e2e/fixture_transport.py deleted file mode 100644 index ce4eec701ca..00000000000 --- a/tests/e2e/fixture_transport.py +++ /dev/null @@ -1,724 +0,0 @@ -"""Record/replay transports behind the same ``Transport`` protocol (LIT-5729). - -``RecordingTransport`` decorates the live transport: every call passes through -unchanged and its request/response pair is appended to the fixture bundle. -``ReplayTransport`` implements the protocol from a recorded bundle alone: no -HTTP, no proxy, no provider spend. Because both fulfil ``Transport``, no test -or client changes shape; ``build_proxy_client`` picks the transport from -``E2E_FIXTURE_MODE`` (live | record | replay, default live). - -Replay matches each call by test node id and canonical content key -(fixture_canonical.py, LIT-5741): volatile headers, credential fields, unique -markers, generated ids, and timestamps are canonicalized out before hashing, so -matching is order-independent across distinct keys, FIFO within a key, and a -miss fails hard (``ReplayMiss``) printing the computed key and the closest -recorded key without ever falling through to a live call. Streaming chunk -fidelity is LIT-5742; scoping record/replay to provider-bound traffic is -LIT-5745. -""" - -from __future__ import annotations - -import difflib -import functools -import hashlib -import os -from collections import deque -from dataclasses import dataclass, field -from datetime import datetime -from itertools import islice -from pathlib import Path -from typing import Final, Literal, assert_never - -from pydantic import BaseModel, JsonValue - -from e2e_http import AuthHeaders, BinaryStream, ProbeResult, Result, StreamingResponse -from fixture_bundle import ( - BundleRecorder, - FreshBundle, - Interaction, - LoadedBundle, - RecordedBinary, - RecordedProbe, - RecordedRequest, - RecordedResponse, - RecordedResult, - RecordedStreaming, - StaleBundle, - UnreadableBundle, - UnsafeBundleDir, - check_freshness, - format_age, - from_result, - interaction_filename, - load_bundle, - prepare_bundle, - slug_for_test, - to_json_value, - to_result, -) -from fixture_canonical import CanonicalRequest, canonicalize, is_secret_field -from transport import Transport - -type FixtureMode = Literal["live", "record", "replay"] - -FIXTURE_MODES: Final[tuple[FixtureMode, ...]] = ("live", "record", "replay") - -SESSION_TEST_KEY: Final = "session" - -REDACTED_HEADER_NAMES: Final[frozenset[str]] = frozenset({"authorization", "x-litellm-api-key"}) -REDACTED_VALUE: Final = "" - - -@dataclass(frozen=True, slots=True) -class InvalidFixtureMode: - value: str - - -def parse_fixture_mode(raw: str) -> FixtureMode | InvalidFixtureMode: - normalized = raw.strip().lower() or "live" - match normalized: - case "live" | "record" | "replay": - return normalized - case _: - return InvalidFixtureMode(value=raw) - - -def current_test_key() -> str: - """The pytest node id of the running test, from the PYTEST_CURRENT_TEST env - var pytest maintains (`` (setup|call|teardown)``); ``session`` for - calls outside any test (e.g. session-finish cleanup).""" - raw = os.environ.get("PYTEST_CURRENT_TEST", "") - if not raw: - return SESSION_TEST_KEY - return raw.rsplit(" (", 1)[0] - - -class ReplayMiss(AssertionError): - """Replay had no recorded interaction for a call the suite made. The test - drifted from the bundle (or the bundle from the suite): re-record.""" - - -_marker_ordinals: Final[dict[str, int]] = {} - - -def deterministic_marker() -> str: - """Stable stand-in for uuid-based unique markers in record and replay modes: - the Nth marker of a test is a pure function of the test's node id and N, so a - replay run regenerates exactly the model names, prompts, and tags the record - run sent and every recorded poll response still satisfies its predicate.""" - test_key = current_test_key() - ordinal = _marker_ordinals.get(test_key, 0) - _marker_ordinals[test_key] = ordinal + 1 - return hashlib.sha1(f"{test_key}#{ordinal}".encode()).hexdigest()[:12] - - -def _dump_flat(model: BaseModel | None) -> dict[str, str]: - if model is None: - return {} - dumped: dict[str, object] = model.model_dump(by_alias=True, exclude_none=True) - return {key: str(value) for key, value in dumped.items()} - - -def _redact(headers: dict[str, str]) -> dict[str, str]: - return { - name: REDACTED_VALUE if name.lower() in REDACTED_HEADER_NAMES else value - for name, value in headers.items() - } - - -def _redact_secret_fields(value: JsonValue) -> JsonValue: - match value: - case dict(): - return { - key: REDACTED_VALUE - if is_secret_field(key) and item is not None - else _redact_secret_fields(item) - for key, item in value.items() - } - case list(): - return [_redact_secret_fields(item) for item in value] - case _: - return value - - -def _redact_flat(fields: dict[str, str]) -> dict[str, str]: - return { - key: REDACTED_VALUE if is_secret_field(key) else value for key, value in fields.items() - } - - -def recorded_request( - method: str, - path: str, - *, - headers: BaseModel, - body: BaseModel | None = None, - params: BaseModel | None = None, - form: BaseModel | None = None, - file_name: str | None = None, - file_content: bytes | None = None, -) -> RecordedRequest: - return RecordedRequest( - method=method, - path=path, - headers=_redact(_dump_flat(headers)), - params=_redact_flat(_dump_flat(params)), - body=None if body is None else _redact_secret_fields(to_json_value(body)), - form=None if form is None else _redact_flat(_dump_flat(form)), - file_name=file_name, - file_sha256=None if file_content is None else hashlib.sha256(file_content).hexdigest(), - file_bytes=None if file_content is None else len(file_content), - ) - - -@dataclass(frozen=True, slots=True) -class RecordingTransport: - """Decorator over the live transport: forwards every call and appends the - interaction to the bundle, so a green live run leaves behind exactly the - traffic replay needs.""" - - inner: Transport - recorder: BundleRecorder - - def _record(self, request: RecordedRequest, response: RecordedResponse) -> None: - self.recorder.record(test_key=current_test_key(), request=request, response=response) - - def bearer(self, key: str) -> AuthHeaders: - return self.inner.bearer(key) - - @property - def master(self) -> AuthHeaders: - return self.inner.master - - def post[R: BaseModel]( - self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] - ) -> Result[R]: - result = self.inner.post(path, headers=headers, json=json, response_type=response_type) - self._record(recorded_request("post", path, headers=headers, body=json), from_result(result)) - return result - - def get[R: BaseModel]( - self, - path: str, - *, - headers: BaseModel, - params: BaseModel, - response_type: type[R], - timeout: float | None = None, - ) -> Result[R]: - result = self.inner.get( - path, headers=headers, params=params, response_type=response_type, timeout=timeout - ) - self._record(recorded_request("get", path, headers=headers, params=params), from_result(result)) - return result - - def delete[R: BaseModel]( - self, - path: str, - *, - headers: BaseModel, - json: BaseModel, - response_type: type[R], - params: BaseModel | None = None, - ) -> Result[R]: - result = self.inner.delete( - path, headers=headers, json=json, response_type=response_type, params=params - ) - self._record( - recorded_request("delete", path, headers=headers, body=json, params=params), - from_result(result), - ) - return result - - def patch[R: BaseModel]( - self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] - ) -> Result[R]: - result = self.inner.patch(path, headers=headers, json=json, response_type=response_type) - self._record(recorded_request("patch", path, headers=headers, body=json), from_result(result)) - return result - - def put[R: BaseModel]( - self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] - ) -> Result[R]: - result = self.inner.put(path, headers=headers, json=json, response_type=response_type) - self._record(recorded_request("put", path, headers=headers, body=json), from_result(result)) - return result - - def stream(self, path: str, *, headers: BaseModel, json: BaseModel) -> StreamingResponse: - response = self.inner.stream(path, headers=headers, json=json) - self._record( - recorded_request("stream", path, headers=headers, body=json), - RecordedStreaming(payload=response), - ) - return response - - def stream_binary( - self, path: str, *, headers: BaseModel, json: BaseModel, chunk_size: int = 8192 - ) -> BinaryStream: - response = self.inner.stream_binary(path, headers=headers, json=json, chunk_size=chunk_size) - self._record( - recorded_request("stream_binary", path, headers=headers, body=json), - RecordedBinary(payload=response), - ) - return response - - def send( - self, - path: str, - *, - headers: BaseModel, - json: BaseModel, - params: BaseModel | None = None, - stream: bool = False, - ) -> StreamingResponse: - response = self.inner.send(path, headers=headers, json=json, params=params, stream=stream) - self._record( - recorded_request("send", path, headers=headers, body=json, params=params), - RecordedStreaming(payload=response), - ) - return response - - def probe(self, path: str, *, params: BaseModel) -> ProbeResult: - response = self.inner.probe(path, params=params) - self._record( - recorded_request("probe", path, headers=self.master, params=params), - RecordedProbe(payload=response), - ) - return response - - def upload[R: BaseModel]( - self, - path: str, - *, - headers: BaseModel, - form: BaseModel, - filename: str, - content: bytes, - file_content_type: str = "application/jsonl", - file_field: str = "file", - params: BaseModel | None = None, - response_type: type[R], - ) -> Result[R]: - result = self.inner.upload( - path, - headers=headers, - form=form, - filename=filename, - content=content, - file_content_type=file_content_type, - file_field=file_field, - params=params, - response_type=response_type, - ) - self._record( - recorded_request( - "upload", - path, - headers=headers, - params=params, - form=form, - file_name=filename, - file_content=content, - ), - from_result(result), - ) - return result - - def download(self, path: str, *, headers: BaseModel) -> StreamingResponse: - response = self.inner.download(path, headers=headers) - self._record( - recorded_request("download", path, headers=headers), - RecordedStreaming(payload=response), - ) - return response - - -def _build_pool(recorded: tuple[Interaction, ...]) -> dict[str, deque[Interaction]]: - keys: Final = tuple(canonicalize(interaction.request).key for interaction in recorded) - return { - key: deque( - interaction - for candidate_key, interaction in zip(keys, recorded, strict=True) - if candidate_key == key - ) - for key in dict.fromkeys(keys) - } - - -def _closest_recorded( - canonical: CanonicalRequest, recorded: tuple[Interaction, ...] -) -> tuple[CanonicalRequest, str]: - candidates: Final = tuple(canonicalize(interaction.request) for interaction in recorded) - ratios: Final = tuple( - difflib.SequenceMatcher( - None, f"{canonical.method} {canonical.path}\n{canonical.content}", - f"{candidate.method} {candidate.path}\n{candidate.content}", - ).ratio() - for candidate in candidates - ) - best: Final = max(range(len(candidates)), key=lambda index: ratios[index]) - return candidates[best], interaction_filename(best, recorded[best].request) - - -def _miss_message(test_key: str, slug: str, canonical: CanonicalRequest, bundle: LoadedBundle) -> str: - recorded: Final = bundle.interactions.get(slug, ()) - if not recorded: - return ( - f"replay miss for {test_key}: computed key {canonical.key} but nothing is recorded " - f"under {slug}; re-record with E2E_FIXTURE_MODE=record" - ) - closest, closest_file = _closest_recorded(canonical, recorded) - diff: Final = "\n".join( - islice( - difflib.unified_diff( - closest.pretty_content().splitlines(), - canonical.pretty_content().splitlines(), - fromfile=f"closest recorded ({closest_file})", - tofile="test made", - lineterm="", - ), - 60, - ) - ) - return ( - f"replay miss for {test_key}: no recorded interaction matches key {canonical.key}; " - f"closest recorded key is {closest.key} ({closest_file})\n{diff}\n" - "re-record with E2E_FIXTURE_MODE=record" - ) - - -@dataclass(slots=True) -class ReplaySource: - """One shared pool per test over a loaded bundle, so every client built in - the session consumes the same recorded interactions. Every pool is built - once at construction and per-key consumption is a single atomic deque pop, - so concurrent replay calls never race. Calls match by canonical content - key: order-independent across distinct keys (concurrent tests interleave - calls nondeterministically), FIFO within one key (a poll loop replays its - recorded responses in recorded order).""" - - bundle: LoadedBundle - _pools: dict[str, dict[str, deque[Interaction]]] = field(init=False) - - def __post_init__(self) -> None: - self._pools = { - slug: _build_pool(recorded) for slug, recorded in self.bundle.interactions.items() - } - - def _pool(self, slug: str) -> dict[str, deque[Interaction]]: - return self._pools.get(slug, {}) - - def next_interaction(self, request: RecordedRequest) -> Interaction: - test_key: Final = current_test_key() - slug: Final = slug_for_test(test_key) - pool: Final = self._pool(slug) - canonical: Final = canonicalize(request) - queue: Final = pool.get(canonical.key) - if queue is None: - raise ReplayMiss(_miss_message(test_key, slug, canonical, self.bundle)) - try: - return queue.popleft() - except IndexError: - raise ReplayMiss( - f"replay exhausted for {test_key}: every recorded interaction for key " - f"{canonical.key} is already consumed; re-record with E2E_FIXTURE_MODE=record" - ) from None - - def leftover_error(self, test_key: str) -> str | None: - """Non-None when the test consumed fewer interactions than were recorded, - meaning a passing replay proved less than the bundle claims.""" - slug: Final = slug_for_test(test_key) - recorded: Final = self.bundle.interactions.get(slug, ()) - if not recorded: - return None - leftover: Final = tuple( - interaction for queue in self._pool(slug).values() for interaction in queue - ) - if not leftover: - return None - return ( - f"replay incomplete for {test_key}: {len(leftover)} of {len(recorded)} recorded " - f"interactions never consumed, e.g. {canonicalize(leftover[0].request).key}; " - "re-record with E2E_FIXTURE_MODE=record" - ) - - -def _expect_result(interaction: Interaction) -> RecordedResult: - match interaction.response: - case RecordedResult() as recorded: - return recorded - case RecordedStreaming() | RecordedBinary() | RecordedProbe(): - raise ReplayMiss( - f"recorded {interaction.request.method} {interaction.request.path} is not a typed result" - ) - - -def _expect_streaming(interaction: Interaction) -> StreamingResponse: - match interaction.response: - case RecordedStreaming(payload=payload): - return payload - case RecordedResult() | RecordedBinary() | RecordedProbe(): - raise ReplayMiss( - f"recorded {interaction.request.method} {interaction.request.path} is not a streaming response" - ) - - -@dataclass(frozen=True, slots=True) -class ReplayTransport: - """A ``Transport`` served entirely from a recorded bundle: never opens a - connection, so a replay run cannot bill a provider.""" - - source: ReplaySource - master_key: str - - def bearer(self, key: str) -> AuthHeaders: - return AuthHeaders(authorization=f"Bearer {key}") - - @property - def master(self) -> AuthHeaders: - return self.bearer(self.master_key) - - def post[R: BaseModel]( - self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] - ) -> Result[R]: - return to_result( - _expect_result( - self.source.next_interaction(recorded_request("post", path, headers=headers, body=json)) - ), - response_type, - ) - - def get[R: BaseModel]( - self, - path: str, - *, - headers: BaseModel, - params: BaseModel, - response_type: type[R], - timeout: float | None = None, - ) -> Result[R]: - return to_result( - _expect_result( - self.source.next_interaction(recorded_request("get", path, headers=headers, params=params)) - ), - response_type, - ) - - def delete[R: BaseModel]( - self, - path: str, - *, - headers: BaseModel, - json: BaseModel, - response_type: type[R], - params: BaseModel | None = None, - ) -> Result[R]: - return to_result( - _expect_result( - self.source.next_interaction( - recorded_request("delete", path, headers=headers, body=json, params=params) - ) - ), - response_type, - ) - - def patch[R: BaseModel]( - self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] - ) -> Result[R]: - return to_result( - _expect_result( - self.source.next_interaction(recorded_request("patch", path, headers=headers, body=json)) - ), - response_type, - ) - - def put[R: BaseModel]( - self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] - ) -> Result[R]: - return to_result( - _expect_result( - self.source.next_interaction(recorded_request("put", path, headers=headers, body=json)) - ), - response_type, - ) - - def stream(self, path: str, *, headers: BaseModel, json: BaseModel) -> StreamingResponse: - return _expect_streaming( - self.source.next_interaction(recorded_request("stream", path, headers=headers, body=json)) - ) - - def stream_binary( - self, path: str, *, headers: BaseModel, json: BaseModel, chunk_size: int = 8192 - ) -> BinaryStream: - interaction = self.source.next_interaction( - recorded_request("stream_binary", path, headers=headers, body=json) - ) - match interaction.response: - case RecordedBinary(payload=payload): - return payload - case RecordedResult() | RecordedStreaming() | RecordedProbe(): - raise ReplayMiss( - f"recorded stream_binary {interaction.request.path} is not a binary stream" - ) - - def send( - self, - path: str, - *, - headers: BaseModel, - json: BaseModel, - params: BaseModel | None = None, - stream: bool = False, - ) -> StreamingResponse: - return _expect_streaming( - self.source.next_interaction( - recorded_request("send", path, headers=headers, body=json, params=params) - ) - ) - - def probe(self, path: str, *, params: BaseModel) -> ProbeResult: - interaction = self.source.next_interaction( - recorded_request("probe", path, headers=self.master, params=params) - ) - match interaction.response: - case RecordedProbe(payload=payload): - return payload - case RecordedResult() | RecordedStreaming() | RecordedBinary(): - raise ReplayMiss(f"recorded probe {interaction.request.path} is not a probe result") - - def upload[R: BaseModel]( - self, - path: str, - *, - headers: BaseModel, - form: BaseModel, - filename: str, - content: bytes, - file_content_type: str = "application/jsonl", - file_field: str = "file", - params: BaseModel | None = None, - response_type: type[R], - ) -> Result[R]: - return to_result( - _expect_result( - self.source.next_interaction( - recorded_request( - "upload", - path, - headers=headers, - params=params, - form=form, - file_name=filename, - file_content=content, - ) - ) - ), - response_type, - ) - - def download(self, path: str, *, headers: BaseModel) -> StreamingResponse: - return _expect_streaming( - self.source.next_interaction(recorded_request("download", path, headers=headers)) - ) - - -@functools.lru_cache(maxsize=8) -def _shared_recorder(root: Path) -> BundleRecorder: - prepared = prepare_bundle(root) - if isinstance(prepared, UnsafeBundleDir): - raise ValueError(f"E2E_FIXTURE_DIR {prepared.path} {prepared.reason}") - return prepared - - -@functools.lru_cache(maxsize=8) -def _shared_replay_source(root: Path) -> ReplaySource: - loaded = load_bundle(root) - if isinstance(loaded, UnreadableBundle): - raise ValueError(f"cannot replay from {root}: {loaded.reason}") - return ReplaySource(bundle=loaded) - - -def replay_leftover_error(*, mode_raw: str, bundle_dir: Path, test_key: str) -> str | None: - """Teardown-time completeness check: in replay mode a passed test with - unconsumed recorded interactions must fail instead of passing against a - recording it no longer matches. Inert in every other mode.""" - if parse_fixture_mode(mode_raw) != "replay": - return None - return _shared_replay_source(bundle_dir).leftover_error(test_key) - - -def select_transport( - live: Transport, *, mode_raw: str, bundle_dir: Path, master_key: str -) -> Transport: - """The one seam every client build goes through: wraps (record), replaces - (replay), or passes through (live) the transport per E2E_FIXTURE_MODE. The - recorder and replay cursors are process-wide singletons per bundle dir, so - every client in a session shares one bundle and one recorded sequence.""" - mode = parse_fixture_mode(mode_raw) - match mode: - case InvalidFixtureMode(value=value): - raise ValueError(f"E2E_FIXTURE_MODE={value!r} is not one of {', '.join(FIXTURE_MODES)}") - case "live": - return live - case "record": - return RecordingTransport(inner=live, recorder=_shared_recorder(bundle_dir)) - case "replay": - return ReplayTransport(source=_shared_replay_source(bundle_dir), master_key=master_key) - case _: - assert_never(mode) - - -def fixture_mode_collection_error(mode_raw: str, bundle_dir: Path, *, now: datetime) -> str | None: - """Session-abort reason for a fixture-mode setup that can never work, or None. - Called at collection time (conftest pytest_sessionstart) so a stale or missing - bundle fails the whole run up front, naming the bundle age, instead of failing - every test individually.""" - mode = parse_fixture_mode(mode_raw) - match mode: - case InvalidFixtureMode(value=value): - return f"E2E_FIXTURE_MODE={value!r} is not one of {', '.join(FIXTURE_MODES)}" - case "live" | "record": - return None - case "replay": - freshness = check_freshness(bundle_dir, now=now) - match freshness: - case FreshBundle(): - return None - case StaleBundle(recorded_at=recorded_at, age=age, limit=limit): - return ( - f"fixture bundle at {bundle_dir} is stale: recorded {recorded_at.isoformat()}, " - f"age {format_age(age)} exceeds the {limit.days}-day limit; " - "re-record with E2E_FIXTURE_MODE=record" - ) - case UnreadableBundle(reason=reason): - return f"E2E_FIXTURE_MODE=replay cannot use bundle at {bundle_dir}: {reason}" - case _: - assert_never(freshness) - case _: - assert_never(mode) - - -def fixture_report_lines(mode_raw: str, bundle_dir: Path, *, now: datetime) -> list[str]: - """pytest report-header lines; empty in live mode so an unset - E2E_FIXTURE_MODE keeps today's output byte-identical.""" - mode = parse_fixture_mode(mode_raw) - match mode: - case InvalidFixtureMode() | "live": - return [] - case "record": - return [f"e2e fixture mode: record -> {bundle_dir}"] - case "replay": - freshness = check_freshness(bundle_dir, now=now) - match freshness: - case FreshBundle(manifest=manifest): - return [ - f"e2e fixture mode: replay <- {bundle_dir} " - f"(recorded {manifest.recorded_at.isoformat()}, harness {manifest.harness_version})" - ] - case StaleBundle() | UnreadableBundle(): - return [f"e2e fixture mode: replay <- {bundle_dir}"] - case _: - assert_never(freshness) - case _: - assert_never(mode) diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py new file mode 100644 index 00000000000..ab0791e6b74 --- /dev/null +++ b/tests/e2e/provider_edge.py @@ -0,0 +1,546 @@ +"""Provider-edge record/replay server for e2e runs (LIT-5745). + +Record and replay scope to provider-bound traffic only: the proxy boots for +real, tests hit it for real, and only the hop from the proxy to the provider +is recorded or served from a bundle. Suites opt in per deployment by pointing +``litellm_params.api_base`` at ``provider_edge_api_base(mount)``, which is an +in-process HTTP server mounting each supported provider under a path prefix +(``http://127.0.0.1:/openai`` forwards to ``https://api.openai.com``). +In record mode the edge relays each request verbatim, stores the interaction, +and serves the proxy the same filtered response replay will serve later; in +replay mode it serves straight from the bundle and never opens a provider +connection, so a green replay run with a fake provider key proves the entire +proxy pipeline (auth, routing, spend logging) without provider spend. + +Request identity reuses fixture_canonical.py: interactions match by canonical +content key, order-independent across keys and FIFO within one. Edge requests +store no headers at all: SDK telemetry headers vary run to run and credential +headers must never touch disk. An unmatched replay call returns HTTP +``REPLAY_MISS_STATUS`` naming the closest recorded interaction, which the +proxy relays as a provider error the failing test surfaces. + +v1 limits: only the mounts in ``EDGE_MOUNTS`` (SigV4 providers like Bedrock +sign the Host header, so a forwarding edge breaks their signatures), JSON and +opaque single-part bodies (multipart boundaries are random per request), +streaming fidelity is LIT-5742, and CI wiring is LIT-5748. Suites that do not +wire the edge keep hitting providers live in every mode. +""" + +from __future__ import annotations + +import base64 +import difflib +import functools +import hashlib +import threading +from collections import deque +from collections.abc import Mapping +from dataclasses import dataclass, field +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from itertools import islice +from pathlib import Path +from types import MappingProxyType +from typing import Final, Literal, assert_never +from urllib.parse import parse_qsl, urlsplit + +from pydantic import JsonValue, TypeAdapter + +from e2e_http import NetworkError, RawResponse, forward +from fixture_bundle import ( + BundleRecorder, + Interaction, + LoadedBundle, + RecordedHttpResponse, + RecordedRequest, + UnreadableBundle, + UnsafeBundleDir, + interaction_filename, + load_bundle, + prepare_bundle, + slug_for_test, +) +from fixture_canonical import CanonicalRequest, canonical_string, canonicalize +from fixture_mode import ( + FIXTURE_MODES, + InvalidFixtureMode, + ReplayMiss, + current_test_key, + parse_fixture_mode, +) + +EDGE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( + { + "openai": "https://api.openai.com", + "anthropic": "https://api.anthropic.com", + } +) + +REPLAY_MISS_STATUS: Final = 599 + +_HOP_BY_HOP_HEADERS: Final[frozenset[str]] = frozenset( + { + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailers", + "transfer-encoding", + "upgrade", + } +) +_REQUEST_DROPPED_HEADERS: Final[frozenset[str]] = _HOP_BY_HOP_HEADERS | { + "host", + "content-length", + "accept-encoding", +} +_RESPONSE_DROPPED_HEADERS: Final[frozenset[str]] = _HOP_BY_HOP_HEADERS | { + "content-encoding", + "content-length", + "set-cookie", +} + +_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) + + +def _edge_request(method: str, path: str, query: str, body: bytes | None) -> RecordedRequest: + """The identity replay matches on: the edge path (mount included), the query + as params, and the body as parsed JSON, or as a canonicalized content digest + when it is not JSON so opaque uploads still match across runs.""" + params: Final = dict(parse_qsl(query, keep_blank_values=True)) + if not body: + return RecordedRequest(method=method.lower(), path=path, headers={}, params=params) + decoded: Final = body.decode("utf-8", errors="replace") + try: + parsed: Final[JsonValue] = _JSON.validate_json(decoded) + except ValueError: + return RecordedRequest( + method=method.lower(), + path=path, + headers={}, + params=params, + file_sha256=hashlib.sha256(canonical_string(decoded).encode()).hexdigest(), + file_bytes=len(body), + ) + return RecordedRequest(method=method.lower(), path=path, headers={}, params=params, body=parsed) + + +def _build_pool(recorded: tuple[Interaction, ...]) -> dict[str, deque[Interaction]]: + keys: Final = tuple(canonicalize(interaction.request).key for interaction in recorded) + return { + key: deque( + interaction + for candidate_key, interaction in zip(keys, recorded, strict=True) + if candidate_key == key + ) + for key in dict.fromkeys(keys) + } + + +def _closest_recorded( + canonical: CanonicalRequest, recorded: tuple[Interaction, ...] +) -> tuple[CanonicalRequest, str]: + candidates: Final = tuple(canonicalize(interaction.request) for interaction in recorded) + ratios: Final = tuple( + difflib.SequenceMatcher( + None, f"{canonical.method} {canonical.path}\n{canonical.content}", + f"{candidate.method} {candidate.path}\n{candidate.content}", + ).ratio() + for candidate in candidates + ) + best: Final = max(range(len(candidates)), key=lambda index: ratios[index]) + return candidates[best], interaction_filename(best, recorded[best].request) + + +def _miss_message(test_key: str, slug: str, canonical: CanonicalRequest, bundle: LoadedBundle) -> str: + recorded: Final = bundle.interactions.get(slug, ()) + if not recorded: + return ( + f"replay miss for {test_key}: computed key {canonical.key} but nothing is recorded " + f"under {slug}; re-record with E2E_FIXTURE_MODE=record" + ) + closest, closest_file = _closest_recorded(canonical, recorded) + diff: Final = "\n".join( + islice( + difflib.unified_diff( + closest.pretty_content().splitlines(), + canonical.pretty_content().splitlines(), + fromfile=f"closest recorded ({closest_file})", + tofile="test made", + lineterm="", + ), + 60, + ) + ) + return ( + f"replay miss for {test_key}: no recorded interaction matches key {canonical.key}; " + f"closest recorded key is {closest.key} ({closest_file})\n{diff}\n" + "re-record with E2E_FIXTURE_MODE=record" + ) + + +@dataclass(slots=True) +class ReplaySource: + """One shared pool per test over a loaded bundle, so every provider call the + proxy makes in the session consumes from the same recorded interactions. + Every pool is built once at construction and per-key consumption is a single + atomic deque pop, so concurrent replay calls never race. Calls match by + canonical content key: order-independent across distinct keys (concurrent + tests interleave calls nondeterministically), FIFO within one key (a retry + or poll loop replays its recorded responses in recorded order).""" + + bundle: LoadedBundle + _pools: dict[str, dict[str, deque[Interaction]]] = field(init=False) + + def __post_init__(self) -> None: + self._pools = { + slug: _build_pool(recorded) for slug, recorded in self.bundle.interactions.items() + } + + def _pool(self, slug: str) -> dict[str, deque[Interaction]]: + return self._pools.get(slug, {}) + + def next_interaction(self, request: RecordedRequest) -> Interaction: + test_key: Final = current_test_key() + slug: Final = slug_for_test(test_key) + pool: Final = self._pool(slug) + canonical: Final = canonicalize(request) + queue: Final = pool.get(canonical.key) + if queue is None: + raise ReplayMiss(_miss_message(test_key, slug, canonical, self.bundle)) + try: + return queue.popleft() + except IndexError: + raise ReplayMiss( + f"replay exhausted for {test_key}: every recorded interaction for key " + f"{canonical.key} is already consumed; re-record with E2E_FIXTURE_MODE=record" + ) from None + + def leftover_error(self, test_key: str) -> str | None: + """Non-None when the test consumed fewer interactions than were recorded, + meaning a passing replay proved less than the bundle claims.""" + slug: Final = slug_for_test(test_key) + recorded: Final = self.bundle.interactions.get(slug, ()) + if not recorded: + return None + leftover: Final = tuple( + interaction for queue in self._pool(slug).values() for interaction in queue + ) + if not leftover: + return None + return ( + f"replay incomplete for {test_key}: {len(leftover)} of {len(recorded)} recorded " + f"interactions never consumed, e.g. {canonicalize(leftover[0].request).key}; " + "re-record with E2E_FIXTURE_MODE=record" + ) + + +@dataclass(frozen=True, slots=True) +class RecordEdge: + """Record backend: forward to the provider, persist, serve the filtered copy. + The lock serializes recorder writes because the edge server handles requests + on concurrent threads.""" + + recorder: BundleRecorder + lock: threading.Lock + + +@dataclass(frozen=True, slots=True) +class ReplayEdge: + source: ReplaySource + + +type EdgeBackend = RecordEdge | ReplayEdge + + +@dataclass(frozen=True, slots=True) +class EdgeReply: + status_code: int + headers: dict[str, str] + body: bytes + + +def _text_reply(status_code: int, message: str) -> EdgeReply: + return EdgeReply( + status_code=status_code, + headers={"content-type": "text/plain; charset=utf-8"}, + body=message.encode(), + ) + + +def _reply_from_recorded(response: RecordedHttpResponse) -> EdgeReply: + return EdgeReply( + status_code=response.status_code, + headers=dict(response.headers), + body=base64.b64decode(response.body_b64), + ) + + +def _recorded_response(outcome: RawResponse | NetworkError) -> RecordedHttpResponse: + match outcome: + case RawResponse(status_code=status_code, headers=headers, body=body): + return RecordedHttpResponse( + status_code=status_code, + headers={ + name: value + for name, value in headers.items() + if name not in _RESPONSE_DROPPED_HEADERS + }, + body_b64=base64.b64encode(body).decode("ascii"), + ) + case NetworkError(message=message): + return RecordedHttpResponse( + status_code=502, + headers={"content-type": "text/plain; charset=utf-8"}, + body_b64=base64.b64encode( + f"provider edge could not reach the provider: {message}".encode() + ).decode("ascii"), + ) + + +def _upstream_url(upstream_base: str, upstream_path: str, query: str) -> str: + url: Final = f"{upstream_base}/{upstream_path}" + return f"{url}?{query}" if query else url + + +def _handle_record( + backend: RecordEdge, + request: RecordedRequest, + *, + method: str, + url: str, + headers: Mapping[str, str], + body: bytes | None, + timeout: float, +) -> EdgeReply: + forwarded: Final = { + name: value for name, value in headers.items() if name.lower() not in _REQUEST_DROPPED_HEADERS + } + outcome: Final = forward(method, url, headers=forwarded, body=body, timeout=timeout) + response: Final = _recorded_response(outcome) + with backend.lock: + backend.recorder.record(test_key=current_test_key(), request=request, response=response) + return _reply_from_recorded(response) + + +def _handle_replay(source: ReplaySource, request: RecordedRequest) -> EdgeReply: + try: + interaction: Final = source.next_interaction(request) + except ReplayMiss as miss: + return _text_reply(REPLAY_MISS_STATUS, str(miss)) + return _reply_from_recorded(interaction.response) + + +def handle_edge_request( + backend: EdgeBackend, + mounts: Mapping[str, str], + method: str, + raw_path: str, + headers: Mapping[str, str], + body: bytes | None, + *, + timeout: float, +) -> EdgeReply: + """The edge's pure core, one HTTP exchange in and out: resolve the mount + prefix, then record (forward + persist) or replay (serve from the bundle). + Socket-free so unit tests exercise every branch without a server.""" + split: Final = urlsplit(raw_path) + mount, _, upstream_path = split.path.lstrip("/").partition("/") + upstream_base: Final = mounts.get(mount) + if upstream_base is None: + return _text_reply( + 404, f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(mounts))}" + ) + request: Final = _edge_request(method, split.path, split.query, body) + match backend: + case RecordEdge(): + return _handle_record( + backend, + request, + method=method, + url=_upstream_url(upstream_base, upstream_path, split.query), + headers=headers, + body=body, + timeout=timeout, + ) + case ReplayEdge(source=source): + return _handle_replay(source, request) + case _: + assert_never(backend) + + +class _EdgeHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_GET(self) -> None: + self._handle() + + def do_POST(self) -> None: + self._handle() + + def do_PUT(self) -> None: + self._handle() + + def do_PATCH(self) -> None: + self._handle() + + def do_DELETE(self) -> None: + self._handle() + + def _handle(self) -> None: + edge_server: Final = self.server + assert isinstance(edge_server, _EdgeHTTPServer) + length: Final = int(self.headers.get("content-length") or "0") + body: Final = self.rfile.read(length) if length else None + reply: Final = handle_edge_request( + edge_server.backend, + edge_server.mounts, + self.command, + self.path, + {name.lower(): value for name, value in self.headers.items()}, + body, + timeout=edge_server.forward_timeout, + ) + self.send_response(reply.status_code) + for name, value in reply.headers.items(): + self.send_header(name, value) + self.send_header("content-length", str(len(reply.body))) + self.end_headers() + self.wfile.write(reply.body) + + def log_message(self, format: str, *args: object) -> None: + """Silence the per-request stderr line BaseHTTPRequestHandler emits.""" + + +class _EdgeHTTPServer(ThreadingHTTPServer): + daemon_threads = True + + def __init__( + self, + bind: tuple[str, int], + *, + backend: EdgeBackend, + mounts: Mapping[str, str], + forward_timeout: float, + ) -> None: + super().__init__(bind, _EdgeHandler) + self.backend: Final = backend + self.mounts: Final = mounts + self.forward_timeout: Final = forward_timeout + + +@dataclass(frozen=True, slots=True) +class ProviderEdge: + port: int + advertise_host: str + + def api_base(self, mount: str) -> str: + return f"http://{self.advertise_host}:{self.port}/{mount}" + + +@dataclass(frozen=True, slots=True) +class RunningEdge: + edge: ProviderEdge + server: _EdgeHTTPServer + + def shutdown(self) -> None: + self.server.shutdown() + self.server.server_close() + + +def start_provider_edge( + backend: EdgeBackend, + *, + mounts: Mapping[str, str] = EDGE_MOUNTS, + bind_host: str = "127.0.0.1", + advertise_host: str | None = None, + forward_timeout: float = 60.0, +) -> RunningEdge: + """Boot an edge server on an OS-assigned port in a daemon thread. + ``advertise_host`` is what api_base URLs name (it differs from the bind + host when the proxy runs in a container and reaches the host machine via + a gateway address like host.docker.internal).""" + server: Final = _EdgeHTTPServer( + (bind_host, 0), backend=backend, mounts=mounts, forward_timeout=forward_timeout + ) + thread: Final = threading.Thread(target=server.serve_forever, name="e2e-provider-edge", daemon=True) + thread.start() + return RunningEdge( + edge=ProviderEdge(port=server.server_address[1], advertise_host=advertise_host or bind_host), + server=server, + ) + + +@functools.lru_cache(maxsize=8) +def _shared_recorder(root: Path) -> BundleRecorder: + prepared = prepare_bundle(root) + if isinstance(prepared, UnsafeBundleDir): + raise ValueError(f"E2E_FIXTURE_DIR {prepared.path} {prepared.reason}") + return prepared + + +@functools.lru_cache(maxsize=8) +def _shared_replay_source(root: Path) -> ReplaySource: + loaded = load_bundle(root) + if isinstance(loaded, UnreadableBundle): + raise ValueError(f"cannot replay from {root}: {loaded.reason}") + return ReplaySource(bundle=loaded) + + +@functools.lru_cache(maxsize=8) +def _shared_edge( + mode: Literal["record", "replay"], + bundle_dir: Path, + bind_host: str, + advertise_host: str, + forward_timeout: float, +) -> ProviderEdge: + backend: Final[EdgeBackend] = ( + RecordEdge(recorder=_shared_recorder(bundle_dir), lock=threading.Lock()) + if mode == "record" + else ReplayEdge(source=_shared_replay_source(bundle_dir)) + ) + return start_provider_edge( + backend, + mounts=EDGE_MOUNTS, + bind_host=bind_host, + advertise_host=advertise_host, + forward_timeout=forward_timeout, + ).edge + + +def replay_leftover_error(*, mode_raw: str, bundle_dir: Path, test_key: str) -> str | None: + """Teardown-time completeness check: in replay mode a passed test with + unconsumed recorded interactions must fail instead of passing against a + recording it no longer matches. Inert in every other mode.""" + if parse_fixture_mode(mode_raw) != "replay": + return None + return _shared_replay_source(bundle_dir).leftover_error(test_key) + + +def provider_edge_api_base( + mount: str, + *, + mode_raw: str, + bundle_dir: Path, + bind_host: str, + advertise_host: str, + forward_timeout: float = 60.0, +) -> str | None: + """The api_base a suite gives an edge-wired deployment: None in live mode + (the deployment keeps its real provider api_base) and the process-wide edge + server's mount URL in record and replay, booting the server on first use.""" + mode: Final = parse_fixture_mode(mode_raw) + match mode: + case InvalidFixtureMode(value=value): + raise ValueError(f"E2E_FIXTURE_MODE={value!r} is not one of {', '.join(FIXTURE_MODES)}") + case "live": + return None + case "record" | "replay": + if mount not in EDGE_MOUNTS: + raise ValueError( + f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(EDGE_MOUNTS))}" + ) + return _shared_edge(mode, bundle_dir, bind_host, advertise_host, forward_timeout).api_base(mount) + case _: + assert_never(mode) diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 3cae337a5ff..6cdd3354bf7 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -65,8 +65,6 @@ from models import ( ) from e2e_config import ( CONTROL_PLANE_BASE_URL, - FIXTURE_DIR, - FIXTURE_MODE_RAW, MASTER_KEY, POLL_INTERVAL, POLL_TIMEOUT, @@ -74,7 +72,6 @@ from e2e_config import ( REQUEST_TIMEOUT, settle_propagation, ) -from fixture_transport import select_transport from transport import HttpTransport, SplitTransport, Transport RowsPredicate = Callable[[list[SpendLogRow]], bool] @@ -547,9 +544,9 @@ def build_proxy_client( pass all three together, since a caller that overrides only the data plane would leave management calls pointed at the env default. - E2E_FIXTURE_MODE wraps (record) or replaces (replay) the transport here, so - every client built from this seam records or replays without changing shape; - unset it stays the plain SplitTransport (see fixture_transport.py).""" + Test-to-proxy traffic always goes over the wire, in every E2E_FIXTURE_MODE: + record and replay scope to the proxy's provider-bound calls via the + provider edge (see provider_edge.py), never to this transport.""" split = SplitTransport( data=HttpTransport( base_url=base_url, @@ -563,12 +560,7 @@ def build_proxy_client( ), ) return ProxyClient( - transport=select_transport( - split, - mode_raw=FIXTURE_MODE_RAW, - bundle_dir=FIXTURE_DIR, - master_key=master_key, - ), + transport=split, poll_timeout=POLL_TIMEOUT, poll_interval=POLL_INTERVAL, ) diff --git a/tests/e2e/quota_management/spend_tracking/test_provider_edge_spend_e2e.py b/tests/e2e/quota_management/spend_tracking/test_provider_edge_spend_e2e.py new file mode 100644 index 00000000000..ced7c819d42 --- /dev/null +++ b/tests/e2e/quota_management/spend_tracking/test_provider_edge_spend_e2e.py @@ -0,0 +1,50 @@ +"""The provider-edge demonstrator: one spend-tracking flow wired through the +record/replay edge (LIT-5745). + +This is the reference for wiring a suite to the edge: register a deployment +whose ``api_base`` comes from ``e2e_config.provider_edge_base``, then exercise +the proxy exactly as a live test would. In live mode the base is None and the +deployment talks to the real provider; in record mode it talks through the +local edge, which forwards to the provider and captures the exchange; in +replay mode the same test drives the REAL proxy and REAL database on the +recorded provider traffic alone, so key auth, routing, and the spend-log +write path are all still under test with zero provider calls. +""" + +import pytest + +from e2e_config import CHEAP_OPENAI_MODEL, provider_edge_base +from lifecycle import ResourceManager +from models import LiteLLMParamsBody +from spend_e2e_client import SpendClient, unique_marker, unwrap + +pytestmark = pytest.mark.e2e + + +@pytest.mark.covers("quota_management.spend_tracking.chat_completions.logs_cost") +def test_edge_wired_chat_writes_nonzero_spend_row( + client: SpendClient, resources: ResourceManager, scoped_key: str +) -> None: + base = provider_edge_base("openai") + model = f"e2e-edge-openai-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody( + model=f"openai/{CHEAP_OPENAI_MODEL}", + api_key="os.environ/OPENAI_API_KEY", + api_base=None if base is None else f"{base}/v1", + ), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + chat = unwrap( + client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16) + ) + assert chat.id + + rows = client.poll_logs_for_key( + scoped_key, predicate=lambda rs: any((r.spend or 0) > 0 for r in rs) + ) + matching = [row for row in rows if row.request_id == chat.id] + assert matching, f"no SpendLogs row for request_id {chat.id}; saw {len(rows)} row(s)" + assert (matching[0].spend or 0) > 0, f"spend row for {chat.id} has zero spend" diff --git a/tests/e2e/test_fixture_bundle.py b/tests/e2e/test_fixture_bundle.py index fd4cca6451f..b49ab565e39 100644 --- a/tests/e2e/test_fixture_bundle.py +++ b/tests/e2e/test_fixture_bundle.py @@ -1,9 +1,9 @@ -"""Harness coverage for the on-disk fixture bundle format (LIT-5729). +"""Harness coverage for the on-disk fixture bundle format (LIT-5729/LIT-5745). No proxy and no ``e2e`` marker: these pin the bundle CONTRACT - the seven-day freshness gate that names the bundle's age, record mode's wipe safety (never delete a directory that is not a bundle), collision-free per-test slugs, and -lossless Result round-trips - so replay can never silently drift from what +grouped-in-order loading - so replay can never silently drift from what record wrote. """ @@ -12,18 +12,6 @@ from __future__ import annotations from datetime import datetime, timedelta, timezone from pathlib import Path -import pytest -from pydantic import BaseModel - -from e2e_http import ( - NetworkError, - RateLimitedError, - Result, - Success, - UnauthorizedError, - UnknownApiError, - ValidationError, -) from fixture_bundle import ( BUNDLE_FORMAT_VERSION, MANIFEST_FILENAME, @@ -32,28 +20,22 @@ from fixture_bundle import ( FreshBundle, LoadedBundle, Manifest, + RecordedHttpResponse, RecordedRequest, - RecordedResult, StaleBundle, UnreadableBundle, UnsafeBundleDir, check_freshness, format_age, - from_result, interaction_filename, load_bundle, prepare_bundle, slug_for_test, - to_result, ) NOW = datetime(2026, 8, 18, 12, 0, 0, tzinfo=timezone.utc) -class Payload(BaseModel): - value: str - - def write_manifest( root: Path, recorded_at: datetime, *, format_version: int = BUNDLE_FORMAT_VERSION ) -> None: @@ -74,20 +56,8 @@ def plain_request(path: str) -> RecordedRequest: return RecordedRequest(method="post", path=path, headers={}) -class TestResultRoundTrip: - @pytest.mark.parametrize( - "result", - [ - Success(status_code=201, data=Payload(value="ok")), - NetworkError(message="connection refused"), - UnauthorizedError(), - RateLimitedError(retry_after_seconds=7, body="slow down"), - ValidationError(message="bad shape"), - UnknownApiError(status_code=502, body="upstream exploded"), - ], - ) - def test_every_result_kind_survives_disk_and_back(self, result: Result[Payload]) -> None: - assert to_result(from_result(result), Payload) == result +def plain_response() -> RecordedHttpResponse: + return RecordedHttpResponse(status_code=401, headers={}, body_b64="") class TestFreshness: @@ -144,7 +114,7 @@ class TestPrepareBundle: prepared(root).record( test_key="old.py::test_old", request=plain_request("/stale"), - response=RecordedResult(kind="unauthorized"), + response=plain_response(), ) assert any(entry.is_dir() for entry in root.iterdir()) prepared(root) @@ -193,7 +163,7 @@ class TestRecordAndLoad: recorder.record( test_key=key, request=plain_request(path), - response=RecordedResult(kind="unauthorized"), + response=plain_response(), ) loaded = load_bundle(root) assert isinstance(loaded, LoadedBundle) @@ -208,7 +178,7 @@ class TestRecordAndLoad: recorder.record( test_key=key, request=plain_request(f"/{key[-3:]}"), - response=RecordedResult(kind="unauthorized"), + response=plain_response(), ) loaded = load_bundle(root) assert isinstance(loaded, LoadedBundle) diff --git a/tests/e2e/test_fixture_mode.py b/tests/e2e/test_fixture_mode.py new file mode 100644 index 00000000000..109bb9e1b11 --- /dev/null +++ b/tests/e2e/test_fixture_mode.py @@ -0,0 +1,114 @@ +"""Harness coverage for fixture-mode selection and determinism (LIT-5729/LIT-5745). + +No proxy and no ``e2e`` marker. Pins the mode parser, the deterministic +per-test marker sequence a replay run must regenerate, the collection-time +gate (including the stale message that names the bundle's age), and the pytest +report header. The provider-edge record/replay behavior itself is pinned in +test_provider_edge.py. +""" + +from __future__ import annotations + +import hashlib +from datetime import datetime, timedelta, timezone +from pathlib import Path + +import pytest + +from fixture_bundle import BUNDLE_FORMAT_VERSION, MANIFEST_FILENAME, Manifest +from fixture_mode import ( + InvalidFixtureMode, + current_test_key, + deterministic_marker, + fixture_mode_collection_error, + fixture_report_lines, + parse_fixture_mode, +) + +NOW = datetime(2026, 8, 18, 12, 0, 0, tzinfo=timezone.utc) + + +def write_manifest(root: Path, recorded_at: datetime) -> None: + root.mkdir(parents=True, exist_ok=True) + manifest = Manifest( + format_version=BUNDLE_FORMAT_VERSION, recorded_at=recorded_at, harness_version="abc1234" + ) + (root / MANIFEST_FILENAME).write_text(manifest.model_dump_json(), encoding="utf-8") + + +class TestParseFixtureMode: + @pytest.mark.parametrize( + ("raw", "expected"), + [("live", "live"), ("record", "record"), ("replay", "replay"), ("", "live"), (" REPLAY ", "replay")], + ) + def test_known_values_normalize(self, raw: str, expected: str) -> None: + assert parse_fixture_mode(raw) == expected + + def test_unknown_value_is_invalid_with_the_original_spelling(self) -> None: + assert parse_fixture_mode("cached") == InvalidFixtureMode(value="cached") + + +class TestDeterministicMarker: + def test_sequence_is_a_pure_function_of_test_and_ordinal(self) -> None: + """A replay process must regenerate exactly the markers the record + process generated, so the Nth marker of a test is pinned to a pure + function of the node id and N.""" + key = current_test_key() + assert deterministic_marker() == hashlib.sha1(f"{key}#0".encode()).hexdigest()[:12] + assert deterministic_marker() == hashlib.sha1(f"{key}#1".encode()).hexdigest()[:12] + + +class TestCurrentTestKey: + def test_names_this_test_and_strips_the_phase(self) -> None: + key = current_test_key() + assert key.endswith("TestCurrentTestKey::test_names_this_test_and_strips_the_phase") + assert "(call)" not in key + + +class TestCollectionGate: + def test_invalid_mode_names_the_value_and_the_choices(self, tmp_path: Path) -> None: + assert ( + fixture_mode_collection_error("cached", tmp_path, now=NOW) + == "E2E_FIXTURE_MODE='cached' is not one of live, record, replay" + ) + + @pytest.mark.parametrize("mode_raw", ["live", "", "record"]) + def test_live_and_record_never_block_collection(self, mode_raw: str, tmp_path: Path) -> None: + assert fixture_mode_collection_error(mode_raw, tmp_path / "missing", now=NOW) is None + + def test_replay_with_no_bundle_says_how_to_record_one(self, tmp_path: Path) -> None: + reason = fixture_mode_collection_error("replay", tmp_path / "missing", now=NOW) + assert reason is not None + assert f"no {MANIFEST_FILENAME}" in reason + assert "E2E_FIXTURE_MODE=record" in reason + + def test_stale_replay_bundle_fails_naming_its_age(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + write_manifest(root, NOW - timedelta(days=9, hours=5)) + reason = fixture_mode_collection_error("replay", root, now=NOW) + assert reason is not None + assert "age 9d5h exceeds the 7-day limit" in reason + assert "re-record with E2E_FIXTURE_MODE=record" in reason + + def test_fresh_replay_bundle_collects(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + write_manifest(root, NOW - timedelta(days=2)) + assert fixture_mode_collection_error("replay", root, now=NOW) is None + + +class TestReportHeader: + def test_live_mode_prints_nothing(self, tmp_path: Path) -> None: + assert fixture_report_lines("live", tmp_path, now=NOW) == [] + assert fixture_report_lines("", tmp_path, now=NOW) == [] + + def test_record_and_replay_name_the_bundle(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + recorded_at = NOW - timedelta(days=1) + write_manifest(root, recorded_at) + assert fixture_report_lines("record", root, now=NOW) == [ + f"e2e fixture mode: record -> {root}" + ] + replay_lines = fixture_report_lines("replay", root, now=NOW) + assert len(replay_lines) == 1 + assert "replay" in replay_lines[0] + assert recorded_at.isoformat() in replay_lines[0] diff --git a/tests/e2e/test_fixture_transport.py b/tests/e2e/test_fixture_transport.py deleted file mode 100644 index e61088d841c..00000000000 --- a/tests/e2e/test_fixture_transport.py +++ /dev/null @@ -1,676 +0,0 @@ -"""Harness coverage for the record/replay transports (LIT-5729). - -No proxy and no ``e2e`` marker. A fake in-memory ``Transport`` stands in for -the live one (dependency injection, no monkeypatching): recording must pass -every value through unchanged while writing one redacted interaction file per -call, and replay must serve identical values from the bundle alone - the -fake's call log proves nothing reaches the inner transport - failing hard -(``ReplayMiss``) on any content drift, printing the computed canonical key and -the closest recorded key (LIT-5741; the pure canonicalizer is pinned in -test_fixture_canonical.py). The collection-time gate and report header are -pinned here too, including the stale message that names the bundle's age. -""" - -from __future__ import annotations - -import hashlib -import sys -import threading -from concurrent.futures import ThreadPoolExecutor -from dataclasses import dataclass, field -from datetime import datetime, timedelta, timezone -from pathlib import Path -from uuid import uuid4 - -import pytest -from pydantic import BaseModel - -from e2e_http import ( - AuthHeaders, - BinaryStream, - ProbeResult, - Result, - StreamingResponse, - Success, -) -from fixture_bundle import ( - BUNDLE_FORMAT_VERSION, - MANIFEST_FILENAME, - BundleRecorder, - Interaction, - LoadedBundle, - Manifest, - RecordedResult, - load_bundle, - prepare_bundle, - slug_for_test, -) -from fixture_canonical import canonicalize -from fixture_transport import ( - InvalidFixtureMode, - RecordingTransport, - ReplayMiss, - ReplaySource, - ReplayTransport, - current_test_key, - deterministic_marker, - fixture_mode_collection_error, - fixture_report_lines, - parse_fixture_mode, - recorded_request, - replay_leftover_error, - select_transport, -) -from transport import Transport - -NOW = datetime(2026, 8, 18, 12, 0, 0, tzinfo=timezone.utc) - - -class Payload(BaseModel): - value: str - - -class Body(BaseModel): - prompt: str - - -class Query(BaseModel): - q: str - - -class DeployParams(BaseModel): - model: str - api_key: str | None = None - aws_secret_access_key: str | None = None - - -class DeployBody(BaseModel): - model_name: str - litellm_params: DeployParams - - -STREAMING = StreamingResponse( - status_code=200, - body="", - content_type="text/event-stream", - chunks=2, - stream_events=["one", "two"], - stream_done=True, -) -BINARY = BinaryStream(status_code=200, content_type="audio/mpeg", chunk_count=3, total_bytes=42) -PROBE = ProbeResult(status_code=200, body="alive") - - -@dataclass -class FakeTransport: - calls: list[str] = field(default_factory=list) - - def bearer(self, key: str) -> AuthHeaders: - return AuthHeaders(authorization=f"Bearer {key}") - - @property - def master(self) -> AuthHeaders: - return self.bearer("sk-fake-master") - - def _success[R: BaseModel](self, response_type: type[R]) -> Result[R]: - return Success(status_code=200, data=response_type.model_validate({"value": "live"})) - - def post[R: BaseModel]( - self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] - ) -> Result[R]: - self.calls.append(f"post {path}") - return self._success(response_type) - - def get[R: BaseModel]( - self, - path: str, - *, - headers: BaseModel, - params: BaseModel, - response_type: type[R], - timeout: float | None = None, - ) -> Result[R]: - self.calls.append(f"get {path}") - return self._success(response_type) - - def delete[R: BaseModel]( - self, - path: str, - *, - headers: BaseModel, - json: BaseModel, - response_type: type[R], - params: BaseModel | None = None, - ) -> Result[R]: - self.calls.append(f"delete {path}") - return self._success(response_type) - - def patch[R: BaseModel]( - self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] - ) -> Result[R]: - self.calls.append(f"patch {path}") - return self._success(response_type) - - def put[R: BaseModel]( - self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] - ) -> Result[R]: - self.calls.append(f"put {path}") - return self._success(response_type) - - def stream(self, path: str, *, headers: BaseModel, json: BaseModel) -> StreamingResponse: - self.calls.append(f"stream {path}") - return STREAMING - - def stream_binary( - self, path: str, *, headers: BaseModel, json: BaseModel, chunk_size: int = 8192 - ) -> BinaryStream: - self.calls.append(f"stream_binary {path}") - return BINARY - - def send( - self, - path: str, - *, - headers: BaseModel, - json: BaseModel, - params: BaseModel | None = None, - stream: bool = False, - ) -> StreamingResponse: - self.calls.append(f"send {path}") - return STREAMING - - def probe(self, path: str, *, params: BaseModel) -> ProbeResult: - self.calls.append(f"probe {path}") - return PROBE - - def upload[R: BaseModel]( - self, - path: str, - *, - headers: BaseModel, - form: BaseModel, - filename: str, - content: bytes, - file_content_type: str = "application/jsonl", - file_field: str = "file", - params: BaseModel | None = None, - response_type: type[R], - ) -> Result[R]: - self.calls.append(f"upload {path}") - return self._success(response_type) - - def download(self, path: str, *, headers: BaseModel) -> StreamingResponse: - self.calls.append(f"download {path}") - return STREAMING - - -def make_recorder(root: Path) -> BundleRecorder: - recorder = prepare_bundle(root) - assert isinstance(recorder, BundleRecorder) - return recorder - - -def replay_source(root: Path) -> ReplaySource: - loaded = load_bundle(root) - assert isinstance(loaded, LoadedBundle) - return ReplaySource(bundle=loaded) - - -def this_tests_files(root: Path) -> list[Path]: - slug_dir = root / slug_for_test(current_test_key()) - return sorted(slug_dir.glob("*.json")) if slug_dir.is_dir() else [] - - -def write_manifest(root: Path, recorded_at: datetime) -> None: - root.mkdir(parents=True, exist_ok=True) - manifest = Manifest( - format_version=BUNDLE_FORMAT_VERSION, recorded_at=recorded_at, harness_version="abc1234" - ) - (root / MANIFEST_FILENAME).write_text(manifest.model_dump_json(), encoding="utf-8") - - -class TestParseFixtureMode: - @pytest.mark.parametrize( - ("raw", "expected"), - [("live", "live"), ("record", "record"), ("replay", "replay"), ("", "live"), (" REPLAY ", "replay")], - ) - def test_known_values_normalize(self, raw: str, expected: str) -> None: - assert parse_fixture_mode(raw) == expected - - def test_unknown_value_is_invalid_with_the_original_spelling(self) -> None: - assert parse_fixture_mode("cached") == InvalidFixtureMode(value="cached") - - -class TestDeterministicMarker: - def test_sequence_is_a_pure_function_of_test_and_ordinal(self) -> None: - """A replay process must regenerate exactly the markers the record - process generated, so the Nth marker of a test is pinned to a pure - function of the node id and N.""" - key = current_test_key() - assert deterministic_marker() == hashlib.sha1(f"{key}#0".encode()).hexdigest()[:12] - assert deterministic_marker() == hashlib.sha1(f"{key}#1".encode()).hexdigest()[:12] - - -class TestCurrentTestKey: - def test_names_this_test_and_strips_the_phase(self) -> None: - key = current_test_key() - assert key.endswith("TestCurrentTestKey::test_names_this_test_and_strips_the_phase") - assert "(call)" not in key - - -class TestRecordingTransport: - def test_passes_the_result_through_and_writes_one_file_per_call(self, tmp_path: Path) -> None: - fake = FakeTransport() - root = tmp_path / "bundle" - recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) - result = recording.post( - "/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload - ) - assert result == Success(status_code=200, data=Payload(value="live")) - assert fake.calls == ["post /model/new"] - files = this_tests_files(root) - assert [file.name for file in files] == ["0000-post-model-new.json"] - interaction = Interaction.model_validate_json(files[0].read_text(encoding="utf-8")) - assert interaction.request.method == "post" - assert interaction.request.path == "/model/new" - - def test_redacts_auth_header_values_in_the_recorded_request(self, tmp_path: Path) -> None: - fake = FakeTransport() - root = tmp_path / "bundle" - recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) - headers = AuthHeaders.model_validate( - {"authorization": "Bearer sk-secret", "x-litellm-api-key": "sk-other"} - ) - recording.post("/key/generate", headers=headers, json=Body(prompt="x"), response_type=Payload) - interaction = Interaction.model_validate_json( - this_tests_files(root)[0].read_text(encoding="utf-8") - ) - assert interaction.request.headers == { - "authorization": "", - "x-litellm-api-key": "", - } - assert "sk-secret" not in this_tests_files(root)[0].read_text(encoding="utf-8") - - def test_redacts_credential_body_fields_in_the_recorded_request(self, tmp_path: Path) -> None: - fake = FakeTransport() - root = tmp_path / "bundle" - recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) - recording.post( - "/model/new", - headers=fake.master, - json=DeployBody( - model_name="m", - litellm_params=DeployParams(model="openai/gpt", api_key="sk-live-provider-secret-123456"), - ), - response_type=Payload, - ) - raw = this_tests_files(root)[0].read_text(encoding="utf-8") - interaction = Interaction.model_validate_json(raw) - assert "sk-live-provider-secret-123456" not in raw - assert isinstance(interaction.request.body, dict) - params = interaction.request.body["litellm_params"] - assert isinstance(params, dict) - assert params["api_key"] == "" - assert params["aws_secret_access_key"] is None - - def test_upload_records_a_content_digest_not_the_bytes(self, tmp_path: Path) -> None: - fake = FakeTransport() - root = tmp_path / "bundle" - recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) - recording.upload( - "/v1/files", - headers=fake.master, - form=Query(q="batch"), - filename="batch.jsonl", - content=b'{"custom_id": "1"}', - response_type=Payload, - ) - interaction = Interaction.model_validate_json( - this_tests_files(root)[0].read_text(encoding="utf-8") - ) - assert interaction.request.file_name == "batch.jsonl" - assert interaction.request.file_bytes == len(b'{"custom_id": "1"}') - assert interaction.request.file_sha256 is not None - assert "custom_id" not in interaction.request.model_dump_json() - - -class TestReplayTransport: - def test_serves_recorded_values_without_touching_the_inner_transport( - self, tmp_path: Path - ) -> None: - fake = FakeTransport() - root = tmp_path / "bundle" - recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) - recorded_post = recording.post( - "/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload - ) - recorded_get = recording.get( - "/v1/models", headers=fake.master, params=Query(q="all"), response_type=Payload - ) - recorded_stream = recording.stream( - "/chat/completions", headers=fake.master, json=Body(prompt="hi") - ) - recorded_probe = recording.probe("/health/liveliness", params=Query(q="1")) - recorded_binary = recording.stream_binary( - "/v1/audio/speech", headers=fake.master, json=Body(prompt="say") - ) - calls_after_record = list(fake.calls) - - replay: Transport = ReplayTransport(source=replay_source(root), master_key="sk-1234") - assert ( - replay.post("/model/new", headers=replay.master, json=Body(prompt="x"), response_type=Payload) - == recorded_post - ) - assert ( - replay.get("/v1/models", headers=replay.master, params=Query(q="all"), response_type=Payload) - == recorded_get - ) - assert ( - replay.stream("/chat/completions", headers=replay.master, json=Body(prompt="hi")) - == recorded_stream - ) - assert replay.probe("/health/liveliness", params=Query(q="1")) == recorded_probe - assert ( - replay.stream_binary("/v1/audio/speech", headers=replay.master, json=Body(prompt="say")) - == recorded_binary - ) - assert fake.calls == calls_after_record - - def test_miss_names_the_computed_key_and_the_closest_recorded_key(self, tmp_path: Path) -> None: - fake = FakeTransport() - root = tmp_path / "bundle" - recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) - recording.post("/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload) - replay: Transport = ReplayTransport(source=replay_source(root), master_key="sk-1234") - with pytest.raises(ReplayMiss) as excinfo: - replay.get("/v1/models", headers=replay.master, params=Query(q="all"), response_type=Payload) - message = str(excinfo.value) - assert "no recorded interaction matches key get /v1/models #" in message - assert "closest recorded key is post /model/new #" in message - assert "0000-post-model-new.json" in message - assert "re-record with E2E_FIXTURE_MODE=record" in message - - def test_content_drift_on_the_same_route_misses_with_no_live_call(self, tmp_path: Path) -> None: - """The naive verb+path match replayed a stale response for a request - whose content had changed, silently passing; a content key must miss, - print both canonical forms' diff, and never reach the inner transport.""" - fake = FakeTransport() - root = tmp_path / "bundle" - recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) - recording.post("/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload) - calls_after_record = list(fake.calls) - replay: Transport = ReplayTransport(source=replay_source(root), master_key="sk-1234") - with pytest.raises(ReplayMiss) as excinfo: - replay.post("/model/new", headers=replay.master, json=Body(prompt="y"), response_type=Payload) - message = str(excinfo.value) - assert "no recorded interaction matches key post /model/new #" in message - assert "closest recorded key is post /model/new #" in message - assert '- "prompt": "x"' in message - assert '+ "prompt": "y"' in message - assert fake.calls == calls_after_record - - def test_exhausted_key_names_the_key(self, tmp_path: Path) -> None: - fake = FakeTransport() - root = tmp_path / "bundle" - recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) - recording.post("/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload) - replay: Transport = ReplayTransport(source=replay_source(root), master_key="sk-1234") - replay.post("/model/new", headers=replay.master, json=Body(prompt="x"), response_type=Payload) - with pytest.raises( - ReplayMiss, match=r"every recorded interaction for key post /model/new #\w{16} is already consumed" - ): - replay.post("/model/new", headers=replay.master, json=Body(prompt="x"), response_type=Payload) - - def test_replays_out_of_recorded_order_across_distinct_keys(self, tmp_path: Path) -> None: - """Concurrent tests interleave independent calls nondeterministically - (e.g. a burst of parallel chat calls), so replay matches by content, - never by recorded position.""" - fake = FakeTransport() - root = tmp_path / "bundle" - recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) - recording.post("/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload) - recording.post("/key/generate", headers=fake.master, json=Body(prompt="k"), response_type=Payload) - source = replay_source(root) - replay: Transport = ReplayTransport(source=source, master_key="sk-1234") - replay.post("/key/generate", headers=replay.master, json=Body(prompt="k"), response_type=Payload) - replay.post("/model/new", headers=replay.master, json=Body(prompt="x"), response_type=Payload) - assert source.leftover_error(current_test_key()) is None - - def test_identical_requests_replay_their_responses_in_recorded_order(self, tmp_path: Path) -> None: - """A poll loop makes the same request repeatedly and asserts on the - progression, so duplicates under one key stay FIFO.""" - root = tmp_path / "bundle" - recorder = make_recorder(root) - recorder.record( - test_key=current_test_key(), - request=recorded_request( - "get", "/v1/models", headers=AuthHeaders(authorization="Bearer sk-x"), params=Query(q="all") - ), - response=RecordedResult(kind="success", status_code=200, data={"value": "first"}), - ) - recorder.record( - test_key=current_test_key(), - request=recorded_request( - "get", "/v1/models", headers=AuthHeaders(authorization="Bearer sk-x"), params=Query(q="all") - ), - response=RecordedResult(kind="success", status_code=200, data={"value": "second"}), - ) - replay: Transport = ReplayTransport(source=replay_source(root), master_key="sk-1234") - first = replay.get("/v1/models", headers=replay.master, params=Query(q="all"), response_type=Payload) - second = replay.get("/v1/models", headers=replay.master, params=Query(q="all"), response_type=Payload) - assert first == Success(status_code=200, data=Payload(value="first")) - assert second == Success(status_code=200, data=Payload(value="second")) - - def test_concurrent_replays_of_one_key_serve_each_recording_exactly_once(self, tmp_path: Path) -> None: - """A burst of parallel identical calls consumes one shared pool: no - response duplicated, none forgotten, nothing left over at teardown. - The tiny switch interval forces thread preemption inside pool setup - and consumption, so a non-atomic pool build or pop fails this test.""" - root = tmp_path / "bundle" - recorder = make_recorder(root) - for ordinal in range(32): - recorder.record( - test_key=current_test_key(), - request=recorded_request( - "get", "/v1/models", headers=AuthHeaders(authorization="Bearer sk-x"), params=Query(q="all") - ), - response=RecordedResult(kind="success", status_code=200, data={"value": f"v{ordinal:02d}"}), - ) - source = replay_source(root) - replay: Transport = ReplayTransport(source=source, master_key="sk-1234") - barrier = threading.Barrier(8) - - def consume_one() -> str: - result = replay.get( - "/v1/models", headers=replay.master, params=Query(q="all"), response_type=Payload - ) - assert isinstance(result, Success) - return result.data.value - - def consume(_: int) -> tuple[str, ...]: - barrier.wait() - return tuple(consume_one() for _call in range(4)) - - previous_interval = sys.getswitchinterval() - sys.setswitchinterval(1e-6) - try: - with ThreadPoolExecutor(max_workers=8) as executor: - served = sorted(value for values in executor.map(consume, range(8)) for value in values) - finally: - sys.setswitchinterval(previous_interval) - assert served == [f"v{ordinal:02d}" for ordinal in range(32)] - assert source.leftover_error(current_test_key()) is None - - -class TestRecordedKeySets: - def test_two_separate_recordings_of_one_flow_produce_identical_key_sets( - self, tmp_path: Path - ) -> None: - """Everything a run randomizes (markers, virtual keys, dates) must - canonicalize out, so separately recorded runs of the same suite agree - on every match key and a bundle recorded elsewhere replays here.""" - - def record_flow(root: Path, run_date: str) -> list[str]: - fake = FakeTransport() - recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) - marker = deterministic_marker() - recording.post( - "/model/new", - headers=fake.master, - json=DeployBody( - model_name=f"e2e-chat-{marker}", - litellm_params=DeployParams(model="openai/gpt", api_key=f"sk-live-{uuid4().hex}"), - ), - response_type=Payload, - ) - recording.post( - "/chat/completions", - headers=recording.bearer(f"sk-{uuid4().hex}"), - json=Body(prompt=f"Reply with the single word ok. {marker}"), - response_type=Payload, - ) - recording.get( - "/spend/logs", headers=fake.master, params=Query(q=run_date), response_type=Payload - ) - loaded = load_bundle(root) - assert isinstance(loaded, LoadedBundle) - return sorted( - canonicalize(interaction.request).key - for interactions in loaded.interactions.values() - for interaction in interactions - ) - - first_keys = record_flow(tmp_path / "one", "2026-08-18") - second_keys = record_flow(tmp_path / "two", "2026-08-19") - assert first_keys == second_keys - assert len(first_keys) == 3 - - -class TestReplayLeftover: - def test_fully_consumed_recording_leaves_nothing(self, tmp_path: Path) -> None: - fake = FakeTransport() - root = tmp_path / "bundle" - recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) - recording.post("/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload) - source = replay_source(root) - replay: Transport = ReplayTransport(source=source, master_key="sk-1234") - replay.post("/model/new", headers=replay.master, json=Body(prompt="x"), response_type=Payload) - assert source.leftover_error(current_test_key()) is None - - def test_unconsumed_trailing_interactions_name_the_next_call(self, tmp_path: Path) -> None: - fake = FakeTransport() - root = tmp_path / "bundle" - recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) - recording.post("/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload) - recording.probe("/health/liveliness", params=Query(q="1")) - source = replay_source(root) - replay: Transport = ReplayTransport(source=source, master_key="sk-1234") - replay.post("/model/new", headers=replay.master, json=Body(prompt="x"), response_type=Payload) - error = source.leftover_error(current_test_key()) - assert error is not None - assert "1 of 2 recorded interactions never consumed" in error - assert "e.g. probe /health/liveliness #" in error - assert "re-record with E2E_FIXTURE_MODE=record" in error - - def test_test_without_recordings_has_no_leftover(self, tmp_path: Path) -> None: - root = tmp_path / "bundle" - make_recorder(root) - assert replay_source(root).leftover_error("suite.py::test_never_recorded") is None - - def test_inert_outside_replay_mode(self, tmp_path: Path) -> None: - missing = tmp_path / "missing" - assert replay_leftover_error(mode_raw="", bundle_dir=missing, test_key="k") is None - assert replay_leftover_error(mode_raw="record", bundle_dir=missing, test_key="k") is None - - def test_replay_mode_reads_the_shared_bundle(self, tmp_path: Path) -> None: - fake = FakeTransport() - root = tmp_path / "bundle" - recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root)) - recording.post("/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload) - error = replay_leftover_error(mode_raw="replay", bundle_dir=root, test_key=current_test_key()) - assert error is not None - assert "1 of 1 recorded interactions never consumed" in error - - -class TestSelectTransport: - def test_live_returns_the_live_transport_untouched(self, tmp_path: Path) -> None: - fake = FakeTransport() - for mode_raw in ("live", ""): - assert ( - select_transport(fake, mode_raw=mode_raw, bundle_dir=tmp_path / "b", master_key="sk") - is fake - ) - - def test_record_wraps_live_and_starts_a_fresh_bundle(self, tmp_path: Path) -> None: - fake = FakeTransport() - root = tmp_path / "bundle" - write_manifest(root, NOW - timedelta(days=30)) - (root / "old-test-slug").mkdir() - (root / "old-test-slug" / "0000-post-old.json").write_text("{}", encoding="utf-8") - selected = select_transport(fake, mode_raw="record", bundle_dir=root, master_key="sk") - assert isinstance(selected, RecordingTransport) - assert selected.inner is fake - assert {entry.name for entry in root.iterdir()} == {MANIFEST_FILENAME} - - def test_replay_builds_a_transport_from_the_bundle_alone(self, tmp_path: Path) -> None: - fake = FakeTransport() - root = tmp_path / "bundle" - make_recorder(root) - selected = select_transport(fake, mode_raw="replay", bundle_dir=root, master_key="sk-master") - assert isinstance(selected, ReplayTransport) - assert selected.master == AuthHeaders(authorization="Bearer sk-master") - - def test_invalid_mode_raises_naming_the_value(self, tmp_path: Path) -> None: - with pytest.raises(ValueError, match="cached"): - select_transport( - FakeTransport(), mode_raw="cached", bundle_dir=tmp_path / "b", master_key="sk" - ) - - -class TestCollectionGate: - def test_invalid_mode_names_the_value_and_the_choices(self, tmp_path: Path) -> None: - assert ( - fixture_mode_collection_error("cached", tmp_path, now=NOW) - == "E2E_FIXTURE_MODE='cached' is not one of live, record, replay" - ) - - @pytest.mark.parametrize("mode_raw", ["live", "", "record"]) - def test_live_and_record_never_block_collection(self, mode_raw: str, tmp_path: Path) -> None: - assert fixture_mode_collection_error(mode_raw, tmp_path / "missing", now=NOW) is None - - def test_replay_with_no_bundle_says_how_to_record_one(self, tmp_path: Path) -> None: - reason = fixture_mode_collection_error("replay", tmp_path / "missing", now=NOW) - assert reason is not None - assert f"no {MANIFEST_FILENAME}" in reason - assert "E2E_FIXTURE_MODE=record" in reason - - def test_stale_replay_bundle_fails_naming_its_age(self, tmp_path: Path) -> None: - root = tmp_path / "bundle" - write_manifest(root, NOW - timedelta(days=9, hours=5)) - reason = fixture_mode_collection_error("replay", root, now=NOW) - assert reason is not None - assert "age 9d5h exceeds the 7-day limit" in reason - assert "re-record with E2E_FIXTURE_MODE=record" in reason - - def test_fresh_replay_bundle_collects(self, tmp_path: Path) -> None: - root = tmp_path / "bundle" - write_manifest(root, NOW - timedelta(days=2)) - assert fixture_mode_collection_error("replay", root, now=NOW) is None - - -class TestReportHeader: - def test_live_mode_prints_nothing(self, tmp_path: Path) -> None: - assert fixture_report_lines("live", tmp_path, now=NOW) == [] - assert fixture_report_lines("", tmp_path, now=NOW) == [] - - def test_record_and_replay_name_the_bundle(self, tmp_path: Path) -> None: - root = tmp_path / "bundle" - recorded_at = NOW - timedelta(days=1) - write_manifest(root, recorded_at) - assert fixture_report_lines("record", root, now=NOW) == [ - f"e2e fixture mode: record -> {root}" - ] - replay_lines = fixture_report_lines("replay", root, now=NOW) - assert len(replay_lines) == 1 - assert "replay" in replay_lines[0] - assert recorded_at.isoformat() in replay_lines[0] diff --git a/tests/e2e/test_provider_edge.py b/tests/e2e/test_provider_edge.py new file mode 100644 index 00000000000..492eee57aaf --- /dev/null +++ b/tests/e2e/test_provider_edge.py @@ -0,0 +1,492 @@ +"""Harness coverage for the provider-edge record/replay server (LIT-5745). + +No proxy and no ``e2e`` marker. A stdlib http.server stands in for the +provider (dependency injection via the mounts mapping, no monkeypatching): +record mode must forward each edge call to it verbatim, persist one +interaction file, and serve the proxy the same filtered response replay will +serve later; replay mode must serve byte-identical responses from the bundle +alone, with the fake provider's hit log proving nothing leaves the process, +and answer any drifted call with HTTP ``REPLAY_MISS_STATUS`` naming the +computed and closest recorded canonical keys (LIT-5741; the pure canonicalizer +is pinned in test_fixture_canonical.py). Requests are made through +``e2e_http.forward`` so the whole HTTP surface of the edge is exercised; the +pure ``handle_edge_request`` core is pinned socket-free alongside. +""" + +from __future__ import annotations + +import base64 +import json +import threading +from collections.abc import Generator, Mapping +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +import pytest +from pydantic import TypeAdapter + +from e2e_http import RawResponse, forward +from fixture_bundle import ( + BundleRecorder, + Interaction, + LoadedBundle, + RecordedHttpResponse, + RecordedRequest, + load_bundle, + prepare_bundle, + slug_for_test, +) +from fixture_mode import current_test_key +from provider_edge import ( + REPLAY_MISS_STATUS, + EdgeBackend, + ProviderEdge, + RecordEdge, + ReplayEdge, + ReplaySource, + handle_edge_request, + provider_edge_api_base, + replay_leftover_error, + start_provider_edge, +) + +CHAT_PATH = "/openai/v1/chat/completions" +REPLAY_MOUNTS = {"openai": "https://replay.invalid"} +JSON_OBJECT = TypeAdapter(dict[str, object]) + + +def json_object(body: bytes) -> dict[str, object]: + return JSON_OBJECT.validate_json(body) + + +class _FakeProvider(ThreadingHTTPServer): + daemon_threads = True + + def __init__(self, bind: tuple[str, int]) -> None: + super().__init__(bind, _FakeProviderHandler) + self.hits: list[str] = [] + + +class _FakeProviderHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_POST(self) -> None: + self._respond() + + def do_GET(self) -> None: + self._respond() + + def _respond(self) -> None: + provider = self.server + assert isinstance(provider, _FakeProvider) + length = int(self.headers.get("content-length") or "0") + body = self.rfile.read(length) if length else b"" + provider.hits.append(f"{self.command} {self.path}") + payload = json.dumps( + {"echo": body.decode("utf-8"), "path": self.path, "hit": len(provider.hits)} + ).encode() + self.send_response(200) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(payload))) + self.send_header("x-upstream", "fake") + self.send_header("set-cookie", "session=fake-cookie") + self.end_headers() + self.wfile.write(payload) + + def log_message(self, format: str, *args: object) -> None: + """Silence the per-request stderr line BaseHTTPRequestHandler emits.""" + + +@contextmanager +def fake_provider() -> Generator[_FakeProvider]: + server = _FakeProvider(("127.0.0.1", 0)) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield server + finally: + server.shutdown() + server.server_close() + + +def provider_url(server: _FakeProvider) -> str: + return f"http://127.0.0.1:{server.server_address[1]}" + + +@contextmanager +def running_edge(backend: EdgeBackend, mounts: Mapping[str, str]) -> Generator[ProviderEdge]: + running = start_provider_edge(backend, mounts=mounts, bind_host="127.0.0.1") + try: + yield running.edge + finally: + running.shutdown() + + +def record_backend(root: Path) -> RecordEdge: + recorder = prepare_bundle(root) + assert isinstance(recorder, BundleRecorder) + return RecordEdge(recorder=recorder, lock=threading.Lock()) + + +def replay_source(root: Path) -> ReplaySource: + loaded = load_bundle(root) + assert isinstance(loaded, LoadedBundle) + return ReplaySource(bundle=loaded) + + +def call_edge( + edge: ProviderEdge, + method: str, + path: str, + *, + body: bytes | None = None, + headers: dict[str, str] | None = None, +) -> RawResponse: + outcome = forward( + method, + f"http://{edge.advertise_host}:{edge.port}{path}", + headers=headers or {}, + body=body, + timeout=10.0, + ) + assert isinstance(outcome, RawResponse) + return outcome + + +def this_tests_files(root: Path) -> list[Path]: + slug_dir = root / slug_for_test(current_test_key()) + return sorted(slug_dir.glob("*.json")) if slug_dir.is_dir() else [] + + +def chat_body(prompt: str) -> bytes: + return json.dumps({"model": "gpt", "messages": [{"role": "user", "content": prompt}]}).encode() + + +class TestRecordMode: + def test_forwards_to_the_provider_and_writes_one_interaction_file(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + with fake_provider() as provider: + with running_edge(record_backend(root), {"openai": provider_url(provider)}) as edge: + reply = call_edge(edge, "POST", CHAT_PATH, body=chat_body("hi")) + assert provider.hits == ["POST /v1/chat/completions"] + assert reply.status_code == 200 + served = json_object(reply.body) + assert served["echo"] == chat_body("hi").decode() + files = this_tests_files(root) + assert [file.name for file in files] == ["0000-post-openai-v1-chat-completions.json"] + interaction = Interaction.model_validate_json(files[0].read_text(encoding="utf-8")) + assert interaction.request.method == "post" + assert interaction.request.path == CHAT_PATH + assert interaction.request.body == json_object(chat_body("hi")) + assert interaction.response.status_code == 200 + + def test_never_stores_headers_so_credentials_never_touch_disk(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + with fake_provider() as provider: + with running_edge(record_backend(root), {"openai": provider_url(provider)}) as edge: + call_edge( + edge, + "POST", + CHAT_PATH, + body=chat_body("hi"), + headers={"authorization": "Bearer sk-live-provider-secret-abc123"}, + ) + raw = this_tests_files(root)[0].read_text(encoding="utf-8") + assert "sk-live-provider-secret-abc123" not in raw + interaction = Interaction.model_validate_json(raw) + assert interaction.request.headers == {} + + def test_strips_volatile_response_headers_and_serves_the_filtered_copy(self, tmp_path: Path) -> None: + """What record serves the proxy must equal what replay will serve later + (record/replay parity), so the filtered stored copy is served in both.""" + root = tmp_path / "bundle" + with fake_provider() as provider: + with running_edge(record_backend(root), {"openai": provider_url(provider)}) as edge: + reply = call_edge(edge, "POST", CHAT_PATH, body=chat_body("hi")) + assert reply.headers.get("x-upstream") == "fake" + assert "set-cookie" not in reply.headers + interaction = Interaction.model_validate_json( + this_tests_files(root)[0].read_text(encoding="utf-8") + ) + assert interaction.response.headers.get("x-upstream") == "fake" + assert "set-cookie" not in interaction.response.headers + assert "content-length" not in interaction.response.headers + + def test_unreachable_provider_records_and_serves_a_502(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + with running_edge(record_backend(root), {"openai": "http://127.0.0.1:9"}) as edge: + reply = call_edge(edge, "POST", CHAT_PATH, body=chat_body("hi")) + assert reply.status_code == 502 + assert b"could not reach the provider" in reply.body + interaction = Interaction.model_validate_json( + this_tests_files(root)[0].read_text(encoding="utf-8") + ) + assert interaction.response.status_code == 502 + + +class TestReplayMode: + def test_serves_recorded_bytes_with_zero_provider_hits(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + with fake_provider() as provider: + with running_edge(record_backend(root), {"openai": provider_url(provider)}) as edge: + recorded = call_edge(edge, "POST", CHAT_PATH, body=chat_body("hi")) + hits_after_record = list(provider.hits) + with running_edge( + ReplayEdge(source=replay_source(root)), {"openai": provider_url(provider)} + ) as edge: + replayed = call_edge(edge, "POST", CHAT_PATH, body=chat_body("hi")) + assert provider.hits == hits_after_record + assert replayed.status_code == recorded.status_code + assert replayed.body == recorded.body + assert replayed.headers.get("x-upstream") == "fake" + + def test_request_identity_ignores_auth_headers(self, tmp_path: Path) -> None: + """The proxy sends different bearer tokens across runs (fresh virtual + keys, rotated provider keys), so headers are no part of the match.""" + root = tmp_path / "bundle" + with fake_provider() as provider: + with running_edge(record_backend(root), {"openai": provider_url(provider)}) as edge: + call_edge( + edge, "POST", CHAT_PATH, body=chat_body("hi"), + headers={"authorization": "Bearer sk-first-run"}, + ) + with running_edge(ReplayEdge(source=replay_source(root)), REPLAY_MOUNTS) as edge: + replayed = call_edge( + edge, "POST", CHAT_PATH, body=chat_body("hi"), + headers={"authorization": "Bearer sk-second-run"}, + ) + assert replayed.status_code == 200 + + def test_content_drift_returns_the_miss_status_naming_both_keys(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + with fake_provider() as provider: + with running_edge(record_backend(root), {"openai": provider_url(provider)}) as edge: + call_edge(edge, "POST", CHAT_PATH, body=chat_body("x")) + with running_edge(ReplayEdge(source=replay_source(root)), REPLAY_MOUNTS) as edge: + missed = call_edge(edge, "POST", CHAT_PATH, body=chat_body("y")) + assert missed.status_code == REPLAY_MISS_STATUS + message = missed.body.decode() + assert f"no recorded interaction matches key post {CHAT_PATH} #" in message + assert f"closest recorded key is post {CHAT_PATH} #" in message + assert '"content": "x"' in message + assert '"content": "y"' in message + assert "re-record with E2E_FIXTURE_MODE=record" in message + + def test_query_params_are_part_of_the_identity(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + with fake_provider() as provider: + with running_edge(record_backend(root), {"openai": provider_url(provider)}) as edge: + call_edge(edge, "GET", "/openai/v1/models?purpose=batch") + assert provider.hits == ["GET /v1/models?purpose=batch"] + with running_edge(ReplayEdge(source=replay_source(root)), REPLAY_MOUNTS) as edge: + missed = call_edge(edge, "GET", "/openai/v1/models?purpose=other") + matched = call_edge(edge, "GET", "/openai/v1/models?purpose=batch") + assert missed.status_code == REPLAY_MISS_STATUS + assert matched.status_code == 200 + + def test_identical_requests_replay_their_responses_in_recorded_order(self, tmp_path: Path) -> None: + """A poll or retry loop repeats the same request and the proxy asserts + on the progression, so duplicates under one key stay FIFO.""" + root = tmp_path / "bundle" + with fake_provider() as provider: + with running_edge(record_backend(root), {"openai": provider_url(provider)}) as edge: + call_edge(edge, "POST", CHAT_PATH, body=chat_body("hi")) + call_edge(edge, "POST", CHAT_PATH, body=chat_body("hi")) + with running_edge(ReplayEdge(source=replay_source(root)), REPLAY_MOUNTS) as edge: + first = json_object(call_edge(edge, "POST", CHAT_PATH, body=chat_body("hi")).body) + second = json_object(call_edge(edge, "POST", CHAT_PATH, body=chat_body("hi")).body) + assert first["hit"] == 1 + assert second["hit"] == 2 + + def test_exhausted_key_returns_the_miss_status(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + with fake_provider() as provider: + with running_edge(record_backend(root), {"openai": provider_url(provider)}) as edge: + call_edge(edge, "POST", CHAT_PATH, body=chat_body("hi")) + with running_edge(ReplayEdge(source=replay_source(root)), REPLAY_MOUNTS) as edge: + call_edge(edge, "POST", CHAT_PATH, body=chat_body("hi")) + exhausted = call_edge(edge, "POST", CHAT_PATH, body=chat_body("hi")) + assert exhausted.status_code == REPLAY_MISS_STATUS + assert b"already consumed" in exhausted.body + + def test_non_json_bodies_match_by_canonical_digest_without_storing_them(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + opaque = b"custom_id one\ncustom_id two\n" + with fake_provider() as provider: + with running_edge(record_backend(root), {"openai": provider_url(provider)}) as edge: + call_edge(edge, "POST", "/openai/v1/files", body=opaque) + raw = this_tests_files(root)[0].read_text(encoding="utf-8") + interaction = Interaction.model_validate_json(raw) + assert interaction.request.body is None + assert interaction.request.file_sha256 is not None + assert interaction.request.file_bytes == len(opaque) + assert "custom_id" not in interaction.request.model_dump_json() + with running_edge(ReplayEdge(source=replay_source(root)), REPLAY_MOUNTS) as edge: + replayed = call_edge(edge, "POST", "/openai/v1/files", body=opaque) + assert replayed.status_code == 200 + + +class TestReplayLeftover: + def test_partially_consumed_recording_names_the_leftover(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + with fake_provider() as provider: + with running_edge(record_backend(root), {"openai": provider_url(provider)}) as edge: + call_edge(edge, "POST", CHAT_PATH, body=chat_body("hi")) + call_edge(edge, "GET", "/openai/v1/models") + source = replay_source(root) + with running_edge(ReplayEdge(source=source), REPLAY_MOUNTS) as edge: + call_edge(edge, "POST", CHAT_PATH, body=chat_body("hi")) + error = source.leftover_error(current_test_key()) + assert error is not None + assert "1 of 2 recorded interactions never consumed" in error + assert "e.g. get /openai/v1/models #" in error + assert "re-record with E2E_FIXTURE_MODE=record" in error + + def test_fully_consumed_recording_leaves_nothing(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + with fake_provider() as provider: + with running_edge(record_backend(root), {"openai": provider_url(provider)}) as edge: + call_edge(edge, "POST", CHAT_PATH, body=chat_body("hi")) + source = replay_source(root) + with running_edge(ReplayEdge(source=source), REPLAY_MOUNTS) as edge: + call_edge(edge, "POST", CHAT_PATH, body=chat_body("hi")) + assert source.leftover_error(current_test_key()) is None + + def test_test_without_recordings_has_no_leftover(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + assert isinstance(prepare_bundle(root), BundleRecorder) + assert replay_source(root).leftover_error("suite.py::test_never_recorded") is None + + def test_inert_outside_replay_mode(self, tmp_path: Path) -> None: + missing = tmp_path / "missing" + assert replay_leftover_error(mode_raw="", bundle_dir=missing, test_key="k") is None + assert replay_leftover_error(mode_raw="record", bundle_dir=missing, test_key="k") is None + + +class TestConcurrentReplay: + def test_parallel_identical_calls_serve_each_recording_exactly_once(self, tmp_path: Path) -> None: + """The edge server handles requests on concurrent threads and a burst + of parallel identical calls consumes one shared pool: no response + duplicated, none forgotten, nothing left over at teardown.""" + root = tmp_path / "bundle" + recorder = prepare_bundle(root) + assert isinstance(recorder, BundleRecorder) + for ordinal in range(32): + recorder.record( + test_key=current_test_key(), + request=RecordedRequest(method="post", path=CHAT_PATH, headers={}, body={"n": "same"}), + response=RecordedHttpResponse( + status_code=200, + headers={"content-type": "application/json"}, + body_b64=base64.b64encode(json.dumps({"value": f"v{ordinal:02d}"}).encode()).decode(), + ), + ) + source = replay_source(root) + body = json.dumps({"n": "same"}).encode() + barrier = threading.Barrier(8) + with running_edge(ReplayEdge(source=source), REPLAY_MOUNTS) as edge: + + def consume(_: int) -> tuple[str, ...]: + barrier.wait() + return tuple( + str(json_object(call_edge(edge, "POST", CHAT_PATH, body=body).body)["value"]) + for _call in range(4) + ) + + with ThreadPoolExecutor(max_workers=8) as executor: + served = sorted(value for values in executor.map(consume, range(8)) for value in values) + assert served == [f"v{ordinal:02d}" for ordinal in range(32)] + assert source.leftover_error(current_test_key()) is None + + +class TestHandleEdgeRequestPure: + def test_unknown_mount_404s_naming_the_known_mounts(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + assert isinstance(prepare_bundle(root), BundleRecorder) + reply = handle_edge_request( + ReplayEdge(source=replay_source(root)), + {"openai": "https://api.openai.com", "anthropic": "https://api.anthropic.com"}, + "POST", + "/bedrock/model/invoke", + {}, + b"{}", + timeout=1.0, + ) + assert reply.status_code == 404 + assert b"unknown provider mount 'bedrock'" in reply.body + assert b"anthropic, openai" in reply.body + + def test_replay_serves_a_directly_recorded_interaction(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + recorder = prepare_bundle(root) + assert isinstance(recorder, BundleRecorder) + recorder.record( + test_key=current_test_key(), + request=RecordedRequest(method="post", path=CHAT_PATH, headers={}, body={"prompt": "x"}), + response=RecordedHttpResponse( + status_code=201, headers={"x-upstream": "fake"}, body_b64=base64.b64encode(b"ok").decode() + ), + ) + reply = handle_edge_request( + ReplayEdge(source=replay_source(root)), + {"openai": "https://api.openai.com"}, + "POST", + CHAT_PATH, + {"authorization": "Bearer sk-anything"}, + json.dumps({"prompt": "x"}).encode(), + timeout=1.0, + ) + assert reply.status_code == 201 + assert reply.body == b"ok" + assert reply.headers == {"x-upstream": "fake"} + + +class TestApiBaseSeam: + def test_live_mode_returns_none(self, tmp_path: Path) -> None: + for mode_raw in ("live", ""): + assert ( + provider_edge_api_base( + "openai", + mode_raw=mode_raw, + bundle_dir=tmp_path / "bundle", + bind_host="127.0.0.1", + advertise_host="127.0.0.1", + ) + is None + ) + + def test_invalid_mode_raises_naming_the_value(self, tmp_path: Path) -> None: + with pytest.raises(ValueError, match="cached"): + provider_edge_api_base( + "openai", + mode_raw="cached", + bundle_dir=tmp_path / "bundle", + bind_host="127.0.0.1", + advertise_host="127.0.0.1", + ) + + def test_unknown_mount_raises_naming_the_known_mounts(self, tmp_path: Path) -> None: + with pytest.raises(ValueError, match="unknown provider mount 'bedrock'"): + provider_edge_api_base( + "bedrock", + mode_raw="record", + bundle_dir=tmp_path / "bundle", + bind_host="127.0.0.1", + advertise_host="127.0.0.1", + ) + + def test_record_mode_boots_one_shared_edge_and_prepares_the_bundle(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + first = provider_edge_api_base( + "openai", mode_raw="record", bundle_dir=root, bind_host="127.0.0.1", advertise_host="127.0.0.1" + ) + second = provider_edge_api_base( + "anthropic", mode_raw="record", bundle_dir=root, bind_host="127.0.0.1", advertise_host="127.0.0.1" + ) + assert first is not None and second is not None + assert first.endswith("/openai") + assert second.endswith("/anthropic") + assert first.rsplit("/", 1)[0] == second.rsplit("/", 1)[0] + assert (root / "manifest.json").is_file() From 0de829d3e44094a808e9c1166d12c4d86b29c6b0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 18:57:35 -0700 Subject: [PATCH 031/281] feat(cli): store the lite login credential in the OS keychain lite login used to write the minted cli-session key in cleartext to ~/.litellm/token.json. The secret material (key plus any JWT) now goes to the OS keychain through the optional keyring package, with the 0600 file kept for non-secret metadata and as the fallback on headless boxes. Legacy plaintext files keep authenticating and are migrated into the keychain, then scrubbed, on first read. A secret still on disk always outranks the keychain entry, so a failed keychain write can never resurrect a stale key. LITELLM_PROXY_API_KEY and --api-key precedence is unchanged, lite logout clears both stores and warns when the keychain will not release the entry, and ~/.litellm is created 0700 (tightened from 0755 where an older CLI left it broader). LITELLM_CLI_DISABLE_KEYRING=1 forces the file fallback. --- basedpyright-code-budget.json | 8 +- .../litellm_proxy_server/cli_token_usage.py | 2 +- litellm/litellm_core_utils/cli_keyring.py | 125 ++++ litellm/litellm_core_utils/cli_token_utils.py | 184 +++++- .../private_json.py | 10 + litellm/proxy/client/README.md | 12 +- litellm/proxy/client/cli/commands/agents.py | 4 +- litellm/proxy/client/cli/commands/auth.py | 148 +++-- .../client/cli/commands/claude_settings.py | 2 +- litellm/proxy/client/cli/commands/config.py | 6 +- litellm/proxy/client/cli/commands/up.py | 22 +- litellm/proxy/client/cli/main.py | 4 +- pyproject.toml | 2 + tests/test_litellm/conftest.py | 65 ++ .../test_cli_token_utils.py | 478 ++++++++++++--- .../proxy/client/cli/test_agents.py | 7 +- .../proxy/client/cli/test_auth_commands.py | 577 +++++++++--------- .../proxy/client/cli/test_claude_settings.py | 15 +- .../proxy/client/cli/test_config_commands.py | 4 +- .../proxy/client/cli/test_up_commands.py | 23 +- type-discipline-budget.json | 8 +- uv.lock | 100 ++- 22 files changed, 1295 insertions(+), 511 deletions(-) create mode 100644 litellm/litellm_core_utils/cli_keyring.py rename litellm/{proxy/client/cli/commands => litellm_core_utils}/private_json.py (64%) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 1ce71c5bd2c..59d56a3f63d 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5663 }, "reportMissingTypeArgument": { - "limit": 15557 + "limit": 15556 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 39043 + "limit": 39042 }, "reportUnknownParameterType": { - "limit": 19887 + "limit": 19886 }, "reportUnknownVariableType": { - "limit": 30574 + "limit": 30571 }, "reportUnnecessaryCast": { "limit": 117 diff --git a/cookbook/litellm_proxy_server/cli_token_usage.py b/cookbook/litellm_proxy_server/cli_token_usage.py index 6306970cdde..e6b3744019c 100644 --- a/cookbook/litellm_proxy_server/cli_token_usage.py +++ b/cookbook/litellm_proxy_server/cli_token_usage.py @@ -60,4 +60,4 @@ if __name__ == "__main__": print("\n💡 Tips:") print("1. Run 'litellm-proxy login' to authenticate first") print("2. Replace 'https://your-proxy.com' with your actual proxy URL") - print("3. The token is stored locally at ~/.litellm/token.json") + print("3. The token is stored in your OS keychain, or in ~/.litellm/token.json when there is none") diff --git a/litellm/litellm_core_utils/cli_keyring.py b/litellm/litellm_core_utils/cli_keyring.py new file mode 100644 index 00000000000..873db64a728 --- /dev/null +++ b/litellm/litellm_core_utils/cli_keyring.py @@ -0,0 +1,125 @@ +""" +CLI Keyring Access + +SDK-level access to the OS keychain (macOS Keychain, Windows Credential Manager, +Linux Secret Service) that holds the credential minted by `lite login`. + +The `keyring` package is optional and imported lazily, so importing this module +never pulls it in. Every failure is returned as a value: a machine with no +keychain, or one whose keychain is locked, must degrade to the token file rather +than break `lite` or the SDK. +""" + +import os +from dataclasses import dataclass +from typing import Final, Protocol, TypeAlias + +KEYRING_SERVICE: Final = "litellm-cli" +KEYRING_ACCOUNT: Final = "credential" +DISABLE_KEYRING_ENV_VAR: Final = "LITELLM_CLI_DISABLE_KEYRING" + +_DISABLED_VALUES: Final = frozenset(("1", "true", "yes", "on")) + + +@dataclass(frozen=True, slots=True) +class SecretFound: + blob: str + + +@dataclass(frozen=True, slots=True) +class SecretMissing: + pass + + +@dataclass(frozen=True, slots=True) +class SecretUnavailable: + pass + + +SecretRead: TypeAlias = SecretFound | SecretMissing | SecretUnavailable + + +class SecretVault(Protocol): + """The single slot holding the CLI credential's secret material.""" + + def read(self) -> SecretRead: ... + + def write(self, blob: str) -> bool: ... + + def erase(self) -> bool: ... + + +class KeyringApi(Protocol): + def get_password(self, service_name: str, username: str) -> str | None: ... + + def set_password(self, service_name: str, username: str, password: str) -> None: ... + + def delete_password(self, service_name: str, username: str) -> None: ... + + +def _keyring_disabled() -> bool: + return os.getenv(DISABLE_KEYRING_ENV_VAR, "").strip().lower() in _DISABLED_VALUES + + +def _import_keyring() -> KeyringApi | None: + try: + import keyring + except ImportError: + return None + return keyring + + +def _keyring_api() -> KeyringApi | None: + return None if _keyring_disabled() else _import_keyring() + + +@dataclass(frozen=True, slots=True) +class KeyringVault: + """The OS keychain, reached through the optional `keyring` package.""" + + def read(self) -> SecretRead: + api: Final = _keyring_api() + if api is None: + return SecretUnavailable() + try: + blob: Final = api.get_password(KEYRING_SERVICE, KEYRING_ACCOUNT) + except Exception: # noqa: BLE001 # backends raise outside keyring.errors; never break the SDK + return SecretUnavailable() + return SecretMissing() if blob is None else SecretFound(blob) + + def write(self, blob: str) -> bool: + api: Final = _keyring_api() + if api is None: + return False + try: + api.set_password(KEYRING_SERVICE, KEYRING_ACCOUNT, blob) + except Exception: # noqa: BLE001 # a keychain that refuses the write falls back to the token file + return False + return True + + def erase(self) -> bool: + if _import_keyring() is None: + return True + if _keyring_disabled(): + # a credential stored before the kill switch was set may still be in the keychain + return False + match self.read(): + case SecretUnavailable(): + return False + case SecretMissing(): + return True + case SecretFound(): + return self._delete() + + def _delete(self) -> bool: + api: Final = _keyring_api() + if api is None: + return False + try: + api.delete_password(KEYRING_SERVICE, KEYRING_ACCOUNT) + except Exception: # noqa: BLE001 # report the failure as a value so `lite logout` can warn + return False + return True + + +SYSTEM_KEYRING: Final[SecretVault] = KeyringVault() diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index a44ce431f4e..9960192180c 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -1,17 +1,68 @@ """ CLI Token Utilities -SDK-level utilities for reading CLI authentication tokens. +SDK-level utilities for reading the credential minted by `lite login`. + +Non-secret metadata lives in ~/.litellm/token.json. The secret material (the +bearer key, plus a JWT when one is issued) lives in the OS keychain when the +machine has one, and in that same 0600 file otherwise. This module hides the +split from callers, and migrates a legacy plaintext file into the keychain the +first time it reads one. + This module has no dependencies on proxy code and can be safely imported at the SDK level. """ -import json -import os +import contextlib import time -from collections.abc import Mapping from pathlib import Path +from types import MappingProxyType from typing import Final +from pydantic import BaseModel, ConfigDict, ValidationError + +from litellm.litellm_core_utils.cli_keyring import ( + SYSTEM_KEYRING, + SecretFound, + SecretMissing, + SecretUnavailable, + SecretVault, +) +from litellm.litellm_core_utils.private_json import ensure_private_dir, write_private_json + + +class CliTokenRecord(BaseModel): + """A stored CLI credential. + + `key is None` means the metadata was found but the secret could not be + produced: the keychain holds nothing for us, or we could not reach it. + """ + + model_config = ConfigDict(frozen=True, extra="allow") + + base_url: str = "" + key: str | None = None + user_id: str = "" + user_email: str = "" + user_role: str = "" + auth_header_name: str = "Authorization" + jwt_token: str = "" + timestamp: float = 0.0 + + +class CliTokenSecret(BaseModel): + """The secret material as stored in the OS keychain. + + `base_url` is duplicated from the metadata file purely as a pairing tag: a + secret minted for one server is never handed to another, even if the + metadata file is edited underneath us. + """ + + model_config = ConfigDict(frozen=True) + + base_url: str + key: str + jwt_token: str = "" + def get_cli_token_file_path() -> str: """Get the path to the CLI token file""" @@ -20,26 +71,39 @@ def get_cli_token_file_path() -> str: return str(config_dir / "token.json") -def load_cli_token() -> dict | None: - """Load CLI token data from file""" - token_file: Final = get_cli_token_file_path() - if not os.path.exists(token_file): +def load_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> CliTokenRecord | None: + """Load the stored CLI credential, or None when this machine has none""" + record: Final = _read_token_file() + if record is None: return None + return _resolve_secret(record, vault) - try: - with open(token_file, "r") as f: - return json.load(f) - except (OSError, json.JSONDecodeError): - return None + +def save_cli_token(record: CliTokenRecord, *, vault: SecretVault = SYSTEM_KEYRING) -> bool: + """Store a freshly minted credential. Returns whether the keychain took the secret""" + if record.key is None or not vault.write(_encode_secret(record.base_url, record.key, record.jwt_token)): + _write_token_file(record) + return False + _write_token_file(_without_secret(record)) + return True + + +def clear_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> bool: + """Remove the credential from both stores. Returns whether the keychain is now free of it""" + erased: Final = vault.erase() + Path(get_cli_token_file_path()).unlink(missing_ok=True) + return erased def get_litellm_gateway_api_key( expected_base_url: str | None = None, + *, + vault: SecretVault = SYSTEM_KEYRING, ) -> str | None: """ Get the stored CLI API key for use with LiteLLM SDK. - This function reads the token file created by `lite login` + This function reads the credential created by `lite login` and returns the API key for use in Python scripts. Args: @@ -47,6 +111,7 @@ def get_litellm_gateway_api_key( originally issued for this URL. Pass the target server URL to prevent credential leakage when the client is pointed at a different (possibly malicious) server. + vault: Where the secret material is stored. Defaults to the OS keychain. Returns: str: The API key if found (and origin matches), None otherwise @@ -62,25 +127,84 @@ def get_litellm_gateway_api_key( >>> base_url="https://your-proxy.com/v1" >>> ) """ - token_data: Final = load_cli_token() - if not token_data or "key" not in token_data: + record: Final = _read_token_file() + if record is None: return None - if expected_base_url is not None: - stored_url: Final = token_data.get("base_url") - if stored_url != expected_base_url.rstrip("/"): - return None - return token_data["key"] + if expected_base_url is not None and record.base_url != expected_base_url.rstrip("/"): + return None + resolved: Final = _resolve_secret(record, vault) + return None if resolved is None else resolved.key -def is_cli_token_fresh(token_data: Mapping[str, object], buffer_hours: float = 0.1) -> bool: - """Check whether a cached CLI token (as stored in token.json) is still - within its expiration window. Used by `lite auth print-token` to fail - fast, without a network round trip, once the cached token is past - `LITELLM_CLI_JWT_EXPIRATION_HOURS`.""" +def is_cli_token_fresh(token_data: CliTokenRecord, buffer_hours: float = 0.1) -> bool: + """Check whether a cached CLI token is still within its expiration window. + Used by `lite auth print-token` to fail fast, without a network round trip, + once the cached token is past `LITELLM_CLI_JWT_EXPIRATION_HOURS`.""" from litellm.constants import CLI_JWT_EXPIRATION_HOURS - timestamp: Final = token_data.get("timestamp") - if not isinstance(timestamp, (int, float)): - return False - age_hours: Final = (time.time() - timestamp) / 3600 + age_hours: Final = (time.time() - token_data.timestamp) / 3600 return age_hours < (CLI_JWT_EXPIRATION_HOURS - buffer_hours) + + +def _read_token_file() -> CliTokenRecord | None: + try: + raw: Final = Path(get_cli_token_file_path()).read_text() + except OSError: + return None + try: + return CliTokenRecord.model_validate_json(raw) + except ValidationError: + return None + + +def _resolve_secret(record: CliTokenRecord, vault: SecretVault) -> CliTokenRecord | None: + match vault.read(): + case SecretFound(blob=blob): + return _apply_vault_secret(record, blob, vault) + case SecretMissing(): + return _migrate_file_secret(record, vault) + case SecretUnavailable(): + return record + + +def _apply_vault_secret(record: CliTokenRecord, blob: str, vault: SecretVault) -> CliTokenRecord | None: + if record.key is not None: + # a secret still on disk means the last keychain write failed: the file outranks the vault + return _migrate_file_secret(record, vault) + try: + secret: Final = CliTokenSecret.model_validate_json(blob) + except ValidationError: + return _migrate_file_secret(record, vault) + if secret.base_url != record.base_url: + return _migrate_file_secret(record, vault) + _scrub_file_secret(record) + return record.model_copy(update=MappingProxyType({"key": secret.key, "jwt_token": secret.jwt_token})) + + +def _migrate_file_secret(record: CliTokenRecord, vault: SecretVault) -> CliTokenRecord | None: + if record.key is None: + return None + if vault.write(_encode_secret(record.base_url, record.key, record.jwt_token)): + _scrub_file_secret(record) + return record + + +def _scrub_file_secret(record: CliTokenRecord) -> None: + if record.key is None and not record.jwt_token: + return + with contextlib.suppress(OSError): + _write_token_file(_without_secret(record)) + + +def _without_secret(record: CliTokenRecord) -> CliTokenRecord: + return record.model_copy(update=MappingProxyType({"key": None, "jwt_token": ""})) + + +def _encode_secret(base_url: str, key: str, jwt_token: str) -> str: + return CliTokenSecret(base_url=base_url, key=key, jwt_token=jwt_token).model_dump_json() + + +def _write_token_file(record: CliTokenRecord) -> None: + path: Final = Path(get_cli_token_file_path()) + ensure_private_dir(path.parent) + write_private_json(str(path), record.model_dump(exclude_none=True)) diff --git a/litellm/proxy/client/cli/commands/private_json.py b/litellm/litellm_core_utils/private_json.py similarity index 64% rename from litellm/proxy/client/cli/commands/private_json.py rename to litellm/litellm_core_utils/private_json.py index 31062e4a799..32bc2e169e2 100644 --- a/litellm/proxy/client/cli/commands/private_json.py +++ b/litellm/litellm_core_utils/private_json.py @@ -1,10 +1,20 @@ import json import os +import stat import tempfile from collections.abc import Mapping from pathlib import Path from typing import Final +PRIVATE_DIR_MODE: Final = 0o700 + + +def ensure_private_dir(directory: Path) -> None: + """Create directory (and parents) owner-only, tightening it if it already exists group/world readable""" + directory.mkdir(mode=PRIVATE_DIR_MODE, parents=True, exist_ok=True) + if stat.S_IMODE(directory.stat().st_mode) & 0o077: + directory.chmod(PRIVATE_DIR_MODE) + def write_private_json(path: str, data: Mapping[str, object]) -> None: """Atomically write JSON to path with owner-only permissions (0600)""" diff --git a/litellm/proxy/client/README.md b/litellm/proxy/client/README.md index 6b28f43ac73..9ece4c2be3d 100644 --- a/litellm/proxy/client/README.md +++ b/litellm/proxy/client/README.md @@ -331,7 +331,7 @@ sequenceDiagram CLI->>Proxy: Poll /sso/cli/poll/login_id with poll_secret header Proxy->>CLI: Return {"status": "ready", "key": "jwt"} - CLI->>CLI: Save key to ~/.litellm/token.json + CLI->>CLI: Save key to the OS keychain (metadata to ~/.litellm/token.json) ``` ### Authentication Commands @@ -352,7 +352,7 @@ The CLI provides these authentication commands: 5. **Callback Processing**: SSO provider redirects back to proxy with state parameter 6. **User Code Verification**: Browser confirms the verification code shown in the CLI 7. **Polling**: CLI polls `/sso/cli/poll/{login_id}` with the polling secret header until the JWT is ready. When `CLI_SSO_CLAIM_MAP` is configured on the proxy, the poll response may include `attribution_metadata` (allowlisted scalar OIDC claims for client attribution). -8. **Token Storage**: CLI saves the authentication token to `~/.litellm/token.json` +8. **Token Storage**: CLI saves the key to the OS keychain and the non-secret session metadata to `~/.litellm/token.json` ### Benefits of This Approach @@ -364,11 +364,11 @@ The CLI provides these authentication commands: ### Token Storage -Authentication tokens are stored in `~/.litellm/token.json` with restricted file permissions (600). The stored token includes: +The key itself goes into the OS keychain (macOS Keychain, Windows Credential Manager, or the Linux Secret Service) under service `litellm-cli`, account `credential`. Only the non-secret session metadata is written to `~/.litellm/token.json`, in a `0700` directory with `0600` file permissions: ```json { - "key": "sk-...", + "base_url": "https://your-proxy.com", "user_id": "cli-user", "user_email": "user@example.com", "user_role": "cli", @@ -377,6 +377,10 @@ Authentication tokens are stored in `~/.litellm/token.json` with restricted file } ``` +Headless boxes and CI runners usually have no keychain. There the key stays in the same `0600` file alongside the metadata, exactly as it did before, and `lite login` tells you which of the two happened. Set `LITELLM_CLI_DISABLE_KEYRING=1` to force the file even where a keychain exists. A `token.json` written by an older `lite` keeps working and is moved into the keychain, and scrubbed from the file, the first time a keychain-capable `lite` reads it. + +`lite logout` clears both stores. If the keychain is locked at that moment it says so, and re-running it once the keychain is unlocked finishes the job. + The stored credential is a short-lived, per-session agent token, not a managed virtual key. It is scoped to the user and team you logged in as and inherits their models and budgets; spend is tracked against the shared team and user budgets rather than a separate per-session cap, so multiple logins or several concurrent agents all draw down the same allowance. It is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); re-run `lite login` to refresh it and pick up your latest team and user settings. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it while fresh and fails once it expires -- there is no silent renewal. It is accepted on a default deployment without `EXPERIMENTAL_UI_LOGIN`, does not appear in the Keys UI, and cannot be rotated or revoked mid-session. For a long-lived, rotatable, Keys-UI-visible credential, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY`. ### Usage diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index ed2bf2be03d..e05e85ae483 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -8,7 +8,7 @@ from typing import Final import click import requests -from .auth import get_stored_api_key, login +from .auth import context_secret_vault, get_stored_api_key, login ANTHROPIC_BASE_URL_ENV: Final = "ANTHROPIC_BASE_URL" ANTHROPIC_AUTH_TOKEN_ENV: Final = "ANTHROPIC_AUTH_TOKEN" @@ -316,7 +316,7 @@ def resolve_api_key(ctx: click.Context) -> str: click.echo("No LiteLLM credentials found; starting login...") ctx.invoke(login) - api_key = get_stored_api_key(expected_base_url=base_url) + api_key = get_stored_api_key(expected_base_url=base_url, vault=context_secret_vault(ctx)) if not api_key: raise click.ClickException("Login did not produce an API key; cannot start the agent.") return api_key diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 0a0bcf80ee5..8b9ef5633da 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -1,9 +1,6 @@ -import json -import os import sys import time import webbrowser -from pathlib import Path from typing import Any, Final from urllib.parse import urlencode @@ -11,10 +8,19 @@ import click import requests from rich.console import Console from rich.table import Table -from typing_extensions import NotRequired, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.constants import CLI_JWT_EXPIRATION_HOURS -from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh +from litellm.litellm_core_utils.cli_keyring import SYSTEM_KEYRING, SecretVault +from litellm.litellm_core_utils.cli_token_utils import ( + CliTokenRecord, + clear_cli_token, + get_cli_token_file_path, + get_litellm_gateway_api_key, + is_cli_token_fresh, + load_cli_token, + save_cli_token, +) from .claude_settings import ( CLAUDE_SETTINGS_PATH, @@ -22,18 +28,6 @@ from .claude_settings import ( ClaudeSettingsError, write_claude_settings, ) -from .private_json import write_private_json - - -class CliTokenData(TypedDict): - base_url: str - key: str - user_id: str - user_email: str - user_role: str - auth_header_name: str - jwt_token: str - timestamp: float class CliTeam(TypedDict, total=False): @@ -46,6 +40,7 @@ class CliTeam(TypedDict, total=False): class CliContextObj(TypedDict): base_url: str base_url_explicit: NotRequired[bool] + secret_vault: NotRequired[ReadOnly[SecretVault]] class CliPollData(TypedDict, total=False): @@ -76,50 +71,32 @@ class CliAuthResult(TypedDict): team_id: str | None +KEYCHAIN_UNREACHABLE_MESSAGE: Final = ( + "Your credential is stored in your OS keychain, which could not be read. Unlock it and retry, or run 'lite login'." +) + + # Token storage utilities -def get_token_file_path() -> str: - """Get the path to store the authentication token""" - home_dir: Final = Path.home() - config_dir: Final = home_dir / ".litellm" - config_dir.mkdir(exist_ok=True) - return str(config_dir / "token.json") +def context_secret_vault(ctx: click.Context) -> SecretVault: + """Where this invocation reads and writes secret material; injectable through ctx.obj for tests""" + ctx_obj: Final[CliContextObj | None] = ctx.obj + if ctx_obj is None: + return SYSTEM_KEYRING + return ctx_obj.get("secret_vault") or SYSTEM_KEYRING -def save_token(token_data: CliTokenData) -> None: - """Save token data to file""" - write_private_json(get_token_file_path(), token_data) - - -def load_token() -> CliTokenData | None: - """Load token data from file""" - token_file: Final = get_token_file_path() - if not os.path.exists(token_file): - return None - - try: - with open(token_file, "r") as f: - return json.load(f) - except (OSError, json.JSONDecodeError): - return None - - -def clear_token() -> None: - """Clear stored token""" - token_file: Final = get_token_file_path() - if os.path.exists(token_file): - os.remove(token_file) - - -def get_stored_api_key(expected_base_url: str | None = None) -> str | None: - """Get the stored API key from token file. +def get_stored_api_key( + expected_base_url: str | None = None, + *, + vault: SecretVault = SYSTEM_KEYRING, +) -> str | None: + """Get the stored API key. If expected_base_url is provided, the key is only returned when it was originally issued for that URL. This prevents credential leakage when the CLI is pointed at a different (possibly malicious) server. """ - from litellm.litellm_core_utils.cli_token_utils import get_litellm_gateway_api_key - - return get_litellm_gateway_api_key(expected_base_url=expected_base_url) + return get_litellm_gateway_api_key(expected_base_url=expected_base_url, vault=vault) # Team selection utilities @@ -689,23 +666,27 @@ def login(ctx: click.Context, config_claude: bool): api_key: Final = auth_result["api_key"] user_id: Final = auth_result["user_id"] - # Save token data. base_url is stored so we can verify origin - # before reusing the key on a subsequent CLI invocation. - save_token( - { - "base_url": base_url.rstrip("/"), - "key": api_key, - "user_id": user_id or "cli-user", - "user_email": "unknown", - "user_role": "cli", - "auth_header_name": "Authorization", - "jwt_token": "", - "timestamp": time.time(), - } + # base_url is stored so we can verify origin before reusing the + # key on a subsequent CLI invocation. + record: Final = CliTokenRecord( + base_url=base_url.rstrip("/"), + key=api_key, + user_id=user_id or "cli-user", + user_email="unknown", + user_role="cli", + auth_header_name="Authorization", + jwt_token="", + timestamp=time.time(), ) + in_keychain: Final = save_cli_token(record, vault=context_secret_vault(ctx)) click.echo("\nLogin successful!") click.echo(f"JWT Token: {api_key[:20]}...") + click.echo( + "Credential stored in your OS keychain." + if in_keychain + else f"No OS keychain available; credential stored in {get_cli_token_file_path()} (owner-only)." + ) click.echo("You can now use the CLI without specifying --api-key") if config_claude: @@ -736,10 +717,14 @@ def login(ctx: click.Context, config_claude: bool): @click.command(name="logout") -def logout(): +@click.pass_context +def logout(ctx: click.Context): """Logout and clear stored authentication""" - clear_token() - click.echo("Logged out successfully. Authentication token cleared.") + if clear_cli_token(vault=context_secret_vault(ctx)): + click.echo("Logged out successfully. Authentication token cleared.") + return + click.echo("Logged out. The local token file is gone, but the OS keychain entry could not be removed.") + click.echo("Unlock your keychain and run 'lite logout' again to clear it.") @click.command(name="print-token") @@ -753,7 +738,7 @@ def print_token(ctx: click.Context): expires after `LITELLM_CLI_JWT_EXPIRATION_HOURS` (default 24h); once expired, run `lite login` again. """ - token_data: Final = load_token() + token_data: Final = load_cli_token(vault=context_secret_vault(ctx)) if not token_data: click.echo("Not authenticated. Run 'lite login'.", err=True) sys.exit(1) @@ -765,7 +750,7 @@ def print_token(ctx: click.Context): ctx_obj: Final[CliContextObj] = ctx.obj if ctx_obj.get("base_url_explicit"): base_url: Final = ctx_obj["base_url"] - if token_data.get("base_url") != base_url.rstrip("/"): + if token_data.base_url != base_url.rstrip("/"): click.echo("Not authenticated for this server. Run 'lite login'.", err=True) sys.exit(1) @@ -773,33 +758,36 @@ def print_token(ctx: click.Context): click.echo("Token expired. Run 'lite login' again.", err=True) sys.exit(1) - api_key: Final = token_data.get("key") + api_key: Final = token_data.key if not api_key: - click.echo("No token available. Run 'lite login'.", err=True) + click.echo(KEYCHAIN_UNREACHABLE_MESSAGE, err=True) sys.exit(1) click.echo(api_key) @click.command(name="whoami") -def whoami(): +@click.pass_context +def whoami(ctx: click.Context): """Show current authentication status""" - token_data: Final = load_token() + token_data: Final = load_cli_token(vault=context_secret_vault(ctx)) if not token_data: click.echo("Not authenticated. Run 'lite login' to authenticate.") return click.echo("Authenticated") - click.echo(f"User Email: {token_data.get('user_email', 'Unknown')}") - click.echo(f"User ID: {token_data.get('user_id', 'Unknown')}") - click.echo(f"User Role: {token_data.get('user_role', 'Unknown')}") + click.echo(f"User Email: {token_data.user_email or 'Unknown'}") + click.echo(f"User ID: {token_data.user_id or 'Unknown'}") + click.echo(f"User Role: {token_data.user_role or 'Unknown'}") # Check if token is still valid (basic timestamp check) - timestamp: Final = token_data.get("timestamp", 0) - age_hours: Final = (time.time() - timestamp) / 3600 + age_hours: Final = (time.time() - token_data.timestamp) / 3600 click.echo(f"Token age: {age_hours:.1f} hours") + if token_data.key is None: + click.echo(KEYCHAIN_UNREACHABLE_MESSAGE) + if age_hours > CLI_JWT_EXPIRATION_HOURS: click.echo(f"Warning: Token is more than {CLI_JWT_EXPIRATION_HOURS} hours old and may have expired.") diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index e9a6a25a064..e18e5b1b7ee 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -15,7 +15,7 @@ from typing import Final from pydantic import JsonValue, TypeAdapter, ValidationError -from .private_json import write_private_json +from litellm.litellm_core_utils.private_json import write_private_json ENV_KEY: Final = "env" API_KEY_HELPER_KEY: Final = "apiKeyHelper" diff --git a/litellm/proxy/client/cli/commands/config.py b/litellm/proxy/client/cli/commands/config.py index 19dd407ba19..2715a0a9a38 100644 --- a/litellm/proxy/client/cli/commands/config.py +++ b/litellm/proxy/client/cli/commands/config.py @@ -10,7 +10,7 @@ from urllib.parse import urlparse import click from pydantic import TypeAdapter -from .private_json import write_private_json +from litellm.litellm_core_utils.private_json import ensure_private_dir, write_private_json HIDDEN_COMMANDS_KEY: Final = "hidden_commands" @@ -42,7 +42,9 @@ def load_config() -> Mapping[str, str]: def save_config(config: Mapping[str, str]) -> None: """Save CLI config to file""" - write_private_json(get_config_file_path(), config) + config_file: Final = Path(get_config_file_path()) + ensure_private_dir(config_file.parent) + write_private_json(str(config_file), config) def get_config_value(key: str) -> str | None: diff --git a/litellm/proxy/client/cli/commands/up.py b/litellm/proxy/client/cli/commands/up.py index dd266b4afa1..a0fd4af8f72 100644 --- a/litellm/proxy/client/cli/commands/up.py +++ b/litellm/proxy/client/cli/commands/up.py @@ -14,10 +14,12 @@ from typing import IO, Final import click from pydantic import JsonValue, TypeAdapter, ValidationError -from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh +from litellm.litellm_core_utils.cli_keyring import SecretVault +from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh, load_cli_token +from litellm.litellm_core_utils.private_json import ensure_private_dir from .agents import AgentRunError, resolve_api_key, verify_proxy_key -from .auth import load_token, login +from .auth import context_secret_vault, login from .claude_settings import ( BACKUP_PATH, CLAUDE_SETTINGS_PATH, @@ -66,7 +68,7 @@ def secure_create(path: Path) -> Iterator[IO[str]]: def write_backup(record: BackupRecord, backup_path: Path | None = None) -> None: path: Final = backup_path if backup_path is not None else BACKUP_PATH - path.parent.mkdir(exist_ok=True) + ensure_private_dir(path.parent) with secure_create(path) as f: json.dump({"existed": record.existed, "content": record.content}, f, indent=2) @@ -103,10 +105,17 @@ def restore_claude_settings(settings_path: Path | None = None, backup_path: Path return record +def _has_fresh_login(base_url: str, vault: SecretVault) -> bool: + token_data: Final = load_cli_token(vault=vault) + if token_data is None or token_data.key is None or token_data.base_url != base_url: + return False + return is_cli_token_fresh(token_data) + + def _ensure_fresh_login(ctx: click.Context) -> None: base_url: Final = ctx.obj["base_url"].rstrip("/") - token_data = load_token() - if token_data and token_data.get("base_url") == base_url and is_cli_token_fresh(token_data): + vault: Final = context_secret_vault(ctx) + if _has_fresh_login(base_url, vault): return if not sys.stdin.isatty(): @@ -117,8 +126,7 @@ def _ensure_fresh_login(ctx: click.Context) -> None: click.echo("No fresh LiteLLM login found for this proxy; starting login...") ctx.invoke(login) - token_data = load_token() - if not token_data or token_data.get("base_url") != base_url or not is_cli_token_fresh(token_data): + if not _has_fresh_login(base_url, vault): raise UpError("Login did not produce a usable token; cannot start `lite up`.") diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index 3a289736c66..664bf5a216c 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -9,7 +9,7 @@ from litellm._version import version as litellm_version from litellm.proxy.client.health import HealthManagementClient from .commands.agents import agent_commands -from .commands.auth import auth_group, get_stored_api_key, login, logout, whoami +from .commands.auth import auth_group, context_secret_vault, get_stored_api_key, login, logout, whoami from .commands.autoroute.commands import autoroute_group from .commands.chat import chat from .commands.config import config_commands, get_config_value, hidden_command_names @@ -94,7 +94,7 @@ def cli(ctx: click.Context, show_version: bool, base_url: str | None, api_key: s # If no API key provided via flag or environment variable, try to load from saved token. # Pass base_url so we only use the stored key when it was issued for this server. if api_key is None: - api_key = get_stored_api_key(expected_base_url=base_url) + api_key = get_stored_api_key(expected_base_url=base_url, vault=context_secret_vault(ctx)) ctx.obj["base_url"] = base_url ctx.obj["api_key"] = api_key diff --git a/pyproject.toml b/pyproject.toml index ffbc96eefb9..32921e14d31 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -85,6 +85,7 @@ cli = [ "pyyaml>=6.0.3,<7.0", "requests>=2.32.0,<3.0", "InquirerPy>=0.3.4,<1.0", + "keyring>=25.6.0,<26.0", ] extra_proxy = [ "prisma>=0.11.0,<1.0", @@ -166,6 +167,7 @@ litellm-proxy = "litellm.proxy.client.cli:cli" dev = [ "diff-cover==9.7.2", "basedpyright==1.39.7", + "keyring==25.7.0", "pytest==9.0.3", "pytest-mock==3.15.1", "pytest-asyncio==1.3.0", diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index c0644c88291..ce0fd197538 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -22,6 +22,12 @@ import litellm from litellm import router as litellm_router_module from litellm import utils as litellm_utils_module from litellm._logging import ALL_LOGGERS +from litellm.litellm_core_utils.cli_keyring import ( + SecretFound, + SecretMissing, + SecretRead, + SecretUnavailable, +) from litellm.litellm_core_utils.prompt_templates import ( image_handling as image_handling_module, ) @@ -106,6 +112,65 @@ def isolate_host_proxy_base_url(monkeypatch): monkeypatch.delenv("PROXY_BASE_URL", raising=False) +@pytest.fixture(scope="function", autouse=True) +def isolate_host_os_keychain(monkeypatch): + """Keep any code path that resolves a CLI credential out of the developer's real OS keychain. + + Tests that exercise keychain behaviour inject their own vault instead. + """ + monkeypatch.setenv("LITELLM_CLI_DISABLE_KEYRING", "1") + + +class FakeSecretVault: + """In-memory stand-in for the OS keychain, injected wherever CLI credential storage is exercised. + + `available=False` models a keychain that is locked or has no backend, `writable=False` one that + refuses to store, and `erasable=False` one that will not release what it already holds. + """ + + def __init__( + self, + blob: str | None = None, + *, + available: bool = True, + writable: bool = True, + erasable: bool = True, + ) -> None: + self.blob: str | None = blob + self.available: bool = available + self.writable: bool = writable + self.erasable: bool = erasable + self.reads: int = 0 + self.writes: list[str] = [] + self.erases: int = 0 + + def read(self) -> SecretRead: + self.reads += 1 + if not self.available: + return SecretUnavailable() + return SecretMissing() if self.blob is None else SecretFound(self.blob) + + def write(self, blob: str) -> bool: + self.writes.append(blob) + if not (self.available and self.writable): + return False + self.blob = blob + return True + + def erase(self) -> bool: + self.erases += 1 + if not (self.available and self.erasable): + return False + self.blob = None + return True + + +@pytest.fixture +def secret_vault_factory(): + """Build FakeSecretVault instances; see its docstring for the failure modes it can model.""" + return FakeSecretVault + + def _run_coroutine_if_needed(result): if not asyncio.iscoroutine(result): return diff --git a/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py b/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py index 27fc5eb4bd0..56ab6bcbfe0 100644 --- a/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py @@ -1,89 +1,429 @@ -""" -Unit tests for CLI token utilities -""" - import json -import os -import tempfile -from pathlib import Path -from unittest.mock import mock_open, patch +import stat +import sys +import time import pytest -from litellm.litellm_core_utils.cli_token_utils import get_litellm_gateway_api_key +from litellm.constants import CLI_JWT_EXPIRATION_HOURS +from litellm.litellm_core_utils.cli_keyring import ( + DISABLE_KEYRING_ENV_VAR, + KEYRING_ACCOUNT, + KEYRING_SERVICE, + KeyringVault, + SecretFound, + SecretMissing, + SecretUnavailable, +) +from litellm.litellm_core_utils.cli_token_utils import ( + CliTokenRecord, + clear_cli_token, + get_cli_token_file_path, + get_litellm_gateway_api_key, + is_cli_token_fresh, + load_cli_token, + save_cli_token, +) + +SERVER = "https://proxy.example.com" +OTHER_SERVER = "https://other-proxy.example.com" -class TestCLITokenUtils: - """Test CLI token utility functions""" +@pytest.fixture +def isolated_home(monkeypatch, tmp_path): + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("USERPROFILE", str(tmp_path)) + return tmp_path - def test_get_litellm_gateway_api_key_success(self): - """Test getting CLI API key when token file exists and is valid""" - token_data = { - "key": "sk-test-cli-key-123", - "user_id": "test-user", - "user_email": "test@example.com", - "timestamp": 1234567890, - } - with ( - patch("os.path.exists", return_value=True), - patch("builtins.open", mock_open(read_data=json.dumps(token_data))), - patch( - "litellm.litellm_core_utils.cli_token_utils.get_cli_token_file_path", - return_value="/test/.litellm/token.json", - ), - ): +def _token_file(home): + return home / ".litellm" / "token.json" - result = get_litellm_gateway_api_key() - assert result == "sk-test-cli-key-123" +def _write_legacy_file(home, **overrides): + payload = { + "base_url": SERVER, + "key": "sk-legacy", + "user_id": "u-1", + "user_email": "user@example.com", + "user_role": "cli", + "timestamp": time.time(), + **overrides, + } + path = _token_file(home) + path.parent.mkdir(exist_ok=True) + path.write_text(json.dumps(payload)) + path.chmod(0o600) + return path - def test_get_litellm_gateway_api_key_no_file(self): - """Test getting CLI API key when token file doesn't exist""" - with ( - patch("os.path.exists", return_value=False), - patch( - "litellm.litellm_core_utils.cli_token_utils.get_cli_token_file_path", - return_value="/test/.litellm/token.json", - ), - ): - result = get_litellm_gateway_api_key() +def _write_metadata_only_file(home): + """What a post-migration token.json looks like: everything except the secret material.""" + path = _token_file(home) + path.parent.mkdir(exist_ok=True) + path.write_text(json.dumps({"base_url": SERVER, "user_id": "u-1", "timestamp": time.time()})) + path.chmod(0o600) + return path - assert result is None - def test_get_litellm_gateway_api_key_invalid_json(self): - """Test getting CLI API key when token file has invalid JSON""" - with ( - patch("os.path.exists", return_value=True), - patch("builtins.open", mock_open(read_data="invalid json")), - patch( - "litellm.litellm_core_utils.cli_token_utils.get_cli_token_file_path", - return_value="/test/.litellm/token.json", - ), - ): +def _blob(base_url=SERVER, key="sk-vault", jwt_token=""): + return json.dumps({"base_url": base_url, "key": key, "jwt_token": jwt_token}) - result = get_litellm_gateway_api_key() - assert result is None +class TestGetCliTokenFilePath: + def test_points_at_the_home_config_file(self, isolated_home): + assert get_cli_token_file_path() == str(isolated_home / ".litellm" / "token.json") - def test_get_litellm_gateway_api_key_no_key_field(self): - """Test getting CLI API key when token file exists but has no key field""" - token_data = { - "user_id": "test-user", - "user_email": "test@example.com", - # Missing 'key' field - } + def test_does_not_create_the_directory(self, isolated_home): + """Merely asking for the path must not leave a directory behind, so an SDK import that + never logs in cannot create a ~/.litellm on someone's machine.""" + get_cli_token_file_path() - with ( - patch("os.path.exists", return_value=True), - patch("builtins.open", mock_open(read_data=json.dumps(token_data))), - patch( - "litellm.litellm_core_utils.cli_token_utils.get_cli_token_file_path", - return_value="/test/.litellm/token.json", - ), - ): + assert not (isolated_home / ".litellm").exists() - result = get_litellm_gateway_api_key() - assert result is None +class TestLoadCliToken: + def test_no_token_file_never_touches_the_keychain(self, isolated_home, secret_vault_factory): + """The SDK calls this on machines that never ran `lite login`; it must not prompt for + keychain access there.""" + vault = secret_vault_factory(blob=_blob()) + + assert load_cli_token(vault=vault) is None + assert vault.reads == 0 + + def test_secret_comes_from_the_vault_when_the_file_holds_only_metadata(self, isolated_home, secret_vault_factory): + _write_metadata_only_file(isolated_home) + vault = secret_vault_factory(blob=_blob(key="sk-from-keychain")) + + record = load_cli_token(vault=vault) + + assert record.key == "sk-from-keychain" + assert "sk-from-keychain" not in _token_file(isolated_home).read_text() + + def test_jwt_token_round_trips_through_the_vault(self, isolated_home, secret_vault_factory): + _write_metadata_only_file(isolated_home) + vault = secret_vault_factory(blob=_blob(key="sk-a", jwt_token="jwt-a")) + + record = load_cli_token(vault=vault) + + assert (record.key, record.jwt_token) == ("sk-a", "jwt-a") + + def test_legacy_plaintext_file_still_authenticates_and_is_migrated(self, isolated_home, secret_vault_factory): + """A token.json written by an older `lite` keeps working, and reading it moves the secret + into the keychain and scrubs it from disk.""" + path = _write_legacy_file(isolated_home) + vault = secret_vault_factory() + + record = load_cli_token(vault=vault) + + assert record.key == "sk-legacy" + assert json.loads(vault.blob)["key"] == "sk-legacy" + on_disk = json.loads(path.read_text()) + assert "key" not in on_disk + assert on_disk["user_email"] == "user@example.com" + assert stat.S_IMODE(path.stat().st_mode) == 0o600 + + def test_legacy_file_survives_a_vault_that_refuses_to_store(self, isolated_home, secret_vault_factory): + """Scrubbing the only copy of the secret after a failed keychain write would log the user + out for good.""" + path = _write_legacy_file(isolated_home) + before = path.read_text() + + record = load_cli_token(vault=secret_vault_factory(writable=False)) + + assert record.key == "sk-legacy" + assert path.read_text() == before + + def test_a_secret_left_on_disk_outranks_a_stale_keychain_entry(self, isolated_home, secret_vault_factory): + """A failed keychain write leaves the fresh secret on disk while the vault still holds the + previous one; the next read must serve the file's secret and move it into the vault, never + resurrect the stale key or scrub the only copy of the fresh one.""" + path = _write_legacy_file(isolated_home, key="sk-fresh") + vault = secret_vault_factory(blob=_blob(key="sk-stale")) + + record = load_cli_token(vault=vault) + + assert record.key == "sk-fresh" + assert json.loads(vault.blob)["key"] == "sk-fresh" + assert "key" not in json.loads(path.read_text()) + + def test_a_disk_secret_survives_when_the_stale_vault_refuses_the_rewrite( + self, isolated_home, secret_vault_factory + ): + path = _write_legacy_file(isolated_home, key="sk-fresh") + before = path.read_text() + + record = load_cli_token(vault=secret_vault_factory(blob=_blob(key="sk-stale"), writable=False)) + + assert record.key == "sk-fresh" + assert path.read_text() == before + + def test_legacy_file_survives_an_unreachable_vault_without_write_attempts( + self, isolated_home, secret_vault_factory + ): + path = _write_legacy_file(isolated_home) + before = path.read_text() + vault = secret_vault_factory(available=False) + + record = load_cli_token(vault=vault) + + assert record.key == "sk-legacy" + assert vault.writes == [] + assert path.read_text() == before + + def test_metadata_only_file_with_an_empty_vault_is_not_a_login(self, isolated_home, secret_vault_factory): + _write_metadata_only_file(isolated_home) + + assert load_cli_token(vault=secret_vault_factory()) is None + + def test_metadata_only_file_with_an_unreachable_vault_reports_a_missing_secret( + self, isolated_home, secret_vault_factory + ): + """The caller needs to tell "never logged in" apart from "locked keychain", so the record + comes back with no key rather than as None.""" + _write_metadata_only_file(isolated_home) + + record = load_cli_token(vault=secret_vault_factory(available=False)) + + assert record.key is None + assert record.user_id == "u-1" + + def test_a_secret_minted_for_another_server_is_never_handed_out(self, isolated_home, secret_vault_factory): + _write_metadata_only_file(isolated_home) + + assert load_cli_token(vault=secret_vault_factory(blob=_blob(base_url=OTHER_SERVER))) is None + + def test_a_secret_minted_for_another_server_loses_to_the_file(self, isolated_home, secret_vault_factory): + _write_legacy_file(isolated_home) + vault = secret_vault_factory(blob=_blob(base_url=OTHER_SERVER, key="sk-elsewhere")) + + record = load_cli_token(vault=vault) + + assert record.key == "sk-legacy" + assert json.loads(vault.blob)["key"] == "sk-legacy" + + def test_unreadable_vault_blob_falls_back_to_the_file_secret(self, isolated_home, secret_vault_factory): + _write_legacy_file(isolated_home) + + record = load_cli_token(vault=secret_vault_factory(blob="not json at all {{{")) + + assert record.key == "sk-legacy" + + def test_corrupt_token_file_is_not_a_login(self, isolated_home, secret_vault_factory): + _token_file(isolated_home).parent.mkdir() + _token_file(isolated_home).write_text("not json at all {{{") + + assert load_cli_token(vault=secret_vault_factory(blob=_blob())) is None + + +class TestGetLitellmGatewayApiKey: + def test_returns_the_vault_secret_when_the_origin_matches(self, isolated_home, secret_vault_factory): + _write_metadata_only_file(isolated_home) + + key = get_litellm_gateway_api_key(expected_base_url=SERVER, vault=secret_vault_factory(blob=_blob())) + + assert key == "sk-vault" + + def test_trailing_slash_on_the_expected_url_is_normalised(self, isolated_home, secret_vault_factory): + _write_metadata_only_file(isolated_home) + + key = get_litellm_gateway_api_key(expected_base_url=SERVER + "/", vault=secret_vault_factory(blob=_blob())) + + assert key == "sk-vault" + + def test_origin_mismatch_returns_nothing_without_reading_the_keychain(self, isolated_home, secret_vault_factory): + """Pointing the SDK at a different server must fail before the keychain is even consulted, + so a hostile base_url cannot provoke an unlock prompt.""" + _write_legacy_file(isolated_home) + vault = secret_vault_factory(blob=_blob()) + + assert get_litellm_gateway_api_key(expected_base_url=OTHER_SERVER, vault=vault) is None + assert vault.reads == 0 + + def test_no_token_file_returns_nothing(self, isolated_home, secret_vault_factory): + assert get_litellm_gateway_api_key(vault=secret_vault_factory(blob=_blob())) is None + + +class TestSaveCliToken: + def test_secret_goes_to_the_keychain_and_never_to_the_file(self, isolated_home, secret_vault_factory): + vault = secret_vault_factory() + + stored = save_cli_token( + CliTokenRecord(base_url=SERVER, key="sk-new", user_id="u-1", timestamp=time.time()), + vault=vault, + ) + + assert stored is True + assert "sk-new" not in _token_file(isolated_home).read_text() + assert json.loads(vault.blob)["key"] == "sk-new" + assert load_cli_token(vault=vault).key == "sk-new" + + def test_falls_back_to_the_owner_only_file_when_there_is_no_keychain(self, isolated_home, secret_vault_factory): + stored = save_cli_token( + CliTokenRecord(base_url=SERVER, key="sk-new", timestamp=time.time()), + vault=secret_vault_factory(available=False), + ) + + path = _token_file(isolated_home) + assert stored is False + assert json.loads(path.read_text())["key"] == "sk-new" + assert stat.S_IMODE(path.stat().st_mode) == 0o600 + assert list(path.parent.glob(".tmp-*")) == [] + + def test_creates_the_config_directory_owner_only(self, isolated_home, secret_vault_factory): + """A 0755 ~/.litellm lets any local process list, and in the fallback case read, the + credential's directory.""" + save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-new"), vault=secret_vault_factory()) + + assert stat.S_IMODE((isolated_home / ".litellm").stat().st_mode) == 0o700 + + def test_tightens_a_directory_left_group_readable_by_an_older_cli(self, isolated_home, secret_vault_factory): + config_dir = isolated_home / ".litellm" + config_dir.mkdir(mode=0o755) + + save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-new"), vault=secret_vault_factory()) + + assert stat.S_IMODE(config_dir.stat().st_mode) == 0o700 + + def test_a_failed_write_leaves_the_previous_credential_intact(self, isolated_home, secret_vault_factory, monkeypatch): + path = _write_legacy_file(isolated_home) + before = path.read_text() + + def _explode(*args, **kwargs): + raise TypeError("not serialisable") + + monkeypatch.setattr("litellm.litellm_core_utils.private_json.json.dump", _explode) + + with pytest.raises(TypeError): + save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-new"), vault=secret_vault_factory(available=False)) + + assert path.read_text() == before + assert list(path.parent.glob(".tmp-*")) == [] + + +class TestClearCliToken: + def test_removes_the_credential_from_both_stores(self, isolated_home, secret_vault_factory): + vault = secret_vault_factory() + save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-new"), vault=vault) + + assert clear_cli_token(vault=vault) is True + assert vault.blob is None + assert not _token_file(isolated_home).exists() + assert load_cli_token(vault=vault) is None + + def test_reports_a_keychain_that_will_not_release_the_secret(self, isolated_home, secret_vault_factory): + _write_legacy_file(isolated_home) + vault = secret_vault_factory(blob=_blob(), erasable=False) + + assert clear_cli_token(vault=vault) is False + assert not _token_file(isolated_home).exists() + + def test_is_safe_when_nothing_was_ever_stored(self, isolated_home, secret_vault_factory): + assert clear_cli_token(vault=secret_vault_factory()) is True + + +class TestIsCliTokenFresh: + def test_a_just_issued_token_is_fresh(self): + assert is_cli_token_fresh(CliTokenRecord(timestamp=time.time())) is True + + def test_a_token_past_its_expiry_is_stale(self): + stale = CliTokenRecord(timestamp=time.time() - (CLI_JWT_EXPIRATION_HOURS + 1) * 3600) + + assert is_cli_token_fresh(stale) is False + + def test_the_buffer_retires_a_token_just_before_it_expires(self): + almost = CliTokenRecord(timestamp=time.time() - (CLI_JWT_EXPIRATION_HOURS * 3600 - 60)) + + assert is_cli_token_fresh(almost, buffer_hours=0.1) is False + + +class _FakeKeyringModule: + def __init__(self, stored=None, *, get_error=None, set_error=None, delete_error=None): + self.stored = stored + self.get_error = get_error + self.set_error = set_error + self.delete_error = delete_error + self.calls = [] + + def get_password(self, service_name, username): + self.calls.append(("get", service_name, username)) + if self.get_error is not None: + raise self.get_error + return self.stored + + def set_password(self, service_name, username, password): + self.calls.append(("set", service_name, username)) + if self.set_error is not None: + raise self.set_error + self.stored = password + + def delete_password(self, service_name, username): + self.calls.append(("delete", service_name, username)) + if self.delete_error is not None: + raise self.delete_error + self.stored = None + + +@pytest.fixture +def install_fake_keyring(monkeypatch): + def _install(fake): + monkeypatch.delenv(DISABLE_KEYRING_ENV_VAR, raising=False) + monkeypatch.setitem(sys.modules, "keyring", fake) + return fake + + return _install + + +class TestKeyringVault: + def test_round_trips_through_the_installed_keyring(self, install_fake_keyring): + fake = install_fake_keyring(_FakeKeyringModule()) + vault = KeyringVault() + + assert vault.write("blob-1") is True + assert vault.read() == SecretFound("blob-1") + assert vault.erase() is True + assert vault.read() == SecretMissing() + assert {call[1:] for call in fake.calls} == {(KEYRING_SERVICE, KEYRING_ACCOUNT)} + + def test_the_kill_switch_reports_no_keychain(self, monkeypatch): + """`LITELLM_CLI_DISABLE_KEYRING` has to work without importing keyring, because keyring + caches its backend on first use and cannot be reconfigured later. Erase still fails: a + credential stored before the switch was set may be in the keychain, and with reads + disabled `lite logout` cannot verify it is gone, so it must warn instead.""" + monkeypatch.setenv(DISABLE_KEYRING_ENV_VAR, "1") + vault = KeyringVault() + + assert vault.read() == SecretUnavailable() + assert vault.write("blob-1") is False + assert vault.erase() is False + + def test_an_uninstalled_keyring_library_degrades_to_the_file(self, monkeypatch): + """keyring is an optional extra, so the SDK must survive its absence rather than raise on + the hot path.""" + monkeypatch.delenv(DISABLE_KEYRING_ENV_VAR, raising=False) + monkeypatch.setitem(sys.modules, "keyring", None) + vault = KeyringVault() + + assert vault.read() == SecretUnavailable() + assert vault.write("blob-1") is False + assert vault.erase() is True + + def test_a_locked_keychain_is_reported_not_raised(self, install_fake_keyring): + install_fake_keyring(_FakeKeyringModule(get_error=RuntimeError("keyring is locked"))) + + assert KeyringVault().read() == SecretUnavailable() + + def test_a_refused_write_is_reported_not_raised(self, install_fake_keyring): + install_fake_keyring(_FakeKeyringModule(set_error=RuntimeError("no backend"))) + + assert KeyringVault().write("blob-1") is False + + def test_a_refused_delete_is_reported_so_logout_can_warn(self, install_fake_keyring): + install_fake_keyring(_FakeKeyringModule(stored="blob-1", delete_error=RuntimeError("locked"))) + + assert KeyringVault().erase() is False + + def test_erasing_a_locked_keychain_is_a_failure(self, install_fake_keyring): + install_fake_keyring(_FakeKeyringModule(get_error=RuntimeError("locked"))) + + assert KeyringVault().erase() is False diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py index a23c573047f..c2858c84c6d 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -672,8 +672,9 @@ class TestAgentCommands: assert "LITELLM_PROXY_API_KEY" in result.output mock_run.assert_not_called() - def test_interactive_without_key_logs_in_then_launches(self): + def test_interactive_without_key_logs_in_then_launches(self, secret_vault_factory): captured = {} + vault = secret_vault_factory() @click.command() def fake_login(): @@ -695,12 +696,12 @@ class TestAgentCommands: result = self.runner.invoke( _agent_command("claude"), [], - obj={"base_url": "http://localhost:4000", "api_key": None}, + obj={"base_url": "http://localhost:4000", "api_key": None, "secret_vault": vault}, ) assert result.exit_code == 0, result.output assert captured["api_key"] == "sk-after-login" - mock_get.assert_called_once_with(expected_base_url="http://localhost:4000") + mock_get.assert_called_once_with(expected_base_url="http://localhost:4000", vault=vault) def test_child_exit_code_reaches_the_shell(self): with patch(f"{AGENTS_MODULE}.run_agent", side_effect=SystemExit(42)): diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 59048067674..e93f05cb4aa 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -4,7 +4,7 @@ import stat import sys import time from pathlib import Path -from unittest.mock import Mock, mock_open, patch +from unittest.mock import Mock, patch sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path @@ -13,21 +13,38 @@ import pytest from click.testing import CliRunner from litellm.constants import CLI_JWT_EXPIRATION_HOURS +from litellm.litellm_core_utils.cli_token_utils import CliTokenRecord, save_cli_token from litellm.proxy.client.cli import cli from litellm.proxy.client.cli.commands.auth import ( - clear_token, + KEYCHAIN_UNREACHABLE_MESSAGE, get_stored_api_key, - get_token_file_path, - load_token, login, logout, print_token, - save_token, whoami, ) from litellm.proxy.client.cli.commands.claude_settings import SettingsFileOwner +@pytest.fixture +def isolated_home(monkeypatch, tmp_path): + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("USERPROFILE", str(tmp_path)) + monkeypatch.delenv("LITELLM_PROXY_URL", raising=False) + monkeypatch.delenv("LITELLM_PROXY_API_KEY", raising=False) + return tmp_path + + +def _write_home_json(home: Path, filename: str, payload: dict[str, object]) -> None: + litellm_dir = home / ".litellm" + litellm_dir.mkdir(exist_ok=True) + (litellm_dir / filename).write_text(json.dumps(payload)) + + +def _secret_blob(base_url: str, key: str) -> str: + return json.dumps({"base_url": base_url, "key": key, "jwt_token": ""}) + + def _mock_cli_sso_start_response( login_id: str = "cli-session-uuid-456", poll_secret: str = "poll-secret", @@ -176,200 +193,50 @@ class TestStartCliSsoFlowErrors: assert "https://unreachable.example.com/sso/cli/start" in message -class TestTokenUtilities: - """Test token file utility functions""" +class TestStoredApiKeyLookup: + """`get_stored_api_key` is what every other `lite` subcommand authenticates with, so the + keychain split and the origin check both have to be invisible to it.""" - def test_get_token_file_path(self): - """Test getting token file path""" - with ( - patch("pathlib.Path.home") as mock_home, - patch("pathlib.Path.mkdir") as mock_mkdir, - ): - mock_home.return_value = Path("/home/user") + def test_returns_the_secret_the_keychain_holds(self, isolated_home, secret_vault_factory): + _write_home_json(isolated_home, "token.json", {"base_url": "https://real-proxy.com", "user_id": "u-1"}) + vault = secret_vault_factory(blob=_secret_blob("https://real-proxy.com", "sk-from-keychain")) - result = get_token_file_path() + assert get_stored_api_key(vault=vault) == "sk-from-keychain" - assert result == "/home/user/.litellm/token.json" - mock_mkdir.assert_called_once_with(exist_ok=True) + def test_returns_a_legacy_plaintext_key(self, isolated_home, secret_vault_factory): + _write_home_json(isolated_home, "token.json", {"base_url": "https://real-proxy.com", "key": "sk-legacy"}) - def test_get_token_file_path_creates_directory(self): - """Test that get_token_file_path creates the config directory""" - with ( - patch("pathlib.Path.home") as mock_home, - patch("pathlib.Path.mkdir") as mock_mkdir, - ): - mock_home.return_value = Path("/home/user") + assert get_stored_api_key(vault=secret_vault_factory()) == "sk-legacy" - get_token_file_path() + def test_no_token_at_all_returns_nothing(self, isolated_home, secret_vault_factory): + assert get_stored_api_key(vault=secret_vault_factory()) is None - mock_mkdir.assert_called_once_with(exist_ok=True) + def test_metadata_without_a_secret_returns_nothing(self, isolated_home, secret_vault_factory): + _write_home_json(isolated_home, "token.json", {"base_url": "https://real-proxy.com", "user_id": "u-1"}) - def test_save_token(self, tmp_path): - """Test saving token data to file""" - token_data = { - "key": "test-key", - "user_id": "test-user", - "timestamp": 1234567890, - } - token_file = tmp_path / "token.json" + assert get_stored_api_key(vault=secret_vault_factory()) is None - with patch("litellm.proxy.client.cli.commands.auth.get_token_file_path") as mock_path: - mock_path.return_value = str(token_file) + def test_matching_base_url_returns_the_key(self, isolated_home, secret_vault_factory): + _write_home_json(isolated_home, "token.json", {"base_url": "https://real-proxy.com", "key": "sk-prod"}) - save_token(token_data) + assert get_stored_api_key("https://real-proxy.com", vault=secret_vault_factory()) == "sk-prod" - assert json.loads(token_file.read_text()) == token_data - assert stat.S_IMODE(token_file.stat().st_mode) == 0o600 + def test_trailing_slash_on_the_expected_url_is_normalised(self, isolated_home, secret_vault_factory): + _write_home_json(isolated_home, "token.json", {"base_url": "https://real-proxy.com", "key": "sk-prod"}) - def test_load_token_success(self): - """Test loading token data from file successfully""" - token_data = { - "key": "test-key", - "user_id": "test-user", - "timestamp": 1234567890, - } + assert get_stored_api_key("https://real-proxy.com/", vault=secret_vault_factory()) == "sk-prod" - with ( - patch("builtins.open", mock_open(read_data=json.dumps(token_data))), - patch("litellm.proxy.client.cli.commands.auth.get_token_file_path") as mock_path, - patch("os.path.exists", return_value=True), - ): - mock_path.return_value = "/test/path/token.json" + def test_mismatched_base_url_withholds_the_key(self, isolated_home, secret_vault_factory): + _write_home_json(isolated_home, "token.json", {"base_url": "https://real-proxy.com", "key": "sk-prod"}) - result = load_token() + assert get_stored_api_key("https://evil.com", vault=secret_vault_factory()) is None - assert result == token_data + def test_old_tokens_without_a_base_url_are_rejected_when_an_origin_is_expected( + self, isolated_home, secret_vault_factory + ): + _write_home_json(isolated_home, "token.json", {"key": "sk-old-token"}) - def test_load_token_file_not_exists(self): - """Test loading token when file doesn't exist""" - with ( - patch("litellm.proxy.client.cli.commands.auth.get_token_file_path") as mock_path, - patch("os.path.exists", return_value=False), - ): - mock_path.return_value = "/test/path/token.json" - - result = load_token() - - assert result is None - - def test_load_token_json_decode_error(self): - """Test loading token with invalid JSON""" - with ( - patch("builtins.open", mock_open(read_data="invalid json")), - patch("litellm.proxy.client.cli.commands.auth.get_token_file_path") as mock_path, - patch("os.path.exists", return_value=True), - ): - mock_path.return_value = "/test/path/token.json" - - result = load_token() - - assert result is None - - def test_load_token_io_error(self): - """Test loading token with IO error""" - with ( - patch("builtins.open", side_effect=OSError("Permission denied")), - patch("litellm.proxy.client.cli.commands.auth.get_token_file_path") as mock_path, - patch("os.path.exists", return_value=True), - ): - mock_path.return_value = "/test/path/token.json" - - result = load_token() - - assert result is None - - def test_clear_token_file_exists(self): - """Test clearing token when file exists""" - with ( - patch("litellm.proxy.client.cli.commands.auth.get_token_file_path") as mock_path, - patch("os.path.exists", return_value=True), - patch("os.remove") as mock_remove, - ): - mock_path.return_value = "/test/path/token.json" - - clear_token() - - mock_remove.assert_called_once_with("/test/path/token.json") - - def test_clear_token_file_not_exists(self): - """Test clearing token when file doesn't exist""" - with ( - patch("litellm.proxy.client.cli.commands.auth.get_token_file_path") as mock_path, - patch("os.path.exists", return_value=False), - patch("os.remove") as mock_remove, - ): - mock_path.return_value = "/test/path/token.json" - - clear_token() - - mock_remove.assert_not_called() - - def test_get_stored_api_key_success(self): - """Test getting stored API key successfully""" - token_data = {"key": "test-api-key-123", "user_id": "test-user"} - - with patch( - "litellm.litellm_core_utils.cli_token_utils.load_cli_token", - return_value=token_data, - ): - result = get_stored_api_key() - assert result == "test-api-key-123" - - def test_get_stored_api_key_no_token(self): - """Test getting stored API key when no token exists""" - with patch( - "litellm.litellm_core_utils.cli_token_utils.load_cli_token", - return_value=None, - ): - result = get_stored_api_key() - assert result is None - - def test_get_stored_api_key_no_key_field(self): - """Test getting stored API key when token has no key field""" - token_data = {"user_id": "test-user"} - - with patch( - "litellm.litellm_core_utils.cli_token_utils.load_cli_token", - return_value=token_data, - ): - result = get_stored_api_key() - assert result is None - - def test_get_stored_api_key_base_url_match(self): - """Stored key is returned when expected_base_url matches stored origin""" - token_data = {"key": "sk-prod", "base_url": "https://real-proxy.com"} - with patch( - "litellm.litellm_core_utils.cli_token_utils.load_cli_token", - return_value=token_data, - ): - assert get_stored_api_key(expected_base_url="https://real-proxy.com") == "sk-prod" - - def test_get_stored_api_key_base_url_match_trailing_slash(self): - """Trailing slash on expected_base_url is normalised before comparison""" - token_data = {"key": "sk-prod", "base_url": "https://real-proxy.com"} - with patch( - "litellm.litellm_core_utils.cli_token_utils.load_cli_token", - return_value=token_data, - ): - assert get_stored_api_key(expected_base_url="https://real-proxy.com/") == "sk-prod" - - def test_get_stored_api_key_base_url_mismatch(self): - """Stored key is NOT returned when expected_base_url differs from stored origin""" - token_data = {"key": "sk-prod", "base_url": "https://real-proxy.com"} - with patch( - "litellm.litellm_core_utils.cli_token_utils.load_cli_token", - return_value=token_data, - ): - assert get_stored_api_key(expected_base_url="https://evil.com") is None - - def test_get_stored_api_key_old_token_no_base_url(self): - """Old tokens without a base_url field are rejected when origin check is requested""" - token_data = {"key": "sk-old-token"} - with patch( - "litellm.litellm_core_utils.cli_token_utils.load_cli_token", - return_value=token_data, - ): - assert get_stored_api_key(expected_base_url="https://real-proxy.com") is None + assert get_stored_api_key("https://real-proxy.com", vault=secret_vault_factory()) is None class TestLoginCommand: @@ -402,7 +269,7 @@ class TestLoginCommand: return_value=_mock_cli_sso_start_response(login_id="cli-test-uuid-123"), ) as mock_post, patch("requests.get", return_value=mock_response) as mock_get, - patch("litellm.proxy.client.cli.commands.auth.save_token") as mock_save, + patch("litellm.proxy.client.cli.commands.auth.save_cli_token") as mock_save, patch("litellm.proxy.client.cli.interface.show_commands") as mock_show_commands, ): result = self.runner.invoke(login, obj=mock_context.obj) @@ -424,8 +291,8 @@ class TestLoginCommand: # Verify JWT was saved mock_save.assert_called_once() saved_data = mock_save.call_args[0][0] - assert saved_data["key"] == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt" - assert saved_data["user_id"] == "test-user-123" + assert saved_data.key == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt" + assert saved_data.user_id == "test-user-123" # Verify commands were shown mock_show_commands.assert_called_once() @@ -557,7 +424,7 @@ class TestLogoutCommand: def test_logout_success(self): """Test successful logout""" - with patch("litellm.proxy.client.cli.commands.auth.clear_token") as mock_clear: + with patch("litellm.proxy.client.cli.commands.auth.clear_cli_token") as mock_clear: result = self.runner.invoke(logout) assert result.exit_code == 0 @@ -574,14 +441,15 @@ class TestWhoamiCommand: def test_whoami_authenticated(self): """Test whoami when user is authenticated""" - token_data = { - "user_email": "test@example.com", - "user_id": "test-user-123", - "user_role": "admin", - "timestamp": time.time() - 3600, # 1 hour ago - } + token_data = CliTokenRecord( + user_email="test@example.com", + user_id="test-user-123", + user_role="admin", + key="sk-live", + timestamp=time.time() - 3600, + ) - with patch("litellm.proxy.client.cli.commands.auth.load_token", return_value=token_data): + with patch("litellm.proxy.client.cli.commands.auth.load_cli_token", return_value=token_data): result = self.runner.invoke(whoami) assert result.exit_code == 0 @@ -593,7 +461,7 @@ class TestWhoamiCommand: def test_whoami_not_authenticated(self): """Test whoami when user is not authenticated""" - with patch("litellm.proxy.client.cli.commands.auth.load_token", return_value=None): + with patch("litellm.proxy.client.cli.commands.auth.load_cli_token", return_value=None): result = self.runner.invoke(whoami) assert result.exit_code == 0 @@ -602,14 +470,15 @@ class TestWhoamiCommand: def test_whoami_old_token(self): """Test whoami with old token showing warning""" - token_data = { - "user_email": "test@example.com", - "user_id": "test-user-123", - "user_role": "admin", - "timestamp": time.time() - (25 * 3600), # 25 hours ago - } + token_data = CliTokenRecord( + user_email="test@example.com", + user_id="test-user-123", + user_role="admin", + key="sk-live", + timestamp=time.time() - (25 * 3600), + ) - with patch("litellm.proxy.client.cli.commands.auth.load_token", return_value=token_data): + with patch("litellm.proxy.client.cli.commands.auth.load_cli_token", return_value=token_data): result = self.runner.invoke(whoami) assert result.exit_code == 0 @@ -618,12 +487,9 @@ class TestWhoamiCommand: def test_whoami_missing_fields(self): """Test whoami with token missing some fields""" - token_data = { - "timestamp": time.time() - 3600 - # Missing user_email, user_id, user_role - } + token_data = CliTokenRecord(key="sk-live", timestamp=time.time() - 3600) - with patch("litellm.proxy.client.cli.commands.auth.load_token", return_value=token_data): + with patch("litellm.proxy.client.cli.commands.auth.load_cli_token", return_value=token_data): result = self.runner.invoke(whoami) assert result.exit_code == 0 @@ -632,16 +498,16 @@ class TestWhoamiCommand: def test_whoami_no_timestamp(self): """Test whoami with token missing timestamp""" - token_data = { - "user_email": "test@example.com", - "user_id": "test-user-123", - "user_role": "admin", - # Missing timestamp - } + token_data = CliTokenRecord( + user_email="test@example.com", + user_id="test-user-123", + user_role="admin", + key="sk-live", + ) with ( patch( - "litellm.proxy.client.cli.commands.auth.load_token", + "litellm.proxy.client.cli.commands.auth.load_cli_token", return_value=token_data, ), patch("time.time", return_value=1000), @@ -701,7 +567,7 @@ class TestCLIKeyRegenerationFlow: return_value=_mock_cli_sso_start_response(login_id="cli-session-uuid-456"), ), patch("requests.get", side_effect=[mock_first_response, mock_second_response]) as mock_get, - patch("litellm.proxy.client.cli.commands.auth.save_token") as mock_save, + patch("litellm.proxy.client.cli.commands.auth.save_cli_token") as mock_save, patch("litellm.proxy.client.cli.interface.show_commands") as mock_show_commands, patch("click.prompt", return_value="2"), ): # User selects index 2 @@ -734,8 +600,8 @@ class TestCLIKeyRegenerationFlow: # Verify JWT was saved mock_save.assert_called_once() saved_data = mock_save.call_args[0][0] - assert saved_data["key"] == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.team-beta.jwt" - assert saved_data["user_id"] == "test-user-456" + assert saved_data.key == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.team-beta.jwt" + assert saved_data.user_id == "test-user-456" mock_show_commands.assert_called_once() @@ -762,7 +628,7 @@ class TestCLIKeyRegenerationFlow: return_value=_mock_cli_sso_start_response(login_id="cli-session-uuid-solo"), ), patch("requests.get", return_value=mock_response), - patch("litellm.proxy.client.cli.commands.auth.save_token") as mock_save, + patch("litellm.proxy.client.cli.commands.auth.save_cli_token") as mock_save, patch("litellm.proxy.client.cli.interface.show_commands"), ): result = self.runner.invoke(login, obj=mock_context.obj) @@ -780,8 +646,8 @@ class TestCLIKeyRegenerationFlow: # Verify JWT was saved mock_save.assert_called_once() saved_data = mock_save.call_args[0][0] - assert saved_data["key"] == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.no-team.jwt" - assert saved_data["user_id"] == "test-user-solo" + assert saved_data.key == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.no-team.jwt" + assert saved_data.user_id == "test-user-solo" class TestPrintTokenCommand: @@ -810,7 +676,7 @@ class TestPrintTokenCommand: self.runner = CliRunner() def test_no_stored_token_fails_cleanly(self): - with patch("litellm.proxy.client.cli.commands.auth.load_token", return_value=None): + with patch("litellm.proxy.client.cli.commands.auth.load_cli_token", return_value=None): result = self.runner.invoke(print_token, obj={}) assert result.exit_code != 0 @@ -822,12 +688,12 @@ class TestPrintTokenCommand: one). Must use token.json's own base_url, not a hardcoded default.""" with ( patch( - "litellm.proxy.client.cli.commands.auth.load_token", - return_value={ - "base_url": "https://litellm-proxy.corp.com", - "key": "sk-prod-fresh", - "timestamp": time.time(), - }, + "litellm.proxy.client.cli.commands.auth.load_cli_token", + return_value=CliTokenRecord( + base_url="https://litellm-proxy.corp.com", + key="sk-prod-fresh", + timestamp=time.time(), + ), ), patch("requests.post") as mock_post, ): @@ -844,12 +710,12 @@ class TestPrintTokenCommand: token minted for proxy A must not reach a helper invocation aimed at proxy B, even though the token itself is otherwise fresh.""" with patch( - "litellm.proxy.client.cli.commands.auth.load_token", - return_value={ - "base_url": "https://other-server.com", - "key": "sk-should-not-print", - "timestamp": time.time(), - }, + "litellm.proxy.client.cli.commands.auth.load_cli_token", + return_value=CliTokenRecord( + base_url="https://other-server.com", + key="sk-should-not-print", + timestamp=time.time(), + ), ): result = self.runner.invoke( print_token, @@ -863,12 +729,12 @@ class TestPrintTokenCommand: """`lite up`'s own bound invocation shape: --base-url matching the token's origin must succeed exactly like the bare/legacy invocation does.""" with patch( - "litellm.proxy.client.cli.commands.auth.load_token", - return_value={ - "base_url": "http://localhost:4000", - "key": "sk-matches", - "timestamp": time.time(), - }, + "litellm.proxy.client.cli.commands.auth.load_cli_token", + return_value=CliTokenRecord( + base_url="http://localhost:4000", + key="sk-matches", + timestamp=time.time(), + ), ): result = self.runner.invoke( print_token, @@ -884,12 +750,12 @@ class TestPrintTokenCommand: frequently).""" with ( patch( - "litellm.proxy.client.cli.commands.auth.load_token", - return_value={ - "base_url": "http://localhost:4000", - "key": "sk-cached-fresh", - "timestamp": time.time(), - }, + "litellm.proxy.client.cli.commands.auth.load_cli_token", + return_value=CliTokenRecord( + base_url="http://localhost:4000", + key="sk-cached-fresh", + timestamp=time.time(), + ), ), patch("requests.post") as mock_post, ): @@ -908,12 +774,12 @@ class TestPrintTokenCommand: with ( patch( - "litellm.proxy.client.cli.commands.auth.load_token", - return_value={ - "base_url": "http://localhost:4000", - "key": "sk-stale-key", - "timestamp": old_timestamp, - }, + "litellm.proxy.client.cli.commands.auth.load_cli_token", + return_value=CliTokenRecord( + base_url="http://localhost:4000", + key="sk-stale-key", + timestamp=old_timestamp, + ), ), patch("requests.post") as mock_post, ): @@ -925,25 +791,11 @@ class TestPrintTokenCommand: mock_post.assert_not_called() -def _write_home_json(home: Path, filename: str, payload: dict[str, object]) -> None: - litellm_dir = home / ".litellm" - litellm_dir.mkdir(exist_ok=True) - (litellm_dir / filename).write_text(json.dumps(payload)) - - class TestPrintTokenWithConfigFile: """A config-file base_url is a drop-in replacement for exporting LITELLM_PROXY_URL, so print-token must treat it as an explicit server choice: a token minted for a different proxy is never handed out.""" - @pytest.fixture - def isolated_home(self, monkeypatch, tmp_path): - monkeypatch.setenv("HOME", str(tmp_path)) - monkeypatch.setenv("USERPROFILE", str(tmp_path)) - monkeypatch.delenv("LITELLM_PROXY_URL", raising=False) - monkeypatch.delenv("LITELLM_PROXY_API_KEY", raising=False) - return tmp_path - def test_config_base_url_mismatch_fails_closed(self, isolated_home): _write_home_json( isolated_home, @@ -1001,37 +853,196 @@ class TestPrintTokenWithConfigFile: assert result.stdout.strip() == "sk-issued-for-a" -class TestSaveTokenPrivateWrite: - """token.json holds the real API key: it must never be world-readable at any - instant, and a failed write must not destroy the previously stored token.""" +class TestFileFallbackStorage: + """On a headless box with no keychain the token file is still the only store, so it has to + stay owner-only and survive a failed write.""" - @pytest.fixture - def isolated_home(self, monkeypatch, tmp_path): - monkeypatch.setenv("HOME", str(tmp_path)) - monkeypatch.setenv("USERPROFILE", str(tmp_path)) - monkeypatch.delenv("LITELLM_PROXY_URL", raising=False) - monkeypatch.delenv("LITELLM_PROXY_API_KEY", raising=False) - return tmp_path - - def test_save_token_owner_only_permissions_and_no_temp_leftovers(self, isolated_home): - save_token({"key": "sk-secret", "user_id": "u-1", "timestamp": 1234567890}) + def test_owner_only_file_and_directory_with_no_temp_leftovers(self, isolated_home, secret_vault_factory): + save_cli_token( + CliTokenRecord(base_url="https://proxy.example.com", key="sk-secret", user_id="u-1", timestamp=1234567890), + vault=secret_vault_factory(available=False), + ) token_file = isolated_home / ".litellm" / "token.json" - assert json.loads(token_file.read_text()) == {"key": "sk-secret", "user_id": "u-1", "timestamp": 1234567890} + assert json.loads(token_file.read_text())["key"] == "sk-secret" assert stat.S_IMODE(token_file.stat().st_mode) == 0o600 + assert stat.S_IMODE(token_file.parent.stat().st_mode) == 0o700 assert list(token_file.parent.glob(".tmp-*")) == [] - def test_save_token_failure_mid_write_preserves_existing_token(self, isolated_home): + def test_a_failed_write_preserves_the_existing_token(self, isolated_home, secret_vault_factory, monkeypatch): _write_home_json(isolated_home, "token.json", {"key": "sk-original", "timestamp": 1234567890}) token_file = isolated_home / ".litellm" / "token.json" + def _explode(*args, **kwargs): + raise TypeError("not serialisable") + + monkeypatch.setattr("litellm.litellm_core_utils.private_json.json.dump", _explode) + with pytest.raises(TypeError): - save_token({"key": object()}) + save_cli_token(CliTokenRecord(key="sk-new"), vault=secret_vault_factory(available=False)) assert json.loads(token_file.read_text()) == {"key": "sk-original", "timestamp": 1234567890} assert list(token_file.parent.glob(".tmp-*")) == [] +class TestKeychainBackedCommands: + """End-to-end through the `lite` commands: the secret lives in the keychain, the file keeps + only metadata, and every command still reads and writes through that split.""" + + def setup_method(self): + self.runner = CliRunner() + + def _login(self, vault, base_url="https://test.example.com"): + poll_response = Mock() + poll_response.status_code = 200 + poll_response.json.return_value = { + "status": "ready", + "key": "sk-minted", + "user_id": "test-user-123", + "team_id": "team-1", + "teams": ["team-1"], + } + with ( + patch("webbrowser.open"), + patch("requests.post", return_value=_mock_cli_sso_start_response()), + patch("requests.get", return_value=poll_response), + patch("litellm.proxy.client.cli.interface.show_commands"), + ): + return self.runner.invoke(login, obj={"base_url": base_url, "secret_vault": vault}) + + def test_login_puts_the_secret_in_the_keychain_and_not_in_the_file(self, isolated_home, secret_vault_factory): + vault = secret_vault_factory() + + result = self._login(vault) + + token_file = isolated_home / ".litellm" / "token.json" + assert result.exit_code == 0 + assert "Credential stored in your OS keychain." in result.output + assert json.loads(vault.blob)["key"] == "sk-minted" + assert "sk-minted" not in token_file.read_text() + assert json.loads(token_file.read_text())["user_id"] == "test-user-123" + + def test_login_without_a_keychain_says_where_the_credential_went(self, isolated_home, secret_vault_factory): + result = self._login(secret_vault_factory(available=False)) + + token_file = isolated_home / ".litellm" / "token.json" + assert result.exit_code == 0 + assert "No OS keychain available" in result.output + assert str(token_file) in result.output + assert json.loads(token_file.read_text())["key"] == "sk-minted" + + def test_whoami_and_print_token_read_through_the_keychain(self, isolated_home, secret_vault_factory): + vault = secret_vault_factory() + self._login(vault) + obj = {"base_url": "https://test.example.com", "secret_vault": vault} + + whoami_result = self.runner.invoke(whoami, obj=obj) + print_result = self.runner.invoke(print_token, obj=obj) + + assert "Authenticated" in whoami_result.output + assert "test-user-123" in whoami_result.output + assert print_result.exit_code == 0 + assert print_result.stdout.strip() == "sk-minted" + + def test_logout_clears_the_keychain_as_well_as_the_file(self, isolated_home, secret_vault_factory): + vault = secret_vault_factory() + self._login(vault) + + result = self.runner.invoke(logout, obj={"base_url": "https://test.example.com", "secret_vault": vault}) + + assert result.exit_code == 0 + assert "Logged out successfully" in result.output + assert vault.blob is None + assert not (isolated_home / ".litellm" / "token.json").exists() + + def test_logout_warns_when_the_keychain_will_not_release_the_secret(self, isolated_home, secret_vault_factory): + """Silently reporting success would leave a live credential in the keychain.""" + vault = secret_vault_factory(erasable=False) + self._login(vault) + + result = self.runner.invoke(logout, obj={"base_url": "https://test.example.com", "secret_vault": vault}) + + assert result.exit_code == 0 + assert "could not be removed" in result.output + assert not (isolated_home / ".litellm" / "token.json").exists() + + def test_print_token_explains_a_locked_keychain_instead_of_printing_nothing( + self, isolated_home, secret_vault_factory + ): + _write_home_json( + isolated_home, + "token.json", + {"base_url": "https://test.example.com", "user_id": "u-1", "timestamp": time.time()}, + ) + obj = {"base_url": "https://test.example.com", "secret_vault": secret_vault_factory(available=False)} + + result = self.runner.invoke(print_token, obj=obj) + + assert result.exit_code == 1 + assert KEYCHAIN_UNREACHABLE_MESSAGE in result.output + + def test_whoami_flags_a_locked_keychain(self, isolated_home, secret_vault_factory): + _write_home_json( + isolated_home, + "token.json", + {"base_url": "https://test.example.com", "user_id": "u-1", "timestamp": time.time()}, + ) + obj = {"base_url": "https://test.example.com", "secret_vault": secret_vault_factory(available=False)} + + result = self.runner.invoke(whoami, obj=obj) + + assert "Authenticated" in result.output + assert KEYCHAIN_UNREACHABLE_MESSAGE in result.output + + +class TestApiKeyPrecedence: + """`LITELLM_PROXY_API_KEY` and `--api-key` outrank the stored credential; moving the secret + into the keychain must not disturb that order.""" + + def _resolved_key(self, args, obj=None): + with patch("litellm.proxy.client.cli.main.print_version") as mock_print_version: + result = CliRunner().invoke(cli, [*args, "version"], obj=obj) + assert result.exit_code == 0, result.output + return mock_print_version.call_args[0][1] + + def test_the_stored_credential_is_the_fallback(self, isolated_home): + _write_home_json( + isolated_home, + "token.json", + {"base_url": "http://localhost:4000", "key": "sk-stored", "timestamp": time.time()}, + ) + + assert self._resolved_key([]) == "sk-stored" + + def test_the_stored_credential_is_read_through_the_injected_keychain(self, isolated_home, secret_vault_factory): + """The vault handed to the CLI through ctx.obj must be the one the group callback reads, + so a keychain-held secret resolves without ever touching the host OS keychain.""" + _write_home_json(isolated_home, "token.json", {"base_url": "http://localhost:4000", "timestamp": time.time()}) + vault = secret_vault_factory(_secret_blob("http://localhost:4000", "sk-keychain")) + + assert self._resolved_key([], obj={"secret_vault": vault}) == "sk-keychain" + + def test_env_var_beats_the_stored_credential(self, isolated_home, monkeypatch): + _write_home_json( + isolated_home, + "token.json", + {"base_url": "http://localhost:4000", "key": "sk-stored", "timestamp": time.time()}, + ) + monkeypatch.setenv("LITELLM_PROXY_API_KEY", "sk-from-env") + + assert self._resolved_key([]) == "sk-from-env" + + def test_explicit_api_key_beats_both(self, isolated_home, monkeypatch): + _write_home_json( + isolated_home, + "token.json", + {"base_url": "http://localhost:4000", "key": "sk-stored", "timestamp": time.time()}, + ) + monkeypatch.setenv("LITELLM_PROXY_API_KEY", "sk-from-env") + + assert self._resolved_key(["--api-key", "sk-explicit"]) == "sk-explicit" + + class TestLoginConfigClaude: """`lite login --config-claude` wiring into ~/.claude/settings.json""" @@ -1054,7 +1065,7 @@ class TestLoginConfigClaude: patch("webbrowser.open"), patch("requests.post", return_value=_mock_cli_sso_start_response()), patch("requests.get", return_value=poll_response), - patch("litellm.proxy.client.cli.commands.auth.save_token"), + patch("litellm.proxy.client.cli.commands.auth.save_cli_token"), patch("litellm.proxy.client.cli.interface.show_commands"), patch("litellm.proxy.client.cli.commands.auth.CLAUDE_SETTINGS_PATH", settings_path), patch( diff --git a/tests/test_litellm/proxy/client/cli/test_claude_settings.py b/tests/test_litellm/proxy/client/cli/test_claude_settings.py index bc9744eb410..9010fb4c022 100644 --- a/tests/test_litellm/proxy/client/cli/test_claude_settings.py +++ b/tests/test_litellm/proxy/client/cli/test_claude_settings.py @@ -7,6 +7,7 @@ from unittest.mock import patch import pytest from click.testing import CliRunner +from litellm.litellm_core_utils.cli_token_utils import CliTokenRecord from litellm.proxy.client.cli import cli from litellm.proxy.client.cli.commands.claude_settings import ( AUTOROUTE_BACKUP_PATH, @@ -181,18 +182,18 @@ class TestApiKeyHelperIsActuallyInvocable: assert result.exit_code != 2 def test_the_generated_command_reaches_print_token(self): - with patch(f"{AUTH_MODULE}.load_token", return_value=None): + with patch(f"{AUTH_MODULE}.load_cli_token", return_value=None): result = CliRunner().invoke(cli, self._helper_args("http://localhost:4000")) assert "Not authenticated" in result.output def test_the_generated_command_carries_the_base_url_through(self): - stale = { - "base_url": "http://other-proxy.example.com", - "key": "sk-stale", - "timestamp": time.time(), - } - with patch(f"{AUTH_MODULE}.load_token", return_value=stale): + stale = CliTokenRecord( + base_url="http://other-proxy.example.com", + key="sk-stale", + timestamp=time.time(), + ) + with patch(f"{AUTH_MODULE}.load_cli_token", return_value=stale): result = CliRunner().invoke(cli, self._helper_args("http://localhost:4000")) assert "Not authenticated for this server" in result.output diff --git a/tests/test_litellm/proxy/client/cli/test_config_commands.py b/tests/test_litellm/proxy/client/cli/test_config_commands.py index d81ee6bd2b1..6f3f4e4b268 100644 --- a/tests/test_litellm/proxy/client/cli/test_config_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_config_commands.py @@ -18,7 +18,7 @@ from litellm.proxy.client.cli.commands.config import ( load_config, save_config, ) -from litellm.proxy.client.cli.commands.private_json import write_private_json +from litellm.litellm_core_utils.private_json import write_private_json from litellm.proxy.client.cli.interface import show_commands @@ -355,7 +355,7 @@ class TestWritePrivateJson: def _interrupt(*args: object, **kwargs: object) -> None: raise KeyboardInterrupt() - monkeypatch.setattr("litellm.proxy.client.cli.commands.private_json.json.dump", _interrupt) + monkeypatch.setattr("litellm.litellm_core_utils.private_json.json.dump", _interrupt) target = tmp_path / "config.json" with pytest.raises(KeyboardInterrupt): diff --git a/tests/test_litellm/proxy/client/cli/test_up_commands.py b/tests/test_litellm/proxy/client/cli/test_up_commands.py index 51de0dcf11d..aebf441f777 100644 --- a/tests/test_litellm/proxy/client/cli/test_up_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_up_commands.py @@ -8,6 +8,7 @@ import click import pytest from click.testing import CliRunner +from litellm.litellm_core_utils.cli_token_utils import CliTokenRecord from litellm.proxy.client.cli.commands import up as up_module from litellm.proxy.client.cli.commands.agents import AgentRunError from litellm.proxy.client.cli.commands.claude_settings import ClaudeSettingsError @@ -220,13 +221,17 @@ def _make_ctx(base_url): return click.Context(click.Command("test"), obj={"base_url": base_url}) +def _token(key, base_url): + return CliTokenRecord(key=key, base_url=base_url) + + class TestEnsureFreshLogin: """A token that is fresh but was issued for a *different* proxy must not be trusted: without this check, a user logged into proxy A who runs `up --base-url proxy-b` would silently get an apiKeyHelper wired up around proxy A's real token, which print-token would then hand to proxy B.""" def test_reuses_a_fresh_token_issued_for_the_same_proxy(self, monkeypatch): - monkeypatch.setattr(up_module, "load_token", lambda: {"key": "sk-a", "base_url": "http://proxy-a:4000"}) + monkeypatch.setattr(up_module, "load_cli_token", lambda **_: _token("sk-a", "http://proxy-a:4000")) monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True) login_calls = [] monkeypatch.setattr(up_module, "login", lambda ctx: login_calls.append(ctx)) @@ -239,11 +244,11 @@ class TestEnsureFreshLogin: monkeypatch.setattr(up_module.sys.stdin, "isatty", lambda: True) tokens = iter( [ - {"key": "sk-a", "base_url": "http://proxy-a:4000"}, - {"key": "sk-b", "base_url": "http://proxy-b:4000"}, + _token("sk-a", "http://proxy-a:4000"), + _token("sk-b", "http://proxy-b:4000"), ] ) - monkeypatch.setattr(up_module, "load_token", lambda: next(tokens)) + monkeypatch.setattr(up_module, "load_cli_token", lambda **_: next(tokens)) monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True) login_calls = [] @@ -259,7 +264,7 @@ class TestEnsureFreshLogin: def test_fails_cleanly_non_interactively_when_only_a_different_proxys_token_is_cached(self, monkeypatch): monkeypatch.setattr(up_module.sys.stdin, "isatty", lambda: False) - monkeypatch.setattr(up_module, "load_token", lambda: {"key": "sk-a", "base_url": "http://proxy-a:4000"}) + monkeypatch.setattr(up_module, "load_cli_token", lambda **_: _token("sk-a", "http://proxy-a:4000")) monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True) with pytest.raises(UpError, match="lite login"): @@ -276,7 +281,7 @@ class TestUpCommand: backup_path.write_text(json.dumps(existing_backup)) with ( - patch(f"{UP_MODULE}.load_token", return_value={"key": "sk-fresh", "base_url": "http://localhost:4000"}), + patch(f"{UP_MODULE}.load_cli_token", return_value=_token("sk-fresh", "http://localhost:4000")), patch(f"{UP_MODULE}.is_cli_token_fresh", return_value=True), patch(f"{UP_MODULE}.resolve_api_key", return_value="sk-fresh"), patch(f"{UP_MODULE}.verify_proxy_key"), @@ -293,7 +298,7 @@ class TestUpCommand: _patch_paths(monkeypatch, tmp_path) monkeypatch.setattr(sys.stdin, "isatty", lambda: False) - with patch(f"{UP_MODULE}.load_token", return_value=None): + with patch(f"{UP_MODULE}.load_cli_token", return_value=None): result = self.runner.invoke(up, obj={"base_url": "http://localhost:4000"}) assert result.exit_code != 0 @@ -303,7 +308,7 @@ class TestUpCommand: _patch_paths(monkeypatch, tmp_path) with ( - patch(f"{UP_MODULE}.load_token", return_value={"key": "sk-fresh", "base_url": "http://localhost:4000"}), + patch(f"{UP_MODULE}.load_cli_token", return_value=_token("sk-fresh", "http://localhost:4000")), patch(f"{UP_MODULE}.is_cli_token_fresh", return_value=True), patch(f"{UP_MODULE}.resolve_api_key", return_value="sk-fresh"), patch( @@ -329,7 +334,7 @@ class TestUpCommand: return True with ( - patch(f"{UP_MODULE}.load_token", return_value={"key": "sk-fresh", "base_url": "http://localhost:4000"}), + patch(f"{UP_MODULE}.load_cli_token", return_value=_token("sk-fresh", "http://localhost:4000")), patch(f"{UP_MODULE}.is_cli_token_fresh", return_value=True), patch(f"{UP_MODULE}.resolve_api_key", return_value="sk-fresh"), patch(f"{UP_MODULE}.verify_proxy_key"), diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 31726acfbaa..a5c5a9f135b 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22806 + "limit": 22805 }, "LIT002": { - "limit": 26878 + "limit": 26877 }, "LIT003": { "limit": 269 @@ -27,12 +27,12 @@ "limit": 0 }, "LIT010": { - "limit": 16695 + "limit": 16693 }, "LIT011": { "limit": 5588 }, "LIT012": { - "limit": 4519 + "limit": 4511 } } diff --git a/uv.lock b/uv.lock index d9e2fb94667..53bb0cb8f82 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-08-16T00:41:08.185444Z" +exclude-newer = "2026-08-17T01:06:38.502388Z" exclude-newer-span = "P3D" [manifest] @@ -710,6 +710,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a0/59/76ab57e3fe74484f48a53f8e337171b4a2349e506eabe136d7e01d059086/backports_asyncio_runner-1.2.0-py3-none-any.whl", hash = "sha256:0da0a936a8aeb554eccb426dc55af3ba63bcdc69fa1a600b5bb305413a4477b5", size = 12313, upload-time = "2025-07-02T02:27:14.263Z" }, ] +[[package]] +name = "backports-tarfile" +version = "1.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/86/72/cd9b395f25e290e633655a100af28cb253e4393396264a98bd5f5951d50f/backports_tarfile-1.2.0.tar.gz", hash = "sha256:d75e02c268746e1b8144c278978b6e98e85de6ad16f8e4b0844a154557eca991", size = 86406, upload-time = "2024-05-28T17:01:54.731Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b9/fa/123043af240e49752f1c4bd24da5053b6bd00cad78c2be53c0d1e8b975bc/backports.tarfile-1.2.0-py3-none-any.whl", hash = "sha256:77e284d754527b01fb1e6fa8a1afe577858ebe4e9dad8919e34c862cb399bc34", size = 30181, upload-time = "2024-05-28T17:01:53.112Z" }, +] + [[package]] name = "basedpyright" version = "1.39.7" @@ -3538,6 +3547,51 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/04/96/92447566d16df59b2a776c0fb82dbc4d9e07cd95062562af01e408583fc4/itsdangerous-2.2.0-py3-none-any.whl", hash = "sha256:c6242fc49e35958c8b15141343aa660db5fc54d4f13a1db01a3f5891b98700ef", size = 16234, upload-time = "2024-04-16T21:28:14.499Z" }, ] +[[package]] +name = "jaraco-classes" +version = "3.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "more-itertools" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/06/c0/ed4a27bc5571b99e3cff68f8a9fa5b56ff7df1c2251cc715a652ddd26402/jaraco.classes-3.4.0.tar.gz", hash = "sha256:47a024b51d0239c0dd8c8540c6c7f484be3b8fcf0b2d85c13825780d3b3f3acd", size = 11780, upload-time = "2024-03-31T07:27:36.643Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7f/66/b15ce62552d84bbfcec9a4873ab79d993a1dd4edb922cbfccae192bd5b5f/jaraco.classes-3.4.0-py3-none-any.whl", hash = "sha256:f662826b6bed8cace05e7ff873ce0f9283b5c924470fe664fff1c2f00f581790", size = 6777, upload-time = "2024-03-31T07:27:34.792Z" }, +] + +[[package]] +name = "jaraco-context" +version = "6.1.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "backports-tarfile", marker = "python_full_version < '3.12'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/af/50/4763cd07e722bb6285316d390a164bc7e479db9d90daa769f22578f698b4/jaraco_context-6.1.2.tar.gz", hash = "sha256:f1a6c9d391e661cc5b8d39861ff077a7dc24dc23833ccee564b234b81c82dfe3", size = 16801, upload-time = "2026-03-20T22:13:33.922Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f2/58/bc8954bda5fcda97bd7c19be11b85f91973d67a706ed4a3aec33e7de22db/jaraco_context-6.1.2-py3-none-any.whl", hash = "sha256:bf8150b79a2d5d91ae48629d8b427a8f7ba0e1097dd6202a9059f29a36379535", size = 7871, upload-time = "2026-03-20T22:13:32.808Z" }, +] + +[[package]] +name = "jaraco-functools" +version = "4.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "more-itertools" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6c/1f/c23395957d41ccf27c4e535c3d334c4051e5395b3752057ba4cbaec35c56/jaraco_functools-4.6.0.tar.gz", hash = "sha256:880c577ec9720b3a052d5bc611fb9f2269b3d87902ef42440df443b88e443280", size = 20837, upload-time = "2026-07-14T01:28:02.544Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/02/36/ecc85bc96c273dc8a11273ed4782272975e6338d4a3e9228621175edf0e3/jaraco_functools-4.6.0-py3-none-any.whl", hash = "sha256:99e3dc0060c5cbe8fcd1cdb36258e2a65ca40f1566b2033b12abb1bb44dd3c30", size = 11677, upload-time = "2026-07-14T01:28:01.59Z" }, +] + +[[package]] +name = "jeepney" +version = "0.9.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7b/6f/357efd7602486741aa73ffc0617fb310a29b588ed0fd69c2399acbb85b0c/jeepney-0.9.0.tar.gz", hash = "sha256:cf0e9e845622b81e4a28df94c40345400256ec608d0e55bb8a3feaa9163f5732", size = 106758, upload-time = "2025-02-27T18:51:01.684Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b2/a3/e137168c9c44d18eff0376253da9f1e9234d0239e0ee230d2fee6cea8e55/jeepney-0.9.0-py3-none-any.whl", hash = "sha256:97e5714520c16fc0a45695e5365a2e11b81ea79bba796e26f9f1d178cb182683", size = 49010, upload-time = "2025-02-27T18:51:00.104Z" }, +] + [[package]] name = "jinja2" version = "3.1.6" @@ -3764,6 +3818,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/41/45/1a4ed80516f02155c51f51e8cedb3c1902296743db0bbc66608a0db2814f/jsonschema_specifications-2025.9.1-py3-none-any.whl", hash = "sha256:98802fee3a11ee76ecaca44429fda8a41bff98b00a0f2838151b113f210cc6fe", size = 18437, upload-time = "2025-09-08T01:34:57.871Z" }, ] +[[package]] +name = "keyring" +version = "25.7.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "importlib-metadata", marker = "python_full_version < '3.12'" }, + { name = "jaraco-classes" }, + { name = "jaraco-context" }, + { name = "jaraco-functools" }, + { name = "jeepney", marker = "sys_platform == 'linux'" }, + { name = "pywin32-ctypes", marker = "sys_platform == 'win32'" }, + { name = "secretstorage", marker = "sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/43/4b/674af6ef2f97d56f0ab5153bf0bfa28ccb6c3ed4d1babf4305449668807b/keyring-25.7.0.tar.gz", hash = "sha256:fe01bd85eb3f8fb3dd0405defdeac9a5b4f6f0439edbb3149577f244a2e8245b", size = 63516, upload-time = "2025-11-16T16:26:09.482Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/81/db/e655086b7f3a705df045bf0933bdd9c2f79bb3c97bfef1384598bb79a217/keyring-25.7.0-py3-none-any.whl", hash = "sha256:be4a0b195f149690c166e850609a477c532ddbfbaed96a404d4e43f8d5e2689f", size = 39160, upload-time = "2025-11-16T16:26:08.402Z" }, +] + [[package]] name = "kiwisolver" version = "1.5.0" @@ -4222,6 +4294,7 @@ caching = [ ] cli = [ { name = "inquirerpy" }, + { name = "keyring" }, { name = "pyyaml" }, { name = "requests" }, { name = "rich" }, @@ -4352,6 +4425,7 @@ dev = [ { name = "diff-cover" }, { name = "fakeredis" }, { name = "fastapi-offline" }, + { name = "keyring" }, { name = "langfuse" }, { name = "openapi-core" }, { name = "opentelemetry-api" }, @@ -4446,6 +4520,7 @@ requires-dist = [ { name = "inquirerpy", marker = "extra == 'proxy'", specifier = ">=0.3.4,<1.0" }, { name = "jinja2", specifier = ">=3.1.6,<4.0" }, { name = "jsonschema", specifier = ">=4.0.0,<5.0" }, + { name = "keyring", marker = "extra == 'cli'", specifier = ">=25.6.0,<26.0" }, { name = "langfuse", marker = "extra == 'proxy-runtime'", specifier = ">=2.59.7,<3.0" }, { name = "litellm-enterprise", marker = "extra == 'proxy'", editable = "enterprise" }, { name = "litellm-proxy-extras", marker = "extra == 'proxy'", editable = "litellm-proxy-extras" }, @@ -4532,6 +4607,7 @@ dev = [ { name = "diff-cover", specifier = "==9.7.2" }, { name = "fakeredis", specifier = "==2.34.1" }, { name = "fastapi-offline", specifier = "==1.7.6" }, + { name = "keyring", specifier = "==25.7.0" }, { name = "langfuse", specifier = "==2.59.7" }, { name = "openapi-core", specifier = "==0.22.0" }, { name = "opentelemetry-api", specifier = "==1.28.0" }, @@ -7790,6 +7866,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c0/d2/21af5c535501a7233e734b8af901574572da66fcc254cb35d0609c9080dd/pywin32-311-cp314-cp314-win_arm64.whl", hash = "sha256:a508e2d9025764a8270f93111a970e1d0fbfc33f4153b388bb649b7eec4f9b42", size = 8932540, upload-time = "2025-07-14T20:13:36.379Z" }, ] +[[package]] +name = "pywin32-ctypes" +version = "0.2.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/85/9f/01a1a99704853cb63f253eea009390c88e7131c67e66a0a02099a8c917cb/pywin32-ctypes-0.2.3.tar.gz", hash = "sha256:d162dc04946d704503b2edc4d55f3dba5c1d539ead017afa00142c38b9885755", size = 29471, upload-time = "2024-08-14T10:15:34.626Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/de/3d/8161f7711c017e01ac9f008dfddd9410dff3674334c233bde66e7ba65bbf/pywin32_ctypes-0.2.3-py3-none-any.whl", hash = "sha256:8a1513379d709975552d202d942d9837758905c8d01eb82b8bcc30918929e7b8", size = 30756, upload-time = "2024-08-14T10:15:33.187Z" }, +] + [[package]] name = "pyyaml" version = "6.0.3" @@ -8632,6 +8717,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/07/39/338d9219c4e87f3e708f18857ecd24d22a0c3094752393319553096b98af/scipy-1.17.1-cp314-cp314t-win_arm64.whl", hash = "sha256:200e1050faffacc162be6a486a984a0497866ec54149a01270adc8a59b7c7d21", size = 25489165, upload-time = "2026-02-23T00:22:29.563Z" }, ] +[[package]] +name = "secretstorage" +version = "3.5.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cryptography" }, + { name = "jeepney" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1c/03/e834bcd866f2f8a49a85eaff47340affa3bfa391ee9912a952a1faa68c7b/secretstorage-3.5.0.tar.gz", hash = "sha256:f04b8e4689cbce351744d5537bf6b1329c6fc68f91fa666f60a380edddcd11be", size = 19884, upload-time = "2025-11-23T19:02:53.191Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/46/f5af3402b579fd5e11573ce652019a67074317e18c1935cc0b4ba9b35552/secretstorage-3.5.0-py3-none-any.whl", hash = "sha256:0ce65888c0725fcb2c5bc0fdb8e5438eece02c523557ea40ce0703c266248137", size = 15554, upload-time = "2025-11-23T19:02:51.545Z" }, +] + [[package]] name = "semantic-router" version = "0.1.15" From c2b3c4b1e49e4ffde1131892e929f6fc9193f751 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Wed, 19 Aug 2026 19:08:14 -0700 Subject: [PATCH 032/281] feat(ptu): accrue flat cost for PTU deployments declared in config.yaml (#37556) * feat(ptu): accrue flat cost for PTU deployments declared in config.yaml The flat-cost rollup reads deployments from LiteLLM_ProxyModelTable, and config.yaml models never reach that table by design, so a PTU deployment declared there accrued no flat cost at all while still billing its traffic per token. The provider bills the reservation whichever file declared it. The rollup now also reads the deployments the router holds that no database row owns, identified by db_model, skipping the per-request credential clones that carry original_model_id and reuse their source's PTU config under a fresh id. Registering such a deployment zeroes its pricing, since reserved capacity already pays for the traffic it serves, and leaving a rate unset falls back to the public cost map, which makes the double charge the default rather than an opt-in. The rules both halves apply now live in one module. The rollup's test for what it will charge and the router's test for what to zero have to agree, or a deployment one accepts and the other declines serves its traffic for free. That module also owns the fields the write endpoints already zero, so the two paths cannot drift: tiered_pricing is emptied rather than zeroed because its tiers outrank the rates beside them, the search context table is written zeroed because an absent one means the provider default, and any further rate the deployment itself declares is zeroed alongside the standing set. The prune is bounded to the deployments a run scanned, but only for a run that priced a config-declared deployment. Deciding a row is garbage on staleness alone stays correct while every run derives its charges from the same table, so a database-only run sweeps exactly as it did before; once one host's charges come from a file the others cannot read, a row it never considered is not evidence of anything. Behaviour change worth calling out: a zeroed deployment sorts ahead of an unpriced sibling in QualityRouter's cost tiebreak, where an unset rate previously sorted last. Reserved capacity really is the cheaper choice, but the ordering moves. * refactor(ptu): drop a Final rebind and two redundant isinstance guards The basedpyright budget rejected reassigning a Final in the datetime coercion and two isinstance calls the router entry's own type already guarantees. Filtering the built records rather than the raw entries removes both guards and leaves _router_deployment as the single validator. --- litellm/litellm_core_utils/ptu_pricing.py | 149 +++++++++ .../model_management_endpoints.py | 22 +- .../proxy/spend_tracking/ptu_feature_flag.py | 22 +- .../spend_tracking/ptu_flat_cost_rollup.py | 219 ++++++++----- litellm/router.py | 12 +- .../litellm_core_utils/test_ptu_pricing.py | 163 ++++++++++ .../test_ptu_flat_cost_rollup.py | 305 +++++++++++++++++- .../test_router_model_cost_isolation.py | 113 +++++++ 8 files changed, 889 insertions(+), 116 deletions(-) create mode 100644 litellm/litellm_core_utils/ptu_pricing.py create mode 100644 tests/test_litellm/litellm_core_utils/test_ptu_pricing.py diff --git a/litellm/litellm_core_utils/ptu_pricing.py b/litellm/litellm_core_utils/ptu_pricing.py new file mode 100644 index 00000000000..a1f8bb36e27 --- /dev/null +++ b/litellm/litellm_core_utils/ptu_pricing.py @@ -0,0 +1,149 @@ +"""Which deployments accrue PTU flat cost, and what that costs them per token. + +Reserved provisioned throughput is billed by the hour whether or not requests are sent, so +a deployment that accrues flat cost must not also bill per token. The two halves live here +together because they have to agree: a deployment the rollup declines to charge but the +router prices at zero serves its traffic for free. +""" + +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import datetime, timezone +from types import MappingProxyType +from typing import Final + +from litellm.secret_managers.main import get_secret_bool +from litellm.types.router import ModelInfo +from litellm.types.utils import CustomPricingLiteLLMParams, MirroredPricingParams + +PTU_COST_ATTRIBUTION_ENV_VAR: Final = "LITELLM_ENABLE_PTU_COST_ATTRIBUTION" + + +def is_ptu_cost_attribution_enabled() -> bool: + """Whether PTU flat-cost attribution is turned on for this process.""" + return get_secret_bool(PTU_COST_ATTRIBUTION_ENV_VAR, False) is True + + +PTU_ZEROED_PRICING_FIELDS: Final = tuple(f for f in MirroredPricingParams.model_fields if f != "tiered_pricing") + ( + "cache_creation_input_token_cost_above_1hr", + "cache_creation_input_token_cost_above_200k_tokens", + "cache_read_input_token_cost_above_200k_tokens", +) +# tiered_pricing is emptied rather than zeroed: its tiers outrank the zeros written beside +# them, so a zero here would leave the cost map's tiers billing the traffic the reserved +# capacity already covers. +PTU_EMPTIED_PRICING_FIELDS: Final = frozenset(("tiered_pricing",)) +# search_context_cost_per_query holds its rates in a table keyed by context size, and an +# absent table means the provider's own default rather than free, so it is zeroed in place +# and written on every PTU deployment rather than only where a table is already stored. +PTU_ZEROED_TABLE_FIELDS: Final = frozenset(("search_context_cost_per_query",)) +SEARCH_CONTEXT_SIZES: Final = ("search_context_size_low", "search_context_size_medium", "search_context_size_high") +# Rate fields only. CustomPricingLiteLLMParams also carries settings that are not charges, +# and zeroing one of those would destroy the deployment's configuration rather than stop a +# charge. +CUSTOM_PRICING_FIELDS: Final = frozenset(f for f in CustomPricingLiteLLMParams.model_fields if "cost" in f) +PTU_ZEROED_PRICING: Final[Mapping[str, float | tuple[()] | Mapping[str, float]]] = MappingProxyType( + { + **dict.fromkeys(PTU_ZEROED_PRICING_FIELDS, 0.0), + **dict.fromkeys(PTU_EMPTIED_PRICING_FIELDS, ()), + **dict.fromkeys(PTU_ZEROED_TABLE_FIELDS, MappingProxyType(dict.fromkeys(SEARCH_CONTEXT_SIZES, 0.0))), + } +) + + +@dataclass(frozen=True, slots=True) +class PTUTerms: + """The reservation a deployment declares, once every field has been validated.""" + + team_id: str + ptu_count: int + cost_per_ptu_per_hour: float + effective_from: datetime + effective_to: datetime | None + + +def _to_utc(parsed: datetime) -> datetime: + """``parsed`` as UTC, reading a naive value as UTC rather than local time.""" + return parsed.replace(tzinfo=timezone.utc) if parsed.tzinfo is None else parsed.astimezone(timezone.utc) + + +def _as_utc(value: object) -> datetime | None: + """A model_info datetime as UTC, parsing an ISO string, else None.""" + if isinstance(value, datetime): + return _to_utc(value) + if not isinstance(value, str): + return None + try: + return _to_utc(datetime.fromisoformat(value.replace("Z", "+00:00"))) + except ValueError: + return None + + +def ptu_terms(model_info: Mapping[str, object]) -> PTUTerms | None: + """The reservation this deployment accrues flat cost for, else None. + + A start is required rather than inferred because flat cost accrues from it, and a + present but unparseable bound would read as no bound and widen the window to the whole + day, so either one leaves the deployment unpriced until the config is fixed. + """ + ptu_count: Final = model_info.get("ptu_count") + cost_per_hour: Final = model_info.get("cost_per_ptu_per_hour") + team_id: Final = model_info.get("team_id") + if ptu_count is None or cost_per_hour is None or not team_id: + return None + try: + ptu_count_int: Final = int(ptu_count) + cost_per_hour_float: Final = float(cost_per_hour) + except (TypeError, ValueError, OverflowError): + return None + if not 0 < ptu_count_int <= ModelInfo.MAX_PTU_COUNT: + return None + if not 0 <= cost_per_hour_float <= ModelInfo.MAX_COST_PER_PTU_PER_HOUR: + return None + + raw_from: Final = model_info.get("ptu_effective_from") + raw_to: Final = model_info.get("ptu_effective_to") + effective_from: Final = _as_utc(raw_from) + effective_to: Final = _as_utc(raw_to) + if effective_from is None or (raw_to is not None and effective_to is None): + return None + if effective_to is not None and effective_to <= effective_from: + return None + return PTUTerms( + team_id=str(team_id), + ptu_count=ptu_count_int, + cost_per_ptu_per_hour=cost_per_hour_float, + effective_from=effective_from, + effective_to=effective_to, + ) + + +def zeroed_ptu_pricing( + model_info: Mapping[str, object], declared: Mapping[str, object] +) -> Mapping[str, float | tuple[()] | Mapping[str, float]] | None: + """The pricing a deployment accruing flat cost must carry, else None. + + Both conditions hold or nothing is zeroed. Without the flag no flat cost accrues, so + zeroing would leave the deployment serving for free with nothing charged in its place, + which is what an SDK user who happens to carry ptu_count would otherwise get. The terms + are checked first only because they are a few dict reads, while the flag can resolve + through a configured secret manager, and this runs for every deployment registered. + + Any further rate the deployment itself declares is zeroed alongside the standing set, + since one left standing bills the traffic the reserved capacity already paid for. + """ + if ptu_terms(model_info) is None: + return None + if not is_ptu_cost_attribution_enabled(): + return None + return MappingProxyType( + { + **PTU_ZEROED_PRICING, + **dict.fromkeys( + CUSTOM_PRICING_FIELDS.intersection(declared) + .difference(PTU_ZEROED_TABLE_FIELDS) + .difference(PTU_EMPTIED_PRICING_FIELDS), + 0.0, + ), + } + ) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index ade24d194d2..1b49e2455e4 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -24,6 +24,13 @@ from pydantic import BaseModel, ConfigDict, Field, ValidationError from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import LITELLM_PROXY_ADMIN_NAME +from litellm.litellm_core_utils.ptu_pricing import ( + CUSTOM_PRICING_FIELDS, + PTU_EMPTIED_PRICING_FIELDS, + PTU_ZEROED_PRICING_FIELDS, + PTU_ZEROED_TABLE_FIELDS, + SEARCH_CONTEXT_SIZES, +) from litellm.proxy._types import ( BlockModelRequest, CommonProxyErrors, @@ -89,7 +96,6 @@ from litellm.types.router import ( ModelInfo, updateDeployment, ) -from litellm.types.utils import CustomPricingLiteLLMParams from litellm.utils import get_utc_datetime router: Final = APIRouter() @@ -346,12 +352,8 @@ def _validate_ptu_model_info(model_info: Mapping[str, object]) -> None: # tiered_pricing is the one mirrored field that is a table of ranges, not a rate, so it is stored # empty (see _PTU_EMPTIED_PRICING_FIELDS): its tiers outrank the zeros written beside them, so # dropping it would leave the cost map's tiers billing the traffic the reserved capacity covers. -_PTU_ZEROED_PRICING_FIELDS: Final = tuple(f for f in SPECIAL_MODEL_INFO_PARAMS if f != "tiered_pricing") + ( - "cache_creation_input_token_cost_above_1hr", - "cache_creation_input_token_cost_above_200k_tokens", - "cache_read_input_token_cost_above_200k_tokens", -) -_PTU_EMPTIED_PRICING_FIELDS: Final = frozenset({"tiered_pricing"}) +_PTU_ZEROED_PRICING_FIELDS: Final = PTU_ZEROED_PRICING_FIELDS +_PTU_EMPTIED_PRICING_FIELDS: Final = PTU_EMPTIED_PRICING_FIELDS _PTU_ZEROED_PRICING: Final[Mapping[str, float | tuple[()]]] = MappingProxyType( { **dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0), @@ -363,13 +365,13 @@ _EMPTY_MODEL_INFO: Final[Mapping[str, object]] = _NO_PRICING_OVERRIDE # Rate fields only. CustomPricingLiteLLMParams also carries settings that are not charges # (an embedding's output_vector_size, the regional uplift multipliers), and zeroing one of # those would destroy the deployment's configuration rather than stop a charge. -_CUSTOM_PRICING_FIELDS: Final = frozenset(f for f in CustomPricingLiteLLMParams.model_fields if "cost" in f) +_CUSTOM_PRICING_FIELDS: Final = CUSTOM_PRICING_FIELDS # search_context_cost_per_query holds its rates in a table keyed by context size, and an absent # table means the provider's own default rate rather than free (litellm/llms/gemini/cost_calculator # falls back to $0.035), so it is zeroed in place rather than emptied like tiered_pricing, and # written on every PTU deployment rather than only where a table is already stored. -_PTU_ZEROED_TABLE_FIELDS: Final = frozenset({"search_context_cost_per_query"}) -_SEARCH_CONTEXT_SIZES: Final = ("search_context_size_low", "search_context_size_medium", "search_context_size_high") +_PTU_ZEROED_TABLE_FIELDS: Final = PTU_ZEROED_TABLE_FIELDS +_SEARCH_CONTEXT_SIZES: Final = SEARCH_CONTEXT_SIZES def _is_nonzero_rate(value: object) -> bool: diff --git a/litellm/proxy/spend_tracking/ptu_feature_flag.py b/litellm/proxy/spend_tracking/ptu_feature_flag.py index 9078079b676..7f52dfa155d 100644 --- a/litellm/proxy/spend_tracking/ptu_feature_flag.py +++ b/litellm/proxy/spend_tracking/ptu_feature_flag.py @@ -1,18 +1,12 @@ -"""Opt-in flag for PTU (provisioned throughput unit) flat-cost attribution. +"""Re-exported from ``litellm.litellm_core_utils.ptu_pricing``. -The whole feature is inert unless an operator sets -``LITELLM_ENABLE_PTU_COST_ATTRIBUTION``: the daily rollup is not scheduled, the -model endpoints reject PTU config, the daily activity read path reports zero flat -cost, and the model form hides the PTU inputs. +The flag lives in core because the router reads it while registering a deployment, and +router code cannot import from the proxy. """ -from typing import Final +from litellm.litellm_core_utils.ptu_pricing import ( + PTU_COST_ATTRIBUTION_ENV_VAR, + is_ptu_cost_attribution_enabled, +) -from litellm.secret_managers.main import get_secret_bool - -PTU_COST_ATTRIBUTION_ENV_VAR: Final = "LITELLM_ENABLE_PTU_COST_ATTRIBUTION" - - -def is_ptu_cost_attribution_enabled() -> bool: - """Report whether this deployment opted into PTU flat-cost attribution.""" - return get_secret_bool(PTU_COST_ATTRIBUTION_ENV_VAR, False) is True +__all__ = ("PTU_COST_ATTRIBUTION_ENV_VAR", "is_ptu_cost_attribution_enabled") diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index 25de9f6d065..381641be96d 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -14,6 +14,7 @@ and share the existing unique constraint. import asyncio import json +import sys from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from datetime import date, datetime, time, timedelta, timezone @@ -29,14 +30,15 @@ from litellm.constants import ( PTU_ROLLUP_MAX_BACKFILL_DAYS, PTU_SENTINEL_API_KEY, ) +from litellm.litellm_core_utils.ptu_pricing import ptu_terms from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled -from litellm.types.router import ModelInfo if TYPE_CHECKING: from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.utils import PrismaClient _HOURS_PER_DAY: Final = 24 +_PRUNE_ID_CHUNK_SIZE: Final = 5_000 _UPSERT_ATTEMPTS: Final = 3 _UPSERT_RETRY_BACKOFF_SECONDS: Final = 0.5 @@ -72,28 +74,6 @@ class PTUModel: effective_to: datetime | None = None -def _parse_utc_datetime(value: object) -> datetime | None: - """Parse a model_info datetime (ISO string or datetime) into a UTC-aware datetime, else None.""" - parsed: Final = _coerce_datetime(value) - if parsed is None: - return None - if parsed.tzinfo is None: - return parsed.replace(tzinfo=timezone.utc) - return parsed.astimezone(timezone.utc) - - -def _coerce_datetime(value: object) -> datetime | None: - """``value`` as a datetime, parsing an ISO string, else None.""" - if isinstance(value, datetime): - return value - if not isinstance(value, str): - return None - try: - return datetime.fromisoformat(value.replace("Z", "+00:00")) - except ValueError: - return None - - def _public_model_name(row: object, model_info: Mapping[str, object]) -> str: """The name an operator recognises for this deployment. @@ -167,46 +147,20 @@ def _parse_ptu_model(row: object) -> PTUModel | None: Valid means model_info has a positive ptu_count, a non-negative cost_per_ptu_per_hour, and a team_id (1 model -> 1 team). """ - raw_model_info: Final = getattr(row, "model_info", None) - model_info: Final = _decode_model_info(raw_model_info) + model_info: Final = _decode_model_info(getattr(row, "model_info", None)) if model_info is None: return None - ptu_count: Final = model_info.get("ptu_count") - cost_per_hour: Final = model_info.get("cost_per_ptu_per_hour") - team_id: Final = model_info.get("team_id") - if ptu_count is None or cost_per_hour is None or not team_id: - return None - try: - ptu_count_int: Final = int(ptu_count) - cost_per_hour_float: Final = float(cost_per_hour) - except (TypeError, ValueError, OverflowError): - return None - if not 0 < ptu_count_int <= ModelInfo.MAX_PTU_COUNT: - return None - if not 0 <= cost_per_hour_float <= ModelInfo.MAX_COST_PER_PTU_PER_HOUR: - return None - if model_info.get("ptu_effective_from") is None: - # The endpoints require a start; a row without one predates that rule or was - # written around them, and inferring one would bill days the deployment did not exist - return None - raw_from: Final = model_info.get("ptu_effective_from") - raw_to: Final = model_info.get("ptu_effective_to") - effective_from: Final = _parse_utc_datetime(raw_from) - effective_to: Final = _parse_utc_datetime(raw_to) - # A present-but-unparseable bound would read as "no bound" and silently widen the - # window to the whole day, so the deployment is skipped until the config is fixed - if (raw_from is not None and effective_from is None) or (raw_to is not None and effective_to is None): - return None - if effective_from is not None and effective_to is not None and effective_to <= effective_from: + terms: Final = ptu_terms(model_info) + if terms is None: return None return PTUModel( model_id=str(getattr(row, "model_id", "") or ""), model_name=_public_model_name(row, model_info), - team_id=str(team_id), - ptu_count=ptu_count_int, - cost_per_ptu_per_hour=cost_per_hour_float, - effective_from=effective_from, - effective_to=effective_to, + team_id=terms.team_id, + ptu_count=terms.ptu_count, + cost_per_ptu_per_hour=terms.cost_per_ptu_per_hour, + effective_from=terms.effective_from, + effective_to=terms.effective_to, ) @@ -358,10 +312,70 @@ async def _upsert_charge_with_retry( return False -async def _load_ptu_models(prisma_client: "PrismaClient") -> tuple[PTUModel, ...]: - """Every model deployment currently carrying valid manual PTU config.""" +@dataclass(frozen=True, slots=True) +class _LoadedDeployments: + """The deployments a run will price, and every deployment id it looked at. + + The id set is deliberately wider than the priced set. A deployment whose PTU config + was removed produces no charge and still has to be prunable, so bounding the prune on + what priced would strand its old rows forever. It is also a guaranteed superset of the + priced set, or a run could write a charge that falls outside its own delete filter. + """ + + models: tuple[PTUModel, ...] + scanned_ids: frozenset[str] + config_sourced: bool + + +def _running_router() -> object | None: + """The proxy's router, or None outside a running proxy. + + Read out of ``sys.modules`` rather than imported, so a rollup driven from a test or a + script does not pull the whole proxy server in behind it. + """ + proxy_server: Final = sys.modules.get("litellm.proxy.proxy_server") + return getattr(proxy_server, "llm_router", None) if proxy_server is not None else None + + +def _config_deployments(router: object | None, *, owned_by_db: frozenset[str]) -> tuple[_PTUDeployment, ...]: + """Deployments the router holds that no ``LiteLLM_ProxyModelTable`` row owns. + + ``db_model`` is forced True on every deployment loaded from that table and defaults to + False on ModelInfo, so the complement is what config.yaml declared. A per-request + credential clone carries ``original_model_id`` and reuses its source's PTU config under + a fresh id, so pricing it would bill one reservation once per distinct client key. + """ + entries: Final = tuple(getattr(router, "model_list", None) or ()) + records: Final = tuple(_router_deployment(entry) for entry in entries) + return tuple( + record + for record in records + if record is not None + and record.model_info.get("db_model") is not True + and record.model_info.get("original_model_id") is None + and record.model_id not in owned_by_db + ) + + +async def _load_ptu_models(prisma_client: "PrismaClient") -> _LoadedDeployments: + """Every deployment carrying valid manual PTU config, and every id the scan saw. + + Reserved capacity is billed by the provider whichever file declared it, so a + deployment the proxy only knows from config.yaml accrues alongside the stored ones. + """ rows: Final = await prisma_client.db.litellm_proxymodeltable.find_many() - return tuple(parsed for parsed in (_parse_ptu_model(row) for row in rows) if parsed is not None) + db_ids: Final = frozenset(model_id for row in rows if (model_id := str(getattr(row, "model_id", "") or ""))) + config_records: Final = _config_deployments(_running_router(), owned_by_db=db_ids) + models: Final = tuple( + parsed for parsed in (_parse_ptu_model(row) for row in (*rows, *config_records)) if parsed is not None + ) + return _LoadedDeployments( + models=models, + config_sourced=bool(config_records), + scanned_ids=db_ids + | frozenset(record.model_id for record in config_records) + | frozenset(model.model_id for model in models), + ) async def run_ptu_flat_cost_rollup( @@ -378,8 +392,10 @@ async def run_ptu_flat_cost_rollup( The prune predicate is ``updated_at < run_started`` rather than "not in the charge set I computed", which matters under concurrency: whether a row is garbage becomes a property of the row instead of one run's in-memory config snapshot, so a run can - never delete a row a concurrent run just wrote. It is still skipped when any charge - failed to write, since a row whose replacement never landed would look unrefreshed. + never delete a row a concurrent run just wrote. It is bounded to the deployments this + run looked at, so a row it cannot account for is out of reach either way. It is still + skipped when any charge failed to write, since a row whose replacement never landed + would look unrefreshed. """ day: Final = target_date or (datetime.now(timezone.utc).date() - timedelta(days=1)) @@ -390,7 +406,8 @@ async def run_ptu_flat_cost_rollup( date_str: Final = day.isoformat() run_started: Final = datetime.now(timezone.utc) - ptu_models: Final = await _load_ptu_models(prisma_client) + loaded: Final = await _load_ptu_models(prisma_client) + ptu_models: Final = loaded.models charges: Final = _aggregate_charges(ptu_models, day) landed: Final = tuple( @@ -415,7 +432,12 @@ async def run_ptu_flat_cost_rollup( date_str, ) else: - await _prune_unrefreshed_sentinel_rows(prisma_client, date_str=date_str, run_started=run_started) + await _prune_unrefreshed_sentinel_rows( + prisma_client, + date_str=date_str, + run_started=run_started, + scanned_ids=loaded.scanned_ids if loaded.config_sourced else None, + ) verbose_proxy_logger.info( "PTU rollup for %s: %d PTU models processed, %d rows written, %d rows failed", @@ -524,7 +546,7 @@ async def run_ptu_flat_cost_backfill( verbose_proxy_logger.warning("PTU backfill: prisma_client is None, skipping") return BackfillResult(start=end, end=end, days_scanned=0, rows_written=0) - ptu_models: Final = await _load_ptu_models(prisma_client) + ptu_models: Final = (await _load_ptu_models(prisma_client)).models days: Final = _backfill_window(ptu_models, end) if not days: @@ -707,26 +729,61 @@ async def _prune_unrefreshed_sentinel_rows( *, date_str: str, run_started: datetime, + scanned_ids: frozenset[str] | None, ) -> None: - """Delete the day's PTU sentinel rows this run did not refresh. + """Delete the day's PTU sentinel rows this run looked at and did not refresh. - Every charge the run wrote bumps ``updated_at`` past ``run_started``, so anything - left below that mark is a (team, model) the current config no longer prices. The mark - is pulled back by ``PTU_PRUNE_SKEW_GRACE_SECONDS`` because the two timestamps come - from different hosts: a stale row is hours old, a concurrently written one is seconds - old, and the grace separates them without waiting on clocks agreeing. The - predicate reads only the row, never the caller's config snapshot, which is what - makes it safe to run twice, out of order, or beside another pod: a row written - after this run began is out of reach of its delete. Mirrors the retention predicate - ``SpendLogCleanup`` deletes by.""" + Two conditions, and a row survives unless it meets both. It must be stale: every + charge the run wrote bumps ``updated_at`` past ``run_started``, so anything left below + that mark is a (team, model) the current config no longer prices. The mark is pulled + back by ``PTU_PRUNE_SKEW_GRACE_SECONDS`` because the two timestamps come from + different hosts, and the grace separates a row that is hours old from one written + seconds ago without waiting on clocks agreeing. + + A run that priced a deployment only its own host declares must also name the + deployments it scanned. Staleness alone is sufficient while every run derives its + charges from the same table, because then any two runs compute the same set, so a + database-only run still sweeps by timestamp exactly as it always has. Once one host's + charges come from a file the others cannot read, a row it never considered is not + evidence of anything, and deleting it drops a charge that host is responsible for. + + Where the bound applies the ids go out in chunks, because each is one bind variable and + the server rejects a statement carrying more than 32767 of them, which a proxy holding + that many deployments would otherwise hit every night with no handler above here. + """ cutoff: Final = run_started - timedelta(seconds=PTU_PRUNE_SKEW_GRACE_SECONDS) - await prisma_client.db.litellm_dailyteamspend.delete_many( - where={ # mutable-ok: prisma delete filter - "date": date_str, - "api_key": PTU_SENTINEL_API_KEY, - "updated_at": {"lt": cutoff}, # mutable-ok: prisma comparison filter - } + unbounded: Final = { # mutable-ok: prisma delete filter + "date": date_str, + "api_key": PTU_SENTINEL_API_KEY, + "updated_at": {"lt": cutoff}, # mutable-ok: prisma comparison filter + } + ordered: Final = () if scanned_ids is None else tuple(sorted(scanned_ids)) + filters: Final = ( + (unbounded,) + if scanned_ids is None + else tuple( + MappingProxyType( + { + **unbounded, + "model": { # mutable-ok: prisma membership filter + "in": ordered[start : start + _PRUNE_ID_CHUNK_SIZE] + }, + } + ) + for start in range(0, len(ordered), _PRUNE_ID_CHUNK_SIZE) + ) ) + deletions: Final = tuple( + [await prisma_client.db.litellm_dailyteamspend.delete_many(where=where) for where in filters] + ) + deleted: Final = sum(deletions) + if deleted: + verbose_proxy_logger.info( + "PTU rollup for %s: pruned %s stale sentinel row(s) across %s deployment(s)", + date_str, + deleted, + "every" if scanned_ids is None else len(scanned_ids), + ) __all__ = ( diff --git a/litellm/router.py b/litellm/router.py index ef04423e3ab..26158ae0a56 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -64,6 +64,7 @@ from litellm.litellm_core_utils.coroutine_checker import coroutine_checker from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.litellm_core_utils.ptu_pricing import zeroed_ptu_pricing from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, ) @@ -7695,7 +7696,16 @@ class Router: - None: If the deployment is not active for the current environment (if 'supported_environments' is set in litellm_params) """ try: - litellm_params: Final[LiteLLM_Params] = LiteLLM_Params(**_litellm_params) + zeroed_pricing: Final = ( + zeroed_ptu_pricing(_model_info, _litellm_params) if _model_info.get("db_model") is not True else None + ) + litellm_params: Final[LiteLLM_Params] = LiteLLM_Params( + **( + _litellm_params + if zeroed_pricing is None + else MappingProxyType({**_litellm_params, **zeroed_pricing}) + ) + ) warn_on_provider_credential_mismatch(model_name=_model_name, litellm_params=_litellm_params) deployment = Deployment( **deployment_info, diff --git a/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py b/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py new file mode 100644 index 00000000000..270c59f595f --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py @@ -0,0 +1,163 @@ +"""Tests for the shared PTU rules: which deployments accrue flat cost, and what that zeroes.""" + +import os +from datetime import datetime, timezone +from unittest.mock import patch + +import pytest + +from litellm.litellm_core_utils.ptu_pricing import ( + CUSTOM_PRICING_FIELDS, + PTU_EMPTIED_PRICING_FIELDS, + PTU_ZEROED_PRICING_FIELDS, + PTU_ZEROED_TABLE_FIELDS, + SEARCH_CONTEXT_SIZES, + ptu_terms, + zeroed_ptu_pricing, +) +from litellm.types.router import ModelInfo + +_VALID = { + "team_id": "team-alpha", + "ptu_count": 100, + "cost_per_ptu_per_hour": 0.02, + "ptu_effective_from": "2026-01-01T00:00:00Z", +} + + +def _with_flag(model_info, declared=None, enabled=True): + with patch.dict(os.environ, {"LITELLM_ENABLE_PTU_COST_ATTRIBUTION": "True" if enabled else ""}, clear=False): + return zeroed_ptu_pricing(model_info, declared or {}) + + +def test_a_complete_reservation_is_accepted(): + terms = ptu_terms(_VALID) + + assert terms is not None + assert terms.team_id == "team-alpha" + assert terms.ptu_count == 100 + assert terms.effective_from == datetime(2026, 1, 1, tzinfo=timezone.utc) + assert terms.effective_to is None + + +@pytest.mark.parametrize( + "override", + [ + {"team_id": None}, + {"team_id": ""}, + {"ptu_count": None}, + {"cost_per_ptu_per_hour": None}, + {"ptu_count": 0}, + {"ptu_count": -1}, + {"ptu_count": ModelInfo.MAX_PTU_COUNT + 1}, + {"cost_per_ptu_per_hour": -0.01}, + {"cost_per_ptu_per_hour": ModelInfo.MAX_COST_PER_PTU_PER_HOUR + 1}, + {"ptu_count": "not-a-number"}, + {"ptu_effective_from": None}, + {"ptu_effective_from": "not-a-date"}, + {"ptu_effective_to": "not-a-date"}, + {"ptu_effective_to": "2025-01-01T00:00:00Z"}, + {"ptu_effective_to": "2026-01-01T00:00:00Z"}, + ], + ids=[ + "no team", + "blank team", + "no count", + "no rate", + "zero count", + "negative count", + "count over the cap", + "negative rate", + "rate over the cap", + "count not a number", + "no start", + "unparseable start", + "unparseable end", + "end before start", + "end equal to start", + ], +) +def test_an_incomplete_reservation_accrues_nothing(override): + """Anything the rollup declines to charge must also decline to be zeroed, or the + deployment serves its traffic for free with nothing charged in its place.""" + assert ptu_terms({**_VALID, **override}) is None + assert _with_flag({**_VALID, **override}) is None + + +def test_a_naive_start_is_read_as_utc(): + """config.yaml is hand-typed, and pydantic hands back a naive datetime for a date with + no offset.""" + terms = ptu_terms({**_VALID, "ptu_effective_from": datetime(2026, 5, 1, 12, 0)}) + + assert terms is not None + assert terms.effective_from == datetime(2026, 5, 1, 12, 0, tzinfo=timezone.utc) + + +def test_an_offset_start_is_converted_rather_than_relabelled(): + terms = ptu_terms({**_VALID, "ptu_effective_from": "2026-05-01T12:00:00-05:00"}) + + assert terms is not None + assert terms.effective_from == datetime(2026, 5, 1, 17, 0, tzinfo=timezone.utc) + + +def test_nothing_is_zeroed_while_the_feature_is_off(): + """No flat cost accrues with the flag off, so zeroing would serve the traffic free.""" + assert _with_flag(_VALID, enabled=False) is None + + +def test_the_standing_rates_are_all_zeroed(): + override = _with_flag(_VALID) + + assert override is not None + assert [field for field in PTU_ZEROED_PRICING_FIELDS if override[field] != 0.0] == [] + + +def test_tiered_pricing_is_emptied_rather_than_zeroed(): + """A tier outranks the flat rates written beside it, so a zero there would leave the + cost map's tiers billing the traffic the reserved capacity already covers.""" + override = _with_flag(_VALID, declared={"tiered_pricing": [{"range": [0, 1000], "input_cost_per_token": 0.003}]}) + + assert override is not None + for field in PTU_EMPTIED_PRICING_FIELDS: + assert override[field] == () + + +def test_the_search_context_table_is_zeroed_in_place_on_every_deployment(): + """An absent table means the provider's own default rather than free, so it is written + even when the deployment never declared one.""" + override = _with_flag(_VALID) + + assert override is not None + for field in PTU_ZEROED_TABLE_FIELDS: + assert dict(override[field]) == dict.fromkeys(SEARCH_CONTEXT_SIZES, 0.0) + + +def test_a_declared_table_does_not_become_a_scalar(): + """Zeroing it as a plain 0.0 would leave the provider's reader without a table to + consult, which is the same as absent.""" + override = _with_flag(_VALID, declared={"search_context_cost_per_query": {"search_context_size_medium": 0.05}}) + + assert override is not None + assert dict(override["search_context_cost_per_query"]) == dict.fromkeys(SEARCH_CONTEXT_SIZES, 0.0) + + +def test_a_rate_the_deployment_declares_itself_is_zeroed_too(): + """The standing set covers the mirrored rates. Anything else the operator wrote would + otherwise survive and bill the traffic the hourly charge already paid for.""" + extra = "input_cost_per_token_above_200k_tokens" + assert extra in CUSTOM_PRICING_FIELDS + assert extra not in PTU_ZEROED_PRICING_FIELDS + + override = _with_flag(_VALID, declared={extra: 9e-06}) + + assert override is not None + assert override[extra] == 0.0 + + +def test_a_setting_that_is_not_a_charge_is_left_alone(): + """CustomPricingLiteLLMParams also carries configuration, and zeroing one of those + would break the deployment rather than stop a charge.""" + override = _with_flag(_VALID, declared={"output_vector_size": 1536}) + + assert override is not None + assert "output_vector_size" not in override diff --git a/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py index f41756d2b87..333fd597b49 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py @@ -210,8 +210,10 @@ async def test_rollup_prunes_stale_row_when_config_is_gone(): where = table.delete_many.await_args.kwargs["where"] assert where["date"] == DAY.isoformat() assert where["api_key"] == PTU_SENTINEL_API_KEY - # the row is garbage because this run did not refresh it, not because of a key list + # the row is garbage because this run did not refresh it, and it is reachable at all + # because the run scanned the deployment it belongs to assert "lt" in where["updated_at"] + assert "model" not in where, "a database-only run has no reason to bound the sweep" @pytest.mark.asyncio @@ -705,13 +707,20 @@ class _FakeSentinelTable: async def delete_many(self, where): self.delete_many_calls.append(where) cutoff = where["updated_at"]["lt"] + # honouring "model" matters: a fake that ignored an unknown clause would delete + # the row the prune-scoping test exists to protect and still report a pass + allowed = where.get("model", {}).get("in") doomed = [ k for k, v in self.rows.items() - if k[1] == where["date"] and k[2] == where["api_key"] and v["updated_at"] < cutoff + if k[1] == where["date"] + and k[2] == where["api_key"] + and v["updated_at"] < cutoff + and (allowed is None or k[3] in allowed) ] for k in doomed: del self.rows[k] + return len(doomed) async def find_many(self, where=None): """Read back sentinel rows the way prisma would, honouring api_key and a date range.""" @@ -785,11 +794,11 @@ async def test_an_older_run_cannot_delete_a_newer_runs_row(): @pytest.mark.asyncio async def test_a_later_clean_run_clears_the_row_the_race_left_behind(): - """The race can leave a charge for a since-removed deployment in place for a day; the - next run, seeing only the current config, must sweep it.""" + """The race can leave a charge for a no-longer-priced deployment in place for a day; + the next run, seeing only the current config, must sweep it.""" table = _FakeSentinelTable() ptu = {"ptu_count": 10, "cost_per_ptu_per_hour": 2.0, "team_id": "t"} - stale_key = ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-removed") + stale_key = ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-retired") table.rows[stale_key] = { "ptu_flat_cost": 480.0, "model_group": "retired", @@ -797,7 +806,14 @@ async def test_a_later_clean_run_clears_the_row_the_race_left_behind(): } await run_ptu_flat_cost_rollup( - _prisma_for([_model_row(model_id="dep-live", model_info=ptu)], table), target_date=DAY + _prisma_for( + [ + _model_row(model_id="dep-live", model_info=ptu), + _model_row(model_id="dep-retired", model_info={"team_id": "t"}), + ], + table, + ), + target_date=DAY, ) assert stale_key not in table.rows @@ -1715,16 +1731,19 @@ async def test_a_run_holding_the_lock_still_prunes(): """Losing the sweep entirely would leave stale charges forever, so the guarded path, which is the normal one, keeps it.""" table = _FakeSentinelTable() - table.seed("t", DAY, "dep-gone", 480.0, updated_at=datetime(2020, 1, 1, tzinfo=timezone.utc)) + table.seed("t", DAY, "dep-unpriced", 480.0, updated_at=datetime(2020, 1, 1, tzinfo=timezone.utc)) prisma = _prisma_for( - [_model_row(model_id="dep-live", model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"})], + [ + _model_row(model_id="dep-live", model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"}), + _model_row(model_id="dep-unpriced", model_info={"team_id": "t"}), + ], table, ) await run_scheduled_ptu_rollup(prisma, pod_lock_manager=_pod_lock(acquired=True), target_date=DAY) assert table.delete_many_calls != [] - assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-gone") not in table.rows + assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-unpriced") not in table.rows assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-live") in table.rows @@ -1737,7 +1756,7 @@ async def test_the_prune_cutoff_allows_for_clock_skew_between_hosts(): just_written = datetime.now(timezone.utc) - timedelta(seconds=30) table.seed("t", DAY, "dep-concurrent", 480.0, updated_at=just_written) table.seed("t", DAY, "dep-stale", 480.0, updated_at=datetime.now(timezone.utc) - timedelta(hours=6)) - prisma = _prisma_for([], table) + prisma = _prisma_for([_model_row(model_id="dep-concurrent"), _model_row(model_id="dep-stale")], table) await run_ptu_flat_cost_rollup(prisma, target_date=DAY) @@ -1747,6 +1766,127 @@ async def test_the_prune_cutoff_allows_for_clock_skew_between_hosts(): assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-stale") not in table.rows +@pytest.mark.asyncio +async def test_a_run_pricing_config_cannot_prune_a_row_it_did_not_scan(monkeypatch): + """Staleness alone stops being evidence once two hosts hold different configuration: a + row this run never considered belongs to a deployment another host is pricing from its + own file, and sweeping it drops that charge.""" + table = _FakeSentinelTable() + table.seed("t", DAY, "dep-elsewhere", 480.0, updated_at=datetime(2020, 1, 1, tzinfo=timezone.utc)) + entry = _router_entry(model_id="cfg-here", model_info=dict(_VALID_PTU)) + monkeypatch.setattr(ptu_rollup, "_running_router", lambda: _router_holding(entry)) + + await run_scheduled_ptu_rollup(_prisma_for([], table), pod_lock_manager=_pod_lock(acquired=True), target_date=DAY) + + assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-elsewhere") in table.rows + assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "cfg-here") in table.rows + assert table.delete_many_calls[-1]["model"]["in"] == ("cfg-here",) + + +@pytest.mark.asyncio +async def test_a_deployment_deleted_from_the_table_keeps_the_day_it_was_charged(monkeypatch): + """The accepted cost of bounding the prune, driven through the sequence that produces + it: charge the day while the deployment exists, remove it, run the day again. Nothing + scans it now, so nothing may judge its row, and the amount it was billed stands.""" + table = _FakeSentinelTable() + ptu = {"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"} + live_row = _model_row(model_id="dep-live", model_info=ptu) + doomed_row = _model_row(model_id="dep-doomed", model_info=ptu) + charged_key = ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-doomed") + monkeypatch.setattr( + ptu_rollup, "_running_router", lambda: _router_holding(_router_entry(model_id="cfg", model_info=dict(ptu))) + ) + + await run_scheduled_ptu_rollup( + _prisma_for([live_row, doomed_row], table), pod_lock_manager=_pod_lock(acquired=True), target_date=DAY + ) + billed = table.rows[charged_key]["ptu_flat_cost"] + table.rows[charged_key]["updated_at"] = datetime(2020, 1, 1, tzinfo=timezone.utc) + + await run_scheduled_ptu_rollup( + _prisma_for([live_row], table), pod_lock_manager=_pod_lock(acquired=True), target_date=DAY + ) + + assert table.rows[charged_key]["ptu_flat_cost"] == billed + assert "dep-doomed" not in table.delete_many_calls[-1]["model"]["in"] + + +@pytest.mark.asyncio +async def test_a_database_only_run_sweeps_exactly_as_it_did_before(): + """The bound exists for charges another host declares. A deployment nobody declares any + more still has its leftover row swept, which is what the table-only sweep always did.""" + table = _FakeSentinelTable() + table.seed("t", DAY, "dep-gone", 480.0, updated_at=datetime(2020, 1, 1, tzinfo=timezone.utc)) + prisma = _prisma_for( + [_model_row(model_id="dep-live", model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"})], + table, + ) + + await run_scheduled_ptu_rollup(prisma, pod_lock_manager=_pod_lock(acquired=True), target_date=DAY) + + assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-gone") not in table.rows + assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-live") in table.rows + + +@pytest.mark.asyncio +async def test_every_deployment_that_prices_is_inside_the_set_that_bounds_the_prune(): + """The bound has to be a superset of what the same run wrote, or a run's own charge + could fall outside its own delete filter and never be reconciled.""" + table = _FakeSentinelTable() + prisma = _prisma_for( + [ + _model_row(model_id="dep-a", model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"}), + _model_row(model_id="dep-b", model_info={"ptu_count": 9, "cost_per_ptu_per_hour": 1.0, "team_id": "u"}), + _model_row(model_id="dep-unpriced", model_info={"team_id": "t"}), + ], + table, + ) + + loaded = await ptu_rollup._load_ptu_models(prisma) + + assert {model.model_id for model in loaded.models} <= loaded.scanned_ids + assert loaded.scanned_ids == {"dep-a", "dep-b", "dep-unpriced"} + + +@pytest.mark.asyncio +async def test_a_priced_deployment_is_in_the_bound_even_with_an_id_the_scan_skips(): + """The bound is built by construction rather than by coincidence. The row scan drops a + falsy id while the parser still prices one, and a charge outside its own run's delete + filter could never be reconciled by any later run.""" + prisma = _prisma_for( + [_model_row(model_id="", model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"})], + _FakeSentinelTable(), + ) + + loaded = await ptu_rollup._load_ptu_models(prisma) + + assert {model.model_id for model in loaded.models} <= loaded.scanned_ids + + +@pytest.mark.asyncio +async def test_the_prune_splits_the_id_set_across_statements(monkeypatch): + """Every id is one bind variable and the server refuses a statement carrying more than + 32767, so a proxy with that many deployments would fail the prune outright, and with it + the rest of the scheduled run.""" + monkeypatch.setattr(ptu_rollup, "_PRUNE_ID_CHUNK_SIZE", 2) + table = _FakeSentinelTable() + ptu = {"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"} + deployments = [_model_row(model_id=f"dep-{n}", model_info=ptu) for n in range(4)] + monkeypatch.setattr( + ptu_rollup, "_running_router", lambda: _router_holding(_router_entry(model_id="dep-4", model_info=dict(ptu))) + ) + table.seed("t", DAY, "dep-3", 480.0, updated_at=datetime(2020, 1, 1, tzinfo=timezone.utc)) + + await run_scheduled_ptu_rollup( + _prisma_for(deployments, table), pod_lock_manager=_pod_lock(acquired=True), target_date=DAY + ) + + chunks = [call["model"]["in"] for call in table.delete_many_calls] + assert len(chunks) == 3 + assert all(len(chunk) <= 2 for chunk in chunks) + assert sorted(i for chunk in chunks for i in chunk) == [f"dep-{n}" for n in range(5)] + + @pytest.mark.asyncio async def test_scheduled_rollup_writes_nothing_when_ptu_attribution_is_disabled(monkeypatch): """Startup already skips scheduling the cron, so this guards the function itself: a @@ -1760,3 +1900,148 @@ async def test_scheduled_rollup_writes_nothing_when_ptu_attribution_is_disabled( assert result is None assert table.rows == {} assert table.upsert_keys == [] + + +# --- config.yaml deployments reach the rollup through the router ---------------- + + +def _router_holding(*entries): + """A stand-in for the proxy's router, carrying whatever model_list is passed.""" + return types.SimpleNamespace(model_list=list(entries)) + + +@pytest.mark.asyncio +async def test_a_config_declared_deployment_is_priced(monkeypatch): + """The whole point. A PTU deployment the proxy only knows from config.yaml is not in + LiteLLM_ProxyModelTable, so a DB-only scan bills the provider's reservation to nobody.""" + entry = _router_entry(model_id="cfg-1", model_name="gpt-4o-ptu", model_info=dict(_VALID_PTU)) + monkeypatch.setattr(ptu_rollup, "_running_router", lambda: _router_holding(entry)) + + loaded = await ptu_rollup._load_ptu_models(_prisma_for([], _FakeSentinelTable())) + + assert [(m.model_id, m.model_name, m.team_id) for m in loaded.models] == [("cfg-1", "gpt-4o-ptu", "t")] + assert "cfg-1" in loaded.scanned_ids + + +@pytest.mark.asyncio +async def test_a_database_backed_router_entry_is_not_counted_twice(monkeypatch): + """Every deployment loaded from the table is also in the router, flagged db_model. Pricing + both copies would write two charges for one reservation.""" + row = _model_row(model_id="db-1", model_info=dict(_VALID_PTU)) + mirrored = _router_entry(model_id="db-1", model_info={**_VALID_PTU, "db_model": True}) + monkeypatch.setattr(ptu_rollup, "_running_router", lambda: _router_holding(mirrored)) + + loaded = await ptu_rollup._load_ptu_models(_prisma_for([row], _FakeSentinelTable())) + + assert [m.model_id for m in loaded.models] == ["db-1"] + + +@pytest.mark.asyncio +async def test_a_router_entry_sharing_an_id_with_the_table_is_priced_once(monkeypatch): + """db_model is data the router carries rather than something this module controls, so the + id anti-join is what actually maps onto the failure: two charges under one id.""" + row = _model_row(model_id="both-1", model_info=dict(_VALID_PTU)) + unflagged = _router_entry(model_id="both-1", model_info=dict(_VALID_PTU)) + monkeypatch.setattr(ptu_rollup, "_running_router", lambda: _router_holding(unflagged)) + + loaded = await ptu_rollup._load_ptu_models(_prisma_for([row], _FakeSentinelTable())) + + assert [m.model_id for m in loaded.models] == ["both-1"] + + +@pytest.mark.asyncio +async def test_a_client_credential_clone_is_not_priced(monkeypatch): + """Supplying an api_key on a request mints a clone of the deployment under a fresh id, + carrying the source's PTU config. Pricing it bills one reservation per distinct caller key.""" + source = _router_entry(model_id="cfg-1", model_info=dict(_VALID_PTU)) + clone = _router_entry(model_id="cfg-1-clone", model_info={**_VALID_PTU, "original_model_id": "cfg-1"}) + monkeypatch.setattr(ptu_rollup, "_running_router", lambda: _router_holding(source, clone)) + + loaded = await ptu_rollup._load_ptu_models(_prisma_for([], _FakeSentinelTable())) + + assert [m.model_id for m in loaded.models] == ["cfg-1"] + + +@pytest.mark.asyncio +async def test_a_config_deployment_without_ptu_config_is_scanned_but_not_priced(monkeypatch): + """It has to stay in the scanned set or its leftover sentinel rows become unprunable.""" + entry = _router_entry(model_id="cfg-plain", model_info={"team_id": "t"}) + monkeypatch.setattr(ptu_rollup, "_running_router", lambda: _router_holding(entry)) + + loaded = await ptu_rollup._load_ptu_models(_prisma_for([], _FakeSentinelTable())) + + assert loaded.models == () + assert "cfg-plain" in loaded.scanned_ids + + +@pytest.mark.asyncio +async def test_no_router_in_the_process_prices_the_database_alone(monkeypatch): + """The rollup is importable and callable outside a running proxy.""" + monkeypatch.setattr(ptu_rollup, "_running_router", lambda: None) + + loaded = await ptu_rollup._load_ptu_models( + _prisma_for([_model_row(model_id="db-1", model_info=dict(_VALID_PTU))], _FakeSentinelTable()) + ) + + assert [m.model_id for m in loaded.models] == ["db-1"] + + +@pytest.mark.asyncio +async def test_a_config_deployment_is_charged_end_to_end(monkeypatch): + """Through the scheduled entry point, so the charge lands in a sentinel row rather than + stopping at the loader.""" + table = _FakeSentinelTable() + entry = _router_entry(model_id="cfg-1", model_name="gpt-4o-ptu", model_info=dict(_VALID_PTU)) + monkeypatch.setattr(ptu_rollup, "_running_router", lambda: _router_holding(entry)) + + await run_scheduled_ptu_rollup(_prisma_for([], table), pod_lock_manager=_pod_lock(acquired=True), target_date=DAY) + + assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "cfg-1") in table.rows + + +@pytest.mark.asyncio +async def test_a_stale_database_backed_router_entry_is_not_treated_as_config(monkeypatch): + """The reconcile can leave a deployment on the router after its row is gone. The id + anti-join cannot see that one, so the flag is what keeps it from being priced as though + config.yaml had declared it.""" + stale = _router_entry(model_id="db-gone", model_info={**_VALID_PTU, "db_model": True}) + monkeypatch.setattr(ptu_rollup, "_running_router", lambda: _router_holding(stale)) + + loaded = await ptu_rollup._load_ptu_models(_prisma_for([], _FakeSentinelTable())) + + assert loaded.models == () + + +def test_the_router_lookup_reads_the_proxys_own_global(): + """Every other config test replaces this helper, so without one test driving the real + body a typo in the module path or the attribute name leaves the whole feature dead in + production with the suite still green.""" + import sys + import types as _types + + assert ptu_rollup._running_router() is None or "litellm.proxy.proxy_server" in sys.modules + + sentinel = object() + stub = _types.SimpleNamespace(llm_router=sentinel) + real = sys.modules.get("litellm.proxy.proxy_server") + sys.modules["litellm.proxy.proxy_server"] = stub + try: + assert ptu_rollup._running_router() is sentinel + del stub.llm_router + assert ptu_rollup._running_router() is None + finally: + if real is None: + del sys.modules["litellm.proxy.proxy_server"] + else: + sys.modules["litellm.proxy.proxy_server"] = real + + +def test_the_router_lookup_returns_none_outside_a_proxy(): + import sys + + real = sys.modules.pop("litellm.proxy.proxy_server", None) + try: + assert ptu_rollup._running_router() is None + finally: + if real is not None: + sys.modules["litellm.proxy.proxy_server"] = real diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index dfe46d54ab8..4674a8b1dfa 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -1584,3 +1584,116 @@ def test_inherit_builtin_tiered_output_rate_leaves_a_user_rate_alone(): ) assert model_info["output_cost_per_token"] == 9e-07 + + +# --- a config.yaml PTU deployment must not also bill per token ------------------ + +_PTU_MODEL_INFO = { + "team_id": "team-alpha", + "ptu_count": 100, + "cost_per_ptu_per_hour": 0.02, + "ptu_effective_from": "2026-01-01T00:00:00Z", +} + + +def _ptu_router(model_info=None, litellm_params=None, ptu_enabled=True): + """A router built the way loading config.yaml builds one.""" + with patch.dict(os.environ, {"LITELLM_ENABLE_PTU_COST_ATTRIBUTION": "True" if ptu_enabled else ""}, clear=False): + return Router( + model_list=[ + { + "model_name": "gpt-4o-ptu", + "litellm_params": { + "model": "anthropic/claude-sonnet-4-5-20250929", + "api_key": "sk-not-used", + **(litellm_params or {}), + }, + "model_info": dict(_PTU_MODEL_INFO if model_info is None else model_info), + } + ] + ) + + +def test_a_config_ptu_deployment_bills_nothing_per_token(): + """Reserved capacity is already billed by the hour, so charging its traffic bills the + same tokens twice. Left unset the rate falls back to the public cost map, which makes + the double charge the default rather than an opt-in.""" + router = _ptu_router(litellm_params={"input_cost_per_token": 5e-06, "output_cost_per_token": 1.5e-05}) + entry = router.model_list[0] + + assert entry["litellm_params"]["input_cost_per_token"] == 0.0 + assert entry["litellm_params"]["output_cost_per_token"] == 0.0 + assert entry["model_info"]["input_cost_per_token"] == 0.0 + assert litellm.model_cost[entry["model_info"]["id"]]["input_cost_per_token"] == 0.0 + + +@pytest.mark.parametrize( + "backend", + ["anthropic/claude-sonnet-4-5-20250929", "azure/gpt-4o", "gemini/gemini-2.5-flash"], +) +def test_a_config_ptu_deployment_imports_no_cache_rate_from_its_backend(backend): + """The cache back-fill runs whenever input_cost_per_token is set, and 0.0 is set, so a + partially zeroed deployment would silently inherit the backend model's real cache rates. + Every backend here publishes non-zero ones, which is what makes the assertion mean + something.""" + cache_fields = ( + "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr", + "cache_creation_input_token_cost_above_200k_tokens", + "cache_read_input_token_cost", + "cache_read_input_token_cost_above_200k_tokens", + ) + builtin = litellm.get_model_info(model=backend) + assert any(builtin.get(field) for field in cache_fields), "backend publishes no cache pricing to leak" + + router = _ptu_router(litellm_params={"model": backend}) + priced = litellm.model_cost[router.model_list[0]["model_info"]["id"]] + + assert [field for field in cache_fields if priced.get(field)] == [] + + +def test_zeroing_a_ptu_deployment_leaves_its_backend_model_priced(): + """A sibling deployment on the same backend must keep billing normally.""" + backend = "anthropic/claude-sonnet-4-5-20250929" + builtin = litellm.get_model_info(model=backend)["input_cost_per_token"] + assert builtin > 0 + + _ptu_router(litellm_params={"model": backend}) + + assert litellm.get_model_info(model=backend)["input_cost_per_token"] == builtin + + +def test_zeroing_does_not_change_the_deployment_id(): + """The id is a hash of the deployment's params and keys its cooldowns, its budget, and + every spend row already written against it.""" + params = {"input_cost_per_token": 5e-06} + priced = _ptu_router(litellm_params=params, ptu_enabled=False).model_list[0]["model_info"]["id"] + zeroed = _ptu_router(litellm_params=params).model_list[0]["model_info"]["id"] + + assert priced == zeroed + + +def test_a_database_backed_deployment_is_left_alone(): + """The write endpoints already zero those, and they answer 400 rather than silently + rewriting a rate the caller sent.""" + entry = _ptu_router(model_info={**_PTU_MODEL_INFO, "db_model": True}).model_list[0] + + assert entry["litellm_params"].get("input_cost_per_token") is None + + +def test_nothing_is_zeroed_while_the_feature_is_off(): + """No flat cost accrues with the flag off, so zeroing would serve the traffic free.""" + entry = _ptu_router(litellm_params={"input_cost_per_token": 5e-06}, ptu_enabled=False).model_list[0] + + assert entry["litellm_params"]["input_cost_per_token"] == 5e-06 + + +@pytest.mark.parametrize("dropped", ["team_id", "ptu_effective_from"], ids=["no team_id", "no ptu_effective_from"]) +def test_a_deployment_the_rollup_will_not_charge_is_not_zeroed(dropped): + """The rollup refuses to price a reservation missing either field, so zeroing on the + looser count-and-rate test alone would leave the deployment serving for free with + nothing charged in its place.""" + incomplete = {k: v for k, v in _PTU_MODEL_INFO.items() if k != dropped} + entry = _ptu_router(model_info=incomplete, litellm_params={"input_cost_per_token": 5e-06}).model_list[0] + + assert entry["litellm_params"]["input_cost_per_token"] == 5e-06 From bd322ed8a7eb6968b2af8bebce28eb9f19251fd3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 19:11:01 -0700 Subject: [PATCH 033/281] refactor(cli): state the credential-store precedence rules as contracts Drop the inline notes on keychain erasure and disk-vs-vault precedence in favour of docstrings on the two functions that own those rules, and remove a stale section header and a field note that the code already says plainly. --- litellm/litellm_core_utils/cli_keyring.py | 6 +++++- litellm/litellm_core_utils/cli_token_utils.py | 6 +++++- litellm/proxy/client/cli/commands/auth.py | 3 --- 3 files changed, 10 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/cli_keyring.py b/litellm/litellm_core_utils/cli_keyring.py index 873db64a728..fcbf5ada55a 100644 --- a/litellm/litellm_core_utils/cli_keyring.py +++ b/litellm/litellm_core_utils/cli_keyring.py @@ -98,10 +98,14 @@ class KeyringVault: return True def erase(self) -> bool: + """Whether the keychain is guaranteed to hold no credential afterwards. + + An uninstalled `keyring` package can never have stored one. A kill switch set after + a credential was stored leaves that entry out of reach, so erasure cannot be promised. + """ if _import_keyring() is None: return True if _keyring_disabled(): - # a credential stored before the kill switch was set may still be in the keychain return False match self.read(): case SecretUnavailable(): diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index 9960192180c..dd263ccd412 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -168,8 +168,12 @@ def _resolve_secret(record: CliTokenRecord, vault: SecretVault) -> CliTokenRecor def _apply_vault_secret(record: CliTokenRecord, blob: str, vault: SecretVault) -> CliTokenRecord | None: + """Resolve the credential when both stores hold one. + + A secret still on disk is the fresher of the two, because it is only left there when the + keychain write that should have removed it failed, so it outranks the vault entry. + """ if record.key is not None: - # a secret still on disk means the last keychain write failed: the file outranks the vault return _migrate_file_secret(record, vault) try: secret: Final = CliTokenSecret.model_validate_json(blob) diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 8b9ef5633da..c2b5b6a620f 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -76,7 +76,6 @@ KEYCHAIN_UNREACHABLE_MESSAGE: Final = ( ) -# Token storage utilities def context_secret_vault(ctx: click.Context) -> SecretVault: """Where this invocation reads and writes secret material; injectable through ctx.obj for tests""" ctx_obj: Final[CliContextObj | None] = ctx.obj @@ -666,8 +665,6 @@ def login(ctx: click.Context, config_claude: bool): api_key: Final = auth_result["api_key"] user_id: Final = auth_result["user_id"] - # base_url is stored so we can verify origin before reusing the - # key on a subsequent CLI invocation. record: Final = CliTokenRecord( base_url=base_url.rstrip("/"), key=api_key, From 2a771caf022c5ad6195190da2f67dbf7b833d481 Mon Sep 17 00:00:00 2001 From: hiraku-miyoshi Date: Wed, 19 Aug 2026 19:12:22 -0700 Subject: [PATCH 034/281] fix(proxy): clamp reservation record TTL so stale records never outlive their counters --- litellm/proxy/hooks/batch_enqueued_tokens.py | 34 ++++++++++++++----- .../proxy/hooks/test_batch_enqueued_tokens.py | 20 +++++++++++ 2 files changed, 46 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/hooks/batch_enqueued_tokens.py b/litellm/proxy/hooks/batch_enqueued_tokens.py index 82d4e7c66fc..32bccca2ab0 100644 --- a/litellm/proxy/hooks/batch_enqueued_tokens.py +++ b/litellm/proxy/hooks/batch_enqueued_tokens.py @@ -9,9 +9,11 @@ the reservation is refunded when the batch reaches a terminal state """ import asyncio +import math +import time import uuid -from collections.abc import Awaitable, Mapping, Sequence -from dataclasses import dataclass +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError @@ -85,6 +87,7 @@ class BatchEnqueuedTokenReservation: scopes: tuple[BatchEnqueuedTokenScope, ...] backend: ReservationBackend = "redis" owner: str = "" + reserved_at_monotonic: float = field(default_factory=time.monotonic, compare=False) @dataclass(frozen=True, slots=True) @@ -180,11 +183,18 @@ class BatchEnqueuedTokenStore: granted them, and in-memory grants also remember the granting worker, so a refund never debits counters the grant did not charge. Everything expires after ``BATCH_ENQUEUED_TOKEN_TTL_SECONDS`` so a crash between submission and the - terminal-state refund can never leak tokens forever. + terminal-state refund can never leak tokens forever, and reservation records + expire no later than the counters they would refund, so a stale record can + never debit an allowance re-granted after its counters expired. """ - def __init__(self, internal_usage_cache: "InternalUsageCache") -> None: + def __init__( + self, + internal_usage_cache: "InternalUsageCache", + monotonic: Callable[[], float] = time.monotonic, + ) -> None: self.internal_usage_cache = internal_usage_cache + self._monotonic: Final = monotonic self._lock = asyncio.Lock() self._owner_token = uuid.uuid4().hex redis_cache = internal_usage_cache.dual_cache.redis_cache @@ -235,6 +245,7 @@ class BatchEnqueuedTokenStore: tokens: int, scopes: tuple[BatchEnqueuedTokenScope, ...], ) -> BatchEnqueuedTokenOutcome: + started: Final = self._monotonic() for index, scope in enumerate(scopes): result = await self._run_reserve_script( reserve_script, @@ -246,7 +257,9 @@ class BatchEnqueuedTokenStore: if result[0] != 1: await self._rollback_partial_reserve(refund_script, tokens=tokens, scopes=scopes[:index]) return BatchEnqueuedTokenOverLimit(scope=scope, enqueued=result[1]) - return BatchEnqueuedTokenReservation(tokens=tokens, scopes=scopes, backend="redis") + return BatchEnqueuedTokenReservation( + tokens=tokens, scopes=scopes, backend="redis", reserved_at_monotonic=started + ) async def _run_reserve_script( self, @@ -295,6 +308,7 @@ class BatchEnqueuedTokenStore: scopes: tuple[BatchEnqueuedTokenScope, ...], span: "Span | None", ) -> BatchEnqueuedTokenOutcome: + started: Final = self._monotonic() async with self._lock: currents: Final = tuple([await self._get_local_counter(scope, span) for scope in scopes]) for scope, current in zip(scopes, currents): @@ -302,7 +316,9 @@ class BatchEnqueuedTokenStore: return BatchEnqueuedTokenOverLimit(scope=scope, enqueued=current) for scope, current in zip(scopes, currents): await self._set_local_counter(scope, current + tokens, span) - return BatchEnqueuedTokenReservation(tokens=tokens, scopes=scopes, backend="memory", owner=self._owner_token) + return BatchEnqueuedTokenReservation( + tokens=tokens, scopes=scopes, backend="memory", owner=self._owner_token, reserved_at_monotonic=started + ) async def refund( self, @@ -349,11 +365,13 @@ class BatchEnqueuedTokenStore: litellm_parent_otel_span: "Span | None" = None, ) -> None: serialized: Final = _RESERVATION_ADAPTER.dump_json(reservation).decode("utf-8") + elapsed: Final = self._monotonic() - reservation.reserved_at_monotonic + ttl: Final = max(1, BATCH_ENQUEUED_TOKEN_TTL_SECONDS - math.ceil(elapsed)) if self._save_script is not None: try: await self._save_script( (self._record_key(batch_id),), - (serialized, BATCH_ENQUEUED_TOKEN_TTL_SECONDS), + (serialized, ttl), ) except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory record verbose_proxy_logger.warning( @@ -364,7 +382,7 @@ class BatchEnqueuedTokenStore: await self.internal_usage_cache.async_set_cache( key=self._record_key(batch_id), value=serialized, - ttl=BATCH_ENQUEUED_TOKEN_TTL_SECONDS, + ttl=ttl, litellm_parent_otel_span=litellm_parent_otel_span, local_only=True, ) diff --git a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py index 440b3050a39..edc921a40a3 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py +++ b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py @@ -16,6 +16,7 @@ from typing import Final import pytest from litellm.caching.caching import DualCache +from litellm.constants import BATCH_ENQUEUED_TOKEN_TTL_SECONDS from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.batch_enqueued_tokens import ( BatchEnqueuedTokenOverLimit, @@ -136,6 +137,7 @@ class _SingleKeyRedisFake: raise_after_landing_save_keys: frozenset[str] = frozenset(), ) -> None: self.script_calls: tuple[tuple[str, tuple[str, ...]], ...] = () + self.save_ttls: tuple[int, ...] = () self.counters: Mapping[str, int] = MappingProxyType({}) self.records: Mapping[str, str] = MappingProxyType({}) self.fail_reserve_keys = fail_reserve_keys @@ -180,6 +182,7 @@ class _SingleKeyRedisFake: if keys[0] in self.fail_save_keys: raise ConnectionError(f"simulated redis failure for {keys[0]}") self.records = MappingProxyType({**self.records, keys[0]: str(args[0])}) + self.save_ttls = (*self.save_ttls, int(args[1])) if keys[0] in self.raise_after_landing_save_keys: raise TimeoutError(f"simulated redis timeout after landing for {keys[0]}") return 1 @@ -303,6 +306,23 @@ async def test_local_ghost_left_by_landed_save_never_refunds_twice(): assert fake.counters[f"batch_enqueued_tokens:{scope.key}:{scope.value}"] == 100 +@pytest.mark.asyncio +async def test_record_ttl_shrinks_by_elapsed_time_so_stale_records_never_outlive_their_counters(): + scope = _scope(limit=100) + fake = _SingleKeyRedisFake() + ticks = iter((1_000.0, 1_030.5)) + store = BatchEnqueuedTokenStore( + internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60)), + monotonic=lambda: next(ticks), + ) + reservation = await store.reserve(tokens=60, scopes=(scope,)) + assert isinstance(reservation, BatchEnqueuedTokenReservation) + assert reservation.reserved_at_monotonic == 1_000.0 + await store.save_reservation("batch_ttl_clamp", reservation) + assert fake.save_ttls == (BATCH_ENQUEUED_TOKEN_TTL_SECONDS - 31,) + assert await store.pop_reservation("batch_ttl_clamp") == reservation + + @pytest.mark.asyncio async def test_memory_refund_skips_reservations_granted_by_another_worker(): store = _in_memory_store() From 01add582982a253c2fff3466b608c1d0409ed1ae Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 19:45:38 -0700 Subject: [PATCH 035/281] fix(cli): name why a login fell back to the token file lite ships with every install of litellm, but the keyring package it needs for keychain storage only ships with the cli extra. Such a user on a Mac was told 'No OS keychain available' about a machine that plainly has one, with nothing pointing at the missing package. The vault now reports which of the three unusable states it is in, so login can point at the install, name the kill switch, or report a genuinely absent keychain. --- litellm/litellm_core_utils/cli_keyring.py | 56 +++++++++++++------ litellm/litellm_core_utils/cli_token_utils.py | 26 +++++---- litellm/proxy/client/README.md | 2 +- litellm/proxy/client/cli/commands/auth.py | 41 +++++++++++--- tests/test_litellm/conftest.py | 18 ++++-- .../test_cli_token_utils.py | 23 ++++---- .../proxy/client/cli/test_auth_commands.py | 28 ++++++++++ 7 files changed, 141 insertions(+), 53 deletions(-) diff --git a/litellm/litellm_core_utils/cli_keyring.py b/litellm/litellm_core_utils/cli_keyring.py index fcbf5ada55a..b19b3d3bc83 100644 --- a/litellm/litellm_core_utils/cli_keyring.py +++ b/litellm/litellm_core_utils/cli_keyring.py @@ -5,9 +5,9 @@ SDK-level access to the OS keychain (macOS Keychain, Windows Credential Manager, Linux Secret Service) that holds the credential minted by `lite login`. The `keyring` package is optional and imported lazily, so importing this module -never pulls it in. Every failure is returned as a value: a machine with no -keychain, or one whose keychain is locked, must degrade to the token file rather -than break `lite` or the SDK. +never pulls it in. Every failure is returned as a value, naming which of the +three ways the keychain can be out of reach applies, so callers can degrade to +the token file and tell the user what to do about it. """ import os @@ -32,11 +32,28 @@ class SecretMissing: @dataclass(frozen=True, slots=True) -class SecretUnavailable: +class SecretStored: pass -SecretRead: TypeAlias = SecretFound | SecretMissing | SecretUnavailable +@dataclass(frozen=True, slots=True) +class KeyringNotInstalled: + pass + + +@dataclass(frozen=True, slots=True) +class KeyringDisabled: + pass + + +@dataclass(frozen=True, slots=True) +class KeyringUnreachable: + pass + + +KeyringUnusable: TypeAlias = KeyringNotInstalled | KeyringDisabled | KeyringUnreachable +SecretRead: TypeAlias = SecretFound | SecretMissing | KeyringUnusable +SecretWrite: TypeAlias = SecretStored | KeyringUnusable class SecretVault(Protocol): @@ -44,7 +61,7 @@ class SecretVault(Protocol): def read(self) -> SecretRead: ... - def write(self, blob: str) -> bool: ... + def write(self, blob: str) -> SecretWrite: ... def erase(self) -> bool: ... @@ -69,8 +86,11 @@ def _import_keyring() -> KeyringApi | None: return keyring -def _keyring_api() -> KeyringApi | None: - return None if _keyring_disabled() else _import_keyring() +def _keyring_api() -> KeyringApi | KeyringNotInstalled | KeyringDisabled: + if _keyring_disabled(): + return KeyringDisabled() + api: Final = _import_keyring() + return KeyringNotInstalled() if api is None else api @dataclass(frozen=True, slots=True) @@ -79,23 +99,23 @@ class KeyringVault: def read(self) -> SecretRead: api: Final = _keyring_api() - if api is None: - return SecretUnavailable() + if isinstance(api, (KeyringNotInstalled, KeyringDisabled)): + return api try: blob: Final = api.get_password(KEYRING_SERVICE, KEYRING_ACCOUNT) except Exception: # noqa: BLE001 # backends raise outside keyring.errors; never break the SDK - return SecretUnavailable() + return KeyringUnreachable() return SecretMissing() if blob is None else SecretFound(blob) - def write(self, blob: str) -> bool: + def write(self, blob: str) -> SecretWrite: api: Final = _keyring_api() - if api is None: - return False + if isinstance(api, (KeyringNotInstalled, KeyringDisabled)): + return api try: api.set_password(KEYRING_SERVICE, KEYRING_ACCOUNT, blob) except Exception: # noqa: BLE001 # a keychain that refuses the write falls back to the token file - return False - return True + return KeyringUnreachable() + return SecretStored() def erase(self) -> bool: """Whether the keychain is guaranteed to hold no credential afterwards. @@ -108,7 +128,7 @@ class KeyringVault: if _keyring_disabled(): return False match self.read(): - case SecretUnavailable(): + case KeyringNotInstalled() | KeyringDisabled() | KeyringUnreachable(): return False case SecretMissing(): return True @@ -117,7 +137,7 @@ class KeyringVault: def _delete(self) -> bool: api: Final = _keyring_api() - if api is None: + if isinstance(api, (KeyringNotInstalled, KeyringDisabled)): return False try: api.delete_password(KEYRING_SERVICE, KEYRING_ACCOUNT) diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index dd263ccd412..40822e5b335 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -22,10 +22,14 @@ from pydantic import BaseModel, ConfigDict, ValidationError from litellm.litellm_core_utils.cli_keyring import ( SYSTEM_KEYRING, + KeyringDisabled, + KeyringNotInstalled, + KeyringUnreachable, SecretFound, SecretMissing, - SecretUnavailable, + SecretStored, SecretVault, + SecretWrite, ) from litellm.litellm_core_utils.private_json import ensure_private_dir, write_private_json @@ -79,13 +83,15 @@ def load_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> CliTokenRecord | N return _resolve_secret(record, vault) -def save_cli_token(record: CliTokenRecord, *, vault: SecretVault = SYSTEM_KEYRING) -> bool: - """Store a freshly minted credential. Returns whether the keychain took the secret""" - if record.key is None or not vault.write(_encode_secret(record.base_url, record.key, record.jwt_token)): - _write_token_file(record) - return False - _write_token_file(_without_secret(record)) - return True +def save_cli_token(record: CliTokenRecord, *, vault: SecretVault = SYSTEM_KEYRING) -> SecretWrite: + """Store a freshly minted credential. Reports whether the keychain took the secret, and why not""" + outcome: Final = ( + SecretStored() + if record.key is None + else vault.write(_encode_secret(record.base_url, record.key, record.jwt_token)) + ) + _write_token_file(_without_secret(record) if isinstance(outcome, SecretStored) else record) + return outcome def clear_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> bool: @@ -163,7 +169,7 @@ def _resolve_secret(record: CliTokenRecord, vault: SecretVault) -> CliTokenRecor return _apply_vault_secret(record, blob, vault) case SecretMissing(): return _migrate_file_secret(record, vault) - case SecretUnavailable(): + case KeyringNotInstalled() | KeyringDisabled() | KeyringUnreachable(): return record @@ -188,7 +194,7 @@ def _apply_vault_secret(record: CliTokenRecord, blob: str, vault: SecretVault) - def _migrate_file_secret(record: CliTokenRecord, vault: SecretVault) -> CliTokenRecord | None: if record.key is None: return None - if vault.write(_encode_secret(record.base_url, record.key, record.jwt_token)): + if isinstance(vault.write(_encode_secret(record.base_url, record.key, record.jwt_token)), SecretStored): _scrub_file_secret(record) return record diff --git a/litellm/proxy/client/README.md b/litellm/proxy/client/README.md index 9ece4c2be3d..d46c7174b4c 100644 --- a/litellm/proxy/client/README.md +++ b/litellm/proxy/client/README.md @@ -377,7 +377,7 @@ The key itself goes into the OS keychain (macOS Keychain, Windows Credential Man } ``` -Headless boxes and CI runners usually have no keychain. There the key stays in the same `0600` file alongside the metadata, exactly as it did before, and `lite login` tells you which of the two happened. Set `LITELLM_CLI_DISABLE_KEYRING=1` to force the file even where a keychain exists. A `token.json` written by an older `lite` keeps working and is moved into the keychain, and scrubbed from the file, the first time a keychain-capable `lite` reads it. +Keychain storage needs the `keyring` package, which ships with `pip install 'litellm[cli]'`. Headless boxes and CI runners usually have no keychain either. In all of those cases the key stays in the same `0600` file alongside the metadata, exactly as it did before, and `lite login` names which one applies: the package is missing, the machine has no keychain, or you set `LITELLM_CLI_DISABLE_KEYRING=1` to force the file even where a keychain exists. A `token.json` written by an older `lite` keeps working and is moved into the keychain, and scrubbed from the file, the first time a keychain-capable `lite` reads it. `lite logout` clears both stores. If the keychain is locked at that moment it says so, and re-running it once the keychain is unlocked finishes the job. diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index c2b5b6a620f..03906f9b6df 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -11,7 +11,16 @@ from rich.table import Table from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.constants import CLI_JWT_EXPIRATION_HOURS -from litellm.litellm_core_utils.cli_keyring import SYSTEM_KEYRING, SecretVault +from litellm.litellm_core_utils.cli_keyring import ( + DISABLE_KEYRING_ENV_VAR, + SYSTEM_KEYRING, + KeyringDisabled, + KeyringNotInstalled, + KeyringUnreachable, + SecretStored, + SecretVault, + SecretWrite, +) from litellm.litellm_core_utils.cli_token_utils import ( CliTokenRecord, clear_cli_token, @@ -72,9 +81,29 @@ class CliAuthResult(TypedDict): KEYCHAIN_UNREACHABLE_MESSAGE: Final = ( - "Your credential is stored in your OS keychain, which could not be read. Unlock it and retry, or run 'lite login'." + "Your credential is stored in your OS keychain, which could not be read. Unlock it and retry, " + "install the keyring package with: pip install 'litellm[cli]', or run 'lite login'." ) +KEYRING_INSTALL_HINT: Final = "pip install 'litellm[cli]'" + + +def storage_notice(outcome: SecretWrite) -> str: + """Tell the user where the credential ended up, and how to get keychain storage if it did not.""" + path: Final = get_cli_token_file_path() + match outcome: + case SecretStored(): + return "Credential stored in your OS keychain." + case KeyringNotInstalled(): + return ( + f"Credential stored in {path} (owner-only). " + f"For OS keychain storage, install the keyring package with: {KEYRING_INSTALL_HINT}" + ) + case KeyringDisabled(): + return f"Keychain storage is off ({DISABLE_KEYRING_ENV_VAR}). Credential stored in {path} (owner-only)." + case KeyringUnreachable(): + return f"No OS keychain available. Credential stored in {path} (owner-only)." + def context_secret_vault(ctx: click.Context) -> SecretVault: """Where this invocation reads and writes secret material; injectable through ctx.obj for tests""" @@ -675,15 +704,11 @@ def login(ctx: click.Context, config_claude: bool): jwt_token="", timestamp=time.time(), ) - in_keychain: Final = save_cli_token(record, vault=context_secret_vault(ctx)) + stored: Final = save_cli_token(record, vault=context_secret_vault(ctx)) click.echo("\nLogin successful!") click.echo(f"JWT Token: {api_key[:20]}...") - click.echo( - "Credential stored in your OS keychain." - if in_keychain - else f"No OS keychain available; credential stored in {get_cli_token_file_path()} (owner-only)." - ) + click.echo(storage_notice(stored)) click.echo("You can now use the CLI without specifying --api-key") if config_claude: diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index ce0fd197538..a716ec0e0aa 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -23,10 +23,13 @@ from litellm import router as litellm_router_module from litellm import utils as litellm_utils_module from litellm._logging import ALL_LOGGERS from litellm.litellm_core_utils.cli_keyring import ( + KeyringUnreachable, + KeyringUnusable, SecretFound, SecretMissing, SecretRead, - SecretUnavailable, + SecretStored, + SecretWrite, ) from litellm.litellm_core_utils.prompt_templates import ( image_handling as image_handling_module, @@ -125,7 +128,8 @@ class FakeSecretVault: """In-memory stand-in for the OS keychain, injected wherever CLI credential storage is exercised. `available=False` models a keychain that is locked or has no backend, `writable=False` one that - refuses to store, and `erasable=False` one that will not release what it already holds. + refuses to store, `erasable=False` one that will not release what it already holds, and `failure` + picks which unusable state those report. """ def __init__( @@ -135,11 +139,13 @@ class FakeSecretVault: available: bool = True, writable: bool = True, erasable: bool = True, + failure: KeyringUnusable = KeyringUnreachable(), ) -> None: self.blob: str | None = blob self.available: bool = available self.writable: bool = writable self.erasable: bool = erasable + self.failure: KeyringUnusable = failure self.reads: int = 0 self.writes: list[str] = [] self.erases: int = 0 @@ -147,15 +153,15 @@ class FakeSecretVault: def read(self) -> SecretRead: self.reads += 1 if not self.available: - return SecretUnavailable() + return self.failure return SecretMissing() if self.blob is None else SecretFound(self.blob) - def write(self, blob: str) -> bool: + def write(self, blob: str) -> SecretWrite: self.writes.append(blob) if not (self.available and self.writable): - return False + return self.failure self.blob = blob - return True + return SecretStored() def erase(self) -> bool: self.erases += 1 diff --git a/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py b/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py index 56ab6bcbfe0..961d986e8f5 100644 --- a/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py @@ -13,7 +13,10 @@ from litellm.litellm_core_utils.cli_keyring import ( KeyringVault, SecretFound, SecretMissing, - SecretUnavailable, + KeyringDisabled, + KeyringNotInstalled, + KeyringUnreachable, + SecretStored, ) from litellm.litellm_core_utils.cli_token_utils import ( CliTokenRecord, @@ -253,7 +256,7 @@ class TestSaveCliToken: vault=vault, ) - assert stored is True + assert stored == SecretStored() assert "sk-new" not in _token_file(isolated_home).read_text() assert json.loads(vault.blob)["key"] == "sk-new" assert load_cli_token(vault=vault).key == "sk-new" @@ -265,7 +268,7 @@ class TestSaveCliToken: ) path = _token_file(isolated_home) - assert stored is False + assert stored == KeyringUnreachable() assert json.loads(path.read_text())["key"] == "sk-new" assert stat.S_IMODE(path.stat().st_mode) == 0o600 assert list(path.parent.glob(".tmp-*")) == [] @@ -379,7 +382,7 @@ class TestKeyringVault: fake = install_fake_keyring(_FakeKeyringModule()) vault = KeyringVault() - assert vault.write("blob-1") is True + assert vault.write("blob-1") == SecretStored() assert vault.read() == SecretFound("blob-1") assert vault.erase() is True assert vault.read() == SecretMissing() @@ -393,8 +396,8 @@ class TestKeyringVault: monkeypatch.setenv(DISABLE_KEYRING_ENV_VAR, "1") vault = KeyringVault() - assert vault.read() == SecretUnavailable() - assert vault.write("blob-1") is False + assert vault.read() == KeyringDisabled() + assert vault.write("blob-1") == KeyringDisabled() assert vault.erase() is False def test_an_uninstalled_keyring_library_degrades_to_the_file(self, monkeypatch): @@ -404,19 +407,19 @@ class TestKeyringVault: monkeypatch.setitem(sys.modules, "keyring", None) vault = KeyringVault() - assert vault.read() == SecretUnavailable() - assert vault.write("blob-1") is False + assert vault.read() == KeyringNotInstalled() + assert vault.write("blob-1") == KeyringNotInstalled() assert vault.erase() is True def test_a_locked_keychain_is_reported_not_raised(self, install_fake_keyring): install_fake_keyring(_FakeKeyringModule(get_error=RuntimeError("keyring is locked"))) - assert KeyringVault().read() == SecretUnavailable() + assert KeyringVault().read() == KeyringUnreachable() def test_a_refused_write_is_reported_not_raised(self, install_fake_keyring): install_fake_keyring(_FakeKeyringModule(set_error=RuntimeError("no backend"))) - assert KeyringVault().write("blob-1") is False + assert KeyringVault().write("blob-1") == KeyringUnreachable() def test_a_refused_delete_is_reported_so_logout_can_warn(self, install_fake_keyring): install_fake_keyring(_FakeKeyringModule(stored="blob-1", delete_error=RuntimeError("locked"))) diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index e93f05cb4aa..1f2d48f6547 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -13,6 +13,11 @@ import pytest from click.testing import CliRunner from litellm.constants import CLI_JWT_EXPIRATION_HOURS +from litellm.litellm_core_utils.cli_keyring import ( + DISABLE_KEYRING_ENV_VAR, + KeyringDisabled, + KeyringNotInstalled, +) from litellm.litellm_core_utils.cli_token_utils import CliTokenRecord, save_cli_token from litellm.proxy.client.cli import cli from litellm.proxy.client.cli.commands.auth import ( @@ -931,6 +936,29 @@ class TestKeychainBackedCommands: assert str(token_file) in result.output assert json.loads(token_file.read_text())["key"] == "sk-minted" + def test_login_points_a_user_missing_the_keyring_package_at_the_install( + self, isolated_home, secret_vault_factory + ): + """`lite` ships with every install, the keyring package only with the cli extra. Telling + that user their machine has no keychain sends them looking for a problem they do not have.""" + result = self._login(secret_vault_factory(available=False, failure=KeyringNotInstalled())) + + token_file = isolated_home / ".litellm" / "token.json" + assert result.exit_code == 0 + assert "pip install 'litellm[cli]'" in result.output + assert "No OS keychain available" not in result.output + assert json.loads(token_file.read_text())["key"] == "sk-minted" + + def test_login_names_the_kill_switch_instead_of_blaming_the_machine( + self, isolated_home, secret_vault_factory + ): + result = self._login(secret_vault_factory(available=False, failure=KeyringDisabled())) + + assert result.exit_code == 0 + assert DISABLE_KEYRING_ENV_VAR in result.output + assert "No OS keychain available" not in result.output + assert json.loads((isolated_home / ".litellm" / "token.json").read_text())["key"] == "sk-minted" + def test_whoami_and_print_token_read_through_the_keychain(self, isolated_home, secret_vault_factory): vault = secret_vault_factory() self._login(vault) From 0b374541bb302697ea66135fdeb04bdfb4d3a44b Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 19 Aug 2026 20:01:07 -0700 Subject: [PATCH 036/281] refactor(ui): migrate the last antd components off antd onto shadcn (#37569) Converts the remaining dashboard components that still imported antd: admin panel, agents, MCP toolsets, policies, prompts, bulk user edit, create user, plugin settings, teams, add model, auto router, cloudzero export, BYOK credentials, credential modal, onboarding link, create key and routing groups. Primitives map onto the house shadcn set: Typography onto semantic tags, Select onto ui/select, SearchSelect or MultiSelect, Input onto ui/input, Tooltip onto SimpleTooltip, Card, Table, Tabs, Switch, Checkbox, Radio, Tag onto Badge, Divider onto Separator, Spin onto UiLoadingSpinner, Modal onto Dialog, message onto toast, and Space, Row, Col, Flex and Layout onto flex containers. --- tests/e2e/ui/helpers/mcp.ts | 18 +- tests/e2e/ui/tests/mcp/mcpServers.spec.ts | 27 +- .../e2e/ui/tests/modelsPage/addModel.spec.ts | 61 +- .../ui/tests/modelsPage/credentials.spec.ts | 2 +- tests/e2e/ui/tests/proxy-admin/keys.spec.ts | 27 +- tests/e2e/ui/tests/proxy-admin/teams.spec.ts | 13 +- .../ui/tests/settings/routerSettings.spec.ts | 26 +- .../e2e/ui/tests/team-admin/teamAdmin.spec.ts | 9 +- tests/e2e/ui/tests/usage/usagePage.spec.ts | 3 +- tests/e2e/ui/tests/users/searchUsers.spec.ts | 2 +- .../ui/tests/users/viewInternalUsers.spec.ts | 2 +- ui/litellm-dashboard/eslint-suppressions.json | 63 -- .../admin-panel/_components/AdminPanel.tsx | 38 +- .../_components/add_agent_form.test.tsx | 15 +- .../agents/_components/add_agent_form.tsx | 178 +++--- .../_components/MCPToolsetsTab.tsx | 40 +- .../policies/_components/add_policy_form.tsx | 5 +- .../prompts/_components/add_prompt_form.tsx | 96 +-- .../users/_components/BulkEditUsers.tsx | 236 +++---- .../src/components/CreateUserButton.test.tsx | 54 +- .../src/components/CreateUserButton.tsx | 206 ++++--- .../PluginSettings/PluginSettings.tsx | 150 +++-- ui/litellm-dashboard/src/components/Teams.tsx | 65 +- .../src/components/add_model/AddModelForm.tsx | 576 +++++++++--------- .../add_model/add_auto_router_tab.tsx | 368 +++++------ ...loudzero_export_modal.integration.test.tsx | 2 +- .../src/components/cloudzero_export_modal.tsx | 23 +- .../mcp_tools/ByokCredentialModal.tsx | 18 +- .../components/model_add/CredentialModal.tsx | 40 +- .../src/components/onboarding_link.tsx | 6 +- .../create_key_button.integration.test.tsx | 38 +- .../organisms/create_key_button.tsx | 428 +++++++------ .../src/components/routing_groups/index.tsx | 102 ++-- .../update_model_credentials_modal.tsx | 7 +- 34 files changed, 1513 insertions(+), 1431 deletions(-) diff --git a/tests/e2e/ui/helpers/mcp.ts b/tests/e2e/ui/helpers/mcp.ts index b41aec59ded..554177e11bc 100644 --- a/tests/e2e/ui/helpers/mcp.ts +++ b/tests/e2e/ui/helpers/mcp.ts @@ -12,23 +12,21 @@ export async function createMcpServer(page: PwPage, url: string): Promise { await expect(discovery).toBeVisible({ timeout: 5_000 }); await discovery.getByRole("button", { name: /Custom Server/i }).click(); - const formModal = page.locator(".ant-modal:visible").filter({ hasText: "MCP Server Name" }); + const formModal = page.getByRole("dialog").filter({ hasText: "MCP Server Name" }); await expect(formModal).toBeVisible({ timeout: 5_000 }); // Name — no spaces or hyphens per validateMCPServerName const uniqueName = `e2e_mcp_${Date.now()}`; createdServerName = uniqueName; - await formModal.locator('input[id="server_name"]').fill(uniqueName); + await formModal.getByLabel("MCP Server Name").fill(uniqueName); - // Transport: Streamable HTTP — the only value the proxy actually accepts is "http" - const transportField = formModal.locator(".ant-form-item", { hasText: "Transport Type" }); - await transportField.locator(".ant-select").click(); - await page.locator(".ant-select-dropdown:visible").getByText("Streamable HTTP").click(); + // Transport: Streamable HTTP — the only value the proxy actually accepts is "http". + // Select popups are portaled to the body, so the option lookup is page-scoped. + await formModal.getByRole("combobox", { name: "Transport Type" }).click(); + await page.getByRole("option", { name: "Streamable HTTP" }).click(); // URL — use a fake URL; the form just persists it, it doesn't have to be reachable - await formModal.locator('input[id="url"]').fill("https://e2e-fake-mcp.test.local/mcp"); + await formModal.getByLabel("MCP Server URL").fill("https://e2e-fake-mcp.test.local/mcp"); - // Authentication: None - // The auth_type Form.Item has no label prop (CreateMCPServer.tsx), so - // it can't be anchored by label text. Scope via the enclosing Collapse - // panel ("Authentication") instead — that anchor is stable even if the - // placeholder copy changes. - const authSection = formModal.locator(".ant-collapse-item", { hasText: /^Authentication/ }); - const authField = authSection.locator(".ant-form-item").first(); - await authField.locator(".ant-select").click(); - await page.locator(".ant-select-dropdown:visible").getByText("None", { exact: true }).click(); + // Authentication: None. "Authentication" is exact so it can't also match the + // "Authentication Value" field that some auth types reveal below it. + await formModal.getByRole("combobox", { name: "Authentication", exact: true }).click(); + await page.getByRole("option", { name: "None", exact: true }).click(); // Submit await formModal.getByRole("button", { name: /^Add MCP Server$/ }).click(); diff --git a/tests/e2e/ui/tests/modelsPage/addModel.spec.ts b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts index bd5373e1569..dad716b4c83 100644 --- a/tests/e2e/ui/tests/modelsPage/addModel.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts @@ -21,15 +21,22 @@ async function findDeploymentByName(page: PlaywrightPage, modelName: string): Pr return body.data.find((row) => row.model_name === modelName); } +/** Anchors a substring match to the whole string, escaping regex metacharacters. */ +const exactly = (text: string): RegExp => new RegExp(`^${text.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}$`); + /** - * Helper to select a provider from the Add Model form dropdown. + * Helper to select a provider from the Add Model form dropdown. The field is a + * searchable combobox: it only opens on click, typing filters the list, and the + * option has to be picked explicitly because nothing is highlighted by default. + * Options are matched on their visible text, not their accessible name, which + * also carries the provider logo's alt text ("Anthropic logo Anthropic"). */ -async function selectProvider(page: any, providerName: string) { - const providerDropdown = page.getByRole("combobox", { name: /Provider/i }); +async function selectProvider(page: PlaywrightPage, providerName: string) { + const providerDropdown = page.getByRole("combobox", { name: "Provider", exact: true }); + await providerDropdown.click(); await providerDropdown.fill(providerName); - await page.waitForTimeout(1000); - await providerDropdown.press("Enter"); - await page.waitForTimeout(2000); + await page.getByRole("option").filter({ hasText: exactly(providerName) }).click(); + await expect(providerDropdown).toHaveValue(providerName); } test.describe("Add Model", () => { @@ -64,11 +71,10 @@ test.describe("Add Model", () => { await selectProvider(page, "Anthropic"); // The model field should be a multi-select dropdown; click to open it - const modelDropdown = page.locator(".ant-select-selection-overflow").first(); - await modelDropdown.click(); + await page.getByRole("combobox", { name: "Select models" }).click(); // Verify provider-specific models are listed - await expect(page.getByTitle("claude-haiku-4-5", { exact: true })).toBeVisible(); + await expect(page.getByRole("option", { name: "claude-haiku-4-5", exact: true })).toBeVisible(); }); test("Edit team model TPM and RPM limits", async ({ page }) => { @@ -156,14 +162,14 @@ test.describe("Add Model", () => { await page.getByRole("tab", { name: "Add Model" }).click(); // Labels come from /public/providers/fields, not the frontend Providers enum, and the two differ. - await selectProvider(page, "OpenAI-Compatible Endpoints"); + await selectProvider(page, "OpenAI-Compatible Endpoints (Together AI, etc.)"); const publicName = `e2e-ui-added-${Date.now()}`; uiAddedModelName = publicName; // The model picker's "custom" entry reveals the free-text name field. - await page.locator(".ant-select-selection-overflow").first().click(); - await page.locator(".ant-select-dropdown:visible").getByText("Custom Model Name (Enter below)").click(); + await page.getByRole("combobox", { name: "Select models" }).click(); + await page.getByRole("option", { name: "Custom Model Name (Enter below)" }).click(); await page.keyboard.press("Escape"); await page.getByPlaceholder("Enter custom model name").fill(publicName); @@ -177,8 +183,8 @@ test.describe("Add Model", () => { await expect(page.getByTestId("connection-success-msg")).toBeVisible({ timeout: 30_000 }); // The modal swallows the Add click. Scope to the footer: the dismiss X is also named "Close". - const resultsModal = page.locator(".ant-modal:visible").filter({ hasText: "Connection Test Results" }); - await resultsModal.locator(".ant-modal-footer").getByRole("button", { name: "Close" }).click(); + const resultsModal = page.getByRole("dialog", { name: "Connection Test Results" }); + await resultsModal.locator('[data-slot="dialog-footer"]').getByRole("button", { name: "Close" }).click(); await expect(resultsModal).toBeHidden({ timeout: 5_000 }); const created = await captureRequestBody(page, { method: "POST", urlIncludes: "/model/new" }, async () => { @@ -213,9 +219,8 @@ test.describe("Add Model", () => { await selectProvider(page, "Anthropic"); // Select model: claude-haiku-4-5 - const modelDropdown = page.locator(".ant-select-selection-overflow").first(); - await modelDropdown.click(); - await page.getByTitle("claude-haiku-4-5", { exact: true }).click(); + await page.getByRole("combobox", { name: "Select models" }).click(); + await page.getByRole("option", { name: "claude-haiku-4-5", exact: true }).click(); await page.keyboard.press("Escape"); // Enter bad API key @@ -239,9 +244,8 @@ test.describe("Add Model", () => { await selectProvider(page, "Anthropic"); // Select model: claude-haiku-4-5 - const modelDropdown = page.locator(".ant-select-selection-overflow").first(); - await modelDropdown.click(); - await page.getByTitle("claude-haiku-4-5", { exact: true }).click(); + await page.getByRole("combobox", { name: "Select models" }).click(); + await page.getByRole("option", { name: "claude-haiku-4-5", exact: true }).click(); await page.keyboard.press("Escape"); // Enter any API key @@ -315,18 +319,15 @@ test.describe("Add Model", () => { await selectProvider(page, "Cohere"); - const modelDropdown = page.locator(".ant-select-selection-overflow").first(); - await modelDropdown.click(); - const wildcardOption = page.getByTitle(/All .* Models \(Wildcard\)/); - await wildcardOption.click(); + await page.getByRole("combobox", { name: "Select models" }).click(); + await page.getByRole("option", { name: /All .* Models \(Wildcard\)/ }).click(); await page.keyboard.press("Escape"); const apiKeyInput = page.locator('input[type="password"]').first(); await apiKeyInput.fill("sk-any-key-for-team-byok-test"); - // Flip the Team-BYOK switch on (Form.Item label "Team-BYOK Model") - const teamByokRow = page.locator(".ant-form-item", { hasText: "Team-BYOK Model" }); - await teamByokRow.getByRole("switch").click(); + // Flip the Team-BYOK switch on; the Switch carries its own aria-label. + await page.getByRole("switch", { name: "Team-BYOK Model" }).click(); // TeamDropdown options show the alias above the team id, so match on the id line by text. const teamDropdown = page.getByTestId("team-dropdown").getByRole("combobox"); @@ -376,10 +377,8 @@ test.describe("Add Model", () => { await selectProvider(page, "Cohere"); // Select All Cohere Models (Wildcard) - const modelDropdown = page.locator(".ant-select-selection-overflow").first(); - await modelDropdown.click(); - const wildcardOption = page.getByTitle(/All .* Models \(Wildcard\)/); - await wildcardOption.click(); + await page.getByRole("combobox", { name: "Select models" }).click(); + await page.getByRole("option", { name: /All .* Models \(Wildcard\)/ }).click(); await page.keyboard.press("Escape"); // Enter any API key diff --git a/tests/e2e/ui/tests/modelsPage/credentials.spec.ts b/tests/e2e/ui/tests/modelsPage/credentials.spec.ts index 7c836068567..ceedc959ccc 100644 --- a/tests/e2e/ui/tests/modelsPage/credentials.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/credentials.spec.ts @@ -41,7 +41,7 @@ test.describe("Edit LLM credential", () => { await row.getByTestId(`credential-actions-${credentialName}`).click(); await page.getByTestId("credential-action-edit").click(); - const modal = page.locator(".ant-modal-content").filter({ hasText: "Edit Credential" }); + const modal = page.getByRole("dialog", { name: "Edit Credential" }); await expect(modal).toBeVisible({ timeout: 10_000 }); const apiKeyField = modal.locator("#api_key"); diff --git a/tests/e2e/ui/tests/proxy-admin/keys.spec.ts b/tests/e2e/ui/tests/proxy-admin/keys.spec.ts index 004fedb3263..0c38641dcc7 100644 --- a/tests/e2e/ui/tests/proxy-admin/keys.spec.ts +++ b/tests/e2e/ui/tests/proxy-admin/keys.spec.ts @@ -45,9 +45,9 @@ test.describe("Proxy Admin - Keys", () => { await page.keyboard.type(E2E_TEAM_CRUD_ALIAS); await page.locator('[data-slot="combobox-content"]:visible').getByText(E2E_TEAM_CRUD_ALIAS).first().click(); - // Select models - await page.locator(".ant-select-selection-overflow").click(); - await page.locator(".ant-select-dropdown:visible").getByText("All Team Models").click(); + // Select models — the popup is portaled to the body, so scope options to the page. + await page.getByRole("combobox", { name: "Select models" }).click(); + await page.getByRole("option", { name: "All Team Models", exact: true }).click(); await page.keyboard.press("Escape"); // Submit @@ -86,7 +86,7 @@ test.describe("Proxy Admin - Keys", () => { // Scope to the modal — the Regenerate button has an icon whose aria-label // ("sync") is concatenated into the button's accessible name, and the // "Regenerate Key" button is still in the DOM behind the modal. - const modal = page.locator(".ant-modal:visible"); + const modal = page.getByRole("dialog", { name: "Regenerate Virtual Key" }); await modal.getByRole("button", { name: /Regenerate/ }).click(); // Success view shows a Copy button in the footer (text varies between modal versions) @@ -198,8 +198,8 @@ test.describe("Proxy Admin - Keys", () => { // Select models — open the multi-select and pick the all-models meta-option. // With no team selected the modal offers "All Proxy Models"; the team-scoped // "All Team Models" option only appears once a team is picked. - await page.locator(".ant-select-selection-overflow").click(); - await page.locator(".ant-select-dropdown:visible").getByText("All Proxy Models").click(); + await page.getByRole("combobox", { name: "Select models" }).click(); + await page.getByRole("option", { name: "All Proxy Models", exact: true }).click(); await page.keyboard.press("Escape"); await page.getByRole("button", { name: "Create Key", exact: true }).click(); @@ -221,17 +221,10 @@ test.describe("Proxy Admin - Keys", () => { const keyName = `e2e-admin-specific-${Date.now()}`; await page.getByLabel(/Key Name/).fill(keyName); - // Open the model multi-select and pick a single specific model. Use - // getByRole("option", ...) to avoid the strict-mode collision between - // the option container and its inner text node. + // Open the model multi-select and pick a single specific model. const modelName = "fake-openai-gpt-4"; - await page.locator(".ant-select-selection-overflow").click(); - const option = page.locator(".ant-select-dropdown:visible").getByRole("option", { name: modelName, exact: true }); - await option.waitFor({ state: "attached" }); - // Dispatch the click via the DOM — antd's dropdown can render the option - // off-viewport during the open animation, which trips Playwright's - // visibility/stability checks. The click handler fires regardless. - await option.evaluate((el: HTMLElement) => el.click()); + await page.getByRole("combobox", { name: "Select models" }).click(); + await page.getByRole("option", { name: modelName, exact: true }).click(); await page.keyboard.press("Escape"); await page.getByRole("button", { name: "Create Key", exact: true }).click(); @@ -242,7 +235,7 @@ test.describe("Proxy Admin - Keys", () => { // verify it can call /chat/completions for the model it was scoped to. // The mock LLM server (fixtures/mock_llm_server/server.py) replies with // a fixed "This is a mock response." body. - const apiKey = (await page.locator(".ant-modal:visible pre").innerText()).trim(); + const apiKey = (await page.getByRole("dialog", { name: "Save your Key" }).locator("pre").innerText()).trim(); expect(apiKey).toMatch(/^sk-/); const response = await page.request.post("/chat/completions", { diff --git a/tests/e2e/ui/tests/proxy-admin/teams.spec.ts b/tests/e2e/ui/tests/proxy-admin/teams.spec.ts index 172fde173df..7383b452162 100644 --- a/tests/e2e/ui/tests/proxy-admin/teams.spec.ts +++ b/tests/e2e/ui/tests/proxy-admin/teams.spec.ts @@ -41,11 +41,12 @@ test.describe("Proxy Admin - Teams", () => { .click(); // Wait for the Create Team modal - const dialog = page.locator(".ant-modal:visible"); + const dialog = page.getByRole("dialog", { name: "Create Team" }); await expect(dialog).toBeVisible({ timeout: 5_000 }); - // Fill Team Name — the input has id="team_alias" - await dialog.locator("#team_alias").fill(uniqueAlias); + // Fill Team Name — FormField derives the control id from React.useId(), so + // the input is only addressable by its label or its test id. + await dialog.getByTestId("team-name-input").fill(uniqueAlias); // Select models — the models multi-select is inside the modal. Its popup is // portaled to the body, so scope the option lookup to the page, not the dialog. @@ -75,7 +76,7 @@ test.describe("Proxy Admin - Teams", () => { await page.getByRole("button", { name: /Add Member/i }).click(); // Wait for Add Team Member modal - const modal = page.locator(".ant-modal:visible"); + const modal = page.getByRole("dialog", { name: "Add Team Member" }); await expect(modal).toBeVisible({ timeout: 5_000 }); // The email field is a Select — type to search, then select from dropdown @@ -112,7 +113,7 @@ test.describe("Proxy Admin - Teams", () => { await page.getByTestId("edit-member").first().click(); - const modal = page.locator(".ant-modal:visible"); + const modal = page.getByRole("dialog", { name: "Edit Member" }); await expect(modal).toBeVisible({ timeout: 5_000 }); await modal.getByRole("button", { name: /Save Changes/i }).click(); @@ -155,7 +156,7 @@ test.describe("Proxy Admin - Teams", () => { await page.getByTestId("edit-member").first().click(); - const modal = page.locator(".ant-modal:visible"); + const modal = page.getByRole("dialog", { name: "Edit Member" }); await expect(modal).toBeVisible({ timeout: 5_000 }); await modal.getByRole("button", { name: /Save Changes/i }).click(); diff --git a/tests/e2e/ui/tests/settings/routerSettings.spec.ts b/tests/e2e/ui/tests/settings/routerSettings.spec.ts index 9784abff040..1188e8f201e 100644 --- a/tests/e2e/ui/tests/settings/routerSettings.spec.ts +++ b/tests/e2e/ui/tests/settings/routerSettings.spec.ts @@ -67,29 +67,25 @@ test.describe("Router Settings - Fallbacks", () => { await page.getByRole("button", { name: /Add Fallbacks/i }).click(); await modelsLoaded; - const modal = page.locator(".ant-modal:visible"); + const modal = page.getByRole("dialog", { name: "Configure Model Fallbacks" }); await expect(modal).toBeVisible({ timeout: 5_000 }); - // FallbackGroupConfig.tsx renders both selects with `showSearch`. The - // most stable interaction is: click to open + focus, type the model name to - // narrow the listbox to a single highlighted option, then press Enter. - // Verify each selection landed by watching the dialog's own state transition - // (the tab title updates to the picked primary; the fallback chain list - // populates) rather than by asserting on the dropdown popup, which sits in - // a custom getPopupContainer and is awkward to scope reliably. - const primarySelect = modal.locator(".ant-select").filter({ hasText: "Select primary model" }); - await primarySelect.click(); + // FallbackGroupConfig.tsx renders both fields as searchable comboboxes: they + // open on click, typing filters the listbox, and the option has to be picked + // explicitly. Verify each selection landed by watching the dialog's own state + // transition (the tab title updates to the picked primary; the fallback chain + // list populates) rather than by asserting on the popup, which is portaled + // out of the dialog. + await modal.getByRole("combobox", { name: /Primary Model/ }).click(); await page.keyboard.type(PRIMARY); - await page.keyboard.press("Enter"); + await page.getByRole("option", { name: PRIMARY, exact: true }).click(); await expect(modal.getByRole("tab", { name: PRIMARY })).toBeVisible({ timeout: 10_000, }); - const fallbackSelect = modal.locator(".ant-select").filter({ hasText: "Select fallback models" }); - await fallbackSelect.click(); + await modal.getByRole("combobox", { name: /Select fallback models/ }).click(); await page.keyboard.type(FALLBACK); - await page.keyboard.press("Enter"); - await page.keyboard.press("Escape"); + await page.getByRole("option", { name: FALLBACK, exact: true }).click(); // The Fallback Chain helper text reads "(N/10 used)"; once it ticks to 1 the // selection has been recorded. await expect(modal.getByText("(1/10 used)")).toBeVisible({ diff --git a/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts index 0ef74e71529..f93cca75347 100644 --- a/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts +++ b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts @@ -61,7 +61,7 @@ test.describe("Team Admin", () => { await page.getByRole("tab", { name: "Members" }).click(); await page.getByRole("button", { name: /Add Member/i }).click(); - const modal = page.locator(".ant-modal:visible"); + const modal = page.getByRole("dialog", { name: "Add Team Member" }); await expect(modal).toBeVisible({ timeout: 5_000 }); // Use a dedicated invitee user so this doesn't race with the proxy-admin @@ -144,9 +144,10 @@ test.describe("Team Admin", () => { await page.keyboard.type(E2E_TEAM_CRUD_ALIAS); await page.locator('[data-slot="combobox-content"]:visible').getByText(E2E_TEAM_CRUD_ALIAS).first().click(); - // Models — pick "All Team Models" - await page.locator(".ant-select-selection-overflow").click(); - await page.locator(".ant-select-dropdown:visible").getByText("All Team Models").click(); + // Models — pick "All Team Models". The popup is portaled to the body, so + // scope the option lookup to the page. + await page.getByRole("combobox", { name: "Select models" }).click(); + await page.getByRole("option", { name: "All Team Models", exact: true }).click(); await page.keyboard.press("Escape"); const generate = await captureRequestBody(page, { method: "POST", urlIncludes: "/key/generate" }, async () => { diff --git a/tests/e2e/ui/tests/usage/usagePage.spec.ts b/tests/e2e/ui/tests/usage/usagePage.spec.ts index 6031aa54055..8fa59beb905 100644 --- a/tests/e2e/ui/tests/usage/usagePage.spec.ts +++ b/tests/e2e/ui/tests/usage/usagePage.spec.ts @@ -22,7 +22,8 @@ async function openUsage(page: PlaywrightPage): Promise { const card = topKeysCard(page); await expect(card).toBeVisible({ timeout: 30_000 }); // Widen past the default top-5 so other keys in the database cannot crowd this one out. - await card.locator(".ant-segmented-item").filter({ hasText: /^50$/ }).click(); + // The radio itself is sr-only and its label covers it, so click the label. + await card.getByRole("radiogroup", { name: "Number of top keys to show" }).getByText("50", { exact: true }).click(); return card; } diff --git a/tests/e2e/ui/tests/users/searchUsers.spec.ts b/tests/e2e/ui/tests/users/searchUsers.spec.ts index a9b0e329a2b..e87218b5a5e 100644 --- a/tests/e2e/ui/tests/users/searchUsers.spec.ts +++ b/tests/e2e/ui/tests/users/searchUsers.spec.ts @@ -11,7 +11,7 @@ test.skip("Internal Users Search", () => { await tab.click(); await expect(page.locator("tbody tr").first()).toBeVisible(); - await expect(page.locator(".ant-skeleton")).toHaveCount(0); + await expect(page.locator('[data-slot="skeleton"]')).toHaveCount(0); } test("can search users by email", async ({ page }) => { diff --git a/tests/e2e/ui/tests/users/viewInternalUsers.spec.ts b/tests/e2e/ui/tests/users/viewInternalUsers.spec.ts index ea61c238c02..614191372d0 100644 --- a/tests/e2e/ui/tests/users/viewInternalUsers.spec.ts +++ b/tests/e2e/ui/tests/users/viewInternalUsers.spec.ts @@ -13,7 +13,7 @@ test.skip("Internal Users Page", () => { const firstRow = page.locator("tbody tr").first(); await expect(firstRow).toBeVisible(); - await expect(page.locator(".ant-skeleton")).toHaveCount(0); + await expect(page.locator('[data-slot="skeleton"]')).toHaveCount(0); } test("renders internal users table correctly", async ({ page }) => { diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 2d68ea2aa68..a9a163f3c78 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -5,9 +5,6 @@ } }, "src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx": { - "no-restricted-imports": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -22,9 +19,6 @@ "no-nested-ternary": { "count": 3 }, - "no-restricted-imports": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 2 } @@ -535,9 +529,6 @@ "no-nested-ternary": { "count": 1 }, - "no-restricted-imports": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -871,9 +862,6 @@ "local/filename-pascal-case": { "count": 1 }, - "no-restricted-imports": { - "count": 1 - }, "react-hooks/immutability": { "count": 2 }, @@ -997,9 +985,6 @@ "src/app/(dashboard)/prompts/_components/add_prompt_form.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 2 } }, "src/app/(dashboard)/prompts/_components/index.tsx": { @@ -1187,9 +1172,6 @@ } }, "src/app/(dashboard)/users/_components/BulkEditUsers.tsx": { - "no-restricted-imports": { - "count": 1 - }, "prefer-const": { "count": 1 } @@ -1323,9 +1305,6 @@ } }, "src/components/CreateUserButton.tsx": { - "no-restricted-imports": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 1 } @@ -1373,11 +1352,6 @@ "count": 1 } }, - "src/components/Settings/AdminSettings/PluginSettings/PluginSettings.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.test.tsx": { "max-nested-callbacks": { "count": 1 @@ -1428,9 +1402,6 @@ "no-nested-ternary": { "count": 2 }, - "no-restricted-imports": { - "count": 1 - }, "prefer-const": { "count": 2 }, @@ -1462,14 +1433,8 @@ } }, "src/components/add_model/AddModelForm.tsx": { - "local/no-complex-jsx-arrow": { - "count": 1 - }, "no-nested-ternary": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/components/add_model/RouterConfigBuilder.tsx": { @@ -1480,9 +1445,6 @@ "src/components/add_model/add_auto_router_tab.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/components/add_model/add_model_modes.tsx": { @@ -1639,9 +1601,6 @@ "local/filename-pascal-case": { "count": 1 }, - "no-restricted-imports": { - "count": 1 - }, "no-restricted-syntax": { "count": 3 }, @@ -1815,11 +1774,6 @@ "count": 1 } }, - "src/components/mcp_tools/ByokCredentialModal.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/mcp_tools/MCPToolArgumentsForm.tsx": { "no-nested-ternary": { "count": 1 @@ -1835,11 +1789,6 @@ "count": 1 } }, - "src/components/model_add/CredentialModal.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/model_add/reuse_credentials.tsx": { "local/filename-pascal-case": { "count": 1 @@ -1907,9 +1856,6 @@ "src/components/onboarding_link.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/components/organisms/create_key_button.tsx": { @@ -1919,9 +1865,6 @@ "max-lines": { "count": 1 }, - "no-restricted-imports": { - "count": 1 - }, "prefer-const": { "count": 2 }, @@ -2010,9 +1953,6 @@ "local/filename-pascal-case": { "count": 1 }, - "no-restricted-imports": { - "count": 1 - }, "react-hooks/preserve-manual-memoization": { "count": 1 } @@ -2377,9 +2317,6 @@ "src/components/update_model_credentials_modal.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "no-restricted-imports": { - "count": 1 } }, "src/components/user_agent_activity.tsx": { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx index 903676d238a..976fb94acea 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx @@ -7,7 +7,7 @@ import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; import { Button } from "@/components/ui/button"; import { Card } from "@/components/ui/card"; import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; -import { Space, Tabs, Typography } from "antd"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { Info, TriangleAlert } from "lucide-react"; import React, { useEffect, useState } from "react"; import NewBadge from "@/components/common_components/NewBadge"; @@ -35,8 +35,6 @@ import { Input } from "@/components/ui/input"; import { useZodForm } from "@/lib/forms/useZodForm"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; -const { Title, Paragraph, Text } = Typography; - const allowedIPSchema = z.object({ ip: z.string().min(1, "Please enter an IP address"), }); @@ -223,7 +221,7 @@ const AdminPanel: React.FC = ({ proxySettings }) => { children: ( <> - ✨ Security Settings +

✨ Security Settings

SSO Configuration Deprecated @@ -329,7 +327,9 @@ const AdminPanel: React.FC = ({ proxySettings }) => { Confirm Delete - Are you sure you want to delete the IP address: {ipToDelete}? + + Are you sure you want to delete the IP address: {ipToDelete}? +
@@ -962,10 +996,10 @@ const AddAgentForm: React.FC = ({ visible, onClose, accessTok

Agent Created!

- + {createdAgentName} - +
{createdKeyValue && (
@@ -998,13 +1032,7 @@ const AddAgentForm: React.FC = ({ visible, onClose, accessTok
- - - - - - - +
event.preventDefault()} className="space-y-4"> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx index d61d6056972..71770a9c945 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx @@ -1,8 +1,7 @@ import React, { useState, useCallback } from "react"; -import { Input, message, Spin } from "antd"; import { z } from "zod/v4"; import { SortingState } from "@tanstack/react-table"; -import { Inbox, Plus } from "lucide-react"; +import { Inbox, Plus, X } from "lucide-react"; import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; import { useQueryClient } from "@tanstack/react-query"; @@ -18,9 +17,11 @@ import { MCPToolset, MCPToolsetTool } from "@/components/mcp_tools/types"; import { FieldGroup } from "@/components/shared/form/field"; import { FormField } from "@/components/shared/form/FormField"; import { Button } from "@/components/ui/button"; -import { Input as ShadcnInput } from "@/components/ui/input"; +import { Input } from "@/components/ui/input"; +import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { useZodForm } from "@/lib/forms/useZodForm"; +import { toast } from "@/lib/toast"; import { displayToolName, getMCPToolsetTableColumns } from "./MCPToolsetTableColumns"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; @@ -97,7 +98,7 @@ function MCPToolList({ serverId, serverName, accessToken, selectedTools, onToggl
{loading ? (
- +
) : tools.length === 0 ? (

No tools found for this server.

@@ -212,10 +213,10 @@ function CreateToolsetModal({ open, onClose, onSave, accessToken, initialToolset event.preventDefault()} className="mt-2"> - {(field) => } + {(field) => } - {(field) => } + {(field) => } @@ -226,13 +227,20 @@ function CreateToolsetModal({ open, onClose, onSave, accessToken, initialToolset

Available Tools

- setServerSearch(e.target.value)} - className="mb-2" - allowClear - /> + + setServerSearch(e.target.value)} + /> + {serverSearch && ( + + setServerSearch("")}> + + + + )} +
{filteredServers.length === 0 ? (

@@ -381,14 +389,14 @@ export function MCPToolsetsTab({ accessToken, userRole }: MCPToolsetsTabProps) { const handleCreate = async (name: string, description: string | undefined, tools: MCPToolsetTool[]) => { if (!accessToken) return; await createMCPToolset(accessToken, { toolset_name: name, description, tools }); - message.success("Toolset created"); + toast.success("Toolset created"); queryClient.invalidateQueries({ queryKey: ["mcpToolsets"] }); }; const handleUpdate = async (name: string, description: string | undefined, tools: MCPToolsetTool[]) => { if (!accessToken || !editToolset) return; await updateMCPToolset(accessToken, { toolset_id: editToolset.toolset_id, toolset_name: name, description, tools }); - message.success("Toolset updated"); + toast.success("Toolset updated"); queryClient.invalidateQueries({ queryKey: ["mcpToolsets"] }); setEditToolset(null); }; @@ -398,7 +406,7 @@ export function MCPToolsetsTab({ accessToken, userRole }: MCPToolsetsTabProps) { setDeleting(true); try { await deleteMCPToolset(accessToken, deleteId); - message.success("Toolset deleted"); + toast.success("Toolset deleted"); queryClient.invalidateQueries({ queryKey: ["mcpToolsets"] }); setDeleteId(null); } finally { 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 b1f00ee64ad..7fc88e9f29f 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 @@ -1,5 +1,4 @@ import React, { useState, useEffect } from "react"; -import { Tag } from "antd"; import { z } from "zod/v4"; import { Policy, PolicyCreateRequest, PolicyUpdateRequest } from "@/components/policies/types"; import { Guardrail } from "@/components/guardrails/types"; @@ -474,9 +473,9 @@ const AddPolicyForm: React.FC = ({

{resolvedGuardrails.map((g) => ( - + {g} - + ))}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/add_prompt_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/add_prompt_form.tsx index 67e0b4e604b..2deda4b30ee 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/add_prompt_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/add_prompt_form.tsx @@ -1,7 +1,5 @@ -import React, { useState } from "react"; -import { Upload } from "antd"; -import type { UploadFile, UploadProps } from "antd"; -import { Upload as UploadIcon } from "lucide-react"; +import React, { useRef, useState } from "react"; +import { Upload as UploadIcon, X } from "lucide-react"; import { z } from "zod/v4"; import { convertPromptFileToJson, createPromptCall } from "@/components/networking"; import { toast } from "@/lib/toast"; @@ -50,25 +48,46 @@ const EMPTY_VALUES: AddPromptFormValues = { prompt_id: "", prompt_integration: " const AddPromptForm: React.FC = ({ visible, onClose, accessToken, onSuccess }) => { const form = useZodForm(addPromptSchema, { defaultValues: EMPTY_VALUES }); const [loading, setLoading] = useState(false); - const [fileList, setFileList] = useState([]); + const [selectedFile, setSelectedFile] = useState(null); + const fileInputRef = useRef(null); const [promptIntegration, setPromptIntegration] = useState("dotprompt"); + const clearSelectedFile = () => { + setSelectedFile(null); + if (fileInputRef.current) { + fileInputRef.current.value = ""; + } + }; + const handleCancel = () => { form.reset(EMPTY_VALUES); - setFileList([]); + clearSelectedFile(); setPromptIntegration("dotprompt"); onClose(); }; + const handleFileChange = (event: React.ChangeEvent) => { + const picked = event.target.files?.[0]; + if (!picked) return; + if (!picked.name.endsWith(".prompt")) { + toast.fromError("Please upload a .prompt file"); + clearSelectedFile(); + return; + } + setSelectedFile(picked); + }; + const handleIntegrationChange = (selected: string | null) => { if (selected === null) return; form.setValue("prompt_integration", selected); setPromptIntegration(selected); }; - const convertUploadedFile = async (token: string, promptId: string): Promise => { - const file = fileList[0].originFileObj as File; - + const convertUploadedFile = async ( + token: string, + promptId: string, + file: File, + ): Promise => { try { const conversionResult = await convertPromptFileToJson(token, file); @@ -98,16 +117,15 @@ const AddPromptForm: React.FC = ({ visible, onClose, accessT const isDotprompt = promptIntegration === "dotprompt"; - if (isDotprompt && fileList.length === 0) { + if (isDotprompt && !selectedFile) { toast.fromError("Please upload a .prompt file"); return; } setLoading(true); - const promptData: CreatePromptRequest | Record | null = isDotprompt - ? await convertUploadedFile(accessToken, values.prompt_id) - : {}; + const promptData: CreatePromptRequest | Record | null = + isDotprompt && selectedFile ? await convertUploadedFile(accessToken, values.prompt_id, selectedFile) : {}; if (promptData === null) { setLoading(false); @@ -127,23 +145,6 @@ const AddPromptForm: React.FC = ({ visible, onClose, accessT } }; - const uploadProps: UploadProps = { - beforeUpload: (file) => { - if (!file.name.endsWith(".prompt")) { - toast.fromError("Please upload a .prompt file"); - return false; - } - return false; // Prevent automatic upload - }, - fileList, - onChange: ({ fileList: newFileList }) => { - setFileList(newFileList.slice(-1)); // Keep only the last file - }, - onRemove: () => { - setFileList([]); - }, - }; - return ( !open && handleCancel()}> @@ -180,14 +181,30 @@ const AddPromptForm: React.FC = ({ visible, onClose, accessT Prompt File - - - - {fileList.length > 0 && ( -
Selected: {fileList[0].name}
+ + + {selectedFile && ( +
+ Selected: {selectedFile.name} + +
)} Upload a .prompt file that follows the Dotprompt specification @@ -196,16 +213,13 @@ const AddPromptForm: React.FC = ({ visible, onClose, accessT - {" "} - , - , ] diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.tsx index 5e540baaaf9..d6d64a4e49f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.tsx @@ -1,12 +1,15 @@ -import React, { useState } from "react"; -import { Typography, Divider, Table, Select, InputNumber, Card, Space, Checkbox } from "antd"; +import React, { useId, useState } from "react"; import { userBulkUpdateUserCall, teamBulkMemberAddCall, Member } from "@/components/networking"; import { UserEditView } from "./user_edit_view"; import { toast } from "@/lib/toast"; import { MoneyCell } from "@/components/shared/table_cells"; +import { MultiSelect } from "@/components/shared/MultiSelect"; +import NumericalInput from "@/components/shared/numerical_input"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { Checkbox } from "@/components/ui/checkbox"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; - -const { Text, Title } = Typography; +import { Separator } from "@/components/ui/separator"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; interface BulkEditUserModalProps { open: boolean; @@ -38,6 +41,10 @@ const BulkEditUserModal: React.FC = ({ const [teamBudget, setTeamBudget] = useState(null); const [addToTeams, setAddToTeams] = useState(false); const [updateAllUsers, setUpdateAllUsers] = useState(false); + const updateAllUsersId = useId(); + const addToTeamsId = useId(); + const selectedTeamsId = useId(); + const teamBudgetId = useId(); const handleCancel = () => { // Reset team management state @@ -210,14 +217,22 @@ const BulkEditUserModal: React.FC = ({ {allowAllUsers && (
- setUpdateAllUsers(e.target.checked)}> - Update ALL users in the system - +
+ setUpdateAllUsers(checked === true)} + aria-label="Update ALL users in the system" + /> + +
{updateAllUsers && ( -
- +
+ ⚠️ This will apply changes to ALL users in the system, not just the selected ones. - +
)}
@@ -225,118 +240,115 @@ const BulkEditUserModal: React.FC = ({ {!updateAllUsers && (
- Selected Users ({selectedUsers.length}): - ( - - {text.length > 20 ? `${text.slice(0, 20)}...` : text} - - ), - }, - { - title: "Email", - dataIndex: "user_email", - key: "user_email", - width: "25%", - render: (text: string) => ( - - {text || "No email"} - - ), - }, - { - title: "Current Role", - dataIndex: "user_role", - key: "user_role", - width: "25%", - render: (role: string) => ( - {possibleUIRoles?.[role]?.ui_label || role} - ), - }, - { - title: "Budget", - dataIndex: "max_budget", - key: "max_budget", - width: "20%", - render: (budget: number | null) => ( - - ), - }, - ]} - /> +
Selected Users ({selectedUsers.length}):
+
+
+ + + User ID + Email + Current Role + Budget + + + + {selectedUsers.map((user) => ( + + + {user.user_id.length > 20 ? `${user.user_id.slice(0, 20)}...` : user.user_id} + + {user.user_email || "No email"} + + {possibleUIRoles?.[user.user_role]?.ui_label || user.user_role} + + + + + + ))} + +
+
)} - +
- +

Instructions: Fill in the fields below with the values you want to apply to all selected users. You can bulk edit: role, budget, models, and metadata. You can also add users to teams. - +

{/* Team Management Section */} - - - setAddToTeams(e.target.checked)}> - Add selected users to teams - + + + Team Management + + +
+
+ setAddToTeams(checked === true)} + aria-label="Add selected users to teams" + /> + +
- {addToTeams && ( - <> -
- Select Teams: - onChange(selected.length === 0 ? undefined : selected)} - > - - - {(selected: string[]) => - selected.length === 0 - ? "Select Organization" - : organizationOptions - .filter((option) => selected.includes(option.value)) - .map((option) => option.label) - .join(", ") - } - - - - {organizationOptions.map((option) => ( - - {option.label} - - ))} - - + !open && handleCancel()}> + + + Invite User + +
+

Create a User who can own keys

+ +
+ +
+ + {userEmailField} + {roleField( + labelWithHint( + "Global Proxy Role", + "This role is independent of any team/org specific roles. Configure Team / Organization Admins in the Settings", + ), )} - + {teamField} - {metadataField} - {sendInviteEmailField} + + {({ id, value, onChange }) => ( + + )} + - - - - Personal Key Creation - - - - {({ value, onChange }) => ( - ({ label: getModelDisplayName(model), value: model })), - ]} - value={value ?? []} - onValueChange={onChange} - placeholder="Select models" - /> - )} - - - - + {metadataField} + {sendInviteEmailField} -
- -
-
- - + + + + Personal Key Creation + + + + {({ value, onChange }) => ( + ({ label: getModelDisplayName(model), value: model })), + ]} + value={value ?? []} + onValueChange={onChange} + placeholder="Select models" + /> + )} + + + + + +
+ +
+ + + + {apiuser && ( {v}, - }, - { title: "Display Name", dataIndex: "display_name", key: "display_name" }, - { - title: "URL", - dataIndex: "url", - key: "url", - render: (v: string) => ( - - {v} - - ), - }, - { - title: "Plugin Key", - dataIndex: "plugin_key", - key: "plugin_key", - render: (v?: string) => (v ? {"•".repeat(8)} : —), - }, - { - title: "Actions", - key: "actions", - render: (_: unknown, plugin: Plugin, idx: number) => ( - - - - - ), - }, - ]; + const renderRows = () => { + if (loading) { + return ( + + + + + + ); + } + + if (plugins.length === 0) { + return ( + + + No data + + + ); + } + + return plugins.map((plugin, idx) => ( + + + {plugin.name} + + {plugin.display_name} + + + {plugin.url} + + + + {plugin.plugin_key ? ( + {"•".repeat(8)} + ) : ( + — + )} + + +
+ + +
+ + + )); + }; return ( - Plugins - - Register external services as plugins. Once added, users can toggle to the plugin from the mode switcher in the - top-left of the sidebar. - - - Each plugin must expose GET /api/plugin-manifest returning nav items and capabilities. - + +

Plugins

+

+ Register external services as plugins. Once added, users can toggle to the plugin from the mode switcher in + the top-left of the sidebar. +

+

+ Each plugin must expose GET /api/plugin-manifest returning nav + items and capabilities. +

+ + + - - - +
+ + + Name + Display Name + URL + Plugin Key + Actions + + + {renderRows()} +
+ !open && setModalOpen(false)}> diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index 8d9fdd7598f..ccd5500622b 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -15,7 +15,7 @@ import { SearchSelect } from "@/components/shared/SearchSelect"; import { labelWithDocsHint, labelWithHint } from "@/components/shared/form/LabelWithHint"; import { useZodForm } from "@/lib/forms/useZodForm"; import { TagsInput } from "@/app/(dashboard)/guardrails/_components/content_filter/TagsInput"; -import { Layout, Tabs } from "antd"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { ChevronDown, Plus, Users } from "lucide-react"; import React, { useEffect, useMemo, useState } from "react"; import { z } from "zod/v4"; @@ -542,8 +542,6 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser return false; }; - const { Content } = Layout; - const tabItems = [ { key: "your-teams", @@ -611,7 +609,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser ]; return ( - +
{selectedTeamId ? ( = ({ accessToken, userID, userRole, premiumUser premiumUser={premiumUser} /> ) : ( - } - title="Teams" - subtitle="Manage teams, members, and their access to models and budgets" - primaryAction={ - canCreateOrManageTeams(userRole, userID, organizations) ? ( - setIsTeamModalVisible(true)} data-testid="create-team-button"> - - Create Team - - ) : undefined - } - tabs={({ leadingControls }) => ( - - )} - /> + + } + title="Teams" + subtitle="Manage teams, members, and their access to models and budgets" + primaryAction={ + canCreateOrManageTeams(userRole, userID, organizations) ? ( + setIsTeamModalVisible(true)} data-testid="create-team-button"> + + Create Team + + ) : undefined + } + tabs={({ leadingControls }) => ( + + {leadingControls} + {tabItems.map((item) => ( + + {item.label} + + ))} + + )} + /> + {tabItems.map((item) => ( + + {item.children} + + ))} + )} {canCreateOrManageTeams(userRole, userID, organizations) && ( @@ -1203,7 +1218,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser
)} - + ); }; diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx index b6dddf43588..83c9e538732 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx @@ -5,7 +5,10 @@ import { all_admin_roles, isUserTeamAdminForAnyTeam } from "@/utils/roles"; import { modelCreationScope } from "@/utils/modelPermissions"; import { Switch } from "@/components/ui/switch"; import { Field, FieldLabel } from "@/components/shared/form/field"; -import { Select as AntdSelect, Card, Col, Row, Tooltip, Typography } from "antd"; +import { Card, CardContent } from "@/components/ui/card"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { SearchSelect, type SearchSelectOption } from "@/components/shared/SearchSelect"; +import { SimpleTooltip } from "@/components/ui/tooltip"; import { Info } from "lucide-react"; import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; import { Button } from "@/components/ui/button"; @@ -24,6 +27,7 @@ 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"; import ConditionalPublicModelName from "./conditional_public_model_name"; import LiteLLMModelNameField from "./litellm_model_name"; @@ -57,8 +61,6 @@ const connectionTestModelName = (values: MountedFormValues): string | undefined return typeof named === "string" ? named : undefined; }; -const { Title, Link } = Typography; - const AddModelForm: React.FC = ({ form, registry, @@ -117,6 +119,34 @@ const AddModelForm: React.FC = ({ return [...providerMetadata].sort((a, b) => a.provider_display_name.localeCompare(b.provider_display_name)); }, [providerMetadata]); + const providerOptions: SearchSelectOption[] = useMemo( + () => + sortedProviderMetadata.map((providerInfo) => ({ + label: providerInfo.provider_display_name, + value: providerInfo.provider, + icon: , + })), + [sortedProviderMetadata], + ); + + const credentialOptions: SearchSelectOption[] = useMemo( + () => [ + { label: "None", value: "" }, + ...credentials.map((credential) => ({ + label: credential.credential_name, + value: credential.credential_name, + })), + ], + [credentials], + ); + + const applyProviderSelection = (provider: Providers) => { + setSelectedProvider(provider); + setProviderModelsFn(provider); + form.setValue("model", []); + form.setValue("model_name", undefined); + }; + const providerMetadataErrorText = providerMetadataError ? providerMetadataError instanceof Error ? providerMetadataError.message @@ -132,311 +162,287 @@ const AddModelForm: React.FC = ({ return ( <> - Add Model +

Add Model

- - -
{ - event.preventDefault(); - void handleOk().then((submitted) => { - if (submitted) { - setTeamAdminSelectedTeam(null); - } - }); - }} - > - <> - {requiresTeamScope && ( - <> - - {(control) => ( - { - control.onChange(value); - setTeamAdminSelectedTeam(value); - }} - /> - )} - - {!teamAdminSelectedTeam && ( - - - Team Selection Required - - As a team admin, you need to select your team first before adding models. - - - )} - - )} - {(isAdmin || (isTeamAdmin && teamAdminSelectedTeam)) && ( - <> - - {(control) => ( - { - control.onChange(value); - setSelectedProvider(value as Providers); - setProviderModelsFn(value as Providers); - form.setValue("model", []); - form.setValue("model_name", undefined); - }} - > - {providerMetadataErrorText && sortedProviderMetadata.length === 0 && ( - - {providerMetadataErrorText} - - )} - {sortedProviderMetadata.map((providerInfo) => { - const displayName = providerInfo.provider_display_name; - const providerKey = providerInfo.provider; - - return ( - -
- - {displayName} -
- - ); - })} - - )} - - - - {/* Conditionally Render "Public Model Name" */} - - - {/* Select Mode */} - - {(control) => ( - { - control.onChange(value); - setTestMode(value); - }} - options={TEST_MODES} - /> - )} - - - - -

- Optional - LiteLLM endpoint to use when health checking this model{" "} - - Learn more - -

- - - - {/* Credentials */} -
- - Either select existing credentials OR enter new provider credentials below - -
- - - {(control) => ( - - (option?.label ?? "").toLowerCase().includes(input.toLowerCase()) - } - value={control.value as string | null | undefined} - onChange={control.onChange} - onBlur={control.onBlur} - options={[ - { value: null, label: "None" }, - ...credentials.map((credential) => ({ - value: credential.credential_name, - label: credential.credential_name, - })), - ]} - allowClear - /> - )} - - - {/* Only show provider specific fields if no credentials selected */} - {!selectedCredentialName && ( - <> -
-
- OR -
-
- - - )} -
-
- Additional Model Info Settings -
-
- {/* Team-only Model Switch - Only show for proxy admins, not team admins */} - {(isAdmin || !isTeamAdmin) && ( - - - {labelWithHint( - "Team-BYOK Model", - "Only use this model + credential combination for this team. Useful when teams want to onboard their own OpenAI keys.", - )} - - - - { - setIsTeamOnly(checked); - if (!checked) { - form.setValue("team_id", undefined); - } - }} - disabled={!premiumUser} - aria-label="Team-BYOK Model" - /> - - - - )} - - {/* Conditional Team Selection */} - {isTeamOnly && !requiresTeamScope && ( + + + + { + event.preventDefault(); + void handleOk().then((submitted) => { + if (submitted) { + setTeamAdminSelectedTeam(null); + } + }); + }} + > + <> + {requiresTeamScope && ( + <> {(control) => ( { + control.onChange(value); + setTeamAdminSelectedTeam(value); + }} /> )} - )} - {isAdmin && ( - <> + {!teamAdminSelectedTeam && ( + + + Team Selection Required + + As a team admin, you need to select your team first before adding models. + + + )} + + )} + {(isAdmin || (isTeamAdmin && teamAdminSelectedTeam)) && ( + <> + + {(control) => ( + { + control.onChange(value); + applyProviderSelection(value as Providers); + }} + /> + )} + + + + {/* Conditionally Render "Public Model Name" */} + + + {/* Select Mode */} + + {(control) => ( + + )} + +
+
+
+

+ Optional - LiteLLM endpoint to use when health checking this model{" "} + + Learn more + +

+
+
+ + {/* Credentials */} +
+ + Either select existing credentials OR enter new provider credentials below + +
+ + + {(control) => ( + control.onChange(value === "" ? null : value)} + /> + )} + + + {/* Only show provider specific fields if no credentials selected */} + {!selectedCredentialName && ( + <> +
+
+ OR +
+
+ + + )} +
+
+ Additional Model Info Settings +
+
+ {/* Team-only Model Switch - Only show for proxy admins, not team admins */} + {(isAdmin || !isTeamAdmin) && ( + + + {labelWithHint( + "Team-BYOK Model", + "Only use this model + credential combination for this team. Useful when teams want to onboard their own OpenAI keys.", + )} + + + + { + setIsTeamOnly(checked); + if (!checked) { + form.setValue("team_id", undefined); + } + }} + disabled={!premiumUser} + aria-label="Team-BYOK Model" + /> + + + + )} + + {/* Conditional Team Selection */} + {isTeamOnly && !requiresTeamScope && ( {(control) => ( - ({ - value: group, - label: group, - }))} - maxTagCount="responsive" - allowClear + disabled={!premiumUser} /> )} - - )} - - - )} -
- - Need Help? - -
- - + )} + {isAdmin && ( + <> + + {(control) => ( + + )} + + + )} + + + )} +
+ + + Need Help? + + +
+ + +
-
- - - - + + + + + {/* Test Connection Results Modal */} 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 64d1519f915..66564c3e007 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 @@ -1,12 +1,12 @@ import React, { useEffect, useState } from "react"; import { useQuery } from "@tanstack/react-query"; import { useWatch } from "react-hook-form"; -import { Card } from "antd"; import { ChevronDown, ChevronRight, CircleHelp } from "lucide-react"; import { z } from "zod/v4"; import { FieldGroup } from "@/components/shared/form/field"; import { FormField } from "@/components/shared/form/FormField"; import { Button } from "@/components/ui/button"; +import { Card, CardContent } from "@/components/ui/card"; import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; @@ -483,205 +483,207 @@ const AddAutoRouterTab: React.FC = ({ return ( -
handleAutoRouterSubmit())} noValidate> - - - {({ ref, ...field }) => } - - -
- - - {modelsUnverifiable && ( -
- Could not load available models.{" "} - -
- )} -
- - {requiresTeamScope && ( + + handleAutoRouterSubmit())} noValidate> + - {({ id, value, onChange }) => } + {({ ref, ...field }) => } - )} -
- +
+ )} +
+ + {requiresTeamScope && ( + - {!detailsExpanded && ( - - {tierConfigSummary(complexityRouterConfig.tiers)} - - )} - - {detailsExpanded && ( -
- -
+ > + {({ id, value, onChange }) => } + )} -
- {isAdmin && ( - + + {detailsExpanded && ( +
+ +
)} - > - {({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => ( - + + {isAdmin && ( + + {({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => ( + + )} + + )} + +
+ + + Need Help? + + } /> - )} - - )} - -
- - Get help on our github + +
+ + + - - - - - + + + +
-
- - + + +
!open && setIsRoutingTestVisible(false)}> diff --git a/ui/litellm-dashboard/src/components/cloudzero_export_modal.integration.test.tsx b/ui/litellm-dashboard/src/components/cloudzero_export_modal.integration.test.tsx index 7f52342ceb3..fbcd9c2fdcc 100644 --- a/ui/litellm-dashboard/src/components/cloudzero_export_modal.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/cloudzero_export_modal.integration.test.tsx @@ -139,7 +139,7 @@ describe("CloudZeroExportModal", () => { await screen.findByLabelText("CloudZero API Key"); await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("Export to CSV")); + await user.click(await screen.findByRole("option", { name: "Export to CSV" })); expect(screen.queryByLabelText("CloudZero API Key")).not.toBeInTheDocument(); await user.click(screen.getByRole("button", { name: "Export CSV" })); diff --git a/ui/litellm-dashboard/src/components/cloudzero_export_modal.tsx b/ui/litellm-dashboard/src/components/cloudzero_export_modal.tsx index c8ef2218093..81d62fc398e 100644 --- a/ui/litellm-dashboard/src/components/cloudzero_export_modal.tsx +++ b/ui/litellm-dashboard/src/components/cloudzero_export_modal.tsx @@ -1,5 +1,4 @@ import React, { useState, useEffect } from "react"; -import { Spin, Select } from "antd"; import { CircleCheck, FileDown } from "lucide-react"; import { z } from "zod/v4"; import { getGlobalLitellmHeaderName } from "@/components/networking"; @@ -10,6 +9,7 @@ import { FieldGroup } from "@/components/shared/form/field"; import { FormField } from "@/components/shared/form/FormField"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { useZodForm } from "@/lib/forms/useZodForm"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; @@ -251,13 +251,18 @@ const CloudZeroExportModal: React.FC = ({ isOpen, onC {/* Export Type Selection */}

Export Destination

- value && setExportType(value)}> + + + + + {exportOptions.map((option) => ( + + {option.label} + + ))} + +
{/* CloudZero Configuration */} @@ -265,7 +270,7 @@ const CloudZeroExportModal: React.FC = ({ isOpen, onC
{settingsLoading ? (
- +
) : ( <> diff --git a/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx b/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx index f21af04d5d5..61fb5a44c66 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx @@ -1,13 +1,14 @@ "use client"; -import React, { useState } from "react"; -import { Input, Switch } from "antd"; +import React, { useId, useState } from "react"; import { toast } from "@/lib/toast"; import { fetchClient } from "@/lib/http/api"; import { ApiError } from "@/lib/http/client"; import { ArrowLeft, ArrowRight, Check, Key, Link2, Lock, X } from "lucide-react"; import { MCPServer } from "./types"; import { Dialog, DialogContent } from "@/components/ui/dialog"; +import { PasswordInput } from "@/components/shared/PasswordInput"; +import { Switch } from "@/components/ui/switch"; const byokSaveErrorMessage = (e: unknown): string => { if (e instanceof ApiError) { @@ -29,6 +30,7 @@ export const ByokCredentialModal: React.FC = ({ server const [apiKey, setApiKey] = useState(""); const [saveKey, setSaveKey] = useState(true); const [loading, setLoading] = useState(false); + const apiKeyInputId = useId(); const serverDisplayName = server.alias || server.server_name || "Service"; const firstLetter = serverDisplayName.charAt(0).toUpperCase(); @@ -169,13 +171,15 @@ export const ByokCredentialModal: React.FC = ({ server

Enter your {serverDisplayName} API key to authorize this connection.

- - + {serverDisplayName} API Key + + setApiKey(e.target.value)} - size="large" - className="rounded-lg" + groupClassName="rounded-lg" /> {server.byok_api_key_help_url && ( = ({ server Save key for future use
- +
{/* Security note */} diff --git a/ui/litellm-dashboard/src/components/model_add/CredentialModal.tsx b/ui/litellm-dashboard/src/components/model_add/CredentialModal.tsx index 0fba7c1cd9f..4f25b2a201a 100644 --- a/ui/litellm-dashboard/src/components/model_add/CredentialModal.tsx +++ b/ui/litellm-dashboard/src/components/model_add/CredentialModal.tsx @@ -1,5 +1,6 @@ import { Input } from "@/components/ui/input"; -import { Select as AntdSelect, Tooltip, Typography } from "antd"; +import { SearchSelect, type SearchSelectOption } from "@/components/shared/SearchSelect"; +import { SimpleTooltip } from "@/components/ui/tooltip"; import { Button } from "@/components/ui/button"; import { useState } from "react"; import { FormProvider, useForm } from "react-hook-form"; @@ -19,7 +20,11 @@ import { Logo } from "@/components/molecules/logo/Logo"; import { resetCredentialFormOnProviderChange } from "./credential_form_helpers"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; -const { Link } = Typography; +const providerOptions: SearchSelectOption[] = Object.entries(Providers).map(([providerEnum, providerDisplayName]) => ({ + label: providerDisplayName, + value: providerEnum, + icon: , +})); interface CredentialModalProps { open: boolean; @@ -122,34 +127,27 @@ export default function CredentialModal({ className="mb-4" > {(control) => ( - { + { control.onChange(value); resetCredentialFormOnProviderChange(formAdapter, value as Providers, setSelectedProvider); }} - > - {Object.entries(Providers).map(([providerEnum, providerDisplayName]) => ( - -
- - {providerDisplayName} -
- - ))} - + /> )}
- - Need Help? - + + + Need Help? + +
- - - Showing {filteredGroups.length} {filteredGroups.length === 1 ? "result" : "results"} - - - +
+ + +
+ + + + + setSearchQuery(e.target.value)} + /> + {searchQuery && ( + + setSearchQuery("")}> + + + + )} + +
+ + + + Showing {filteredGroups.length} {filteredGroups.length === 1 ? "result" : "results"} + +
+
- setDeletingGroup(g)} - proxyBaseUrl={proxySettings.LITELLM_UI_API_DOC_BASE_URL?.trim() || proxySettings.PROXY_BASE_URL || ""} - /> + setDeletingGroup(g)} + proxyBaseUrl={proxySettings.LITELLM_UI_API_DOC_BASE_URL?.trim() || proxySettings.PROXY_BASE_URL || ""} + /> + { Delete routing group? - - Models in {deletingGroup?.group_name} will fall back to the proxy's top-level - routing strategy. This cannot be undone. - +

+ Models in {deletingGroup?.group_name} will fall back to the + proxy's top-level routing strategy. This cannot be undone. +

- +
); }; diff --git a/ui/litellm-dashboard/src/components/update_model_credentials_modal.tsx b/ui/litellm-dashboard/src/components/update_model_credentials_modal.tsx index 4db8c83ba9d..1198e7510ed 100644 --- a/ui/litellm-dashboard/src/components/update_model_credentials_modal.tsx +++ b/ui/litellm-dashboard/src/components/update_model_credentials_modal.tsx @@ -1,4 +1,3 @@ -import { Typography } from "antd"; import { TriangleAlert } from "lucide-react"; import { useState } from "react"; import { z } from "zod/v4"; @@ -13,8 +12,6 @@ import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/u import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { useZodForm } from "@/lib/forms/useZodForm"; -const { Text } = Typography; - const updateCredentialsSchema = z.object({ api_key: z.string().min(1, "Enter a new API key"), }); @@ -77,10 +74,10 @@ export default function UpdateModelCredentialsModal({ Update API Key - + Update this model's API key. Only the new key is sent; the rest of the deployment configuration is left untouched. - + From 1750893a6905f45d7fa9cc65167856ae6c182208 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 20:26:41 -0700 Subject: [PATCH 037/281] fix(cli): never report success while a credential is still readable Migration moved the secret into the keychain and then suppressed any OSError from rewriting token.json, so a file that could not be rewritten kept the credential in cleartext while every command reported success. That file is now removed instead: signing in again costs one command, a stranded live credential costs the credential `lite logout` also reported a clean logout whenever the keyring package was missing, on the reasoning that an install without it could never have stored anything. The entry belongs to the OS, so a keychain-backed login survives a logout run from a venv without the cli extra. erase() now reports which keychain state applies, and logout warns with the advice that fixes each one, staying quiet for file-backed logins whose token file still carries its own secret Also pins the migration path's tightening of a world-readable legacy token.json, and moves the logout tests off patch() onto the injected vault --- litellm/litellm_core_utils/cli_keyring.py | 41 ++++++---- litellm/litellm_core_utils/cli_token_utils.py | 49 ++++++++++-- litellm/proxy/client/cli/commands/auth.py | 31 +++++--- tests/test_litellm/conftest.py | 13 +++- .../test_cli_token_utils.py | 77 ++++++++++++++++--- .../proxy/client/cli/test_auth_commands.py | 66 ++++++++++++++-- 6 files changed, 226 insertions(+), 51 deletions(-) diff --git a/litellm/litellm_core_utils/cli_keyring.py b/litellm/litellm_core_utils/cli_keyring.py index b19b3d3bc83..0497991c2a0 100644 --- a/litellm/litellm_core_utils/cli_keyring.py +++ b/litellm/litellm_core_utils/cli_keyring.py @@ -36,6 +36,16 @@ class SecretStored: pass +@dataclass(frozen=True, slots=True) +class SecretErased: + pass + + +@dataclass(frozen=True, slots=True) +class SecretStranded: + pass + + @dataclass(frozen=True, slots=True) class KeyringNotInstalled: pass @@ -54,6 +64,7 @@ class KeyringUnreachable: KeyringUnusable: TypeAlias = KeyringNotInstalled | KeyringDisabled | KeyringUnreachable SecretRead: TypeAlias = SecretFound | SecretMissing | KeyringUnusable SecretWrite: TypeAlias = SecretStored | KeyringUnusable +SecretErase: TypeAlias = SecretErased | SecretStranded | KeyringUnusable class SecretVault(Protocol): @@ -63,7 +74,7 @@ class SecretVault(Protocol): def write(self, blob: str) -> SecretWrite: ... - def erase(self) -> bool: ... + def erase(self) -> SecretErase: ... class KeyringApi(Protocol): @@ -117,33 +128,31 @@ class KeyringVault: return KeyringUnreachable() return SecretStored() - def erase(self) -> bool: - """Whether the keychain is guaranteed to hold no credential afterwards. + def erase(self) -> SecretErase: + """Remove our entry, reporting whether the keychain is guaranteed to be free of it. - An uninstalled `keyring` package can never have stored one. A kill switch set after - a credential was stored leaves that entry out of reach, so erasure cannot be promised. + A keychain out of reach is never an erasure: the entry belongs to the OS, not to this + install, so it outlives an uninstalled `keyring` package and a kill switch set after login. + Those cases are reported apart from a confirmed entry that would not delete, because only + the caller knows whether this machine ever put a secret in a keychain. """ - if _import_keyring() is None: - return True - if _keyring_disabled(): - return False match self.read(): - case KeyringNotInstalled() | KeyringDisabled() | KeyringUnreachable(): - return False + case KeyringNotInstalled() | KeyringDisabled() | KeyringUnreachable() as unusable: + return unusable case SecretMissing(): - return True + return SecretErased() case SecretFound(): return self._delete() - def _delete(self) -> bool: + def _delete(self) -> SecretErase: api: Final = _keyring_api() if isinstance(api, (KeyringNotInstalled, KeyringDisabled)): - return False + return api try: api.delete_password(KEYRING_SERVICE, KEYRING_ACCOUNT) except Exception: # noqa: BLE001 # report the failure as a value so `lite logout` can warn - return False - return True + return SecretStranded() + return SecretErased() SYSTEM_KEYRING: Final[SecretVault] = KeyringVault() diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index 40822e5b335..cd69f9470a3 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -25,9 +25,12 @@ from litellm.litellm_core_utils.cli_keyring import ( KeyringDisabled, KeyringNotInstalled, KeyringUnreachable, + SecretErase, + SecretErased, SecretFound, SecretMissing, SecretStored, + SecretStranded, SecretVault, SecretWrite, ) @@ -84,7 +87,7 @@ def load_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> CliTokenRecord | N def save_cli_token(record: CliTokenRecord, *, vault: SecretVault = SYSTEM_KEYRING) -> SecretWrite: - """Store a freshly minted credential. Reports whether the keychain took the secret, and why not""" + """Store a freshly minted credential. Reports where its secret material ended up, and why""" outcome: Final = ( SecretStored() if record.key is None @@ -94,11 +97,33 @@ def save_cli_token(record: CliTokenRecord, *, vault: SecretVault = SYSTEM_KEYRIN return outcome -def clear_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> bool: - """Remove the credential from both stores. Returns whether the keychain is now free of it""" - erased: Final = vault.erase() +def clear_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> SecretErase: + """Remove the credential from both stores. Reports whether the keychain is now free of it""" + outcome: Final = vault.erase() + settled: Final = _nothing_left_behind(outcome) Path(get_cli_token_file_path()).unlink(missing_ok=True) - return erased + return SecretErased() if settled else outcome + + +def _nothing_left_behind(outcome: SecretErase) -> bool: + """Whether the keychain can be trusted to hold no credential of ours once the file is gone""" + match outcome: + case SecretErased(): + return True + case SecretStranded(): + return False + case KeyringNotInstalled() | KeyringDisabled() | KeyringUnreachable(): + return not _secret_lives_in_keychain() + + +def _secret_lives_in_keychain() -> bool: + """Whether the token file is the metadata half of a pair whose secret half went to a keychain. + + A file that still carries its own secret rules one out, which keeps `lite logout` quiet on the + machines that never had a keychain to begin with. + """ + record: Final = _read_token_file() + return record is not None and record.key is None and not record.jwt_token def get_litellm_gateway_api_key( @@ -200,10 +225,22 @@ def _migrate_file_secret(record: CliTokenRecord, vault: SecretVault) -> CliToken def _scrub_file_secret(record: CliTokenRecord) -> None: + """Leave no secret material in the token file once the vault holds it. + + A file that cannot be rewritten without the secret is removed instead. Signing in again costs + the user one command; a live credential left behind in cleartext costs them the credential. + """ if record.key is None and not record.jwt_token: return - with contextlib.suppress(OSError): + try: _write_token_file(_without_secret(record)) + except OSError: + _discard_token_file() + + +def _discard_token_file() -> None: + with contextlib.suppress(OSError): + Path(get_cli_token_file_path()).unlink(missing_ok=True) def _without_secret(record: CliTokenRecord) -> CliTokenRecord: diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 03906f9b6df..4cf18435473 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -17,7 +17,9 @@ from litellm.litellm_core_utils.cli_keyring import ( KeyringDisabled, KeyringNotInstalled, KeyringUnreachable, + SecretErased, SecretStored, + SecretStranded, SecretVault, SecretWrite, ) @@ -80,12 +82,16 @@ class CliAuthResult(TypedDict): team_id: str | None -KEYCHAIN_UNREACHABLE_MESSAGE: Final = ( - "Your credential is stored in your OS keychain, which could not be read. Unlock it and retry, " - "install the keyring package with: pip install 'litellm[cli]', or run 'lite login'." +KEYRING_INSTALL_HINT: Final = "pip install 'litellm[cli]'" + +STRANDED_CREDENTIAL_MESSAGE: Final = ( + "Logged out locally, but your credential is still in the OS keychain and could not be removed." ) -KEYRING_INSTALL_HINT: Final = "pip install 'litellm[cli]'" +KEYCHAIN_UNREACHABLE_MESSAGE: Final = ( + "Your credential is stored in your OS keychain, which could not be read. Unlock it, or install " + f"the keyring package with: {KEYRING_INSTALL_HINT}. Run 'lite login' to start over." +) def storage_notice(outcome: SecretWrite) -> str: @@ -742,11 +748,18 @@ def login(ctx: click.Context, config_claude: bool): @click.pass_context def logout(ctx: click.Context): """Logout and clear stored authentication""" - if clear_cli_token(vault=context_secret_vault(ctx)): - click.echo("Logged out successfully. Authentication token cleared.") - return - click.echo("Logged out. The local token file is gone, but the OS keychain entry could not be removed.") - click.echo("Unlock your keychain and run 'lite logout' again to clear it.") + match clear_cli_token(vault=context_secret_vault(ctx)): + case SecretErased(): + click.echo("Logged out successfully. Authentication token cleared.") + case KeyringNotInstalled(): + click.echo(STRANDED_CREDENTIAL_MESSAGE) + click.echo(f"Install the keyring package with: {KEYRING_INSTALL_HINT}, then run 'lite logout' again.") + case KeyringDisabled(): + click.echo(STRANDED_CREDENTIAL_MESSAGE) + click.echo(f"Unset {DISABLE_KEYRING_ENV_VAR} and run 'lite logout' again to clear it.") + case SecretStranded() | KeyringUnreachable(): + click.echo(STRANDED_CREDENTIAL_MESSAGE) + click.echo("Unlock your keychain and run 'lite logout' again to clear it.") @click.command(name="print-token") diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index a716ec0e0aa..b42355fa045 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -25,10 +25,13 @@ from litellm._logging import ALL_LOGGERS from litellm.litellm_core_utils.cli_keyring import ( KeyringUnreachable, KeyringUnusable, + SecretErase, + SecretErased, SecretFound, SecretMissing, SecretRead, SecretStored, + SecretStranded, SecretWrite, ) from litellm.litellm_core_utils.prompt_templates import ( @@ -163,12 +166,14 @@ class FakeSecretVault: self.blob = blob return SecretStored() - def erase(self) -> bool: + def erase(self) -> SecretErase: self.erases += 1 - if not (self.available and self.erasable): - return False + if not self.available: + return self.failure + if not self.erasable: + return SecretStranded() if self.blob is not None else SecretErased() self.blob = None - return True + return SecretErased() @pytest.fixture diff --git a/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py b/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py index 961d986e8f5..925b440dfb7 100644 --- a/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py @@ -16,7 +16,9 @@ from litellm.litellm_core_utils.cli_keyring import ( KeyringDisabled, KeyringNotInstalled, KeyringUnreachable, + SecretErased, SecretStored, + SecretStranded, ) from litellm.litellm_core_utils.cli_token_utils import ( CliTokenRecord, @@ -126,6 +128,16 @@ class TestLoadCliToken: assert on_disk["user_email"] == "user@example.com" assert stat.S_IMODE(path.stat().st_mode) == 0o600 + def test_migration_tightens_a_world_readable_legacy_file(self, isolated_home, secret_vault_factory): + """An older `lite`, a loose umask, or a restored backup can leave token.json readable by + every account on the box. Migrating it must not preserve those permissions.""" + path = _write_legacy_file(isolated_home) + path.chmod(0o644) + + load_cli_token(vault=secret_vault_factory()) + + assert stat.S_IMODE(path.stat().st_mode) == 0o600 + def test_legacy_file_survives_a_vault_that_refuses_to_store(self, isolated_home, secret_vault_factory): """Scrubbing the only copy of the secret after a failed keychain write would log the user out for good.""" @@ -304,12 +316,35 @@ class TestSaveCliToken: assert list(path.parent.glob(".tmp-*")) == [] +class TestScrubFailure: + """A keychain that took the secret while the file kept it is the worst of both stores: the + credential is live, it is in cleartext on disk, and every command reports success.""" + + def test_a_file_that_cannot_be_rewritten_is_removed_instead( + self, isolated_home, secret_vault_factory, monkeypatch + ): + path = _write_legacy_file(isolated_home) + vault = secret_vault_factory() + + def _explode(*args, **kwargs): + raise OSError("no space left on device") + + monkeypatch.setattr("litellm.litellm_core_utils.private_json.json.dump", _explode) + + record = load_cli_token(vault=vault) + + assert record.key == "sk-legacy" + assert json.loads(vault.blob)["key"] == "sk-legacy" + assert not path.exists() + assert list(path.parent.glob(".tmp-*")) == [] + + class TestClearCliToken: def test_removes_the_credential_from_both_stores(self, isolated_home, secret_vault_factory): vault = secret_vault_factory() save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-new"), vault=vault) - assert clear_cli_token(vault=vault) is True + assert clear_cli_token(vault=vault) == SecretErased() assert vault.blob is None assert not _token_file(isolated_home).exists() assert load_cli_token(vault=vault) is None @@ -318,11 +353,32 @@ class TestClearCliToken: _write_legacy_file(isolated_home) vault = secret_vault_factory(blob=_blob(), erasable=False) - assert clear_cli_token(vault=vault) is False + assert clear_cli_token(vault=vault) == SecretStranded() + assert not _token_file(isolated_home).exists() + + def test_logout_from_an_install_without_keyring_does_not_claim_the_keychain_is_clear( + self, isolated_home, secret_vault_factory + ): + """Log in where `litellm[cli]` is installed and the secret goes to the OS keychain; log out + from a venv without it and the entry survives, because it belongs to the OS rather than to + the package. Reporting a clean logout there leaves a live credential the user thinks is gone.""" + _write_metadata_only_file(isolated_home) + vault = secret_vault_factory(available=False, failure=KeyringNotInstalled()) + + assert clear_cli_token(vault=vault) == KeyringNotInstalled() + assert not _token_file(isolated_home).exists() + + def test_a_file_backed_login_logs_out_quietly_without_keyring(self, isolated_home, secret_vault_factory): + """The complement: a user who never had a keychain keeps their whole credential in the file, + so removing it is a complete logout and must not warn about an entry that cannot exist.""" + _write_legacy_file(isolated_home) + vault = secret_vault_factory(available=False, failure=KeyringNotInstalled()) + + assert clear_cli_token(vault=vault) == SecretErased() assert not _token_file(isolated_home).exists() def test_is_safe_when_nothing_was_ever_stored(self, isolated_home, secret_vault_factory): - assert clear_cli_token(vault=secret_vault_factory()) is True + assert clear_cli_token(vault=secret_vault_factory()) == SecretErased() class TestIsCliTokenFresh: @@ -384,7 +440,7 @@ class TestKeyringVault: assert vault.write("blob-1") == SecretStored() assert vault.read() == SecretFound("blob-1") - assert vault.erase() is True + assert vault.erase() == SecretErased() assert vault.read() == SecretMissing() assert {call[1:] for call in fake.calls} == {(KEYRING_SERVICE, KEYRING_ACCOUNT)} @@ -392,24 +448,25 @@ class TestKeyringVault: """`LITELLM_CLI_DISABLE_KEYRING` has to work without importing keyring, because keyring caches its backend on first use and cannot be reconfigured later. Erase still fails: a credential stored before the switch was set may be in the keychain, and with reads - disabled `lite logout` cannot verify it is gone, so it must warn instead.""" + disabled `lite logout` cannot verify it is gone, so it must say so instead.""" monkeypatch.setenv(DISABLE_KEYRING_ENV_VAR, "1") vault = KeyringVault() assert vault.read() == KeyringDisabled() assert vault.write("blob-1") == KeyringDisabled() - assert vault.erase() is False + assert vault.erase() == KeyringDisabled() def test_an_uninstalled_keyring_library_degrades_to_the_file(self, monkeypatch): """keyring is an optional extra, so the SDK must survive its absence rather than raise on - the hot path.""" + the hot path. Erase cannot succeed: the entry belongs to the OS and outlives the package, + so an install without it is not evidence that the keychain is empty.""" monkeypatch.delenv(DISABLE_KEYRING_ENV_VAR, raising=False) monkeypatch.setitem(sys.modules, "keyring", None) vault = KeyringVault() assert vault.read() == KeyringNotInstalled() assert vault.write("blob-1") == KeyringNotInstalled() - assert vault.erase() is True + assert vault.erase() == KeyringNotInstalled() def test_a_locked_keychain_is_reported_not_raised(self, install_fake_keyring): install_fake_keyring(_FakeKeyringModule(get_error=RuntimeError("keyring is locked"))) @@ -424,9 +481,9 @@ class TestKeyringVault: def test_a_refused_delete_is_reported_so_logout_can_warn(self, install_fake_keyring): install_fake_keyring(_FakeKeyringModule(stored="blob-1", delete_error=RuntimeError("locked"))) - assert KeyringVault().erase() is False + assert KeyringVault().erase() == SecretStranded() def test_erasing_a_locked_keychain_is_a_failure(self, install_fake_keyring): install_fake_keyring(_FakeKeyringModule(get_error=RuntimeError("locked"))) - assert KeyringVault().erase() is False + assert KeyringVault().erase() == KeyringUnreachable() diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 1f2d48f6547..33bd8307c21 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -46,6 +46,12 @@ def _write_home_json(home: Path, filename: str, payload: dict[str, object]) -> N (litellm_dir / filename).write_text(json.dumps(payload)) +def _write_token_file(home: Path, *, key: str | None) -> None: + """A stored login: `key=None` is the metadata half of a keychain-backed pair, a key is a file-backed one.""" + payload: dict[str, object] = {"base_url": "https://test.example.com", "user_id": "u-1", "timestamp": time.time()} + _write_home_json(home, "token.json", payload if key is None else {**payload, "key": key}) + + def _secret_blob(base_url: str, key: str) -> str: return json.dumps({"base_url": base_url, "key": key, "jwt_token": ""}) @@ -427,14 +433,62 @@ class TestLogoutCommand: """Setup for each test""" self.runner = CliRunner() - def test_logout_success(self): + def test_logout_success(self, isolated_home, secret_vault_factory): """Test successful logout""" - with patch("litellm.proxy.client.cli.commands.auth.clear_cli_token") as mock_clear: - result = self.runner.invoke(logout) + vault = secret_vault_factory(blob=_secret_blob("https://test.example.com", "sk-stored")) + _write_token_file(isolated_home, key=None) - assert result.exit_code == 0 - assert "Logged out successfully" in result.output - mock_clear.assert_called_once() + result = self.runner.invoke(logout, obj={"secret_vault": vault}) + + assert result.exit_code == 0 + assert "Logged out successfully" in result.output + assert vault.blob is None + assert not (isolated_home / ".litellm" / "token.json").exists() + + def test_logout_without_the_keyring_package_does_not_claim_the_keychain_is_clear( + self, isolated_home, secret_vault_factory + ): + """Logging out from an install without the cli extra cannot touch an entry a keychain-backed + login left behind, so it must point at the package rather than report a clean logout.""" + _write_token_file(isolated_home, key=None) + + result = self.runner.invoke( + logout, obj={"secret_vault": secret_vault_factory(available=False, failure=KeyringNotInstalled())} + ) + + assert result.exit_code == 0 + assert "Logged out successfully" not in result.output + assert "still in the OS keychain" in result.output + assert "pip install 'litellm[cli]'" in result.output + + def test_logout_warns_when_the_keychain_refuses_to_release_the_entry( + self, isolated_home, secret_vault_factory + ): + """A locked keychain leaves a live credential behind that the user believes is gone.""" + vault = secret_vault_factory( + blob=_secret_blob("https://test.example.com", "sk-stored"), erasable=False + ) + _write_token_file(isolated_home, key=None) + + result = self.runner.invoke(logout, obj={"secret_vault": vault}) + + assert result.exit_code == 0 + assert "Logged out successfully" not in result.output + assert "still in the OS keychain" in result.output + assert "Unlock your keychain" in result.output + + def test_logout_from_a_file_only_login_stays_quiet(self, isolated_home, secret_vault_factory): + """The credential never went to a keychain, so removing the file is the whole logout and + warning about a keychain entry would send the user chasing one that cannot exist.""" + _write_token_file(isolated_home, key="sk-in-file") + + result = self.runner.invoke( + logout, obj={"secret_vault": secret_vault_factory(available=False, failure=KeyringNotInstalled())} + ) + + assert result.exit_code == 0 + assert "Logged out successfully" in result.output + assert "still in the OS keychain" not in result.output class TestWhoamiCommand: From 424e74ba9ffc9f33757d3c68f90d0fef2cda4079 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 20:31:43 -0700 Subject: [PATCH 038/281] fix(cli): roll the keychain write back when the plaintext copy cannot be removed Removing the file when it could not be rewritten covered a full disk, but not a ~/.litellm that permits neither the rewrite nor the delete, which is what a `sudo lite login` leaves behind. There the secret was copied into the keychain and kept in cleartext on disk, so migration widened exposure instead of narrowing it Migration now only keeps the vault copy if the file's copy is gone. When it is not, the write is rolled back and the user is left exactly as they were, logged in with one copy of the credential --- litellm/litellm_core_utils/cli_token_utils.py | 26 +++++++++++++------ .../test_cli_token_utils.py | 20 ++++++++++++++ 2 files changed, 38 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index cd69f9470a3..53289770c1a 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -12,7 +12,6 @@ first time it reads one. This module has no dependencies on proxy code and can be safely imported at the SDK level. """ -import contextlib import time from pathlib import Path from types import MappingProxyType @@ -217,30 +216,41 @@ def _apply_vault_secret(record: CliTokenRecord, blob: str, vault: SecretVault) - def _migrate_file_secret(record: CliTokenRecord, vault: SecretVault) -> CliTokenRecord | None: + """Move a file-held secret into the vault, but only if the file's copy can be taken away. + + Migrating without scrubbing would leave the credential live in two stores instead of one, so a + file that will not give its copy up rolls the vault write back rather than widening exposure. + """ if record.key is None: return None - if isinstance(vault.write(_encode_secret(record.base_url, record.key, record.jwt_token)), SecretStored): - _scrub_file_secret(record) + if not isinstance(vault.write(_encode_secret(record.base_url, record.key, record.jwt_token)), SecretStored): + return record + if not _scrub_file_secret(record): + vault.erase() return record -def _scrub_file_secret(record: CliTokenRecord) -> None: +def _scrub_file_secret(record: CliTokenRecord) -> bool: """Leave no secret material in the token file once the vault holds it. A file that cannot be rewritten without the secret is removed instead. Signing in again costs the user one command; a live credential left behind in cleartext costs them the credential. """ if record.key is None and not record.jwt_token: - return + return True try: _write_token_file(_without_secret(record)) except OSError: - _discard_token_file() + return _discard_token_file() + return True -def _discard_token_file() -> None: - with contextlib.suppress(OSError): +def _discard_token_file() -> bool: + try: Path(get_cli_token_file_path()).unlink(missing_ok=True) + except OSError: + return False + return True def _without_secret(record: CliTokenRecord) -> CliTokenRecord: diff --git a/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py b/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py index 925b440dfb7..cbc93bbde71 100644 --- a/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py @@ -1,4 +1,5 @@ import json +import os import stat import sys import time @@ -320,6 +321,25 @@ class TestScrubFailure: """A keychain that took the secret while the file kept it is the worst of both stores: the credential is live, it is in cleartext on disk, and every command reports success.""" + @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores directory permissions") + def test_a_file_that_will_not_give_its_copy_up_rolls_the_vault_write_back( + self, isolated_home, secret_vault_factory + ): + """Handing the keychain a copy without taking the file's away leaves the credential live in + two stores instead of one. A directory that permits neither the rewrite nor the delete, a + root-owned ~/.litellm left behind by a `sudo lite login`, must widen nothing.""" + path = _write_legacy_file(isolated_home) + vault = secret_vault_factory() + path.parent.chmod(0o500) + try: + record = load_cli_token(vault=vault) + finally: + path.parent.chmod(0o700) + + assert record.key == "sk-legacy" + assert json.loads(path.read_text())["key"] == "sk-legacy" + assert vault.blob is None + def test_a_file_that_cannot_be_rewritten_is_removed_instead( self, isolated_home, secret_vault_factory, monkeypatch ): From 7b574b9df6ea2bb4269bc5421ab1059c1bbab711 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 19 Aug 2026 20:44:27 -0700 Subject: [PATCH 039/281] chore(ui): drop the antd dependency and its leftovers (#37574) Nothing in the dashboard renders antd any more, so the package and the scaffolding around it can go. This removes `antd` and `@ant-design/cssinjs` from package.json, deletes the global StyleProvider the root layout wrapped every page in, drops the `antd` cascade layer and the z-index override that lifted Base UI popups over an antd Modal, and retires the lint rules that policed antd imports and antd class selectors in tests. Fifteen test files still carried `vi.mock("antd", ...)` factories for components that stopped importing antd during the migration. They were inert, and they resolve the real module, so they would have broken the moment the package left node_modules. The compatibility shims keep their behaviour and lose the antd name: `antdRules`/`antdRequired` become `validatorRules`/`requiredRule`, `isAntdUrl` becomes `isValidUrl`, and `ABOVE_ANTD_MODAL` becomes `NESTED_DIALOG_LAYER`. Comments that explain why a contract looks the way it does still name antd, because that history is the reason. --- ui/litellm-dashboard/CLAUDE.md | 4 +- ui/litellm-dashboard/eslint-suppressions.json | 5 - ui/litellm-dashboard/eslint.config.mjs | 10 - ui/litellm-dashboard/package-lock.json | 992 ------------------ ui/litellm-dashboard/package.json | 2 - .../scripts/eslint-rules/index.mjs | 2 - .../eslint-rules/no-antd-class-selectors.mjs | 45 - .../_components/TeamGuardrailsTab.tsx | 4 +- .../content_filter/CustomPatternModal.tsx | 4 +- .../content_filter/KeywordModal.tsx | 4 +- .../content_filter/PatternModal.tsx | 4 +- .../content_filter/dialog_layering.ts | 2 +- .../_components/AwsSigV4Fields.tsx | 4 +- .../_components/CreateMCPServer.tsx | 14 +- .../_components/EnvVarsSection.tsx | 4 +- .../_components/IdJagFormFields.tsx | 4 +- .../_components/MCPPermissionManagement.tsx | 6 +- .../_components/OAuthFormFields.tsx | 4 +- .../_components/OpenAPIFormSection.tsx | 4 +- .../_components/StdioConfiguration.tsx | 4 +- .../_components/TokenExchangeFormFields.tsx | 6 +- .../_components/mcp_server_edit.tsx | 18 +- .../_components/impact_popover.test.tsx | 47 - .../policies/_components/index.test.tsx | 10 - .../_components/SearchToolTester.test.tsx | 12 - .../users/_components/user_edit_view.test.tsx | 63 -- ui/litellm-dashboard/src/app/globals.css | 12 +- ui/litellm-dashboard/src/app/layout.tsx | 7 +- .../CloudZeroCreateModal.test.tsx | 11 - .../CloudZeroIntegrationSettings.test.tsx | 12 - .../CloudZeroUpdateModal.test.tsx | 11 - .../AdminSettings/PluginSettings/schema.ts | 4 +- .../src/components/add_model/AddModelForm.tsx | 8 +- .../add_model/advanced_settings.tsx | 24 +- .../conditional_public_model_name.tsx | 4 +- .../add_model/litellm_model_name.tsx | 6 +- .../add_model/provider_specific_fields.tsx | 4 +- .../src/components/chat/design.md | 2 +- .../{antdFormRules.ts => formRules.ts} | 12 +- .../edit_auto_router_modal.test.tsx | 2 - .../components/model_add/CredentialModal.tsx | 6 +- .../src/components/navbar.test.tsx | 5 +- .../KeyInfoView.handleKeyUpdate.test.tsx | 37 - .../templates/estimatedOutputTokens.ts | 6 +- .../LogDetailsDrawer/InputCard.test.tsx | 10 - .../LogDetailsDrawer/OutputCard.test.tsx | 10 - .../PrettyMessagesView.test.tsx | 12 +- .../src/contexts/AntdGlobalProvider.tsx | 8 - ...{antdUrl.test.ts => urlValidation.test.ts} | 18 +- .../forms/{antdUrl.ts => urlValidation.ts} | 6 +- 50 files changed, 105 insertions(+), 1410 deletions(-) delete mode 100644 ui/litellm-dashboard/scripts/eslint-rules/no-antd-class-selectors.mjs rename ui/litellm-dashboard/src/components/common_components/{antdFormRules.ts => formRules.ts} (80%) delete mode 100644 ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx rename ui/litellm-dashboard/src/lib/forms/{antdUrl.test.ts => urlValidation.test.ts} (82%) rename ui/litellm-dashboard/src/lib/forms/{antdUrl.ts => urlValidation.ts} (84%) diff --git a/ui/litellm-dashboard/CLAUDE.md b/ui/litellm-dashboard/CLAUDE.md index d7e0a11aaa1..f79258600c1 100644 --- a/ui/litellm-dashboard/CLAUDE.md +++ b/ui/litellm-dashboard/CLAUDE.md @@ -12,13 +12,13 @@ Most of the suite predates this split and is not yet classified, so an unsuffixe Assert something the user could perceive, and assert it precisely enough that the test fails when the behaviour breaks. `eslint-plugin-testing-library` and `eslint-plugin-jest-dom` enforce the mechanical part of that. Two of the enabled rules exist because the failure they catch is silent rather than cosmetic: `await-async-queries` catches an unawaited `findBy*`, whose returned Promise is always truthy and makes the whole assertion vacuous, and `no-wait-for-side-effects` catches work inside a `waitFor` callback, which is retried on every poll. Prefer `findBy*` over `waitFor` wrapped around `getBy*`, and keep a `waitFor` callback to a single assertion -Do not trust `eslint --fix` for these two plugins. Fixing the suite in bulk produced seven distinct kinds of broken output. Four fail loudly: `no-wait-for-side-effects` and `no-wait-for-multiple-assertions` hoist a statement out of the `waitFor` callback while leaving the `const` it reads inside, `prefer-enabled-disabled` drops a closing paren when the subject carries a type assertion, `prefer-presence-queries` swaps in a query it never destructures, and `prefer-in-document` collapses `getAllBy*` to `getBy*` on a value still indexed as an array. Two fail quietly, which is worse: `prefer-checked` swaps the `checked` attribute for the `.checked` property, and antd radios set one without the other, and `prefer-to-have-text-content` wraps arbitrary strings in `new RegExp()` without escaping, so `toContain("100K+ requests")` becomes a pattern meaning "100 followed by one-or-more K". That last one compiles, lints clean, and keeps passing while no longer asserting what it says. Pass a plain string to `toHaveTextContent`, which is already a substring match. Run the fixer on a handful of files at a time and read the diff +Do not trust `eslint --fix` for these two plugins. Fixing the suite in bulk produced seven distinct kinds of broken output. Four fail loudly: `no-wait-for-side-effects` and `no-wait-for-multiple-assertions` hoist a statement out of the `waitFor` callback while leaving the `const` it reads inside, `prefer-enabled-disabled` drops a closing paren when the subject carries a type assertion, `prefer-presence-queries` swaps in a query it never destructures, and `prefer-in-document` collapses `getAllBy*` to `getBy*` on a value still indexed as an array. Two fail quietly, which is worse: `prefer-checked` swaps the `checked` attribute for the `.checked` property, and a radio can set one without the other, and `prefer-to-have-text-content` wraps arbitrary strings in `new RegExp()` without escaping, so `toContain("100K+ requests")` becomes a pattern meaning "100 followed by one-or-more K". That last one compiles, lints clean, and keeps passing while no longer asserting what it says. Pass a plain string to `toHaveTextContent`, which is already a substring match. Run the fixer on a handful of files at a time and read the diff `jest-dom/prefer-to-have-value` stays off because its fixer is wrong here, not merely noisy. It matches any attribute whose name contains "value", so it rewrites `toHaveAttribute("aria-valuenow", n)` into `toHaveValue(n)`, and jest-dom's `toHaveValue` only supports form controls, so the assertion fails on the `role="meter"` elements the dashboard renders. Assert ARIA value attributes with `toHaveAttribute` Reach for `fireEvent.change` rather than `user.type` when a test only needs a field to hold a value. `user.type` dispatches one event per character and re-renders the whole form each time, which is why a single form test could burn seven seconds. Keep `user.type` where the typing itself is the behaviour under test: an autocomplete that filters per keystroke, a debounce, a key handler, or any Base UI combobox, whose filter state is driven by real keyboard input and does not react to a raw change event -A test may reach for a component library's own CSS class only when that library exposes no role, label, title or ARIA state to query instead, and then the line carries a suppression naming the rule and the reason. Check first: antd icons render as `role="img"` with an `aria-label`, and antd `Form.Item` associates its label with the control, so both are reachable accessibly. When a label does not resolve, suspect the control rather than the test, since a custom wrapper that destructures props without spreading them drops the `id` antd injects and leaves the rendered label pointing at nothing +A test may reach for a component library's own CSS class only when that library exposes no role, label, title or ARIA state to query instead. Check first: the shadcn primitives forward roles and `aria-label`, and the shared form field associates its label with the control, so both are reachable accessibly. When nothing accessible identifies the element, prefer its `data-slot` attribute, which the primitives set deliberately and treat as stable. When a label does not resolve, suspect the control rather than the test, since a custom wrapper that destructures props without spreading them drops the `id` the field generates and leaves the rendered label pointing at nothing Rules beyond the enabled set were measured against the whole suite and left off rather than recorded in a budget file, because a ceiling that permits a violation anywhere is worse than an honest gap. `no-node-access` and `no-container` are the ones worth revisiting first, since they catch the DOM archaeology the rules above only discourage. `prefer-implicit-assert` and `prefer-explicit-assert` contradict each other, so neither is enabled diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index a9a163f3c78..3bc93b7ebc4 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1181,11 +1181,6 @@ "count": 1 } }, - "src/app/(dashboard)/users/_components/user_edit_view.test.tsx": { - "no-nested-ternary": { - "count": 1 - } - }, "src/app/(dashboard)/users/_components/user_edit_view.tsx": { "local/filename-pascal-case": { "count": 1 diff --git a/ui/litellm-dashboard/eslint.config.mjs b/ui/litellm-dashboard/eslint.config.mjs index 3cbb93f9a48..df10ca1befa 100644 --- a/ui/litellm-dashboard/eslint.config.mjs +++ b/ui/litellm-dashboard/eslint.config.mjs @@ -57,15 +57,6 @@ const eslintConfig = [ message: "@tremor/react is being phased out; build new UI with shadcn/ui primitives instead of adding tremor imports.", }, - { - group: ["antd", "antd/*"], - message: - "antd is being phased out; build new UI with shadcn/ui primitives instead of adding antd imports.", - }, - { - group: ["@ant-design/icons", "@ant-design/icons/*"], - message: "@ant-design/icons is gone from the dashboard; use lucide-react instead.", - }, ], }, ], @@ -94,7 +85,6 @@ const eslintConfig = [ files: ["src/**/*.test.{ts,tsx}", "tests/**/*.{ts,tsx}"], plugins: { "testing-library": testingLibrary, "jest-dom": jestDom }, rules: { - "local/no-antd-class-selectors": "error", "testing-library/await-async-queries": "error", "testing-library/no-wait-for-multiple-assertions": "error", "testing-library/no-wait-for-side-effects": "error", diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 154a19da7f3..7ea6aa5b2a4 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -8,7 +8,6 @@ "name": "litellm-dashboard", "version": "0.1.0", "dependencies": { - "@ant-design/cssinjs": "1.24.0", "@anthropic-ai/sdk": "0.92.0", "@base-ui/react": "^1.6.0", "@headlessui/tailwindcss": "0.2.2", @@ -18,7 +17,6 @@ "@tanstack/react-query": "5.100.7", "@tanstack/react-table": "8.21.3", "@types/papaparse": "5.5.2", - "antd": "5.29.3", "cva": "1.0.0-beta.4", "date-fns": "^4.4.0", "dayjs": "1.11.19", @@ -123,103 +121,6 @@ "node": ">=6.0.0" } }, - "node_modules/@ant-design/colors": { - "version": "7.2.1", - "resolved": "https://registry.npmjs.org/@ant-design/colors/-/colors-7.2.1.tgz", - "integrity": "sha512-lCHDcEzieu4GA3n8ELeZ5VQ8pKQAWcGGLRTQ50aQM2iqPpq2evTxER84jfdPvsPAtEcZ7m44NI45edFMo8oOYQ==", - "license": "MIT", - "dependencies": { - "@ant-design/fast-color": "^2.0.6" - } - }, - "node_modules/@ant-design/cssinjs": { - "version": "1.24.0", - "resolved": "https://registry.npmjs.org/@ant-design/cssinjs/-/cssinjs-1.24.0.tgz", - "integrity": "sha512-K4cYrJBsgvL+IoozUXYjbT6LHHNt+19a9zkvpBPxLjFHas1UpPM2A5MlhROb0BT8N8WoavM5VsP9MeSeNK/3mg==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.11.1", - "@emotion/hash": "^0.8.0", - "@emotion/unitless": "^0.7.5", - "classnames": "^2.3.1", - "csstype": "^3.1.3", - "rc-util": "^5.35.0", - "stylis": "^4.3.4" - }, - "peerDependencies": { - "react": ">=16.0.0", - "react-dom": ">=16.0.0" - } - }, - "node_modules/@ant-design/cssinjs-utils": { - "version": "1.1.3", - "resolved": "https://registry.npmjs.org/@ant-design/cssinjs-utils/-/cssinjs-utils-1.1.3.tgz", - "integrity": "sha512-nOoQMLW1l+xR1Co8NFVYiP8pZp3VjIIzqV6D6ShYF2ljtdwWJn5WSsH+7kvCktXL/yhEtWURKOfH5Xz/gzlwsg==", - "license": "MIT", - "dependencies": { - "@ant-design/cssinjs": "^1.21.0", - "@babel/runtime": "^7.23.2", - "rc-util": "^5.38.0" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/@ant-design/fast-color": { - "version": "2.0.6", - "resolved": "https://registry.npmjs.org/@ant-design/fast-color/-/fast-color-2.0.6.tgz", - "integrity": "sha512-y2217gk4NqL35giHl72o6Zzqji9O7vHh9YmhUVkPtAOpoTCH4uWxo/pr4VE8t0+ChEPs0qo4eJRC5Q1eXWo3vA==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.24.7" - }, - "engines": { - "node": ">=8.x" - } - }, - "node_modules/@ant-design/icons": { - "version": "5.6.1", - "resolved": "https://registry.npmjs.org/@ant-design/icons/-/icons-5.6.1.tgz", - "integrity": "sha512-0/xS39c91WjPAZOWsvi1//zjx6kAp4kxWwctR6kuU6p133w8RU0D2dSCvZC19uQyharg/sAvYxGYWl01BbZZfg==", - "license": "MIT", - "dependencies": { - "@ant-design/colors": "^7.0.0", - "@ant-design/icons-svg": "^4.4.0", - "@babel/runtime": "^7.24.8", - "classnames": "^2.2.6", - "rc-util": "^5.31.1" - }, - "engines": { - "node": ">=8" - }, - "peerDependencies": { - "react": ">=16.0.0", - "react-dom": ">=16.0.0" - } - }, - "node_modules/@ant-design/icons-svg": { - "version": "4.5.0", - "resolved": "https://registry.npmjs.org/@ant-design/icons-svg/-/icons-svg-4.5.0.tgz", - "integrity": "sha512-1BTUFyKPTBZ53MuTP8s0k5SFEXL7o3VHEOwLgzaoWKwnBeqIcqUtVshc4SKzhI6uACfqhJqBwBUE9FsWR3uULA==", - "license": "MIT" - }, - "node_modules/@ant-design/react-slick": { - "version": "1.1.2", - "resolved": "https://registry.npmjs.org/@ant-design/react-slick/-/react-slick-1.1.2.tgz", - "integrity": "sha512-EzlvzE6xQUBrZuuhSAFTdsr4P2bBBHGZwKFemEfq8gIGyIQCxalYfZW/T2ORbtQx5rU69o+WycP3exY/7T1hGA==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.10.4", - "classnames": "^2.2.5", - "json2mq": "^0.2.0", - "resize-observer-polyfill": "^1.5.1", - "throttle-debounce": "^5.0.0" - }, - "peerDependencies": { - "react": ">=16.9.0" - } - }, "node_modules/@anthropic-ai/sdk": { "version": "0.92.0", "resolved": "https://registry.npmjs.org/@anthropic-ai/sdk/-/sdk-0.92.0.tgz", @@ -810,18 +711,6 @@ "tslib": "^2.4.0" } }, - "node_modules/@emotion/hash": { - "version": "0.8.0", - "resolved": "https://registry.npmjs.org/@emotion/hash/-/hash-0.8.0.tgz", - "integrity": "sha512-kBJtf7PH6aWwZ6fka3zQ0p6SBYzx4fl1LoZXE2RrnYST9Xljm7WfKJrU4g/Xr3Beg72MLrp1AWNUmuYJTL7Cow==", - "license": "MIT" - }, - "node_modules/@emotion/unitless": { - "version": "0.7.5", - "resolved": "https://registry.npmjs.org/@emotion/unitless/-/unitless-0.7.5.tgz", - "integrity": "sha512-OWORNpfjMsSSUBVrRBVGECkhWcULOAJz9ZW8uK9qgxD+87M7jHRcvh/A96XXNhXTLmKcoYSQtBEX7lHMO7YRwg==", - "license": "MIT" - }, "node_modules/@esbuild/aix-ppc64": { "version": "0.28.1", "resolved": "https://registry.npmjs.org/@esbuild/aix-ppc64/-/aix-ppc64-0.28.1.tgz", @@ -2625,153 +2514,6 @@ "dev": true, "license": "MIT" }, - "node_modules/@rc-component/async-validator": { - "version": "5.1.2", - "resolved": "https://registry.npmjs.org/@rc-component/async-validator/-/async-validator-5.1.2.tgz", - "integrity": "sha512-WYbrZSjzznU1ekD0qFq2qRxt309VoS61MTG5npnFQlKYcoy9IzU8T+ZCIhq5bGAXRbXysABFWTspicMfmWFwow==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.24.4" - }, - "engines": { - "node": ">=14.x" - } - }, - "node_modules/@rc-component/color-picker": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/@rc-component/color-picker/-/color-picker-2.0.1.tgz", - "integrity": "sha512-WcZYwAThV/b2GISQ8F+7650r5ZZJ043E57aVBFkQ+kSY4C6wdofXgB0hBx+GPGpIU0Z81eETNoDUJMr7oy/P8Q==", - "license": "MIT", - "dependencies": { - "@ant-design/fast-color": "^2.0.6", - "@babel/runtime": "^7.23.6", - "classnames": "^2.2.6", - "rc-util": "^5.38.1" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/@rc-component/context": { - "version": "1.4.0", - "resolved": "https://registry.npmjs.org/@rc-component/context/-/context-1.4.0.tgz", - "integrity": "sha512-kFcNxg9oLRMoL3qki0OMxK+7g5mypjgaaJp/pkOis/6rVxma9nJBF/8kCIuTYHUQNr0ii7MxqE33wirPZLJQ2w==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.10.1", - "rc-util": "^5.27.0" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/@rc-component/mini-decimal": { - "version": "1.1.4", - "resolved": "https://registry.npmjs.org/@rc-component/mini-decimal/-/mini-decimal-1.1.4.tgz", - "integrity": "sha512-xiuXcaCwyOWpD8a8scdExFl+bntNphAW8XeenL1ig2en0AAZY0Pcp4pC0dI22qJ+NvxKn9RoNIoRdqYU3BLH4w==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.18.0" - }, - "engines": { - "node": ">=8.x" - } - }, - "node_modules/@rc-component/mutate-observer": { - "version": "1.1.0", - "resolved": "https://registry.npmjs.org/@rc-component/mutate-observer/-/mutate-observer-1.1.0.tgz", - "integrity": "sha512-QjrOsDXQusNwGZPf4/qRQasg7UFEj06XiCJ8iuiq/Io7CrHrgVi6Uuetw60WAMG1799v+aM8kyc+1L/GBbHSlw==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.18.0", - "classnames": "^2.3.2", - "rc-util": "^5.24.4" - }, - "engines": { - "node": ">=8.x" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/@rc-component/portal": { - "version": "1.1.2", - "resolved": "https://registry.npmjs.org/@rc-component/portal/-/portal-1.1.2.tgz", - "integrity": "sha512-6f813C0IsasTZms08kfA8kPAGxbbkYToa8ALaiDIGGECU4i9hj8Plgbx0sNJDrey3EtHO30hmdaxtT0138xZcg==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.18.0", - "classnames": "^2.3.2", - "rc-util": "^5.24.4" - }, - "engines": { - "node": ">=8.x" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/@rc-component/qrcode": { - "version": "1.1.3", - "resolved": "https://registry.npmjs.org/@rc-component/qrcode/-/qrcode-1.1.3.tgz", - "integrity": "sha512-aGv6alnn4HbDEsURzKP+jv13rbi1VxmAYfBNZr5GKF1iohMNWy5tAVoJ1E3cOvzMB1kbUPvCXchM6zSFlRGPhA==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.24.7" - }, - "engines": { - "node": ">=8.x" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/@rc-component/tour": { - "version": "1.15.1", - "resolved": "https://registry.npmjs.org/@rc-component/tour/-/tour-1.15.1.tgz", - "integrity": "sha512-Tr2t7J1DKZUpfJuDZWHxyxWpfmj8EZrqSgyMZ+BCdvKZ6r1UDsfU46M/iWAAFBy961Ssfom2kv5f3UcjIL2CmQ==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.18.0", - "@rc-component/portal": "^1.0.0-9", - "@rc-component/trigger": "^2.0.0", - "classnames": "^2.3.2", - "rc-util": "^5.24.4" - }, - "engines": { - "node": ">=8.x" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/@rc-component/trigger": { - "version": "2.3.1", - "resolved": "https://registry.npmjs.org/@rc-component/trigger/-/trigger-2.3.1.tgz", - "integrity": "sha512-ORENF39PeXTzM+gQEshuk460Z8N4+6DkjpxlpE7Q3gYy1iBpLrx0FOJz3h62ryrJZ/3zCAUIkT1Pb/8hHWpb3A==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.23.2", - "@rc-component/portal": "^1.1.0", - "classnames": "^2.3.2", - "rc-motion": "^2.0.0", - "rc-resize-observer": "^1.3.1", - "rc-util": "^5.44.0" - }, - "engines": { - "node": ">=8.x" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, "node_modules/@redocly/ajv": { "version": "8.11.2", "resolved": "https://registry.npmjs.org/@redocly/ajv/-/ajv-8.11.2.tgz", @@ -4861,71 +4603,6 @@ "url": "https://github.com/chalk/ansi-styles?sponsor=1" } }, - "node_modules/antd": { - "version": "5.29.3", - "resolved": "https://registry.npmjs.org/antd/-/antd-5.29.3.tgz", - "integrity": "sha512-3DdbGCa9tWAJGcCJ6rzR8EJFsv2CtyEbkVabZE14pfgUHfCicWCj0/QzQVLDYg8CPfQk9BH7fHCoTXHTy7MP/A==", - "license": "MIT", - "dependencies": { - "@ant-design/colors": "^7.2.1", - "@ant-design/cssinjs": "^1.23.0", - "@ant-design/cssinjs-utils": "^1.1.3", - "@ant-design/fast-color": "^2.0.6", - "@ant-design/icons": "^5.6.1", - "@ant-design/react-slick": "~1.1.2", - "@babel/runtime": "^7.26.0", - "@rc-component/color-picker": "~2.0.1", - "@rc-component/mutate-observer": "^1.1.0", - "@rc-component/qrcode": "~1.1.0", - "@rc-component/tour": "~1.15.1", - "@rc-component/trigger": "^2.3.0", - "classnames": "^2.5.1", - "copy-to-clipboard": "^3.3.3", - "dayjs": "^1.11.11", - "rc-cascader": "~3.34.0", - "rc-checkbox": "~3.5.0", - "rc-collapse": "~3.9.0", - "rc-dialog": "~9.6.0", - "rc-drawer": "~7.3.0", - "rc-dropdown": "~4.2.1", - "rc-field-form": "~2.7.1", - "rc-image": "~7.12.0", - "rc-input": "~1.8.0", - "rc-input-number": "~9.5.0", - "rc-mentions": "~2.20.0", - "rc-menu": "~9.16.1", - "rc-motion": "^2.9.5", - "rc-notification": "~5.6.4", - "rc-pagination": "~5.1.0", - "rc-picker": "~4.11.3", - "rc-progress": "~4.0.0", - "rc-rate": "~2.13.1", - "rc-resize-observer": "^1.4.3", - "rc-segmented": "~2.7.0", - "rc-select": "~14.16.8", - "rc-slider": "~11.1.9", - "rc-steps": "~6.0.1", - "rc-switch": "~4.1.0", - "rc-table": "~7.54.0", - "rc-tabs": "~15.7.0", - "rc-textarea": "~1.10.2", - "rc-tooltip": "~6.4.0", - "rc-tree": "~5.13.1", - "rc-tree-select": "~5.27.0", - "rc-upload": "~4.11.0", - "rc-util": "^5.44.4", - "scroll-into-view-if-needed": "^3.1.0", - "throttle-debounce": "^5.0.2" - }, - "funding": { - "type": "opencollective", - "url": "https://opencollective.com/ant-design" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, "node_modules/argparse": { "version": "2.0.1", "resolved": "https://registry.npmjs.org/argparse/-/argparse-2.0.1.tgz", @@ -5483,12 +5160,6 @@ "node": ">= 16" } }, - "node_modules/classnames": { - "version": "2.5.1", - "resolved": "https://registry.npmjs.org/classnames/-/classnames-2.5.1.tgz", - "integrity": "sha512-saHYOzhIQs6wy2sVxTM6bUDsQO4F50V9RQ22qBpEdCW+I+/Wmke2HOl6lS6dTpdxVhb88/I6+Hs+438c3lfUow==", - "license": "MIT" - }, "node_modules/client-only": { "version": "0.0.1", "resolved": "https://registry.npmjs.org/client-only/-/client-only-0.0.1.tgz", @@ -5553,12 +5224,6 @@ "url": "https://github.com/sponsors/wooorm" } }, - "node_modules/compute-scroll-into-view": { - "version": "3.1.1", - "resolved": "https://registry.npmjs.org/compute-scroll-into-view/-/compute-scroll-into-view-3.1.1.tgz", - "integrity": "sha512-VRhuHOLoKYOy4UbilLbUzbYg93XLjv2PncJC50EuTWPA3gaja1UjBsUP/D/9/juV3vQFr6XBEzn9KCAHdUvOHw==", - "license": "MIT" - }, "node_modules/convert-source-map": { "version": "2.0.0", "resolved": "https://registry.npmjs.org/convert-source-map/-/convert-source-map-2.0.0.tgz", @@ -8465,15 +8130,6 @@ "dev": true, "license": "MIT" }, - "node_modules/json2mq": { - "version": "0.2.0", - "resolved": "https://registry.npmjs.org/json2mq/-/json2mq-0.2.0.tgz", - "integrity": "sha512-SzoRg7ux5DWTII9J2qkrZrqV1gt+rTaoufMxEzXbS26Uid0NwaJd123HcoB80TgubEppxxIGdNxCx50fEoEWQA==", - "license": "MIT", - "dependencies": { - "string-convert": "^0.2.0" - } - }, "node_modules/json5": { "version": "1.0.2", "resolved": "https://registry.npmjs.org/json5/-/json5-1.0.2.tgz", @@ -10937,618 +10593,6 @@ ], "license": "MIT" }, - "node_modules/rc-cascader": { - "version": "3.34.0", - "resolved": "https://registry.npmjs.org/rc-cascader/-/rc-cascader-3.34.0.tgz", - "integrity": "sha512-KpXypcvju9ptjW9FaN2NFcA2QH9E9LHKq169Y0eWtH4e/wHQ5Wh5qZakAgvb8EKZ736WZ3B0zLLOBsrsja5Dag==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.25.7", - "classnames": "^2.3.1", - "rc-select": "~14.16.2", - "rc-tree": "~5.13.0", - "rc-util": "^5.43.0" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-checkbox": { - "version": "3.5.0", - "resolved": "https://registry.npmjs.org/rc-checkbox/-/rc-checkbox-3.5.0.tgz", - "integrity": "sha512-aOAQc3E98HteIIsSqm6Xk2FPKIER6+5vyEFMZfo73TqM+VVAIqOkHoPjgKLqSNtVLWScoaM7vY2ZrGEheI79yg==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.10.1", - "classnames": "^2.3.2", - "rc-util": "^5.25.2" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-collapse": { - "version": "3.9.0", - "resolved": "https://registry.npmjs.org/rc-collapse/-/rc-collapse-3.9.0.tgz", - "integrity": "sha512-swDdz4QZ4dFTo4RAUMLL50qP0EY62N2kvmk2We5xYdRwcRn8WcYtuetCJpwpaCbUfUt5+huLpVxhvmnK+PHrkA==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.10.1", - "classnames": "2.x", - "rc-motion": "^2.3.4", - "rc-util": "^5.27.0" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-dialog": { - "version": "9.6.0", - "resolved": "https://registry.npmjs.org/rc-dialog/-/rc-dialog-9.6.0.tgz", - "integrity": "sha512-ApoVi9Z8PaCQg6FsUzS8yvBEQy0ZL2PkuvAgrmohPkN3okps5WZ5WQWPc1RNuiOKaAYv8B97ACdsFU5LizzCqg==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.10.1", - "@rc-component/portal": "^1.0.0-8", - "classnames": "^2.2.6", - "rc-motion": "^2.3.0", - "rc-util": "^5.21.0" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-drawer": { - "version": "7.3.0", - "resolved": "https://registry.npmjs.org/rc-drawer/-/rc-drawer-7.3.0.tgz", - "integrity": "sha512-DX6CIgiBWNpJIMGFO8BAISFkxiuKitoizooj4BDyee8/SnBn0zwO2FHrNDpqqepj0E/TFTDpmEBCyFuTgC7MOg==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.23.9", - "@rc-component/portal": "^1.1.1", - "classnames": "^2.2.6", - "rc-motion": "^2.6.1", - "rc-util": "^5.38.1" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-dropdown": { - "version": "4.2.1", - "resolved": "https://registry.npmjs.org/rc-dropdown/-/rc-dropdown-4.2.1.tgz", - "integrity": "sha512-YDAlXsPv3I1n42dv1JpdM7wJ+gSUBfeyPK59ZpBD9jQhK9jVuxpjj3NmWQHOBceA1zEPVX84T2wbdb2SD0UjmA==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.18.3", - "@rc-component/trigger": "^2.0.0", - "classnames": "^2.2.6", - "rc-util": "^5.44.1" - }, - "peerDependencies": { - "react": ">=16.11.0", - "react-dom": ">=16.11.0" - } - }, - "node_modules/rc-field-form": { - "version": "2.7.1", - "resolved": "https://registry.npmjs.org/rc-field-form/-/rc-field-form-2.7.1.tgz", - "integrity": "sha512-vKeSifSJ6HoLaAB+B8aq/Qgm8a3dyxROzCtKNCsBQgiverpc4kWDQihoUwzUj+zNWJOykwSY4dNX3QrGwtVb9A==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.18.0", - "@rc-component/async-validator": "^5.0.3", - "rc-util": "^5.32.2" - }, - "engines": { - "node": ">=8.x" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-image": { - "version": "7.12.0", - "resolved": "https://registry.npmjs.org/rc-image/-/rc-image-7.12.0.tgz", - "integrity": "sha512-cZ3HTyyckPnNnUb9/DRqduqzLfrQRyi+CdHjdqgsyDpI3Ln5UX1kXnAhPBSJj9pVRzwRFgqkN7p9b6HBDjmu/Q==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.11.2", - "@rc-component/portal": "^1.0.2", - "classnames": "^2.2.6", - "rc-dialog": "~9.6.0", - "rc-motion": "^2.6.2", - "rc-util": "^5.34.1" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-input": { - "version": "1.8.0", - "resolved": "https://registry.npmjs.org/rc-input/-/rc-input-1.8.0.tgz", - "integrity": "sha512-KXvaTbX+7ha8a/k+eg6SYRVERK0NddX8QX7a7AnRvUa/rEH0CNMlpcBzBkhI0wp2C8C4HlMoYl8TImSN+fuHKA==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.11.1", - "classnames": "^2.2.1", - "rc-util": "^5.18.1" - }, - "peerDependencies": { - "react": ">=16.0.0", - "react-dom": ">=16.0.0" - } - }, - "node_modules/rc-input-number": { - "version": "9.5.0", - "resolved": "https://registry.npmjs.org/rc-input-number/-/rc-input-number-9.5.0.tgz", - "integrity": "sha512-bKaEvB5tHebUURAEXw35LDcnRZLq3x1k7GxfAqBMzmpHkDGzjAtnUL8y4y5N15rIFIg5IJgwr211jInl3cipag==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.10.1", - "@rc-component/mini-decimal": "^1.0.1", - "classnames": "^2.2.5", - "rc-input": "~1.8.0", - "rc-util": "^5.40.1" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-mentions": { - "version": "2.20.0", - "resolved": "https://registry.npmjs.org/rc-mentions/-/rc-mentions-2.20.0.tgz", - "integrity": "sha512-w8HCMZEh3f0nR8ZEd466ATqmXFCMGMN5UFCzEUL0bM/nGw/wOS2GgRzKBcm19K++jDyuWCOJOdgcKGXU3fXfbQ==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.22.5", - "@rc-component/trigger": "^2.0.0", - "classnames": "^2.2.6", - "rc-input": "~1.8.0", - "rc-menu": "~9.16.0", - "rc-textarea": "~1.10.0", - "rc-util": "^5.34.1" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-menu": { - "version": "9.16.1", - "resolved": "https://registry.npmjs.org/rc-menu/-/rc-menu-9.16.1.tgz", - "integrity": "sha512-ghHx6/6Dvp+fw8CJhDUHFHDJ84hJE3BXNCzSgLdmNiFErWSOaZNsihDAsKq9ByTALo/xkNIwtDFGIl6r+RPXBg==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.10.1", - "@rc-component/trigger": "^2.0.0", - "classnames": "2.x", - "rc-motion": "^2.4.3", - "rc-overflow": "^1.3.1", - "rc-util": "^5.27.0" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-motion": { - "version": "2.9.5", - "resolved": "https://registry.npmjs.org/rc-motion/-/rc-motion-2.9.5.tgz", - "integrity": "sha512-w+XTUrfh7ArbYEd2582uDrEhmBHwK1ZENJiSJVb7uRxdE7qJSYjbO2eksRXmndqyKqKoYPc9ClpPh5242mV1vA==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.11.1", - "classnames": "^2.2.1", - "rc-util": "^5.44.0" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-notification": { - "version": "5.6.4", - "resolved": "https://registry.npmjs.org/rc-notification/-/rc-notification-5.6.4.tgz", - "integrity": "sha512-KcS4O6B4qzM3KH7lkwOB7ooLPZ4b6J+VMmQgT51VZCeEcmghdeR4IrMcFq0LG+RPdnbe/ArT086tGM8Snimgiw==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.10.1", - "classnames": "2.x", - "rc-motion": "^2.9.0", - "rc-util": "^5.20.1" - }, - "engines": { - "node": ">=8.x" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-overflow": { - "version": "1.5.0", - "resolved": "https://registry.npmjs.org/rc-overflow/-/rc-overflow-1.5.0.tgz", - "integrity": "sha512-Lm/v9h0LymeUYJf0x39OveU52InkdRXqnn2aYXfWmo8WdOonIKB2kfau+GF0fWq6jPgtdO9yMqveGcK6aIhJmg==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.11.1", - "classnames": "^2.2.1", - "rc-resize-observer": "^1.0.0", - "rc-util": "^5.37.0" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-pagination": { - "version": "5.1.0", - "resolved": "https://registry.npmjs.org/rc-pagination/-/rc-pagination-5.1.0.tgz", - "integrity": "sha512-8416Yip/+eclTFdHXLKTxZvn70duYVGTvUUWbckCCZoIl3jagqke3GLsFrMs0bsQBikiYpZLD9206Ej4SOdOXQ==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.10.1", - "classnames": "^2.3.2", - "rc-util": "^5.38.0" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-picker": { - "version": "4.11.3", - "resolved": "https://registry.npmjs.org/rc-picker/-/rc-picker-4.11.3.tgz", - "integrity": "sha512-MJ5teb7FlNE0NFHTncxXQ62Y5lytq6sh5nUw0iH8OkHL/TjARSEvSHpr940pWgjGANpjCwyMdvsEV55l5tYNSg==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.24.7", - "@rc-component/trigger": "^2.0.0", - "classnames": "^2.2.1", - "rc-overflow": "^1.3.2", - "rc-resize-observer": "^1.4.0", - "rc-util": "^5.43.0" - }, - "engines": { - "node": ">=8.x" - }, - "peerDependencies": { - "date-fns": ">= 2.x", - "dayjs": ">= 1.x", - "luxon": ">= 3.x", - "moment": ">= 2.x", - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - }, - "peerDependenciesMeta": { - "date-fns": { - "optional": true - }, - "dayjs": { - "optional": true - }, - "luxon": { - "optional": true - }, - "moment": { - "optional": true - } - } - }, - "node_modules/rc-progress": { - "version": "4.0.0", - "resolved": "https://registry.npmjs.org/rc-progress/-/rc-progress-4.0.0.tgz", - "integrity": "sha512-oofVMMafOCokIUIBnZLNcOZFsABaUw8PPrf1/y0ZBvKZNpOiu5h4AO9vv11Sw0p4Hb3D0yGWuEattcQGtNJ/aw==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.10.1", - "classnames": "^2.2.6", - "rc-util": "^5.16.1" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-rate": { - "version": "2.13.1", - "resolved": "https://registry.npmjs.org/rc-rate/-/rc-rate-2.13.1.tgz", - "integrity": "sha512-QUhQ9ivQ8Gy7mtMZPAjLbxBt5y9GRp65VcUyGUMF3N3fhiftivPHdpuDIaWIMOTEprAjZPC08bls1dQB+I1F2Q==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.10.1", - "classnames": "^2.2.5", - "rc-util": "^5.0.1" - }, - "engines": { - "node": ">=8.x" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-resize-observer": { - "version": "1.4.3", - "resolved": "https://registry.npmjs.org/rc-resize-observer/-/rc-resize-observer-1.4.3.tgz", - "integrity": "sha512-YZLjUbyIWox8E9i9C3Tm7ia+W7euPItNWSPX5sCcQTYbnwDb5uNpnLHQCG1f22oZWUhLw4Mv2tFmeWe68CDQRQ==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.20.7", - "classnames": "^2.2.1", - "rc-util": "^5.44.1", - "resize-observer-polyfill": "^1.5.1" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-segmented": { - "version": "2.7.1", - "resolved": "https://registry.npmjs.org/rc-segmented/-/rc-segmented-2.7.1.tgz", - "integrity": "sha512-izj1Nw/Dw2Vb7EVr+D/E9lUTkBe+kKC+SAFSU9zqr7WV2W5Ktaa9Gc7cB2jTqgk8GROJayltaec+DBlYKc6d+g==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.11.1", - "classnames": "^2.2.1", - "rc-motion": "^2.4.4", - "rc-util": "^5.17.0" - }, - "peerDependencies": { - "react": ">=16.0.0", - "react-dom": ">=16.0.0" - } - }, - "node_modules/rc-select": { - "version": "14.16.8", - "resolved": "https://registry.npmjs.org/rc-select/-/rc-select-14.16.8.tgz", - "integrity": "sha512-NOV5BZa1wZrsdkKaiK7LHRuo5ZjZYMDxPP6/1+09+FB4KoNi8jcG1ZqLE3AVCxEsYMBe65OBx71wFoHRTP3LRg==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.10.1", - "@rc-component/trigger": "^2.1.1", - "classnames": "2.x", - "rc-motion": "^2.0.1", - "rc-overflow": "^1.3.1", - "rc-util": "^5.16.1", - "rc-virtual-list": "^3.5.2" - }, - "engines": { - "node": ">=8.x" - }, - "peerDependencies": { - "react": "*", - "react-dom": "*" - } - }, - "node_modules/rc-slider": { - "version": "11.1.9", - "resolved": "https://registry.npmjs.org/rc-slider/-/rc-slider-11.1.9.tgz", - "integrity": "sha512-h8IknhzSh3FEM9u8ivkskh+Ef4Yo4JRIY2nj7MrH6GQmrwV6mcpJf5/4KgH5JaVI1H3E52yCdpOlVyGZIeph5A==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.10.1", - "classnames": "^2.2.5", - "rc-util": "^5.36.0" - }, - "engines": { - "node": ">=8.x" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-steps": { - "version": "6.0.1", - "resolved": "https://registry.npmjs.org/rc-steps/-/rc-steps-6.0.1.tgz", - "integrity": "sha512-lKHL+Sny0SeHkQKKDJlAjV5oZ8DwCdS2hFhAkIjuQt1/pB81M0cA0ErVFdHq9+jmPmFw1vJB2F5NBzFXLJxV+g==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.16.7", - "classnames": "^2.2.3", - "rc-util": "^5.16.1" - }, - "engines": { - "node": ">=8.x" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-switch": { - "version": "4.1.0", - "resolved": "https://registry.npmjs.org/rc-switch/-/rc-switch-4.1.0.tgz", - "integrity": "sha512-TI8ufP2Az9oEbvyCeVE4+90PDSljGyuwix3fV58p7HV2o4wBnVToEyomJRVyTaZeqNPAp+vqeo4Wnj5u0ZZQBg==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.21.0", - "classnames": "^2.2.1", - "rc-util": "^5.30.0" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-table": { - "version": "7.54.0", - "resolved": "https://registry.npmjs.org/rc-table/-/rc-table-7.54.0.tgz", - "integrity": "sha512-/wDTkki6wBTjwylwAGjpLKYklKo9YgjZwAU77+7ME5mBoS32Q4nAwoqhA2lSge6fobLW3Tap6uc5xfwaL2p0Sw==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.10.1", - "@rc-component/context": "^1.4.0", - "classnames": "^2.2.5", - "rc-resize-observer": "^1.1.0", - "rc-util": "^5.44.3", - "rc-virtual-list": "^3.14.2" - }, - "engines": { - "node": ">=8.x" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-tabs": { - "version": "15.7.0", - "resolved": "https://registry.npmjs.org/rc-tabs/-/rc-tabs-15.7.0.tgz", - "integrity": "sha512-ZepiE+6fmozYdWf/9gVp7k56PKHB1YYoDsKeQA1CBlJ/POIhjkcYiv0AGP0w2Jhzftd3AVvZP/K+V+Lpi2ankA==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.11.2", - "classnames": "2.x", - "rc-dropdown": "~4.2.0", - "rc-menu": "~9.16.0", - "rc-motion": "^2.6.2", - "rc-resize-observer": "^1.0.0", - "rc-util": "^5.34.1" - }, - "engines": { - "node": ">=8.x" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-textarea": { - "version": "1.10.2", - "resolved": "https://registry.npmjs.org/rc-textarea/-/rc-textarea-1.10.2.tgz", - "integrity": "sha512-HfaeXiaSlpiSp0I/pvWpecFEHpVysZ9tpDLNkxQbMvMz6gsr7aVZ7FpWP9kt4t7DB+jJXesYS0us1uPZnlRnwQ==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.10.1", - "classnames": "^2.2.1", - "rc-input": "~1.8.0", - "rc-resize-observer": "^1.0.0", - "rc-util": "^5.27.0" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-tooltip": { - "version": "6.4.0", - "resolved": "https://registry.npmjs.org/rc-tooltip/-/rc-tooltip-6.4.0.tgz", - "integrity": "sha512-kqyivim5cp8I5RkHmpsp1Nn/Wk+1oeloMv9c7LXNgDxUpGm+RbXJGL+OPvDlcRnx9DBeOe4wyOIl4OKUERyH1g==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.11.2", - "@rc-component/trigger": "^2.0.0", - "classnames": "^2.3.1", - "rc-util": "^5.44.3" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-tree": { - "version": "5.13.1", - "resolved": "https://registry.npmjs.org/rc-tree/-/rc-tree-5.13.1.tgz", - "integrity": "sha512-FNhIefhftobCdUJshO7M8uZTA9F4OPGVXqGfZkkD/5soDeOhwO06T/aKTrg0WD8gRg/pyfq+ql3aMymLHCTC4A==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.10.1", - "classnames": "2.x", - "rc-motion": "^2.0.1", - "rc-util": "^5.16.1", - "rc-virtual-list": "^3.5.1" - }, - "engines": { - "node": ">=10.x" - }, - "peerDependencies": { - "react": "*", - "react-dom": "*" - } - }, - "node_modules/rc-tree-select": { - "version": "5.27.0", - "resolved": "https://registry.npmjs.org/rc-tree-select/-/rc-tree-select-5.27.0.tgz", - "integrity": "sha512-2qTBTzwIT7LRI1o7zLyrCzmo5tQanmyGbSaGTIf7sYimCklAToVVfpMC6OAldSKolcnjorBYPNSKQqJmN3TCww==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.25.7", - "classnames": "2.x", - "rc-select": "~14.16.2", - "rc-tree": "~5.13.0", - "rc-util": "^5.43.0" - }, - "peerDependencies": { - "react": "*", - "react-dom": "*" - } - }, - "node_modules/rc-upload": { - "version": "4.11.0", - "resolved": "https://registry.npmjs.org/rc-upload/-/rc-upload-4.11.0.tgz", - "integrity": "sha512-ZUyT//2JAehfHzjWowqROcwYJKnZkIUGWaTE/VogVrepSl7AFNbQf4+zGfX4zl9Vrj/Jm8scLO0R6UlPDKK4wA==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.18.3", - "classnames": "^2.2.5", - "rc-util": "^5.2.0" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-util": { - "version": "5.44.4", - "resolved": "https://registry.npmjs.org/rc-util/-/rc-util-5.44.4.tgz", - "integrity": "sha512-resueRJzmHG9Q6rI/DfK6Kdv9/Lfls05vzMs1Sk3M2P+3cJa+MakaZyWY8IPfehVuhPJFKrIY1IK4GqbiaiY5w==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.18.3", - "react-is": "^18.2.0" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, - "node_modules/rc-util/node_modules/react-is": { - "version": "18.3.1", - "resolved": "https://registry.npmjs.org/react-is/-/react-is-18.3.1.tgz", - "integrity": "sha512-/LLMVyas0ljjAtoYiPqYiL8VWXzUUdThrmU5+n20DZv+a+ClRoevUzw5JxU+Ieh5/c87ytoTBV9G1FiKfNJdmg==", - "license": "MIT" - }, - "node_modules/rc-virtual-list": { - "version": "3.19.2", - "resolved": "https://registry.npmjs.org/rc-virtual-list/-/rc-virtual-list-3.19.2.tgz", - "integrity": "sha512-Ys6NcjwGkuwkeaWBDqfI3xWuZ7rDiQXlH1o2zLfFzATfEgXcqpk8CkgMfbJD81McqjcJVez25a3kPxCR807evA==", - "license": "MIT", - "dependencies": { - "@babel/runtime": "^7.20.0", - "classnames": "^2.2.6", - "rc-resize-observer": "^1.0.0", - "rc-util": "^5.36.0" - }, - "engines": { - "node": ">=8.x" - }, - "peerDependencies": { - "react": ">=16.9.0", - "react-dom": ">=16.9.0" - } - }, "node_modules/react": { "version": "19.2.8", "resolved": "https://registry.npmjs.org/react/-/react-19.2.8.tgz", @@ -11986,12 +11030,6 @@ "integrity": "sha512-AgZ3UOZm3YndfrJ4OYjgrT7bmCm/1iqkjvEfH/oYjzh6PD2qw4QuT3jjnXIrpdt4MTpMXclMT3lXbmRY+XRakw==", "license": "MIT" }, - "node_modules/resize-observer-polyfill": { - "version": "1.5.1", - "resolved": "https://registry.npmjs.org/resize-observer-polyfill/-/resize-observer-polyfill-1.5.1.tgz", - "integrity": "sha512-LwZrotdHOo12nQuZlHEmtuXdqGoOD0OhaxopaNFxWzInpEgaLWoVuAMbTzixuosCx2nEG58ngzW3vxdWoxIgdg==", - "license": "MIT" - }, "node_modules/resolve": { "version": "2.0.0-next.6", "resolved": "https://registry.npmjs.org/resolve/-/resolve-2.0.0-next.6.tgz", @@ -12190,15 +11228,6 @@ "integrity": "sha512-eNv+WrVbKu1f3vbYJT/xtiF5syA5HPIMtf9IgY/nKg0sWqzAUEvqY/xm7OcZc/qafLx/iO9FgOmeSAp4v5ti/Q==", "license": "MIT" }, - "node_modules/scroll-into-view-if-needed": { - "version": "3.1.0", - "resolved": "https://registry.npmjs.org/scroll-into-view-if-needed/-/scroll-into-view-if-needed-3.1.0.tgz", - "integrity": "sha512-49oNpRjWRvnU8NyGVmUaYG4jtTkNonFZI86MmGRDqBphEK2EXT9gdEUoQPZhuBM8yWHxCWbobltqYO5M4XrUvQ==", - "license": "MIT", - "dependencies": { - "compute-scroll-into-view": "^3.0.2" - } - }, "node_modules/semver": { "version": "7.8.5", "resolved": "https://registry.npmjs.org/semver/-/semver-7.8.5.tgz", @@ -12515,12 +11544,6 @@ "node": ">= 0.4" } }, - "node_modules/string-convert": { - "version": "0.2.1", - "resolved": "https://registry.npmjs.org/string-convert/-/string-convert-0.2.1.tgz", - "integrity": "sha512-u/1tdPl4yQnPBjnVrmdLo9gtuLvELKsAoRapekWggdiQNvvvum+jYF329d84NAa660KQw7pB2n36KrIKVoXa3A==", - "license": "MIT" - }, "node_modules/string.prototype.includes": { "version": "2.0.1", "resolved": "https://registry.npmjs.org/string.prototype.includes/-/string.prototype.includes-2.0.1.tgz", @@ -12745,12 +11768,6 @@ } } }, - "node_modules/stylis": { - "version": "4.4.0", - "resolved": "https://registry.npmjs.org/stylis/-/stylis-4.4.0.tgz", - "integrity": "sha512-5Z9ZpRzfuH6l/UAvCPAPUo3665Nk2wLaZU3x+TLHKVzIz33+sbJqbtrYoC3KD4/uVOr2Zp+L0LySezP9OHV9yA==", - "license": "MIT" - }, "node_modules/supports-color": { "version": "7.2.0", "resolved": "https://registry.npmjs.org/supports-color/-/supports-color-7.2.0.tgz", @@ -12829,15 +11846,6 @@ "node": ">=18" } }, - "node_modules/throttle-debounce": { - "version": "5.0.2", - "resolved": "https://registry.npmjs.org/throttle-debounce/-/throttle-debounce-5.0.2.tgz", - "integrity": "sha512-B71/4oyj61iNH0KeCamLuE2rmKuTO5byTOSVwECM5FA7TiAiAW+UqTKZ9ERueC4qvgSttUhdmq1mXC3kJqGX7A==", - "license": "MIT", - "engines": { - "node": ">=12.22" - } - }, "node_modules/tiny-invariant": { "version": "1.3.3", "resolved": "https://registry.npmjs.org/tiny-invariant/-/tiny-invariant-1.3.3.tgz", diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index 3f9c29d2afe..6b2d4e6106b 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -24,7 +24,6 @@ "gen:api": "node scripts/gen-api-types.mjs" }, "dependencies": { - "@ant-design/cssinjs": "1.24.0", "@anthropic-ai/sdk": "0.92.0", "@base-ui/react": "^1.6.0", "@headlessui/tailwindcss": "0.2.2", @@ -34,7 +33,6 @@ "@tanstack/react-query": "5.100.7", "@tanstack/react-table": "8.21.3", "@types/papaparse": "5.5.2", - "antd": "5.29.3", "cva": "1.0.0-beta.4", "date-fns": "^4.4.0", "dayjs": "1.11.19", diff --git a/ui/litellm-dashboard/scripts/eslint-rules/index.mjs b/ui/litellm-dashboard/scripts/eslint-rules/index.mjs index 52db03a2c10..9e9f901a6df 100644 --- a/ui/litellm-dashboard/scripts/eslint-rules/index.mjs +++ b/ui/litellm-dashboard/scripts/eslint-rules/index.mjs @@ -2,7 +2,6 @@ import noLargeInlineObjectArg from "./no-large-inline-object-arg.mjs"; import noLongConditionChain from "./no-long-condition-chain.mjs"; import noComplexJsxArrow from "./no-complex-jsx-arrow.mjs"; import filenamePascalCase from "./filename-pascal-case.mjs"; -import noAntdClassSelectors from "./no-antd-class-selectors.mjs"; const plugin = { rules: { @@ -10,7 +9,6 @@ const plugin = { "no-long-condition-chain": noLongConditionChain, "no-complex-jsx-arrow": noComplexJsxArrow, "filename-pascal-case": filenamePascalCase, - "no-antd-class-selectors": noAntdClassSelectors, }, }; diff --git a/ui/litellm-dashboard/scripts/eslint-rules/no-antd-class-selectors.mjs b/ui/litellm-dashboard/scripts/eslint-rules/no-antd-class-selectors.mjs deleted file mode 100644 index c95bdb0467f..00000000000 --- a/ui/litellm-dashboard/scripts/eslint-rules/no-antd-class-selectors.mjs +++ /dev/null @@ -1,45 +0,0 @@ -const SELECTOR_REFERENCE = /\.(?:ant|anticon)-[a-z0-9-]+/; -const BARE_CLASS_REFERENCE = /^(?:ant|anticon)-[a-z0-9-]+$/; -const CLASS_ASSERTION_CALLEES = new Set(["toHaveClass", "contains", "toContain"]); - -const isClassAssertionArgument = (node) => { - const call = node.parent; - if (call?.type !== "CallExpression" || !call.arguments.includes(node)) return false; - const callee = call.callee; - return callee?.type === "MemberExpression" && CLASS_ASSERTION_CALLEES.has(callee.property?.name); -}; - -const rule = { - meta: { - type: "problem", - docs: { - description: - "Disallow locating or asserting on antd's internal CSS classes in tests; query by role, label or text instead.", - }, - schema: [], - messages: { - antdClass: - 'Test depends on antd internal class "{{value}}". Query by role, label or text (getByLabelText, getByRole("combobox"), getByTitle) so the test survives the shadcn migration.', - }, - }, - create(context) { - const report = (node, value) => { - if (typeof value !== "string") return; - const matches = - SELECTOR_REFERENCE.test(value) || (BARE_CLASS_REFERENCE.test(value) && isClassAssertionArgument(node)); - if (!matches) return; - context.report({ node, messageId: "antdClass", data: { value } }); - }; - - return { - Literal(node) { - report(node, node.value); - }, - TemplateElement(node) { - report(node, node.value.cooked); - }, - }; - }, -}; - -export default rule; 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 146e0fa86e6..eb880068d12 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx @@ -36,7 +36,7 @@ import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Textarea } from "@/components/ui/textarea"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; -import { isAntdUrl } from "@/lib/forms/antdUrl"; +import { isValidUrl } from "@/lib/forms/urlValidation"; import { useZodForm } from "@/lib/forms/useZodForm"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; import { Button } from "@/components/ui/button"; @@ -51,7 +51,7 @@ const submitGuardrailSchema = z.object({ team_id: z.string().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(isAntdUrl, "Must be a valid URL"), + api_base: z.string().min(1, "Enter the API base URL").refine(isValidUrl, "Must be a valid URL"), extra_litellm_params: z.string().superRefine((value, ctx) => { if (!value) return; try { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/CustomPatternModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/CustomPatternModal.tsx index 20bdddd3b74..64441c90e5a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/CustomPatternModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/CustomPatternModal.tsx @@ -4,7 +4,7 @@ import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from " import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { ACTION_ITEMS } from "./action_options"; -import { ABOVE_ANTD_MODAL } from "./dialog_layering"; +import { NESTED_DIALOG_LAYER } from "./dialog_layering"; interface CustomPatternModalProps { visible: boolean; @@ -31,7 +31,7 @@ const CustomPatternModal: React.FC = ({ }) => { return ( !open && onCancel()}> - + Add custom regex pattern diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/KeywordModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/KeywordModal.tsx index 45627a2f101..177c0d2fac6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/KeywordModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/KeywordModal.tsx @@ -5,7 +5,7 @@ import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Textarea } from "@/components/ui/textarea"; import { ACTION_ITEMS } from "./action_options"; -import { ABOVE_ANTD_MODAL } from "./dialog_layering"; +import { NESTED_DIALOG_LAYER } from "./dialog_layering"; interface KeywordModalProps { visible: boolean; @@ -32,7 +32,7 @@ const KeywordModal: React.FC = ({ }) => { return ( !open && onCancel()}> - + Add blocked keyword diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.tsx index 2d98b67cb4a..e703711a03a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.tsx @@ -14,7 +14,7 @@ import { import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { ACTION_ITEMS } from "./action_options"; -import { ABOVE_ANTD_MODAL } from "./dialog_layering"; +import { NESTED_DIALOG_LAYER } from "./dialog_layering"; interface PrebuiltPattern { name: string; @@ -66,7 +66,7 @@ const PatternModal: React.FC = ({ return ( !open && onCancel()}> - + Add prebuilt pattern diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/dialog_layering.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/dialog_layering.ts index f31ab788beb..0e29ffb6250 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/dialog_layering.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/dialog_layering.ts @@ -1 +1 @@ -export const ABOVE_ANTD_MODAL = "z-[1100]"; +export const NESTED_DIALOG_LAYER = "z-[1100]"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/AwsSigV4Fields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/AwsSigV4Fields.tsx index 77a77596bed..6249614749d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/AwsSigV4Fields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/AwsSigV4Fields.tsx @@ -3,7 +3,7 @@ import React from "react"; import { SimpleTooltip } from "@/components/ui/tooltip"; import { MountedFormField } from "@/components/common_components/MountedFormField"; -import { antdRequired } from "@/components/common_components/antdFormRules"; +import { requiredRule } from "@/components/common_components/formRules"; import { PasswordInput } from "@/components/shared/PasswordInput"; import { Input } from "@/components/ui/input"; import { requiredWhenSiblingSet, textControl } from "./mcpFieldRules"; @@ -39,7 +39,7 @@ const AwsSigV4Fields: React.FC = () => ( label={} name={["credentials", "aws_region_name"]} required - rules={{ validate: { required: antdRequired("AWS region is required for SigV4 auth") } }} + rules={{ validate: { required: requiredRule("AWS region is required for SigV4 auth") } }} > {(control) => } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.tsx index 75680d2e03f..49694d90ae4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.tsx @@ -63,7 +63,7 @@ import { useMountRegistry, type MountedFormValues, } from "@/components/common_components/MountedFormField"; -import { antdRequired, antdRules } from "@/components/common_components/antdFormRules"; +import { requiredRule, validatorRules } from "@/components/common_components/formRules"; import { allFieldsValue, mountedPaths, resetFields, setFieldsValue, singleBranchChange } from "./mcpFormStore"; import { numberControl, notOnlyWhitespace, selectControl, selectTriggerControl, textControl } from "./mcpFieldRules"; import mcpLogo from "../../../../../public/assets/logos/mcp_logo.png"; @@ -657,7 +657,7 @@ const CreateMCPServer: React.FC = ({ } name="server_name" - rules={{ validate: antdRules({ validator: (_, value) => validateMCPServerName(value) }) }} + rules={{ validate: validatorRules({ validator: (_, value) => validateMCPServerName(value) }) }} > {(control) => ( = ({ } name="alias" - rules={{ validate: antdRules({ validator: (_, value) => validateMCPServerName(value) }) }} + rules={{ validate: validatorRules({ validator: (_, value) => validateMCPServerName(value) }) }} > {(control) => ( = ({ label={Transport Type} name="transport" required - rules={{ validate: { required: antdRequired("Please select a transport type") } }} + rules={{ validate: { required: requiredRule("Please select a transport type") } }} > {(control) => ( (control)} items={AUTH_TYPE_ITEMS}> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/EnvVarsSection.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/EnvVarsSection.tsx index f6f724b666d..b2822761e29 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/EnvVarsSection.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/EnvVarsSection.tsx @@ -12,7 +12,7 @@ import { useMountedName, type MountedFormValues, } from "@/components/common_components/MountedFormField"; -import { antdRequired } from "@/components/common_components/antdFormRules"; +import { requiredRule } from "@/components/common_components/formRules"; import { matchesPattern, selectControl, selectTriggerControl, textControl } from "./mcpFieldRules"; import { listControl } from "./mcpFormStore"; @@ -80,7 +80,7 @@ const EnvVarsSection: React.FC = () => { className="mb-0 flex-1" rules={{ validate: { - required: antdRequired("Variable name is required"), + required: requiredRule("Variable name is required"), pattern: matchesPattern( VARIABLE_NAME_PATTERN, "Use letters, digits, underscores; cannot start with a digit.", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/IdJagFormFields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/IdJagFormFields.tsx index d4760a65e6c..7745e1d6586 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/IdJagFormFields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/IdJagFormFields.tsx @@ -3,7 +3,7 @@ import React from "react"; import { SimpleTooltip } from "@/components/ui/tooltip"; import { MountedFormField } from "@/components/common_components/MountedFormField"; -import { antdRequired } from "@/components/common_components/antdFormRules"; +import { requiredRule } from "@/components/common_components/formRules"; import { MultiSelect } from "@/components/shared/MultiSelect"; import { PasswordInput } from "@/components/shared/PasswordInput"; import { Input } from "@/components/ui/input"; @@ -30,7 +30,7 @@ const PRIVATE_KEY_PATH = ["credentials", "client_private_key"] as const; const IdJagFormFields: React.FC = ({ isEditing = false }) => { const placeholderSuffix = isEditing ? " (leave blank to keep existing)" : ""; const requiredWhenCreating = (message: string) => - isEditing ? undefined : { validate: { required: antdRequired(message) } }; + isEditing ? undefined : { validate: { required: requiredRule(message) } }; return ( <> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.tsx index c1da3f105db..14c82fbb0d9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.tsx @@ -15,7 +15,7 @@ import { type MountedFieldControlProps, type MountedFormValues, } from "@/components/common_components/MountedFormField"; -import { antdRequired } from "@/components/common_components/antdFormRules"; +import { requiredRule } from "@/components/common_components/formRules"; import { Field, FieldLabel } from "@/components/shared/form/field"; import { invertedSwitchControl, switchControl, tagsControl, textControl } from "./mcpFieldRules"; import { listControl } from "./mcpFormStore"; @@ -64,7 +64,7 @@ const StaticHeadersFieldArray: React.FC = () => { {(headerControl) => ( { {(valueControl) => ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.tsx index e88c78271d3..60830281293 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.tsx @@ -9,7 +9,7 @@ import { Input } from "@/components/ui/input"; import { Textarea } from "@/components/ui/textarea"; import { OAUTH_FLOW } from "@/components/mcp_tools/types"; import { MountedFormField } from "@/components/common_components/MountedFormField"; -import { antdRequired } from "@/components/common_components/antdFormRules"; +import { requiredRule } from "@/components/common_components/formRules"; import TokenEndpointAuthMethodField from "./TokenEndpointAuthMethodField"; import { numberControl, @@ -79,7 +79,7 @@ const OAuthFormFields: React.FC = ({ }) => { const placeholderSuffix = isEditing ? " (leave blank to keep existing)" : ""; const requiredWhenCreating = (message: string) => - isEditing ? undefined : { validate: { required: antdRequired(message) } }; + isEditing ? undefined : { validate: { required: requiredRule(message) } }; return ( <> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OpenAPIFormSection.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OpenAPIFormSection.tsx index 8907595d92a..56dce02ba6f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OpenAPIFormSection.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OpenAPIFormSection.tsx @@ -4,7 +4,7 @@ import { SimpleTooltip } from "@/components/ui/tooltip"; import { Input } from "@/components/ui/input"; import { AUTH_TYPE, OAUTH_FLOW } from "@/components/mcp_tools/types"; import { MountedFormField } from "@/components/common_components/MountedFormField"; -import { antdRequired } from "@/components/common_components/antdFormRules"; +import { requiredRule } from "@/components/common_components/formRules"; import OpenAPIQuickPicker, { OpenAPIRegistryEntry, OpenAPIKeyTool } from "./OpenAPIQuickPicker"; import { McpForm, resetFields, setFieldsValue } from "./mcpFormStore"; import { textControl } from "./mcpFieldRules"; @@ -76,7 +76,7 @@ const OpenAPIFormSection: React.FC = ({ } name="spec_path" required - rules={{ validate: { required: antdRequired("Please enter an OpenAPI spec URL") } }} + rules={{ validate: { required: requiredRule("Please enter an OpenAPI spec URL") } }} > {(control) => ( = ({ isVisible, requ required={required} rules={{ validate: { - ...(required ? { required: antdRequired("Please enter stdio configuration") } : {}), + ...(required ? { required: requiredRule("Please enter stdio configuration") } : {}), json: parsesAsJson("Please enter valid JSON"), }, }} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/TokenExchangeFormFields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/TokenExchangeFormFields.tsx index 5aaa013aba4..cb7a7d57f7c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/TokenExchangeFormFields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/TokenExchangeFormFields.tsx @@ -6,7 +6,7 @@ import { SimpleTooltip } from "@/components/ui/tooltip"; import { useWatch } from "react-hook-form"; import { MountedFormField } from "@/components/common_components/MountedFormField"; -import { antdRequired } from "@/components/common_components/antdFormRules"; +import { requiredRule } from "@/components/common_components/formRules"; import { PasswordInput } from "@/components/shared/PasswordInput"; import { Input } from "@/components/ui/input"; import { selectControl, selectTriggerControl, tagsControl, textControl } from "./mcpFieldRules"; @@ -35,7 +35,7 @@ const TokenExchangeFormFields: React.FC = ({ isEdi const placeholderSuffix = isEditing ? " (leave blank to keep existing)" : ""; const isEntraObo = useWatch({ name: "token_exchange_profile" }) === "entra_obo"; const requiredWhenCreating = (message: string) => - isEditing ? undefined : { validate: { required: antdRequired(message) } }; + isEditing ? undefined : { validate: { required: requiredRule(message) } }; return ( <> @@ -170,7 +170,7 @@ const TokenExchangeFormFields: React.FC = ({ isEdi isEntraObo ? { validate: { - required: antdRequired("Microsoft Entra OBO requires a scope, e.g. api:///.default"), + required: requiredRule("Microsoft Entra OBO requires a scope, e.g. api:///.default"), }, } : undefined diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx index b42d9ec5f40..4fbdd37240c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx @@ -60,7 +60,7 @@ import { useMountRegistry, type MountedFormValues, } from "@/components/common_components/MountedFormField"; -import { antdRequired, antdRules } from "@/components/common_components/antdFormRules"; +import { requiredRule, validatorRules } from "@/components/common_components/formRules"; import { allFieldsValue, mountedPaths, @@ -774,7 +774,7 @@ const MCPServerEdit: React.FC = ({ validateMCPServerName(value) }) }} + rules={{ validate: validatorRules({ validator: (_, value) => validateMCPServerName(value) }) }} > {(control) => ( = ({ validateMCPServerName(value) }) }} + rules={{ validate: validatorRules({ validator: (_, value) => validateMCPServerName(value) }) }} > {(control) => ( = ({ label="Transport Type" name="transport" required - rules={{ validate: { required: antdRequired("Transport Type is required") } }} + rules={{ validate: { required: requiredRule("Transport Type is required") } }} > {(control) => ( = ({ label="Authentication" name="auth_type" required - rules={{ validate: { required: antdRequired("Authentication is required") } }} + rules={{ validate: { required: requiredRule("Authentication is required") } }} > {(control) => ( { - icon: React.ComponentType; -} - -interface LegacyPopoverProps { - children: React.ReactElement>; - content: React.ReactNode; - onOpenChange?: (open: boolean) => void; - title?: React.ReactNode; -} - -interface LegacyTooltipProps { - children: React.ReactElement>; - title?: React.ReactNode; -} - vi.mock("@heroicons/react/outline", () => ({ EyeIcon: function EyeIcon() { return null; }, })); -vi.mock("antd", async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - Popover: ({ children, content, onOpenChange, title }: LegacyPopoverProps) => { - const [open, setOpen] = React.useState(false); - return ( - <> - {React.cloneElement(children, { - onClick: () => { - const nextOpen = !open; - setOpen(nextOpen); - onOpenChange?.(nextOpen); - }, - })} - {open && ( -
-

{title}

- {content} -
- )} - - ); - }, - Tooltip: ({ children, title }: LegacyTooltipProps) => - React.cloneElement(children, { "aria-label": typeof title === "string" ? title : undefined }), - Spin: () =>