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/542] 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 f08742c6f4ef09ab22ac07f96a691bab6c864d25 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 14 Jul 2026 17:29:48 -0700 Subject: [PATCH 002/542] fix(interactions): track cost and spend for Google Interactions API requests --- litellm/cost_calculator.py | 7 + litellm/litellm_core_utils/litellm_logging.py | 54 +++++-- .../usage_object_transformation.py | 109 ++++++++++++- .../test_usage_object_transformation.py | 121 ++++++++++++++ .../test_litellm_logging.py | 147 ++++++++++++++++++ tests/test_litellm/test_cost_calculator.py | 33 ++++ 6 files changed, 460 insertions(+), 11 deletions(-) create mode 100644 tests/test_litellm/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index a40a8e1389c..3c541605caa 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -19,6 +19,7 @@ from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import ( + InteractionsUsageObjectTransformation, TranscriptionUsageObjectTransformation, ) from litellm.litellm_core_utils.llm_cost_calc.utils import ( @@ -899,6 +900,8 @@ def _get_usage_object( usage_obj, ) ) + elif isinstance(usage_obj, dict) and InteractionsUsageObjectTransformation.is_interactions_usage_object(usage_obj): + return InteractionsUsageObjectTransformation.transform_interactions_usage_object(usage_obj) elif isinstance(usage_obj, dict): return Usage(**usage_obj) elif isinstance(usage_obj, BaseModel): @@ -1267,6 +1270,10 @@ def completion_cost( ) if tr_usage is not None: _usage = tr_usage.model_dump() + elif InteractionsUsageObjectTransformation.is_interactions_usage_object(_usage): + _usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object( + _usage + ).model_dump() else: _usage = _usage diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 461ab62b815..0f6e0ccf9c1 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -68,6 +68,9 @@ from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) +from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import ( + InteractionsUsageObjectTransformation, +) from litellm.litellm_core_utils.logging_utils import truncate_base64_in_messages from litellm.litellm_core_utils.model_param_helper import ModelParamHelper from litellm.litellm_core_utils.redact_messages import ( @@ -93,6 +96,10 @@ from litellm.types.llms.openai import ( ResponseIncompleteEvent, ResponsesAPIResponse, ) +from litellm.types.interactions import ( + InteractionsAPIResponse, + InteractionsAPIStreamingResponse, +) from litellm.types.mcp import MCPPostCallResponseObject from litellm.types.prompts.init_prompts import PromptSpec from litellm.types.rerank import RerankResponse @@ -1888,6 +1895,7 @@ class Logging(LiteLLMLoggingBaseClass): or isinstance(logging_result, FineTuningJob) or isinstance(logging_result, LiteLLMBatch) or isinstance(logging_result, ResponsesAPIResponse) + or isinstance(logging_result, InteractionsAPIResponse) or isinstance(logging_result, OpenAIFileObject) or isinstance(logging_result, LiteLLMRealtimeStreamLoggingObject) or isinstance(logging_result, OpenAIModerationResponse) @@ -1973,7 +1981,7 @@ class Logging(LiteLLMLoggingBaseClass): try: ## BUILD COMPLETE STREAMED RESPONSE complete_streaming_response: Optional[ - Union[ModelResponse, TextCompletionResponse, ResponsesAPIResponse] + Union[ModelResponse, TextCompletionResponse, ResponsesAPIResponse, InteractionsAPIResponse] ] = None if "complete_streaming_response" in self.model_call_details: return # break out of this. @@ -2428,14 +2436,14 @@ class Logging(LiteLLMLoggingBaseClass): ## BUILD COMPLETE STREAMED RESPONSE if "async_complete_streaming_response" in self.model_call_details: return # break out of this. - complete_streaming_response: Optional[Union[ModelResponse, TextCompletionResponse, ResponsesAPIResponse]] = ( - self._get_assembled_streaming_response( - result=result, - start_time=start_time, - end_time=end_time, - is_async=True, - streaming_chunks=self.streaming_chunks, - ) + complete_streaming_response: ( + Union[ModelResponse, TextCompletionResponse, ResponsesAPIResponse, InteractionsAPIResponse] | None + ) = self._get_assembled_streaming_response( + result=result, + start_time=start_time, + end_time=end_time, + is_async=True, + streaming_chunks=self.streaming_chunks, ) if complete_streaming_response is not None: @@ -3153,7 +3161,7 @@ class Logging(LiteLLMLoggingBaseClass): end_time: datetime.datetime, is_async: bool, streaming_chunks: List[Any], - ) -> Optional[Union[ModelResponse, TextCompletionResponse, ResponsesAPIResponse]]: + ) -> Union[ModelResponse, TextCompletionResponse, ResponsesAPIResponse, InteractionsAPIResponse] | None: if self.stream is not True: return None if isinstance(result, ModelResponse): @@ -3180,9 +3188,31 @@ class Logging(LiteLLMLoggingBaseClass): ), ) return result.response + elif isinstance(result, InteractionsAPIStreamingResponse): + return self._assemble_completed_interaction_response(result) else: return None + @staticmethod + def _assemble_completed_interaction_response( + result: InteractionsAPIStreamingResponse, + ) -> InteractionsAPIResponse | None: + """ + The Interactions API streaming iterator hands the terminal event to the + success handlers: the new schema (Api-Revision: 2026-05-20) emits + ``interaction.completed`` carrying the full interaction object, the + legacy schema (2026-05-07) emits a chunk with ``status="completed"`` + and usage on the chunk itself. Build the equivalent non-streaming + response so cost calculation and spend tracking see one shape. + """ + if result.event_type == "interaction.completed" and result.interaction is not None: + return InteractionsAPIResponse(**result.interaction) + if result.status == "completed": + return InteractionsAPIResponse( + **result.model_dump(exclude={"event_type", "delta", "index", "step", "interaction_id", "interaction"}) + ) + return None + def _handle_anthropic_messages_response_logging(self, result: Any) -> ModelResponse: """ Handles logging for Anthropic messages responses. @@ -4687,6 +4717,8 @@ class StandardLoggingPayloadSetup: elif isinstance(usage, dict): if ResponseAPILoggingUtils._is_response_api_usage(usage): return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + if InteractionsUsageObjectTransformation.is_interactions_usage_object(usage): + return InteractionsUsageObjectTransformation.transform_interactions_usage_object(usage) return Usage(**usage) raise ValueError(f"usage is required, got={usage} of type {type(usage)}") @@ -4713,6 +4745,8 @@ class StandardLoggingPayloadSetup: if isinstance(_raw, dict): if ResponseAPILoggingUtils._is_response_api_usage(_raw): return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(_raw).model_dump() + if InteractionsUsageObjectTransformation.is_interactions_usage_object(_raw): + return InteractionsUsageObjectTransformation.transform_interactions_usage_object(_raw).model_dump() return _raw if isinstance(_raw, Usage): return _raw.model_dump() diff --git a/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py b/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py index 1c6adbec174..3f591180b77 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py +++ b/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py @@ -1,6 +1,7 @@ -from typing import Any, Optional, Union +from typing import Any, Mapping, Optional, Sequence, Union from litellm.types.utils import ( + CompletionTokensDetailsWrapper, PromptTokensDetailsWrapper, TranscriptionUsageDurationObject, TranscriptionUsageTokensObject, @@ -34,3 +35,109 @@ class TranscriptionUsageObjectTransformation: ), ) return None + + +_INTERACTIONS_MODALITY_FIELDS: Mapping[str, str] = { + "text": "text_tokens", + "audio": "audio_tokens", + "image": "image_tokens", + "video": "video_tokens", + "document": "text_tokens", +} + + +def _modality_field(entry: Mapping[str, Any]) -> str | None: + return _INTERACTIONS_MODALITY_FIELDS.get(str(entry.get("modality", "")).lower()) + + +def _token_count(value: Any) -> int: + return value if isinstance(value, int) else 0 + + +def _modality_token_sums(entries: Sequence[Mapping[str, Any]]) -> Mapping[str, int]: + fields = {field for entry in entries if (field := _modality_field(entry)) is not None} + return { + field: sum(_token_count(entry.get("tokens")) for entry in entries if _modality_field(entry) == field) + for field in fields + } + + +def _subtract_cached_from_input( + input_sums: Mapping[str, int], + cached_sums: Mapping[str, int], + total_cached_tokens: int, +) -> Mapping[str, int]: + if cached_sums: + return {field: max(0, tokens - cached_sums.get(field, 0)) for field, tokens in input_sums.items()} + if total_cached_tokens and "text_tokens" in input_sums: + return { + **input_sums, + "text_tokens": max(0, input_sums["text_tokens"] - total_cached_tokens), + } + return input_sums + + +class InteractionsUsageObjectTransformation: + """ + Maps the Google Interactions API usage block (total_input_tokens, + output_tokens_by_modality, ...) into LiteLLM's chat-format ``Usage`` so the + generic cost calculator and spend tracking can bill it. + """ + + @staticmethod + def is_interactions_usage_object(usage_object: Any) -> bool: + if not isinstance(usage_object, dict): + return False + if "prompt_tokens" in usage_object or "input_tokens" in usage_object: + return False + return "total_input_tokens" in usage_object or "total_output_tokens" in usage_object + + @staticmethod + def transform_interactions_usage_object(usage_object: Mapping[str, Any]) -> Usage: + input_entries = tuple(usage_object.get("input_tokens_by_modality") or ()) + tuple( + usage_object.get("tool_use_tokens_by_modality") or () + ) + cached_sums = _modality_token_sums(tuple(usage_object.get("cached_tokens_by_modality") or ())) + output_sums = _modality_token_sums(tuple(usage_object.get("output_tokens_by_modality") or ())) + + total_cached_tokens = _token_count(usage_object.get("total_cached_tokens")) + input_sums = _subtract_cached_from_input( + input_sums=_modality_token_sums(input_entries), + cached_sums=cached_sums, + total_cached_tokens=total_cached_tokens, + ) + + reasoning_tokens = _token_count(usage_object.get("total_reasoning_tokens")) or _token_count( + usage_object.get("total_thought_tokens") + ) + prompt_tokens = _token_count(usage_object.get("total_input_tokens")) + _token_count( + usage_object.get("total_tool_use_tokens") + ) + completion_tokens = _token_count(usage_object.get("total_output_tokens")) + reasoning_tokens + total_tokens = _token_count(usage_object.get("total_tokens")) or (prompt_tokens + completion_tokens) + + prompt_tokens_details = ( + PromptTokensDetailsWrapper( + cached_tokens=total_cached_tokens or None, + **input_sums, + ) + if input_sums or total_cached_tokens + else None + ) + completion_tokens_details = ( + CompletionTokensDetailsWrapper( + reasoning_tokens=reasoning_tokens or None, + **output_sums, + ) + if output_sums or reasoning_tokens + else None + ) + + return Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=total_tokens, + prompt_tokens_details=prompt_tokens_details, + completion_tokens_details=completion_tokens_details, + cache_read_input_tokens=total_cached_tokens or None, + ) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py new file mode 100644 index 00000000000..01241cf260d --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py @@ -0,0 +1,121 @@ +from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import ( + InteractionsUsageObjectTransformation, +) +from litellm.types.utils import Usage + +OMNI_VIDEO_USAGE = { + "total_tokens": 18247, + "total_input_tokens": 16, + "input_tokens_by_modality": [{"modality": "text", "tokens": 16}], + "total_cached_tokens": 0, + "total_output_tokens": 17937, + "output_tokens_by_modality": [{"modality": "video", "tokens": 17376}], + "total_tool_use_tokens": 0, + "total_thought_tokens": 294, +} + + +def test_detects_interactions_usage_object(): + assert InteractionsUsageObjectTransformation.is_interactions_usage_object(OMNI_VIDEO_USAGE) is True + + +def test_rejects_chat_and_responses_api_usage_objects(): + chat_usage = {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30} + responses_api_usage = {"input_tokens": 10, "output_tokens": 20, "total_tokens": 30} + assert InteractionsUsageObjectTransformation.is_interactions_usage_object(chat_usage) is False + assert InteractionsUsageObjectTransformation.is_interactions_usage_object(responses_api_usage) is False + assert InteractionsUsageObjectTransformation.is_interactions_usage_object(None) is False + assert InteractionsUsageObjectTransformation.is_interactions_usage_object("usage") is False + + +def test_transforms_real_omni_video_usage_block(): + usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object(OMNI_VIDEO_USAGE) + + assert isinstance(usage, Usage) + assert usage.prompt_tokens == 16 + assert usage.completion_tokens == 17937 + 294 + assert usage.total_tokens == 18247 + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.text_tokens == 16 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.video_tokens == 17376 + assert usage.completion_tokens_details.reasoning_tokens == 294 + + +def test_transforms_reasoning_tokens_spec_field_name(): + usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object( + { + "total_input_tokens": 10, + "total_output_tokens": 20, + "total_reasoning_tokens": 5, + } + ) + assert usage.completion_tokens == 25 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 5 + assert usage.total_tokens == 35 + + +def test_cached_tokens_subtracted_from_text_input(): + usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object( + { + "total_input_tokens": 1000, + "input_tokens_by_modality": [{"modality": "text", "tokens": 1000}], + "total_cached_tokens": 400, + "total_output_tokens": 50, + } + ) + assert usage.prompt_tokens == 1000 + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.text_tokens == 600 + assert usage.prompt_tokens_details.cached_tokens == 400 + assert usage._cache_read_input_tokens == 400 + + +def test_cached_tokens_subtracted_per_modality_when_breakdown_present(): + usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object( + { + "total_input_tokens": 1500, + "input_tokens_by_modality": [ + {"modality": "text", "tokens": 1000}, + {"modality": "audio", "tokens": 500}, + ], + "total_cached_tokens": 300, + "cached_tokens_by_modality": [{"modality": "audio", "tokens": 300}], + "total_output_tokens": 50, + } + ) + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.text_tokens == 1000 + assert usage.prompt_tokens_details.audio_tokens == 200 + assert usage.prompt_tokens_details.cached_tokens == 300 + + +def test_tool_use_tokens_billed_as_input(): + usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object( + { + "total_input_tokens": 100, + "input_tokens_by_modality": [{"modality": "text", "tokens": 100}], + "total_tool_use_tokens": 40, + "tool_use_tokens_by_modality": [{"modality": "text", "tokens": 40}], + "total_output_tokens": 10, + } + ) + assert usage.prompt_tokens == 140 + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.text_tokens == 140 + + +def test_document_modality_folds_into_text(): + usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object( + { + "total_input_tokens": 80, + "input_tokens_by_modality": [ + {"modality": "text", "tokens": 30}, + {"modality": "document", "tokens": 50}, + ], + "total_output_tokens": 10, + } + ) + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.text_tokens == 80 diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index ade2c677745..3769c656b0c 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -3773,3 +3773,150 @@ def test_zero_token_video_usage_preserves_duration_seconds(logging_obj): assert payload["metadata"]["usage_object"]["duration_seconds"] == 4.0 assert payload["total_tokens"] == 0 assert payload["completion_tokens"] == 0 + + +INTERACTIONS_USAGE_BLOCK = { + "total_tokens": 175, + "total_input_tokens": 100, + "input_tokens_by_modality": [{"modality": "text", "tokens": 100}], + "total_cached_tokens": 0, + "total_output_tokens": 50, + "output_tokens_by_modality": [{"modality": "text", "tokens": 50}], + "total_tool_use_tokens": 0, + "total_thought_tokens": 25, +} + + +def _interactions_logging_obj(stream: bool): + logging_obj = LitellmLogging( + model="gemini-2.5-flash", + messages=[], + stream=stream, + call_type="acreate", + start_time=time.time(), + litellm_call_id="interactions-call-id", + function_id="interactions-fn-id", + ) + logging_obj.update_environment_variables( + litellm_params={}, + optional_params={}, + model="gemini-2.5-flash", + custom_llm_provider="gemini", + input="hi", + ) + return logging_obj + + +def test_interactions_response_is_recognized_for_logging(): + from litellm.types.interactions import InteractionsAPIResponse + + logging_obj = _interactions_logging_obj(stream=False) + response = InteractionsAPIResponse(id="interactions/abc", model="gemini-2.5-flash", status="completed") + assert logging_obj._is_recognized_call_type_for_logging(logging_result=response) is True + + +def test_non_streaming_interactions_success_sets_response_cost_and_usage(): + import datetime as dt + + from litellm.types.interactions import InteractionsAPIResponse + + logging_obj = _interactions_logging_obj(stream=False) + response = InteractionsAPIResponse( + id="interactions/abc", + model="gemini-2.5-flash", + status="completed", + steps=[], + usage=dict(INTERACTIONS_USAGE_BLOCK), + ) + + logging_obj._success_handler_helper_fn( + result=response, + start_time=dt.datetime.now(), + end_time=dt.datetime.now(), + cache_hit=False, + ) + + assert logging_obj.model_call_details["response_cost"] > 0 + standard_logging_object = logging_obj.model_call_details["standard_logging_object"] + assert standard_logging_object["prompt_tokens"] == 100 + assert standard_logging_object["completion_tokens"] == 75 + assert standard_logging_object["total_tokens"] == 175 + assert standard_logging_object["response_cost"] == logging_obj.model_call_details["response_cost"] + + +def test_assembled_streaming_response_from_completed_interaction_event(): + import datetime as dt + + from litellm.types.interactions import ( + InteractionsAPIResponse, + InteractionsAPIStreamingResponse, + ) + + logging_obj = _interactions_logging_obj(stream=True) + completed_event = InteractionsAPIStreamingResponse( + event_type="interaction.completed", + interaction={ + "id": "interactions/abc", + "model": "gemini-2.5-flash", + "status": "completed", + "steps": [], + "usage": dict(INTERACTIONS_USAGE_BLOCK), + }, + ) + + assembled = logging_obj._get_assembled_streaming_response( + result=completed_event, + start_time=dt.datetime.now(), + end_time=dt.datetime.now(), + is_async=True, + streaming_chunks=[], + ) + + assert isinstance(assembled, InteractionsAPIResponse) + assert assembled.usage == INTERACTIONS_USAGE_BLOCK + + in_progress_event = InteractionsAPIStreamingResponse(event_type="interaction.in_progress") + assert ( + logging_obj._get_assembled_streaming_response( + result=in_progress_event, + start_time=dt.datetime.now(), + end_time=dt.datetime.now(), + is_async=True, + streaming_chunks=[], + ) + is None + ) + + +def test_assembled_streaming_response_from_legacy_completed_chunk(): + from litellm.types.interactions import ( + InteractionsAPIResponse, + InteractionsAPIStreamingResponse, + ) + + legacy_chunk = InteractionsAPIStreamingResponse( + event_type="interaction.complete", + id="interactions/legacy", + model="gemini-2.5-flash", + status="completed", + outputs=[], + usage=dict(INTERACTIONS_USAGE_BLOCK), + ) + + assembled = LitellmLogging._assemble_completed_interaction_response(legacy_chunk) + + assert isinstance(assembled, InteractionsAPIResponse) + assert assembled.id == "interactions/legacy" + assert assembled.usage == INTERACTIONS_USAGE_BLOCK + + +def test_standard_logging_payload_maps_interactions_usage(): + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + usage = StandardLoggingPayloadSetup.get_usage_from_response_obj( + response_obj={"usage": dict(INTERACTIONS_USAGE_BLOCK)} + ) + + assert usage.prompt_tokens == 100 + assert usage.completion_tokens == 75 + assert usage.total_tokens == 175 diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 9636db4f4cd..c76439e5ad1 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -3479,3 +3479,36 @@ def test_batch_cost_calculator_cache_creation_falls_back_to_input_rate(): ) assert prompt_cost == pytest.approx((1000 * 3e-6 + 8000 * 3e-7 + 2000 * 3e-6) / 2) + + +def test_completion_cost_bills_interactions_api_response(): + from litellm.types.interactions import InteractionsAPIResponse + + model_info = litellm.get_model_info(model="gemini-2.5-flash", custom_llm_provider="gemini") + response = InteractionsAPIResponse( + id="interactions/abc123", + model="gemini-2.5-flash", + status="completed", + steps=[], + usage={ + "total_tokens": 175, + "total_input_tokens": 100, + "input_tokens_by_modality": [{"modality": "text", "tokens": 100}], + "total_cached_tokens": 0, + "total_output_tokens": 50, + "output_tokens_by_modality": [{"modality": "text", "tokens": 50}], + "total_tool_use_tokens": 0, + "total_thought_tokens": 25, + }, + ) + + cost = completion_cost(completion_response=response, custom_llm_provider="gemini") + + reasoning_rate = model_info.get("output_cost_per_reasoning_token") or model_info["output_cost_per_token"] + expected = ( + 100 * model_info["input_cost_per_token"] + + 50 * model_info["output_cost_per_token"] + + 25 * reasoning_rate + ) + assert cost == pytest.approx(expected) + assert cost > 0 From f1a5054a16a503c5e087121abf87c0246aa8b6d0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 14 Jul 2026 18:11:20 -0700 Subject: [PATCH 003/542] fix(interactions): bill only interaction creation, not GET polls --- litellm/litellm_core_utils/litellm_logging.py | 18 ++++++++- litellm/types/utils.py | 6 +++ .../test_litellm_logging.py | 40 +++++++++++++++++-- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 4 files changed, 60 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 0f6e0ccf9c1..94ee595a169 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1895,7 +1895,7 @@ class Logging(LiteLLMLoggingBaseClass): or isinstance(logging_result, FineTuningJob) or isinstance(logging_result, LiteLLMBatch) or isinstance(logging_result, ResponsesAPIResponse) - or isinstance(logging_result, InteractionsAPIResponse) + or (isinstance(logging_result, InteractionsAPIResponse) and self._is_interactions_create_call_type()) or isinstance(logging_result, OpenAIFileObject) or isinstance(logging_result, LiteLLMRealtimeStreamLoggingObject) or isinstance(logging_result, OpenAIModerationResponse) @@ -1913,6 +1913,22 @@ class Logging(LiteLLMLoggingBaseClass): return True return False + def _is_interactions_create_call_type(self) -> bool: + """ + Only interaction creation is billable. GET polls, deletes, and cancels + also return an ``InteractionsAPIResponse`` (with usage once completed), + so recognizing those would write spend on every poll of a background + interaction. The proxy sets ``call_type`` from its route_type + (``create_interaction``/``acreate_interaction``); the SDK sets it from + the decorated function name (``create``/``acreate``). + """ + return self.call_type in ( + CallTypes.create_interaction.value, + CallTypes.acreate_interaction.value, + "create", + "acreate", + ) + def _flush_passthrough_collected_chunks_helper( self, raw_bytes: List[bytes], diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 90ea99ceb23..0fafd4a27fe 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -396,6 +396,12 @@ class CallTypes(str, Enum): vector_store_search = "vector_store_search" avector_store_search = "avector_store_search" + ######################################################### + # Google Interactions API Call Types + ######################################################### + create_interaction = "create_interaction" + acreate_interaction = "acreate_interaction" + ######################################################### # Container Call Types ######################################################### diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 3769c656b0c..2bf9f54d7c9 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -3787,12 +3787,12 @@ INTERACTIONS_USAGE_BLOCK = { } -def _interactions_logging_obj(stream: bool): +def _interactions_logging_obj(stream: bool, call_type: str = "acreate"): logging_obj = LitellmLogging( model="gemini-2.5-flash", messages=[], stream=stream, - call_type="acreate", + call_type=call_type, start_time=time.time(), litellm_call_id="interactions-call-id", function_id="interactions-fn-id", @@ -3807,14 +3807,46 @@ def _interactions_logging_obj(stream: bool): return logging_obj -def test_interactions_response_is_recognized_for_logging(): +@pytest.mark.parametrize("call_type", ["create", "acreate", "create_interaction", "acreate_interaction"]) +def test_interactions_response_is_recognized_for_logging(call_type): from litellm.types.interactions import InteractionsAPIResponse - logging_obj = _interactions_logging_obj(stream=False) + logging_obj = _interactions_logging_obj(stream=False, call_type=call_type) response = InteractionsAPIResponse(id="interactions/abc", model="gemini-2.5-flash", status="completed") assert logging_obj._is_recognized_call_type_for_logging(logging_result=response) is True +@pytest.mark.parametrize( + "call_type", + ["aget", "get", "aget_interaction", "adelete_interaction", "acancel_interaction"], +) +def test_interactions_get_poll_is_not_billed(call_type): + import datetime as dt + + from litellm.types.interactions import InteractionsAPIResponse + + logging_obj = _interactions_logging_obj(stream=False, call_type=call_type) + response = InteractionsAPIResponse( + id="interactions/abc", + model="gemini-2.5-flash", + status="completed", + steps=[], + usage=dict(INTERACTIONS_USAGE_BLOCK), + ) + + assert logging_obj._is_recognized_call_type_for_logging(logging_result=response) is False + + logging_obj._success_handler_helper_fn( + result=response, + start_time=dt.datetime.now(), + end_time=dt.datetime.now(), + cache_hit=False, + ) + + assert logging_obj.model_call_details.get("response_cost") is None + assert logging_obj.model_call_details.get("standard_logging_object") is None + + def test_non_streaming_interactions_success_sets_response_cost_and_usage(): import datetime as dt diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 5629cd6c8d4..59c0532ccbe 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -21685,7 +21685,7 @@ export interface components { * CallTypes * @enum {string} */ - CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "create_sandbox" | "acreate_sandbox" | "delete_sandbox" | "adelete_sandbox" | "run_code" | "arun_code" | "code_interpreter_tool" | "acode_interpreter_tool" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill"; + CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "create_interaction" | "acreate_interaction" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "create_sandbox" | "acreate_sandbox" | "delete_sandbox" | "adelete_sandbox" | "run_code" | "arun_code" | "code_interpreter_tool" | "acode_interpreter_tool" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill"; /** CallbackDelete */ CallbackDelete: { /** Callback Name */ From 59d4e52a3d2b5e088af0e894e34bbd0ed5fbfe9a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 14 Jul 2026 20:05:33 -0700 Subject: [PATCH 004/542] fix(interactions): bill background interactions once completed via cost polling --- litellm/constants.py | 11 ++ .../interactions/background_cost_polling.py | 131 +++++++++++++ litellm/interactions/main.py | 9 + litellm/litellm_core_utils/litellm_logging.py | 27 ++- .../proxy/hooks/proxy_track_cost_callback.py | 12 ++ .../credential_migration.py | 73 ++----- .../test_background_cost_polling.py | 184 ++++++++++++++++++ .../test_litellm_logging.py | 59 +++++- .../hooks/test_proxy_track_cost_callback.py | 43 ++++ tests/test_litellm/test_cost_calculator.py | 29 +++ 10 files changed, 519 insertions(+), 59 deletions(-) create mode 100644 litellm/interactions/background_cost_polling.py create mode 100644 tests/test_litellm/interactions/test_background_cost_polling.py diff --git a/litellm/constants.py b/litellm/constants.py index 715d57e594d..be880b72854 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1457,6 +1457,17 @@ STALE_OBJECT_CLEANUP_BATCH_SIZE = max(1, int(os.getenv("STALE_OBJECT_CLEANUP_BAT # installations with large numbers of stale managed objects). _batch_polling_env = os.getenv("PROXY_BATCH_POLLING_ENABLED", "true").lower() PROXY_BATCH_POLLING_ENABLED = _batch_polling_env == "true" +BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS = float( + os.getenv("BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS", 5) +) +BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS = float( + os.getenv("BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS", 60) +) +BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS = float( + os.getenv("BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS", 3600) +) +_background_interaction_cost_polling_env = os.getenv("BACKGROUND_INTERACTION_COST_POLLING_ENABLED", "true").lower() +BACKGROUND_INTERACTION_COST_POLLING_ENABLED = _background_interaction_cost_polling_env == "true" PROXY_BUDGET_RESCHEDULER_MAX_TIME = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605)) PROXY_BATCH_WRITE_AT = int(os.getenv("PROXY_BATCH_WRITE_AT", 10)) # in seconds, increased from 10 diff --git a/litellm/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py new file mode 100644 index 00000000000..fefbf8715cb --- /dev/null +++ b/litellm/interactions/background_cost_polling.py @@ -0,0 +1,131 @@ +""" +Cost tracking for background interactions. + +A create request with ``background=true`` returns ``in_progress`` with no +usage block, and GET polls are deliberately never billed (billing them would +double-charge every poll; the GET response also does not echo ``background``, +so a poll cannot be told apart from a re-fetch of an already-billed +interaction). The create call is therefore the only place that can own +billing: it schedules a poll task that fetches the interaction until it +reaches a terminal status and logs the final usage as a single success event +attributed to the original request. +""" + +import asyncio +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Awaitable, Callable, Iterator, Optional + +from litellm._logging import verbose_logger +from litellm.constants import ( + BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS, + BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS, + BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS, + BACKGROUND_INTERACTION_COST_POLLING_ENABLED, +) +from litellm.types.interactions import InteractionsAPIResponse + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +_TERMINAL_STATUSES = frozenset({"completed", "failed", "cancelled", "incomplete", "budget_exceeded"}) + + +@dataclass(frozen=True, slots=True) +class BackgroundInteractionPollContext: + interaction_id: str + custom_llm_provider: str + logging_obj: "LiteLLMLoggingObj" + api_key: Optional[str] = None + api_base: Optional[str] = None + initial_interval_seconds: float = BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS + max_interval_seconds: float = BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS + timeout_seconds: float = BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS + + +FetchInteraction = Callable[[BackgroundInteractionPollContext], Awaitable[InteractionsAPIResponse]] + + +async def _fetch_interaction(context: BackgroundInteractionPollContext) -> InteractionsAPIResponse: + from litellm.interactions import aget + + return await aget( + interaction_id=context.interaction_id, + custom_llm_provider=context.custom_llm_provider, + **{"api_key": context.api_key, "api_base": context.api_base, "no-log": True}, + ) + + +def _poll_intervals(initial: float, maximum: float, timeout: float) -> Iterator[float]: + elapsed = 0.0 + interval = initial + while elapsed + interval <= timeout: + yield interval + elapsed += interval + interval = min(interval * 2, maximum) + + +async def poll_and_log_background_interaction_cost( + context: BackgroundInteractionPollContext, + fetch_interaction: FetchInteraction = _fetch_interaction, +) -> None: + for interval in _poll_intervals( + initial=context.initial_interval_seconds, + maximum=context.max_interval_seconds, + timeout=context.timeout_seconds, + ): + await asyncio.sleep(interval) + try: + response = await fetch_interaction(context) + except Exception as e: # noqa: BLE001 # any fetch error must not kill the billing poll loop + verbose_logger.debug( + "Background interaction cost poll for %s failed, will retry: %s", + context.interaction_id, + e, + ) + continue + if response.status not in _TERMINAL_STATUSES: + continue + if response.usage is not None: + await context.logging_obj.async_log_background_interaction_completion(result=response) + return + verbose_logger.warning( + "Gave up cost polling for background interaction %s after %ss; its usage will not be tracked", + context.interaction_id, + context.timeout_seconds, + ) + + +_ACTIVE_POLL_TASKS: set["asyncio.Task[None]"] = set() # mutable-ok: asyncio requires strong refs to running tasks + + +def maybe_schedule_background_interaction_cost_polling( + response: Any, + create_kwargs: dict[str, Any], + custom_llm_provider: str, +) -> Optional["asyncio.Task[None]"]: + from litellm.litellm_core_utils.litellm_logging import Logging + + if not BACKGROUND_INTERACTION_COST_POLLING_ENABLED: + return None + if not isinstance(response, InteractionsAPIResponse): + return None + if response.status != "in_progress" or not response.id: + return None + logging_obj = create_kwargs.get("litellm_logging_obj") + if not isinstance(logging_obj, Logging): + return None + try: + asyncio.get_running_loop() + except RuntimeError: + return None + context = BackgroundInteractionPollContext( + interaction_id=response.id, + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + api_key=create_kwargs.get("api_key"), + api_base=create_kwargs.get("api_base"), + ) + task = asyncio.create_task(poll_and_log_background_interaction_cost(context)) + _ACTIVE_POLL_TASKS.add(task) + task.add_done_callback(_ACTIVE_POLL_TASKS.discard) + return task diff --git a/litellm/interactions/main.py b/litellm/interactions/main.py index 8634269ee94..5a0ea9280b2 100644 --- a/litellm/interactions/main.py +++ b/litellm/interactions/main.py @@ -39,6 +39,9 @@ from typing import Any, AsyncIterator, Coroutine, Dict, Iterator, List, Optional import httpx import litellm +from litellm.interactions.background_cost_polling import ( + maybe_schedule_background_interaction_cost_polling, +) from litellm.interactions.http_handler import interactions_http_handler from litellm.interactions.utils import ( InteractionsAPIRequestUtils, @@ -170,6 +173,12 @@ async def acreate( else: response = init_response + maybe_schedule_background_interaction_cost_polling( + response=response, + create_kwargs=kwargs, + custom_llm_provider=custom_llm_provider, + ) + return response # type: ignore except Exception as e: raise litellm.exception_type( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 94ee595a169..ac6334bea76 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1895,7 +1895,11 @@ class Logging(LiteLLMLoggingBaseClass): or isinstance(logging_result, FineTuningJob) or isinstance(logging_result, LiteLLMBatch) or isinstance(logging_result, ResponsesAPIResponse) - or (isinstance(logging_result, InteractionsAPIResponse) and self._is_interactions_create_call_type()) + or ( + isinstance(logging_result, InteractionsAPIResponse) + and logging_result.usage is not None + and self._is_interactions_create_call_type() + ) or isinstance(logging_result, OpenAIFileObject) or isinstance(logging_result, LiteLLMRealtimeStreamLoggingObject) or isinstance(logging_result, OpenAIModerationResponse) @@ -1921,6 +1925,13 @@ class Logging(LiteLLMLoggingBaseClass): interaction. The proxy sets ``call_type`` from its route_type (``create_interaction``/``acreate_interaction``); the SDK sets it from the decorated function name (``create``/``acreate``). + + Recognition additionally requires a usage block (checked at the call + site): a ``background=true`` create returns ``in_progress`` without + usage, and billing it would write a $0 spend log under the interaction + id that collides with the row the background poll task writes once the + interaction completes (see + ``litellm.interactions.background_cost_polling``). """ return self.call_type in ( CallTypes.create_interaction.value, @@ -1929,6 +1940,20 @@ class Logging(LiteLLMLoggingBaseClass): "acreate", ) + async def async_log_background_interaction_completion( + self, + result: InteractionsAPIResponse, + ) -> None: + """ + Log the terminal result of a background interaction as a fresh success + event. The create request already ran success logging for its + ``in_progress`` response (no usage, so no cost was tracked); clearing + the dedup flag lets the completed result flow through cost calculation + and spend tracking exactly once, spanning create to completion. + """ + self.model_call_details.pop("has_logged_async_success", None) + await self.async_success_handler(result=result) + def _flush_passthrough_collected_chunks_helper( self, raw_bytes: List[bytes], diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index b839426fcda..6359eae0fa4 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -279,6 +279,12 @@ class _ProxyDBLogger(CustomLogger): await _release_budget_reservation(budget_reservation=budget_reservation) else: await _release_budget_reservation(budget_reservation=budget_reservation) + if _is_unbilled_in_progress_interaction(completion_response): + verbose_proxy_logger.debug( + "Cost tracking deferred for in-progress background interaction; " + "a poll task logs the final usage once it completes" + ) + return # Non-model call types (health checks, afile_delete) have no model or standard_logging_object. # Use .get() for "stream" to avoid KeyError on health checks. # WS session wrappers (_aresponses_websocket, _arealtime) also reach here with @@ -418,6 +424,12 @@ def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None: bucket[key] = value +def _is_unbilled_in_progress_interaction(completion_response: Any) -> bool: + from litellm.types.interactions import InteractionsAPIResponse + + return isinstance(completion_response, InteractionsAPIResponse) and completion_response.usage is None + + def _should_track_cost_callback( user_api_key: Optional[str], user_id: Optional[str], diff --git a/litellm/proxy/management_endpoints/credential_migration.py b/litellm/proxy/management_endpoints/credential_migration.py index 4d51295f8dc..6f79a39c883 100644 --- a/litellm/proxy/management_endpoints/credential_migration.py +++ b/litellm/proxy/management_endpoints/credential_migration.py @@ -130,9 +130,7 @@ def classify_value(value: object, key: str = "scan") -> ValueClass: return "plaintext" if value.startswith(_V2_GCM_PREFIX): return "migrated" - decrypted = decrypt_value_helper( - value=value, key=key, exception_type="debug", return_original_value=False - ) + decrypted = decrypt_value_helper(value=value, key=key, exception_type="debug", return_original_value=False) if decrypted is None: # Did not decrypt under nacl and has no v2 marker: legacy plaintext. return "plaintext" @@ -151,9 +149,7 @@ def reencrypt_value(value: object, key: str = "migrate") -> object: return value if value.startswith(_V2_GCM_PREFIX): return value # idempotent: already migrated - decrypted = decrypt_value_helper( - value=value, key=key, exception_type="debug", return_original_value=False - ) + decrypted = decrypt_value_helper(value=value, key=key, exception_type="debug", return_original_value=False) if decrypted is None: # Either legacy plaintext (no ciphertext to migrate) or corrupt. Either # way, do not overwrite — preserve the value as stored. @@ -161,9 +157,7 @@ def reencrypt_value(value: object, key: str = "migrate") -> object: return encrypt_value_helper(decrypted) -def reencrypt_selective_dict( - data: dict[str, object], sensitive_keys: list[str] -) -> dict[str, object]: +def reencrypt_selective_dict(data: dict[str, object], sensitive_keys: list[str]) -> dict[str, object]: """Return a copy of ``data`` with only ``sensitive_keys`` re-encrypted. Non-sensitive fields (e.g. ``base_url``, ``connection_id``) are left as-is. @@ -212,9 +206,7 @@ async def _migrate_config_settings_row( dict with selected sensitive fields (vantage_settings / cloudzero_settings). """ report = LocationReport(location=param_name) - record = await prisma_client.db.litellm_config.find_unique( - where={"param_name": param_name} - ) + record = await prisma_client.db.litellm_config.find_unique(where={"param_name": param_name}) if record is None or record.param_value is None: return report @@ -266,9 +258,7 @@ async def _migrate_sso_config(prisma_client: object, dry_run: bool) -> LocationR every present string field. """ report = LocationReport(location="sso_config") - record = await prisma_client.db.litellm_ssoconfig.find_unique( - where={"id": "sso_config"} - ) + record = await prisma_client.db.litellm_ssoconfig.find_unique(where={"id": "sso_config"}) if record is None or record.sso_settings is None: return report @@ -344,9 +334,7 @@ async def _migrate_callback_vars_table( rows = await table.find_many() for row in rows or []: metadata = getattr(row, "metadata", None) - if not isinstance(metadata, dict) or ( - "logging" not in metadata and "callback_settings" not in metadata - ): + if not isinstance(metadata, dict) or ("logging" not in metadata and "callback_settings" not in metadata): continue # Classify every callback-var value directly (strip the litellm_enc:: @@ -534,9 +522,7 @@ async def _scan_config_env_vars(prisma_client: object) -> LocationReport: """Scan the ``environment_variables`` config row (``param_value`` dict).""" report = LocationReport(location="config_environment_variables") try: - record = await prisma_client.db.litellm_config.find_unique( - where={"param_name": "environment_variables"} - ) + record = await prisma_client.db.litellm_config.find_unique(where={"param_name": "environment_variables"}) except Exception as e: # pragma: no cover - defensive verbose_proxy_logger.debug("scan: config env vars unavailable: %s", str(e)) return report @@ -557,11 +543,7 @@ async def _scan_covered_tables(prisma_client: object) -> list[LocationReport]: """Read-only classification of every rotation-covered table. No writes.""" reports: list[LocationReport] = [] for location, db_attr, json_cols, scalar_cols in _COVERED_TABLE_SPECS: - reports.append( - await _scan_one_table( - prisma_client, location, db_attr, json_cols, scalar_cols - ) - ) + reports.append(await _scan_one_table(prisma_client, location, db_attr, json_cols, scalar_cols)) reports.append(await _scan_config_env_vars(prisma_client)) return reports @@ -575,9 +557,7 @@ _VANTAGE_SENSITIVE = ["api_key", "integration_token"] _CLOUDZERO_SENSITIVE = ["api_key"] -async def _migrate_covered_tables( - prisma_client: object, user_api_key_dict: object -) -> list[LocationReport]: +async def _migrate_covered_tables(prisma_client: object, user_api_key_dict: object) -> list[LocationReport]: """Re-encrypt the tables already covered by ``_rotate_master_key`` (model table, credentials, MCP credential/env tables, config environment_variables) by running that orchestrator in *same-key* mode. With the AES gate on, the @@ -597,8 +577,7 @@ async def _migrate_covered_tables( current_key = _get_salt_key() if current_key is None: raise RuntimeError( - "Cannot migrate covered tables: no salt key / master key is set. " - "Set LITELLM_SALT_KEY before migrating." + "Cannot migrate covered tables: no salt key / master key is set. Set LITELLM_SALT_KEY before migrating." ) await _rotate_master_key( prisma_client=cast("PrismaClient", prisma_client), @@ -648,19 +627,9 @@ async def migrate_encryption( # Net-new walkers (items 3, 4, 11, 12, 13). report.add(await _migrate_callback_vars_table(prisma_client, "team", dry_run)) - report.add( - await _migrate_callback_vars_table(prisma_client, "verification_token", dry_run) - ) - report.add( - await _migrate_config_settings_row( - prisma_client, "vantage_settings", _VANTAGE_SENSITIVE, dry_run - ) - ) - report.add( - await _migrate_config_settings_row( - prisma_client, "cloudzero_settings", _CLOUDZERO_SENSITIVE, dry_run - ) - ) + report.add(await _migrate_callback_vars_table(prisma_client, "verification_token", dry_run)) + report.add(await _migrate_config_settings_row(prisma_client, "vantage_settings", _VANTAGE_SENSITIVE, dry_run)) + report.add(await _migrate_config_settings_row(prisma_client, "cloudzero_settings", _CLOUDZERO_SENSITIVE, dry_run)) report.add(await _migrate_sso_config(prisma_client, dry_run)) return report @@ -683,20 +652,10 @@ async def check_encryption(prisma_client: object) -> MigrationReport: # Net-new walker locations, in dry-run (read-only) mode. report.add(await _migrate_callback_vars_table(prisma_client, "team", dry_run=True)) + report.add(await _migrate_callback_vars_table(prisma_client, "verification_token", dry_run=True)) + report.add(await _migrate_config_settings_row(prisma_client, "vantage_settings", _VANTAGE_SENSITIVE, dry_run=True)) report.add( - await _migrate_callback_vars_table( - prisma_client, "verification_token", dry_run=True - ) - ) - report.add( - await _migrate_config_settings_row( - prisma_client, "vantage_settings", _VANTAGE_SENSITIVE, dry_run=True - ) - ) - report.add( - await _migrate_config_settings_row( - prisma_client, "cloudzero_settings", _CLOUDZERO_SENSITIVE, dry_run=True - ) + await _migrate_config_settings_row(prisma_client, "cloudzero_settings", _CLOUDZERO_SENSITIVE, dry_run=True) ) report.add(await _migrate_sso_config(prisma_client, dry_run=True)) return report diff --git a/tests/test_litellm/interactions/test_background_cost_polling.py b/tests/test_litellm/interactions/test_background_cost_polling.py new file mode 100644 index 00000000000..aac97826b02 --- /dev/null +++ b/tests/test_litellm/interactions/test_background_cost_polling.py @@ -0,0 +1,184 @@ +import asyncio +import time + +import pytest + +from litellm.interactions.background_cost_polling import ( + BackgroundInteractionPollContext, + maybe_schedule_background_interaction_cost_polling, + poll_and_log_background_interaction_cost, +) +from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging +from litellm.types.interactions import InteractionsAPIResponse + +USAGE_BLOCK = { + "total_tokens": 175, + "total_input_tokens": 100, + "input_tokens_by_modality": [{"modality": "text", "tokens": 100}], + "total_cached_tokens": 0, + "total_output_tokens": 50, + "output_tokens_by_modality": [{"modality": "text", "tokens": 50}], + "total_tool_use_tokens": 0, + "total_thought_tokens": 25, +} + + +def _logging_obj(call_type: str = "acreate_interaction") -> LitellmLogging: + logging_obj = LitellmLogging( + model="gemini-2.5-flash", + messages=[], + stream=False, + call_type=call_type, + start_time=time.time(), + litellm_call_id="bg-interactions-call-id", + function_id="bg-interactions-fn-id", + ) + logging_obj.update_environment_variables( + litellm_params={}, + optional_params={}, + model="gemini-2.5-flash", + custom_llm_provider="gemini", + input="hi", + ) + return logging_obj + + +def _context(logging_obj: LitellmLogging, timeout_seconds: float = 1.0) -> BackgroundInteractionPollContext: + return BackgroundInteractionPollContext( + interaction_id="interactions/bg-abc", + custom_llm_provider="gemini", + logging_obj=logging_obj, + initial_interval_seconds=0.001, + max_interval_seconds=0.002, + timeout_seconds=timeout_seconds, + ) + + +def _response(status: str, with_usage: bool) -> InteractionsAPIResponse: + return InteractionsAPIResponse( + id="interactions/bg-abc", + model="gemini-2.5-flash", + status=status, + steps=[], + usage=dict(USAGE_BLOCK) if with_usage else None, + ) + + +def _fetch_sequence(*responses): + remaining = list(responses) + calls = [] + + async def fetch(context): + calls.append(context.interaction_id) + item = remaining.pop(0) if len(remaining) > 1 else remaining[0] + if isinstance(item, Exception): + raise item + return item + + return fetch, calls + + +@pytest.mark.asyncio +async def test_poller_bills_once_when_interaction_completes(): + logging_obj = _logging_obj() + fetch, calls = _fetch_sequence( + _response("in_progress", with_usage=False), + _response("completed", with_usage=True), + ) + + await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) + + assert len(calls) == 2 + assert logging_obj.model_call_details["response_cost"] > 0 + assert logging_obj.model_call_details["standard_logging_object"]["total_tokens"] == 175 + + +@pytest.mark.asyncio +async def test_poller_stops_without_billing_on_terminal_status_without_usage(): + logging_obj = _logging_obj() + fetch, calls = _fetch_sequence(_response("failed", with_usage=False)) + + await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) + + assert len(calls) == 1 + assert logging_obj.model_call_details.get("response_cost") is None + + +@pytest.mark.asyncio +async def test_poller_gives_up_after_timeout_without_billing(): + logging_obj = _logging_obj() + fetch, calls = _fetch_sequence(_response("in_progress", with_usage=False)) + + await poll_and_log_background_interaction_cost( + _context(logging_obj, timeout_seconds=0.01), + fetch_interaction=fetch, + ) + + assert len(calls) >= 2 + assert logging_obj.model_call_details.get("response_cost") is None + + +@pytest.mark.asyncio +async def test_poller_retries_after_fetch_error_and_still_bills(): + logging_obj = _logging_obj() + fetch, calls = _fetch_sequence( + RuntimeError("transient network error"), + _response("completed", with_usage=True), + ) + + await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) + + assert len(calls) == 2 + assert logging_obj.model_call_details["response_cost"] > 0 + + +@pytest.mark.asyncio +async def test_schedule_creates_poll_task_for_in_progress_create(): + logging_obj = _logging_obj() + task = maybe_schedule_background_interaction_cost_polling( + response=_response("in_progress", with_usage=False), + create_kwargs={"litellm_logging_obj": logging_obj}, + custom_llm_provider="gemini", + ) + + assert isinstance(task, asyncio.Task) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response,create_kwargs", + [ + (_response("completed", with_usage=True), {"litellm_logging_obj": "placeholder"}), + (_response("in_progress", with_usage=False), {}), + ("not a response", {"litellm_logging_obj": "placeholder"}), + ], +) +async def test_schedule_skips_non_pollable_results(response, create_kwargs): + if create_kwargs.get("litellm_logging_obj") == "placeholder": + create_kwargs = {"litellm_logging_obj": _logging_obj()} + + task = maybe_schedule_background_interaction_cost_polling( + response=response, + create_kwargs=create_kwargs, + custom_llm_provider="gemini", + ) + + assert task is None + + +@pytest.mark.asyncio +async def test_schedule_respects_kill_switch(monkeypatch): + import litellm.interactions.background_cost_polling as module + + monkeypatch.setattr(module, "BACKGROUND_INTERACTION_COST_POLLING_ENABLED", False) + + task = maybe_schedule_background_interaction_cost_polling( + response=_response("in_progress", with_usage=False), + create_kwargs={"litellm_logging_obj": _logging_obj()}, + custom_llm_provider="gemini", + ) + + assert task is None diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 2bf9f54d7c9..4d8b19150f7 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -3812,10 +3812,67 @@ def test_interactions_response_is_recognized_for_logging(call_type): from litellm.types.interactions import InteractionsAPIResponse logging_obj = _interactions_logging_obj(stream=False, call_type=call_type) - response = InteractionsAPIResponse(id="interactions/abc", model="gemini-2.5-flash", status="completed") + response = InteractionsAPIResponse( + id="interactions/abc", + model="gemini-2.5-flash", + status="completed", + usage=dict(INTERACTIONS_USAGE_BLOCK), + ) assert logging_obj._is_recognized_call_type_for_logging(logging_result=response) is True +@pytest.mark.parametrize("call_type", ["acreate", "acreate_interaction"]) +def test_in_progress_background_create_is_not_billed(call_type): + import datetime as dt + + from litellm.types.interactions import InteractionsAPIResponse + + logging_obj = _interactions_logging_obj(stream=False, call_type=call_type) + response = InteractionsAPIResponse(id="interactions/abc", model="gemini-2.5-flash", status="in_progress") + + assert logging_obj._is_recognized_call_type_for_logging(logging_result=response) is False + + logging_obj._success_handler_helper_fn( + result=response, + start_time=dt.datetime.now(), + end_time=dt.datetime.now(), + cache_hit=False, + ) + + assert logging_obj.model_call_details.get("response_cost") is None + assert logging_obj.model_call_details.get("standard_logging_object") is None + + +@pytest.mark.asyncio +async def test_background_interaction_completion_rebills_after_in_progress_success(): + import datetime as dt + + from litellm.types.interactions import InteractionsAPIResponse + + logging_obj = _interactions_logging_obj(stream=False) + in_progress = InteractionsAPIResponse(id="interactions/abc", model="gemini-2.5-flash", status="in_progress") + await logging_obj.async_success_handler( + result=in_progress, + start_time=dt.datetime.now(), + end_time=dt.datetime.now(), + ) + + assert logging_obj.model_call_details.get("response_cost") is None + assert logging_obj.should_run_logging(event_type="async_success") is False + + completed = InteractionsAPIResponse( + id="interactions/abc", + model="gemini-2.5-flash", + status="completed", + steps=[], + usage=dict(INTERACTIONS_USAGE_BLOCK), + ) + await logging_obj.async_log_background_interaction_completion(result=completed) + + assert logging_obj.model_call_details["response_cost"] > 0 + assert logging_obj.model_call_details["standard_logging_object"]["total_tokens"] == 175 + + @pytest.mark.parametrize( "call_type", ["aget", "get", "aget_interaction", "adelete_interaction", "acancel_interaction"], diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index f289148101a..02a586a3dcc 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -604,6 +604,49 @@ async def test_track_cost_callback_skips_when_no_standard_logging_object(): mock_proxy_logging.failed_tracking_alert.assert_not_called() +@pytest.mark.asyncio +async def test_track_cost_callback_defers_in_progress_background_interaction(): + """ + A background=true interaction create returns in_progress with no usage + block, so its success event has a model but no standard_logging_object. + The callback must skip quietly (billing happens later via the background + poll task) instead of raising 'Cost tracking failed' and alerting. + """ + from litellm.types.interactions import InteractionsAPIResponse + + logger = _ProxyDBLogger() + + kwargs = { + "call_type": "acreate_interaction", + "model": "gemini/gemini-3-flash-preview", + "litellm_call_id": "test-call-id", + "litellm_params": {}, + "stream": False, + } + in_progress_response = InteractionsAPIResponse( + id="interactions/bg-abc", + model="gemini-3-flash-preview", + status="in_progress", + ) + + with patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging: + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=in_progress_response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + mock_proxy_logging.db_spend_update_writer.update_database.assert_not_called() + mock_proxy_logging.failed_tracking_alert.assert_not_called() + + @pytest.mark.asyncio async def test_async_post_call_failure_hook_propagates_trace_id_from_logging_obj(): """ diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index c76439e5ad1..42fef833944 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -3512,3 +3512,32 @@ def test_completion_cost_bills_interactions_api_response(): ) assert cost == pytest.approx(expected) assert cost > 0 + + +def test_completion_cost_bills_interactions_video_output_at_video_rate(): + from litellm.types.interactions import InteractionsAPIResponse + + model_info = litellm.get_model_info(model="gemini-omni-flash-preview", custom_llm_provider="gemini") + video_tokens = 5792 * 8 + response = InteractionsAPIResponse( + id="interactions/video123", + model="gemini-omni-flash-preview", + status="completed", + steps=[], + usage={ + "total_tokens": 10 + video_tokens, + "total_input_tokens": 10, + "input_tokens_by_modality": [{"modality": "text", "tokens": 10}], + "total_cached_tokens": 0, + "total_output_tokens": video_tokens, + "output_tokens_by_modality": [{"modality": "video", "tokens": video_tokens}], + "total_tool_use_tokens": 0, + "total_thought_tokens": 0, + }, + ) + + cost = completion_cost(completion_response=response, custom_llm_provider="gemini") + + expected = 10 * model_info["input_cost_per_token"] + video_tokens * model_info["output_cost_per_video_token"] + assert model_info["output_cost_per_video_token"] != model_info["output_cost_per_token"] + assert cost == pytest.approx(expected) From bbb156577245a9a27ae529276bcaa771e3d06266 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 15 Jul 2026 06:32:16 -0700 Subject: [PATCH 005/542] fix(interactions): hold budget reservation open until background interaction completes --- .../interactions/background_cost_polling.py | 26 ++++++ .../proxy/hooks/proxy_track_cost_callback.py | 14 ++- .../test_background_cost_polling.py | 56 +++++++++++- .../hooks/test_proxy_track_cost_callback.py | 85 +++++++++++++++++++ 4 files changed, 176 insertions(+), 5 deletions(-) diff --git a/litellm/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py index fefbf8715cb..846b26f28b9 100644 --- a/litellm/interactions/background_cost_polling.py +++ b/litellm/interactions/background_cost_polling.py @@ -22,6 +22,7 @@ from litellm.constants import ( BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS, BACKGROUND_INTERACTION_COST_POLLING_ENABLED, ) +from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs from litellm.types.interactions import InteractionsAPIResponse if TYPE_CHECKING: @@ -87,12 +88,37 @@ async def poll_and_log_background_interaction_cost( continue if response.usage is not None: await context.logging_obj.async_log_background_interaction_completion(result=response) + else: + await _release_open_budget_reservation(logging_obj=context.logging_obj) return verbose_logger.warning( "Gave up cost polling for background interaction %s after %ss; its usage will not be tracked", context.interaction_id, context.timeout_seconds, ) + await _release_open_budget_reservation(logging_obj=context.logging_obj) + + +async def _release_open_budget_reservation(logging_obj: "LiteLLMLoggingObj") -> None: + """ + The proxy keeps the pre-call budget reservation open for an in-progress + background interaction so concurrent creates cannot stack past the budget. + The completion success event reconciles it to the actual cost; when the + interaction terminates without billable usage (or polling gives up), no + such event fires, so the poller must release the reservation here or the + spend counters stay pinned at the estimated cost. + """ + metadata = get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details) + budget_reservation = metadata.get("user_api_key_budget_reservation") + if not isinstance(budget_reservation, dict): + return + + from litellm.proxy.spend_tracking.budget_reservation import release_budget_reservation + + try: + await release_budget_reservation(budget_reservation=budget_reservation) + except Exception: # noqa: BLE001 # a failed release must not crash the poll task; counters expire via TTL + verbose_logger.exception("Failed to release budget reservation for an unbilled background interaction") _ACTIVE_POLL_TASKS: set["asyncio.Task[None]"] = set() # mutable-ok: asyncio requires strong refs to running tasks diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 6359eae0fa4..a645bb9929f 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -5,6 +5,7 @@ from typing import Any, List, Optional, Union, cast import litellm from litellm._logging import verbose_proxy_logger +from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, @@ -278,13 +279,20 @@ class _ProxyDBLogger(CustomLogger): elif budget_reservation is not None: await _release_budget_reservation(budget_reservation=budget_reservation) else: - await _release_budget_reservation(budget_reservation=budget_reservation) if _is_unbilled_in_progress_interaction(completion_response): + if BACKGROUND_INTERACTION_COST_POLLING_ENABLED: + verbose_proxy_logger.debug( + "Cost tracking deferred for in-progress background interaction; " + "the budget reservation stays open until the poll task logs the final usage" + ) + return + await _release_budget_reservation(budget_reservation=budget_reservation) verbose_proxy_logger.debug( - "Cost tracking deferred for in-progress background interaction; " - "a poll task logs the final usage once it completes" + "Background interaction cost polling is disabled; released the budget " + "reservation for an in-progress interaction that will not be billed" ) return + await _release_budget_reservation(budget_reservation=budget_reservation) # Non-model call types (health checks, afile_delete) have no model or standard_logging_object. # Use .get() for "stream" to avoid KeyError on health checks. # WS session wrappers (_aresponses_websocket, _arealtime) also reach here with diff --git a/tests/test_litellm/interactions/test_background_cost_polling.py b/tests/test_litellm/interactions/test_background_cost_polling.py index aac97826b02..24abeec1053 100644 --- a/tests/test_litellm/interactions/test_background_cost_polling.py +++ b/tests/test_litellm/interactions/test_background_cost_polling.py @@ -1,5 +1,6 @@ import asyncio import time +from typing import Optional import pytest @@ -23,7 +24,10 @@ USAGE_BLOCK = { } -def _logging_obj(call_type: str = "acreate_interaction") -> LitellmLogging: +def _logging_obj( + call_type: str = "acreate_interaction", + litellm_params: Optional[dict] = None, +) -> LitellmLogging: logging_obj = LitellmLogging( model="gemini-2.5-flash", messages=[], @@ -34,7 +38,7 @@ def _logging_obj(call_type: str = "acreate_interaction") -> LitellmLogging: function_id="bg-interactions-fn-id", ) logging_obj.update_environment_variables( - litellm_params={}, + litellm_params=litellm_params or {}, optional_params={}, model="gemini-2.5-flash", custom_llm_provider="gemini", @@ -43,6 +47,14 @@ def _logging_obj(call_type: str = "acreate_interaction") -> LitellmLogging: return logging_obj +def _reservation() -> dict: + return {"reserved_cost": 0.05, "entries": [], "finalized": False, "input_cost": 0.001} + + +def _logging_obj_with_reservation(reservation: dict) -> LitellmLogging: + return _logging_obj(litellm_params={"metadata": {"user_api_key_budget_reservation": reservation}}) + + def _context(logging_obj: LitellmLogging, timeout_seconds: float = 1.0) -> BackgroundInteractionPollContext: return BackgroundInteractionPollContext( interaction_id="interactions/bg-abc", @@ -118,6 +130,46 @@ async def test_poller_gives_up_after_timeout_without_billing(): assert logging_obj.model_call_details.get("response_cost") is None +@pytest.mark.asyncio +async def test_poller_releases_budget_reservation_when_interaction_ends_without_usage(): + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + fetch, _ = _fetch_sequence(_response("failed", with_usage=False)) + + await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) + + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_poller_releases_budget_reservation_on_timeout_give_up(): + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False)) + + await poll_and_log_background_interaction_cost( + _context(logging_obj, timeout_seconds=0.01), + fetch_interaction=fetch, + ) + + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_poller_leaves_reservation_reconciliation_to_the_completion_event(): + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + fetch, _ = _fetch_sequence( + _response("in_progress", with_usage=False), + _response("completed", with_usage=True), + ) + + await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) + + assert logging_obj.model_call_details["response_cost"] > 0 + assert reservation["finalized"] is False + + @pytest.mark.asyncio async def test_poller_retries_after_fetch_error_and_still_bills(): logging_obj = _logging_obj() diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 02a586a3dcc..725e7f22f7d 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -647,6 +647,91 @@ async def test_track_cost_callback_defers_in_progress_background_interaction(): mock_proxy_logging.failed_tracking_alert.assert_not_called() +def _in_progress_interaction_kwargs(reservation: dict) -> dict: + return { + "call_type": "acreate_interaction", + "model": "gemini/gemini-3-flash-preview", + "litellm_call_id": "test-call-id", + "litellm_params": {"metadata": {"user_api_key_budget_reservation": reservation}}, + "stream": False, + } + + +@pytest.mark.asyncio +async def test_track_cost_callback_keeps_reservation_open_for_in_progress_background_interaction(): + """ + The pre-call budget reservation must stay open while a background + interaction is in flight, so concurrent creates cannot stack past the + budget; the poll task's completion event reconciles it to the actual cost. + """ + from litellm.types.interactions import InteractionsAPIResponse + + logger = _ProxyDBLogger() + reservation = {"reserved_cost": 0.05, "entries": [], "finalized": False} + in_progress_response = InteractionsAPIResponse( + id="interactions/bg-abc", + model="gemini-3-flash-preview", + status="in_progress", + ) + + with patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging: + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=_in_progress_interaction_kwargs(reservation), + completion_response=in_progress_response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert reservation["finalized"] is False + mock_proxy_logging.failed_tracking_alert.assert_not_called() + + +@pytest.mark.asyncio +async def test_track_cost_callback_releases_reservation_for_in_progress_interaction_when_polling_disabled( + monkeypatch, +): + """ + With the poll task kill switch off nothing will ever reconcile the + reservation, so the callback must release it or the spend counters stay + pinned at the estimated cost forever. + """ + import litellm.proxy.hooks.proxy_track_cost_callback as callback_module + from litellm.types.interactions import InteractionsAPIResponse + + monkeypatch.setattr(callback_module, "BACKGROUND_INTERACTION_COST_POLLING_ENABLED", False) + + logger = _ProxyDBLogger() + reservation = {"reserved_cost": 0.05, "entries": [], "finalized": False} + in_progress_response = InteractionsAPIResponse( + id="interactions/bg-abc", + model="gemini-3-flash-preview", + status="in_progress", + ) + + with patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging: + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=_in_progress_interaction_kwargs(reservation), + completion_response=in_progress_response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert reservation["finalized"] is True + mock_proxy_logging.failed_tracking_alert.assert_not_called() + + @pytest.mark.asyncio async def test_async_post_call_failure_hook_propagates_trace_id_from_logging_obj(): """ From 194dca7dd0d67ef23795fc13668092af1e2a86ef Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 15 Jul 2026 07:15:50 -0700 Subject: [PATCH 006/542] fix(interactions): settle pending background interaction billing before delete --- .../interactions/background_cost_polling.py | 91 ++++++++++++- litellm/interactions/main.py | 3 + .../test_background_cost_polling.py | 123 ++++++++++++++++++ 3 files changed, 211 insertions(+), 6 deletions(-) diff --git a/litellm/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py index 846b26f28b9..ed30d2573bd 100644 --- a/litellm/interactions/background_cost_polling.py +++ b/litellm/interactions/background_cost_polling.py @@ -9,6 +9,16 @@ interaction). The create call is therefore the only place that can own billing: it schedules a poll task that fetches the interaction until it reaches a terminal status and logs the final usage as a single success event attributed to the original request. + +Deleting an interaction makes every subsequent poll fail, which would let a +caller retrieve the completed output themselves and then delete it before the +poll task settles, leaving the work unbilled and the budget reservation +refunded at the poll timeout. ``adelete`` therefore settles any pending poll +for the interaction before dispatching the delete: it fetches the current +state with the create's credentials, bills it if it is terminal with usage, +and releases the reservation otherwise. A settlement gate on the create's +logging object makes the poll task and the delete path mutually exclusive, so +the interaction is billed exactly once no matter who settles first. """ import asyncio @@ -65,6 +75,25 @@ def _poll_intervals(initial: float, maximum: float, timeout: float) -> Iterator[ interval = min(interval * 2, maximum) +_SETTLED_KEY = "background_interaction_settled" + + +def _is_settled(logging_obj: "LiteLLMLoggingObj") -> bool: + return logging_obj.model_call_details.get(_SETTLED_KEY) is True + + +def _claim_settlement(logging_obj: "LiteLLMLoggingObj") -> bool: + """ + Exactly-once gate between the poll task and the delete-time settlement: + both run on the same event loop and neither awaits between reading and + setting the flag, so whichever claims first owns billing or release. + """ + if _is_settled(logging_obj): + return False + logging_obj.model_call_details[_SETTLED_KEY] = True + return True + + async def poll_and_log_background_interaction_cost( context: BackgroundInteractionPollContext, fetch_interaction: FetchInteraction = _fetch_interaction, @@ -75,6 +104,8 @@ async def poll_and_log_background_interaction_cost( timeout=context.timeout_seconds, ): await asyncio.sleep(interval) + if _is_settled(context.logging_obj): + return try: response = await fetch_interaction(context) except Exception as e: # noqa: BLE001 # any fetch error must not kill the billing poll loop @@ -86,11 +117,15 @@ async def poll_and_log_background_interaction_cost( continue if response.status not in _TERMINAL_STATUSES: continue + if not _claim_settlement(context.logging_obj): + return if response.usage is not None: await context.logging_obj.async_log_background_interaction_completion(result=response) else: await _release_open_budget_reservation(logging_obj=context.logging_obj) return + if not _claim_settlement(context.logging_obj): + return verbose_logger.warning( "Gave up cost polling for background interaction %s after %ss; its usage will not be tracked", context.interaction_id, @@ -104,9 +139,10 @@ async def _release_open_budget_reservation(logging_obj: "LiteLLMLoggingObj") -> The proxy keeps the pre-call budget reservation open for an in-progress background interaction so concurrent creates cannot stack past the budget. The completion success event reconciles it to the actual cost; when the - interaction terminates without billable usage (or polling gives up), no - such event fires, so the poller must release the reservation here or the - spend counters stay pinned at the estimated cost. + interaction terminates without billable usage (or polling gives up, or it + is deleted before settling), no such event fires, so whoever claims the + settlement must release the reservation here or the spend counters stay + pinned at the estimated cost. """ metadata = get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details) budget_reservation = metadata.get("user_api_key_budget_reservation") @@ -121,7 +157,21 @@ async def _release_open_budget_reservation(logging_obj: "LiteLLMLoggingObj") -> verbose_logger.exception("Failed to release budget reservation for an unbilled background interaction") -_ACTIVE_POLL_TASKS: set["asyncio.Task[None]"] = set() # mutable-ok: asyncio requires strong refs to running tasks +@dataclass(frozen=True, slots=True) +class _ActiveBackgroundPoll: + task: "asyncio.Task[None]" + context: BackgroundInteractionPollContext + + +_ACTIVE_POLLS: dict[ + str, _ActiveBackgroundPoll +] = {} # mutable-ok: asyncio requires strong refs to running tasks, and delete settlement looks polls up by interaction id + + +def _discard_poll(interaction_id: str, task: "asyncio.Task[None]") -> None: + entry = _ACTIVE_POLLS.get(interaction_id) + if entry is not None and entry.task is task: + del _ACTIVE_POLLS[interaction_id] def maybe_schedule_background_interaction_cost_polling( @@ -152,6 +202,35 @@ def maybe_schedule_background_interaction_cost_polling( api_base=create_kwargs.get("api_base"), ) task = asyncio.create_task(poll_and_log_background_interaction_cost(context)) - _ACTIVE_POLL_TASKS.add(task) - task.add_done_callback(_ACTIVE_POLL_TASKS.discard) + _ACTIVE_POLLS[context.interaction_id] = _ActiveBackgroundPoll(task=task, context=context) + task.add_done_callback( + lambda finished, interaction_id=context.interaction_id: _discard_poll(interaction_id, finished) + ) return task + + +async def maybe_settle_background_interaction_before_delete( + interaction_id: str, + fetch_interaction: FetchInteraction = _fetch_interaction, +) -> None: + entry = _ACTIVE_POLLS.get(interaction_id) + if entry is None: + return + context = entry.context + try: + response = await fetch_interaction(context) + except Exception as e: # noqa: BLE001 # unfetchable pre-delete state settles by releasing the reservation + verbose_logger.debug( + "Could not fetch background interaction %s before delete, releasing its reservation: %s", + interaction_id, + e, + ) + if _claim_settlement(context.logging_obj): + await _release_open_budget_reservation(logging_obj=context.logging_obj) + return + if not _claim_settlement(context.logging_obj): + return + if response.status in _TERMINAL_STATUSES and response.usage is not None: + await context.logging_obj.async_log_background_interaction_completion(result=response) + return + await _release_open_budget_reservation(logging_obj=context.logging_obj) diff --git a/litellm/interactions/main.py b/litellm/interactions/main.py index 5a0ea9280b2..985c60c2cc2 100644 --- a/litellm/interactions/main.py +++ b/litellm/interactions/main.py @@ -41,6 +41,7 @@ import httpx import litellm from litellm.interactions.background_cost_polling import ( maybe_schedule_background_interaction_cost_polling, + maybe_settle_background_interaction_before_delete, ) from litellm.interactions.http_handler import interactions_http_handler from litellm.interactions.utils import ( @@ -474,6 +475,8 @@ async def adelete( loop = asyncio.get_event_loop() kwargs["adelete_interaction"] = True + await maybe_settle_background_interaction_before_delete(interaction_id=interaction_id) + func = partial( delete, interaction_id=interaction_id, diff --git a/tests/test_litellm/interactions/test_background_cost_polling.py b/tests/test_litellm/interactions/test_background_cost_polling.py index 24abeec1053..64436ebe4a7 100644 --- a/tests/test_litellm/interactions/test_background_cost_polling.py +++ b/tests/test_litellm/interactions/test_background_cost_polling.py @@ -5,8 +5,10 @@ from typing import Optional import pytest from litellm.interactions.background_cost_polling import ( + _SETTLED_KEY, BackgroundInteractionPollContext, maybe_schedule_background_interaction_cost_polling, + maybe_settle_background_interaction_before_delete, poll_and_log_background_interaction_cost, ) from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging @@ -221,6 +223,127 @@ async def test_schedule_skips_non_pollable_results(response, create_kwargs): assert task is None +def _register_poll(logging_obj: LitellmLogging, poll_fetch=None) -> asyncio.Task: + import litellm.interactions.background_cost_polling as bg + + if poll_fetch is None: + poll_fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False)) + context = _context(logging_obj) + task = asyncio.create_task(poll_and_log_background_interaction_cost(context, fetch_interaction=poll_fetch)) + bg._ACTIVE_POLLS[context.interaction_id] = bg._ActiveBackgroundPoll(task=task, context=context) + task.add_done_callback(lambda finished: bg._discard_poll(context.interaction_id, finished)) + return task + + +@pytest.mark.asyncio +async def test_delete_settlement_bills_pending_background_interaction(): + logging_obj = _logging_obj() + task = _register_poll(logging_obj) + fetch, calls = _fetch_sequence(_response("completed", with_usage=True)) + + await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", + fetch_interaction=fetch, + ) + + assert len(calls) == 1 + assert logging_obj.model_call_details["response_cost"] > 0 + assert logging_obj.model_call_details["standard_logging_object"]["total_tokens"] == 175 + await asyncio.wait_for(task, timeout=5) + + +@pytest.mark.asyncio +async def test_delete_settlement_releases_reservation_when_still_in_progress(): + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + task = _register_poll(logging_obj) + fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False)) + + await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", + fetch_interaction=fetch, + ) + + assert reservation["finalized"] is True + assert logging_obj.model_call_details.get("response_cost") is None + await asyncio.wait_for(task, timeout=5) + + +@pytest.mark.asyncio +async def test_delete_settlement_releases_reservation_when_prefetch_fails(): + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + task = _register_poll(logging_obj) + fetch, _ = _fetch_sequence(RuntimeError("interaction already deleted")) + + await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", + fetch_interaction=fetch, + ) + + assert reservation["finalized"] is True + assert logging_obj.model_call_details.get("response_cost") is None + await asyncio.wait_for(task, timeout=5) + + +@pytest.mark.asyncio +async def test_delete_settlement_ignores_interactions_without_pending_poll(): + fetch, calls = _fetch_sequence(_response("completed", with_usage=True)) + + await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/never-polled", + fetch_interaction=fetch, + ) + + assert calls == [] + + +@pytest.mark.asyncio +async def test_delete_settlement_noop_after_poll_task_finished(): + logging_obj = _logging_obj() + poll_fetch, _ = _fetch_sequence(_response("completed", with_usage=True)) + task = _register_poll(logging_obj, poll_fetch=poll_fetch) + await asyncio.wait_for(task, timeout=5) + assert logging_obj.model_call_details["response_cost"] > 0 + + settle_fetch, settle_calls = _fetch_sequence(_response("completed", with_usage=True)) + await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", + fetch_interaction=settle_fetch, + ) + + assert settle_calls == [] + + +@pytest.mark.asyncio +async def test_delete_settlement_does_not_rebill_when_gate_already_claimed(): + logging_obj = _logging_obj() + logging_obj.model_call_details[_SETTLED_KEY] = True + task = _register_poll(logging_obj) + fetch, calls = _fetch_sequence(_response("completed", with_usage=True)) + + await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", + fetch_interaction=fetch, + ) + + assert len(calls) == 1 + assert logging_obj.model_call_details.get("response_cost") is None + await asyncio.wait_for(task, timeout=5) + + +@pytest.mark.asyncio +async def test_poller_exits_without_billing_once_settled_elsewhere(): + logging_obj = _logging_obj() + logging_obj.model_call_details[_SETTLED_KEY] = True + fetch, calls = _fetch_sequence(_response("completed", with_usage=True)) + + await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) + + assert calls == [] + assert logging_obj.model_call_details.get("response_cost") is None + + @pytest.mark.asyncio async def test_schedule_respects_kill_switch(monkeypatch): import litellm.interactions.background_cost_polling as module From e458aa1230a1e23a06aa6ebff88561f5fa510dda Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 15 Jul 2026 08:12:49 -0700 Subject: [PATCH 007/542] fix(interactions): bill google_search grounding queries per query --- .../usage_object_transformation.py | 12 ++++++- .../test_usage_object_transformation.py | 29 +++++++++++++++ tests/test_litellm/test_cost_calculator.py | 36 +++++++++++++++++++ 3 files changed, 76 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py b/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py index 3f591180b77..23fafc4132d 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py +++ b/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py @@ -62,6 +62,14 @@ def _modality_token_sums(entries: Sequence[Mapping[str, Any]]) -> Mapping[str, i } +def _google_search_query_count(usage_object: Mapping[str, Any]) -> int: + return sum( + _token_count(entry.get("count")) + for entry in tuple(usage_object.get("grounding_tool_count") or ()) + if isinstance(entry, Mapping) and entry.get("type") == "google_search" + ) + + def _subtract_cached_from_input( input_sums: Mapping[str, int], cached_sums: Mapping[str, int], @@ -116,12 +124,14 @@ class InteractionsUsageObjectTransformation: completion_tokens = _token_count(usage_object.get("total_output_tokens")) + reasoning_tokens total_tokens = _token_count(usage_object.get("total_tokens")) or (prompt_tokens + completion_tokens) + web_search_requests = _google_search_query_count(usage_object) prompt_tokens_details = ( PromptTokensDetailsWrapper( cached_tokens=total_cached_tokens or None, + web_search_requests=web_search_requests or None, **input_sums, ) - if input_sums or total_cached_tokens + if input_sums or total_cached_tokens or web_search_requests else None ) completion_tokens_details = ( diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py index 01241cf260d..2d8092959ce 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py @@ -106,6 +106,35 @@ def test_tool_use_tokens_billed_as_input(): assert usage.prompt_tokens_details.text_tokens == 140 +def test_google_search_grounding_count_maps_to_web_search_requests(): + usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object( + { + "total_input_tokens": 103, + "input_tokens_by_modality": [{"modality": "text", "tokens": 103}], + "total_output_tokens": 226, + "total_thought_tokens": 351, + "grounding_tool_count": [ + {"type": "google_search", "count": 3}, + {"type": "url_context", "count": 2}, + ], + } + ) + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.web_search_requests == 3 + + +def test_no_grounding_leaves_web_search_requests_unset(): + usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object( + { + "total_input_tokens": 10, + "input_tokens_by_modality": [{"modality": "text", "tokens": 10}], + "total_output_tokens": 5, + } + ) + assert usage.prompt_tokens_details is not None + assert getattr(usage.prompt_tokens_details, "web_search_requests", None) is None + + def test_document_modality_folds_into_text(): usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object( { diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 42fef833944..b144c614474 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -3514,6 +3514,42 @@ def test_completion_cost_bills_interactions_api_response(): assert cost > 0 +def test_completion_cost_bills_interactions_google_search_per_query(): + from litellm.types.interactions import InteractionsAPIResponse + + model_info = litellm.get_model_info(model="gemini-3-flash-preview", custom_llm_provider="gemini") + response = InteractionsAPIResponse( + id="interactions/search123", + model="gemini-3-flash-preview", + status="completed", + steps=[], + usage={ + "total_tokens": 680, + "total_input_tokens": 103, + "input_tokens_by_modality": [{"modality": "text", "tokens": 103}], + "total_cached_tokens": 0, + "total_output_tokens": 226, + "total_tool_use_tokens": 0, + "total_thought_tokens": 351, + "grounding_tool_count": [{"type": "google_search", "count": 3}], + }, + ) + + cost = completion_cost(completion_response=response, custom_llm_provider="gemini") + + per_query_cost = model_info["search_context_cost_per_query"]["search_context_size_medium"] + reasoning_rate = model_info.get("output_cost_per_reasoning_token") or model_info["output_cost_per_token"] + expected = ( + 103 * model_info["input_cost_per_token"] + + 226 * model_info["output_cost_per_token"] + + 351 * reasoning_rate + + 3 * per_query_cost + ) + assert model_info.get("web_search_billing_unit") == "per_query" + assert cost == pytest.approx(expected) + assert cost > 3 * per_query_cost + + def test_completion_cost_bills_interactions_video_output_at_video_rate(): from litellm.types.interactions import InteractionsAPIResponse From b5c59b67874c2212d708f3f2e3d2b0ee6fead017 Mon Sep 17 00:00:00 2001 From: bhuvan2134686 Date: Mon, 27 Jul 2026 16:51:18 +1000 Subject: [PATCH 008/542] feat(providers): add SCX.ai as OpenAI-compatible provider --- litellm/constants.py | 2 + .../get_llm_provider_logic.py | 3 + litellm/llms/openai_like/providers.json | 12 ++ litellm/types/utils.py | 1 + .../llms/openai_like/test_scx_ai_provider.py | 130 ++++++++++++++++++ 5 files changed, 148 insertions(+) create mode 100644 tests/test_litellm/llms/openai_like/test_scx_ai_provider.py diff --git a/litellm/constants.py b/litellm/constants.py index a9edf135731..cc6d9e2b1e0 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -728,6 +728,7 @@ openai_compatible_endpoints: List = [ "https://api.libertai.io/v1", "https://pinstripes.io/v1", "https://api.meta.ai/v1", + "https://api.scx.ai/v1", ] @@ -795,6 +796,7 @@ openai_compatible_providers: List = [ "pinstripes", # Pinstripes - JSON-configured provider "darkbloom", "meta", # Meta Model API (Muse Spark) - JSON-configured provider + "scx-ai", ] openai_text_completion_compatible_providers: List = [ # providers that support `/v1/completions` "together_ai", diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 487a7b7e25f..52a4f866fbd 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -349,6 +349,9 @@ def get_llm_provider( elif endpoint == "https://api.meta.ai/v1": custom_llm_provider = "meta" dynamic_api_key = get_secret_str("META_API_KEY") + elif endpoint == "https://api.scx.ai/v1": + custom_llm_provider = "scx-ai" + dynamic_api_key = get_secret_str("SCX_API_KEY") if api_base is not None and not isinstance(api_base, str): raise Exception("api base needs to be a string. api_base={}".format(api_base)) diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 164100d4194..d796d140878 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -183,5 +183,17 @@ "max_completion_tokens": "max_tokens" }, "supported_endpoints": ["/v1/chat/completions", "/v1/responses", "/v1/embeddings"] + }, + "scx-ai": { + "base_url": "https://api.scx.ai/v1", + "api_key_env": "SCX_API_KEY", + "api_base_env": "SCX_API_BASE", + "param_mappings": { + "max_completion_tokens": "max_tokens" + }, + "constraints": { + "temperature_max": 1.0 + }, + "supported_endpoints": ["/v1/chat/completions"] } } diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e4dfac48141..feafce9ace1 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3508,6 +3508,7 @@ class LlmProviders(str, Enum): TENSORMESH = "tensormesh" LIBERTAI = "libertai" PINSTRIPES = "pinstripes" + SCX_AI = "scx-ai" DARKBLOOM = "darkbloom" META = "meta" LITELLM_AGENT = "litellm_agent" diff --git a/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py b/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py new file mode 100644 index 00000000000..bebd079646e --- /dev/null +++ b/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py @@ -0,0 +1,130 @@ +""" +Tests for SCX.ai provider configuration and integration. +""" + +import litellm + + +class TestSCXAIProviderConfig: + def test_scx_ai_in_provider_list(self): + from litellm import LlmProviders + + assert hasattr(LlmProviders, "SCX_AI") + assert LlmProviders.SCX_AI.value == "scx-ai" + assert "scx-ai" in litellm.provider_list + + def test_scx_ai_json_config_exists(self): + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + assert JSONProviderRegistry.exists("scx-ai") + + scx = JSONProviderRegistry.get("scx-ai") + assert scx is not None + assert scx.base_url == "https://api.scx.ai/v1" + assert scx.api_key_env == "SCX_API_KEY" + assert scx.param_mappings.get("max_completion_tokens") == "max_tokens" + assert scx.constraints.get("temperature_max") == 1.0 + + def test_scx_ai_in_openai_compatible_providers(self): + from litellm.constants import openai_compatible_providers + + assert "scx-ai" in openai_compatible_providers + + def test_scx_ai_provider_resolution(self): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, api_key, api_base = get_llm_provider( + model="scx-ai/gpt-oss-120b", + custom_llm_provider=None, + api_base=None, + api_key=None, + ) + + assert model == "gpt-oss-120b" + assert provider == "scx-ai" + assert api_base == "https://api.scx.ai/v1" + + def test_scx_ai_api_base_override(self): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, api_key, api_base = get_llm_provider( + model="scx-ai/gpt-oss-120b", + custom_llm_provider=None, + api_base="https://custom.scx.ai/v1", + api_key="sk-test", + ) + + assert provider == "scx-ai" + assert api_base == "https://custom.scx.ai/v1" + assert api_key == "sk-test" + + def test_scx_ai_url_autodetection(self): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, api_key, api_base = get_llm_provider( + model="gpt-oss-120b", + custom_llm_provider=None, + api_base="https://api.scx.ai/v1", + api_key=None, + ) + assert provider == "scx-ai" + assert api_base == "https://api.scx.ai/v1" + + def test_scx_ai_temperature_clamped_to_max(self): + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("scx-ai") + assert provider is not None + config = create_config_class(provider)() + + optional_params = config.map_openai_params( + non_default_params={"temperature": 1.7}, + optional_params={}, + model="gpt-oss-120b", + drop_params=False, + ) + assert optional_params["temperature"] == 1.0 + + optional_params = config.map_openai_params( + non_default_params={"temperature": 0.4}, + optional_params={}, + model="gpt-oss-120b", + drop_params=False, + ) + assert optional_params["temperature"] == 0.4 + + def test_scx_ai_max_completion_tokens_mapped(self): + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("scx-ai") + assert provider is not None + config = create_config_class(provider)() + + optional_params = config.map_openai_params( + non_default_params={"max_completion_tokens": 256}, + optional_params={}, + model="gpt-oss-120b", + drop_params=False, + ) + assert optional_params["max_tokens"] == 256 + assert "max_completion_tokens" not in optional_params + + def test_scx_ai_router_config(self): + from litellm import Router + + router = Router( + model_list=[ + { + "model_name": "scx-chat", + "litellm_params": { + "model": "scx-ai/gpt-oss-120b", + "api_key": "test-key", + }, + } + ] + ) + + assert len(router.model_list) == 1 + assert router.model_list[0]["model_name"] == "scx-chat" From 912c7c7f42edecd099076cd8f279cc5621ad32e9 Mon Sep 17 00:00:00 2001 From: bhuvan2134686 Date: Mon, 27 Jul 2026 16:51:18 +1000 Subject: [PATCH 009/542] feat(providers): register scx-ai in the endpoint support matrix --- litellm/provider_endpoints_support_backup.json | 17 +++++++++++++++++ provider_endpoints_support.json | 17 +++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index dd7712aabca..a7ab8187bd3 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -2010,6 +2010,23 @@ "interactions": true } }, + "scx-ai": { + "display_name": "SCX.ai (`scx-ai`)", + "url": "https://docs.litellm.ai/docs/providers/scx-ai", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, "snowflake": { "display_name": "Snowflake (`snowflake`)", "url": "https://docs.litellm.ai/docs/providers/snowflake", diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 65db63dc045..53b6f6e9fa0 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -2244,6 +2244,23 @@ "interactions": true } }, + "scx-ai": { + "display_name": "SCX.ai (`scx-ai`)", + "url": "https://docs.litellm.ai/docs/providers/scx-ai", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, "snowflake": { "display_name": "Snowflake (`snowflake`)", "url": "https://docs.litellm.ai/docs/providers/snowflake", From 8f61073af80ded2e637ed0fdb8fd39bf5e73c607 Mon Sep 17 00:00:00 2001 From: bhuvan2134686 Date: Mon, 27 Jul 2026 16:51:18 +1000 Subject: [PATCH 010/542] feat(ui): add SCX.ai to the dashboard provider list with logo --- ui/litellm-dashboard/public/assets/logos/scx_ai.svg | 1 + .../src/components/provider_info_helpers.test.tsx | 11 +++++++++++ .../src/components/provider_info_helpers.tsx | 4 ++++ 3 files changed, 16 insertions(+) create mode 100644 ui/litellm-dashboard/public/assets/logos/scx_ai.svg diff --git a/ui/litellm-dashboard/public/assets/logos/scx_ai.svg b/ui/litellm-dashboard/public/assets/logos/scx_ai.svg new file mode 100644 index 00000000000..545176a945b --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/scx_ai.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index 777cdc62987..55983b20bf4 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -62,6 +62,17 @@ describe("provider_info_helpers", () => { expect(result.logo).toBe(providerLogoMap[Providers.Groq]); }); + it("should map scx-ai slug and SCX_AI enum key to the SCX.ai display name and logo", () => { + const fromSlug = getProviderLogoAndName("scx-ai"); + expect(fromSlug.displayName).toBe(Providers.SCX_AI); + expect(fromSlug.logo).toBe(providerLogoMap[Providers.SCX_AI]); + expect(fromSlug.logo).toBeTruthy(); + + const fromEnumKey = getProviderLogoAndName("SCX_AI"); + expect(fromEnumKey.displayName).toBe(Providers.SCX_AI); + expect(fromEnumKey.logo).toBe(providerLogoMap[Providers.SCX_AI]); + }); + it("should map bedrock_mantle slug to Bedrock Mantle display name and logo", () => { const result = getProviderLogoAndName("bedrock_mantle"); expect(result.displayName).toBe(Providers.BedrockMantle); diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index fa6b3c79230..4f5c0b6000b 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -50,6 +50,7 @@ import replicateLogo from "../../public/assets/logos/replicate.svg"; import runwayLogo from "../../public/assets/logos/runway.png"; import sambanovaLogo from "../../public/assets/logos/sambanova.svg"; import sapLogo from "../../public/assets/logos/sap.png"; +import scxAiLogo from "../../public/assets/logos/scx_ai.svg"; import snowflakeLogo from "../../public/assets/logos/snowflake.svg"; import sonioxLogo from "../../public/assets/logos/soniox.svg"; import togetheraiLogo from "../../public/assets/logos/togetherai.svg"; @@ -151,6 +152,7 @@ export enum Providers { SAGEMAKER_LEGACY = "Sagemaker", Sambanova = "Sambanova", SAP = "SAP Generative AI Hub", + SCX_AI = "SCX.ai", Snowflake = "Snowflake", Soniox = "Soniox", TEXT_COMPLETION_CODESTRAL = "Text-Completion-Codestral", @@ -260,6 +262,7 @@ export const provider_map: Record = { SageMaker: "sagemaker_chat", Sambanova: "sambanova", SAP: "sap", + SCX_AI: "scx-ai", Snowflake: "snowflake", Soniox: "soniox", TEXT_COMPLETION_CODESTRAL: "text-completion-codestral", @@ -351,6 +354,7 @@ export const providerLogoMap: Partial> = { [Providers.SAGEMAKER_LEGACY]: bedrockLogo.src, [Providers.Sambanova]: sambanovaLogo.src, [Providers.SAP]: sapLogo.src, + [Providers.SCX_AI]: scxAiLogo.src, [Providers.Snowflake]: snowflakeLogo.src, [Providers.Soniox]: sonioxLogo.src, [Providers.TEXT_COMPLETION_CODESTRAL]: mistralLogo.src, From 7f48431e22b4664d5fa0035779393e4242c3dc8e Mon Sep 17 00:00:00 2001 From: bhuvan2134686 Date: Mon, 27 Jul 2026 17:21:58 +1000 Subject: [PATCH 011/542] feat(models): add pricing and metadata for 5 scx-ai models --- ...odel_prices_and_context_window_backup.json | 66 +++++++++++++++++++ model_prices_and_context_window.json | 66 +++++++++++++++++++ .../llms/openai_like/test_scx_ai_provider.py | 40 +++++++++++ 3 files changed, 172 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2cef600ea32..2f47643da50 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -33546,6 +33546,72 @@ "supports_vision": true, "source": "https://cloud.sambanova.ai/plans/pricing" }, + "scx-ai/gemma-4-31B-it": { + "input_cost_per_token": 3e-07, + "litellm_provider": "scx-ai", + "max_input_tokens": 131000, + "max_tokens": 131000, + "mode": "chat", + "output_cost_per_token": 9.1e-07, + "source": "https://scx.ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "scx-ai/gpt-oss-120b": { + "input_cost_per_token": 1.7e-07, + "litellm_provider": "scx-ai", + "max_input_tokens": 131000, + "max_tokens": 131000, + "mode": "chat", + "output_cost_per_token": 5.5e-07, + "source": "https://scx.ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "scx-ai/Llama-4-Maverick-17B-128E-Instruct": { + "input_cost_per_token": 5.3e-07, + "litellm_provider": "scx-ai", + "max_input_tokens": 131000, + "max_tokens": 131000, + "mode": "chat", + "output_cost_per_token": 1.62e-06, + "source": "https://scx.ai/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "scx-ai/MiniMax-M2.7": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 4.8e-07, + "litellm_provider": "scx-ai", + "max_input_tokens": 192000, + "max_tokens": 192000, + "mode": "chat", + "output_cost_per_token": 1.79e-06, + "source": "https://scx.ai/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "scx-ai/Qwen3-32B": { + "input_cost_per_token": 3.6e-07, + "litellm_provider": "scx-ai", + "max_input_tokens": 33000, + "max_tokens": 33000, + "mode": "chat", + "output_cost_per_token": 8.7e-07, + "source": "https://scx.ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "snowflake/claude-3-5-sonnet": { "litellm_provider": "snowflake", "max_input_tokens": 200000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index db28118d52b..8464aec769f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -33637,6 +33637,72 @@ "supports_vision": true, "source": "https://cloud.sambanova.ai/plans/pricing" }, + "scx-ai/gemma-4-31B-it": { + "input_cost_per_token": 3e-07, + "litellm_provider": "scx-ai", + "max_input_tokens": 131000, + "max_tokens": 131000, + "mode": "chat", + "output_cost_per_token": 9.1e-07, + "source": "https://scx.ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "scx-ai/gpt-oss-120b": { + "input_cost_per_token": 1.7e-07, + "litellm_provider": "scx-ai", + "max_input_tokens": 131000, + "max_tokens": 131000, + "mode": "chat", + "output_cost_per_token": 5.5e-07, + "source": "https://scx.ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "scx-ai/Llama-4-Maverick-17B-128E-Instruct": { + "input_cost_per_token": 5.3e-07, + "litellm_provider": "scx-ai", + "max_input_tokens": 131000, + "max_tokens": 131000, + "mode": "chat", + "output_cost_per_token": 1.62e-06, + "source": "https://scx.ai/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "scx-ai/MiniMax-M2.7": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 4.8e-07, + "litellm_provider": "scx-ai", + "max_input_tokens": 192000, + "max_tokens": 192000, + "mode": "chat", + "output_cost_per_token": 1.79e-06, + "source": "https://scx.ai/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "scx-ai/Qwen3-32B": { + "input_cost_per_token": 3.6e-07, + "litellm_provider": "scx-ai", + "max_input_tokens": 33000, + "max_tokens": 33000, + "mode": "chat", + "output_cost_per_token": 8.7e-07, + "source": "https://scx.ai/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "snowflake/claude-3-5-sonnet": { "litellm_provider": "snowflake", "max_input_tokens": 200000, diff --git a/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py b/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py index bebd079646e..7994a133a83 100644 --- a/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py +++ b/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py @@ -128,3 +128,43 @@ class TestSCXAIProviderConfig: assert len(router.model_list) == 1 assert router.model_list[0]["model_name"] == "scx-chat" + + +class TestSCXAIModelMetadata: + SCX_MODELS = ( + "scx-ai/Llama-4-Maverick-17B-128E-Instruct", + "scx-ai/gemma-4-31B-it", + "scx-ai/Qwen3-32B", + "scx-ai/MiniMax-M2.7", + "scx-ai/gpt-oss-120b", + ) + VISION_MODELS = ("scx-ai/Llama-4-Maverick-17B-128E-Instruct", "scx-ai/gemma-4-31B-it") + + @staticmethod + def _load(path_parts): + import json + from pathlib import Path + + json_path = Path(__file__).parents[4].joinpath(*path_parts) + with open(json_path) as f: + return json.load(f) + + def test_scx_ai_models_registered_with_correct_metadata(self): + model_cost = self._load(("model_prices_and_context_window.json",)) + for model in self.SCX_MODELS: + info = model_cost.get(model) + assert info is not None, f"{model} missing from model_prices_and_context_window.json" + assert info["litellm_provider"] == "scx-ai" + assert info["mode"] == "chat" + assert info["input_cost_per_token"] > 0 + assert info["output_cost_per_token"] > 0 + assert info["supports_function_calling"] is True + assert info["supports_tool_choice"] is True + assert info.get("supports_vision", False) is (model in self.VISION_MODELS) + + def test_scx_ai_models_synced_to_backup(self): + model_cost = self._load(("model_prices_and_context_window.json",)) + backup = self._load(("litellm", "model_prices_and_context_window_backup.json")) + for model in self.SCX_MODELS: + assert model in backup, f"{model} missing from backup json" + assert backup[model] == model_cost[model], f"{model} differs between root and backup json" From b8e2848fcd7e3be397de758de9a0d2061e6a10d5 Mon Sep 17 00:00:00 2001 From: bhuvan2134686 Date: Mon, 27 Jul 2026 17:21:58 +1000 Subject: [PATCH 012/542] fix(ui): make SCX.ai selectable in the Add Model provider dropdown --- .../provider_create_fields.json | 28 +++++++++++++++++++ .../llms/openai_like/test_scx_ai_provider.py | 27 ++++++++++++++++++ 2 files changed, 55 insertions(+) diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index fcc6aac1c14..36db02d814e 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -2660,6 +2660,34 @@ ], "default_model_placeholder": "sap/gpt-4" }, + { + "provider": "SCX_AI", + "provider_display_name": "SCX.ai", + "litellm_provider": "scx-ai", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://api.scx.ai/v1", + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "scx-ai/gpt-oss-120b" + }, { "provider": "Snowflake", "provider_display_name": "Snowflake", diff --git a/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py b/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py index 7994a133a83..92921067728 100644 --- a/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py +++ b/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py @@ -168,3 +168,30 @@ class TestSCXAIModelMetadata: for model in self.SCX_MODELS: assert model in backup, f"{model} missing from backup json" assert backup[model] == model_cost[model], f"{model} differs between root and backup json" + + +class TestSCXAIDashboardRegistration: + @staticmethod + def _provider_create_fields(): + import json + from pathlib import Path + + import litellm + + path = Path(litellm.__file__).parent / "proxy" / "public_endpoints" / "provider_create_fields.json" + with open(path) as f: + return json.load(f) + + def test_scx_ai_is_selectable_in_the_add_model_form(self): + entries = [e for e in self._provider_create_fields() if e["litellm_provider"] == "scx-ai"] + assert len(entries) == 1, "scx-ai must appear exactly once in provider_create_fields.json" + + entry = entries[0] + assert entry["provider"] == "SCX_AI" + assert entry["provider_display_name"] == "SCX.ai" + assert entry["default_model_placeholder"].startswith("scx-ai/") + + fields = {f["key"]: f for f in entry["credential_fields"]} + assert fields["api_key"]["required"] is True + assert fields["api_key"]["field_type"] == "password" + assert fields["api_base"]["required"] is False From cabbc7ebfb9e3aa84649ab184bbab5c692f40baf Mon Sep 17 00:00:00 2001 From: bhuvan2134686 Date: Tue, 28 Jul 2026 09:07:01 +1000 Subject: [PATCH 013/542] feat(ui): default the SCX.ai Add Model placeholder to MiniMax-M2.7 --- litellm/proxy/public_endpoints/provider_create_fields.json | 2 +- .../src/components/provider_info_helpers.test.tsx | 4 ++++ ui/litellm-dashboard/src/components/provider_info_helpers.tsx | 2 ++ 3 files changed, 7 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 36db02d814e..4cfe9e4ef5e 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -2686,7 +2686,7 @@ "default_value": null } ], - "default_model_placeholder": "scx-ai/gpt-oss-120b" + "default_model_placeholder": "scx-ai/MiniMax-M2.7" }, { "provider": "Snowflake", diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index 55983b20bf4..8a4c9dfd24d 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -172,6 +172,10 @@ describe("provider_info_helpers", () => { expect(getPlaceholder(Providers.Vertex_AI)).toBe("gemini-pro"); }); + it("should return an scx-ai model placeholder for SCX_AI provider", () => { + expect(getPlaceholder(Providers.SCX_AI)).toBe("scx-ai/MiniMax-M2.7"); + }); + it("should return claude-3-opus placeholder for Anthropic provider", () => { expect(getPlaceholder(Providers.Anthropic)).toBe("claude-3-opus"); }); diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index 4f5c0b6000b..fa1f9c11eb3 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -443,6 +443,8 @@ export const getPlaceholder = (selectedProvider: string): string => { return "cursor/claude-4-sonnet"; } else if (selectedProvider === Providers.ZAI) { return "zai/glm-4.5"; + } else if (selectedProvider === Providers.SCX_AI) { + return "scx-ai/MiniMax-M2.7"; } else { return "gpt-3.5-turbo"; } From 7e1f44f0cc15e639d8fe95fc8831c5a0b25cc76b Mon Sep 17 00:00:00 2001 From: mateo Date: Thu, 30 Jul 2026 02:38:28 +0000 Subject: [PATCH 014/542] 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 015/542] 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 016/542] 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 abf7dab0c2822edf8c3b2bc78618e62e5e6941f8 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 31 Jul 2026 21:14:18 +0000 Subject: [PATCH 017/542] feat(azure_ai): support entra id / oauth auth on every azure ai foundry route Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/images/main.py | 27 ++-- litellm/llms/azure/common_utils.py | 29 +++- litellm/llms/azure_ai/common_utils.py | 47 +++++- .../image_edit/flux2_transformation.py | 18 +-- .../azure_ai/image_edit/mai_transformation.py | 19 +-- .../azure_ai/image_edit/transformation.py | 22 ++- .../document_intelligence/transformation.py | 16 +- litellm/llms/azure_ai/ocr/transformation.py | 10 +- .../llms/azure_ai/rerank/transformation.py | 8 +- .../llms/base_llm/rerank/transformation.py | 2 + litellm/llms/cohere/rerank/transformation.py | 2 + litellm/llms/custom_httpx/llm_http_handler.py | 1 + .../llms/dashscope/rerank/transformation.py | 2 + .../llms/deepinfra/rerank/transformation.py | 2 + .../fireworks_ai/rerank/transformation.py | 2 + .../llms/hosted_vllm/rerank/transformation.py | 2 + .../llms/huggingface/rerank/transformation.py | 2 + .../llms/infinity/rerank/transformation.py | 4 +- litellm/llms/jina_ai/rerank/transformation.py | 2 + .../llms/nvidia_nim/rerank/transformation.py | 2 + .../llms/vertex_ai/rerank/transformation.py | 8 +- litellm/llms/voyage/rerank/transformation.py | 2 + litellm/llms/watsonx/rerank/transformation.py | 2 + litellm/main.py | 4 +- .../llms/azure/test_azure_common_utils.py | 63 +++++++- ...test_azure_ai_image_edit_transformation.py | 33 ++++ .../test_mai_image_edit_transformation.py | 14 ++ .../test_azure_ai_rerank_transformation.py | 24 +++ .../llms/azure_ai/test_azure_ai_entra_auth.py | 153 ++++++++++++++++++ ...ocument_intelligence_ocr_transformation.py | 27 ++++ 30 files changed, 480 insertions(+), 69 deletions(-) create mode 100644 tests/test_litellm/llms/azure_ai/test_azure_ai_entra_auth.py diff --git a/litellm/images/main.py b/litellm/images/main.py index 17ea9aa177b..4c88eb52cd8 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -430,24 +430,31 @@ def image_generation( aimg_generation=aimg_generation, ) elif custom_llm_provider == "azure_ai": - from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo + from litellm.llms.azure_ai.common_utils import ( + AzureFoundryModelInfo, + get_azure_ai_auth_headers, + ) api_base = AzureFoundryModelInfo.get_api_base(api_base) api_key = AzureFoundryModelInfo.get_api_key(api_key) if extra_headers is not None: optional_params["extra_headers"] = extra_headers - default_headers = { + caller_set_auth = "api-key" in headers or "Authorization" in headers + auth_headers = ( + headers + if caller_set_auth + else get_azure_ai_auth_headers( + api_key=api_key, + litellm_params=litellm_params_dict, + api_key_header="api-key", + ) + ) + headers = { "Content-Type": "application/json", + **auth_headers, + **headers, } - # Only add api-key header if api_key is not None - # Azure AD authentication will use Authorization header instead - if api_key is not None: - default_headers["api-key"] = api_key - - for k, v in default_headers.items(): - if k not in headers: - headers[k] = v model_response = azure_chat_completions.image_generation( model=model, diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 91f5793e269..85100e595e6 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -2,6 +2,7 @@ import asyncio import hashlib import json import os +from functools import lru_cache from typing import Any, Callable, Dict, Literal, NamedTuple, Optional, Union, cast import httpx @@ -57,6 +58,24 @@ def process_azure_headers(headers: Union[httpx.Headers, dict]) -> dict: return {**llm_response_headers, **openai_headers} +@lru_cache(maxsize=128) +def _cached_entra_id_token_provider( + tenant_id: str, + client_id: str, + client_secret: str, + scope: str, +) -> Callable[[], str]: + """Build (once per credential set) a bearer token provider backed by a `ClientSecretCredential`. + + The credential caches the access token internally and only talks to Entra ID when it is close + to expiry, so reusing the provider keeps one AAD round trip per token lifetime instead of one + per request. + """ + from azure.identity import ClientSecretCredential, get_bearer_token_provider + + return get_bearer_token_provider(ClientSecretCredential(tenant_id, client_id, client_secret), scope) + + def get_azure_ad_token_from_entra_id( tenant_id: str, client_id: str, @@ -75,8 +94,6 @@ def get_azure_ad_token_from_entra_id( Returns: callable that returns a bearer token. """ - from azure.identity import ClientSecretCredential, get_bearer_token_provider - verbose_logger.debug("Getting Azure AD Token from Entra ID") if tenant_id.startswith("os.environ/"): @@ -102,9 +119,13 @@ def get_azure_ad_token_from_entra_id( ) if _tenant_id is None or _client_id is None or _client_secret is None: raise ValueError("tenant_id, client_id, and client_secret must be provided") - credential = ClientSecretCredential(_tenant_id, _client_id, _client_secret) - token_provider = get_bearer_token_provider(credential, scope) + token_provider = _cached_entra_id_token_provider( + tenant_id=_tenant_id, + client_id=_client_id, + client_secret=_client_secret, + scope=scope, + ) verbose_logger.debug("token_provider %s", token_provider) diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index 9965aa693c3..5dd5f5c78cc 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -1,9 +1,54 @@ +from collections.abc import Mapping from typing import List, Literal, Optional import litellm from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues +from litellm.types.router import GenericLiteLLMParams + +AzureAIApiKeyHeader = Literal["Authorization", "api-key", "Api-Key", "Ocp-Apim-Subscription-Key"] + + +def get_azure_ai_entra_token(litellm_params: Mapping[str, object] | None = None) -> str | None: + """ + Resolve an Entra ID / OAuth access token for an Azure AI Foundry deployment. + + Accepts the same credential set as the `azure` provider: service principal + (`tenant_id` / `client_id` / `client_secret`), a pre-fetched `azure_ad_token`, an OIDC + federated token, username/password, or `DefaultAzureCredential` / managed identity. + """ + from litellm.llms.azure.common_utils import get_azure_ad_token + + params = GenericLiteLLMParams.model_validate(litellm_params) if litellm_params else GenericLiteLLMParams() + + return get_azure_ad_token(params) + + +def get_azure_ai_auth_headers( + api_key: str | None, + litellm_params: Mapping[str, object] | None = None, + api_key_header: AzureAIApiKeyHeader = "Authorization", + api_key_env_var: str = "AZURE_AI_API_KEY", +) -> dict[str, str]: + """ + Build the auth headers for an Azure AI Foundry route. + + Prefers the API key when one is configured, and otherwise falls back to Entra ID / OAuth, + sending the access token as a bearer token. + """ + if api_key: + return {api_key_header: f"Bearer {api_key}" if api_key_header == "Authorization" else api_key} + + azure_ad_token = get_azure_ai_entra_token(litellm_params=litellm_params) + if azure_ad_token: + return {"Authorization": f"Bearer {azure_ad_token}"} + + raise ValueError( + f"Missing Azure AI credentials - set an API key (`api_key` or {api_key_env_var}), or Entra ID / OAuth " + "credentials (`tenant_id` + `client_id` + `client_secret`, `azure_ad_token`, an OIDC token, or a managed " + "identity with `litellm.enable_azure_ad_token_refresh = True`)" + ) class AzureFoundryModelInfo(BaseLLMModelInfo): @@ -43,7 +88,7 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): @staticmethod def get_api_key(api_key: Optional[str] = None) -> Optional[str]: - return api_key or litellm.api_key or litellm.openai_key or get_secret_str("AZURE_AI_API_KEY") + return api_key or litellm.api_key or get_secret_str("AZURE_AI_API_KEY") @property def api_version(self, api_version: Optional[str] = None) -> Optional[str]: diff --git a/litellm/llms/azure_ai/image_edit/flux2_transformation.py b/litellm/llms/azure_ai/image_edit/flux2_transformation.py index 1bc3bdcddc1..db429b85082 100644 --- a/litellm/llms/azure_ai/image_edit/flux2_transformation.py +++ b/litellm/llms/azure_ai/image_edit/flux2_transformation.py @@ -5,7 +5,10 @@ from typing import Any, Dict, Optional, Tuple from httpx._types import RequestFiles import litellm -from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo +from litellm.llms.azure_ai.common_utils import ( + AzureFoundryModelInfo, + get_azure_ai_auth_headers, +) from litellm.llms.azure_ai.image_generation.flux_transformation import ( AzureFoundryFluxImageGenerationConfig, ) @@ -71,16 +74,13 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): """ Validate Azure AI Foundry environment and set up authentication """ - api_key = AzureFoundryModelInfo.get_api_key(api_key) - - if not api_key: - raise ValueError( - f"Azure AI API key is required for model {model}. Set AZURE_AI_API_KEY environment variable or pass api_key parameter." - ) - headers.update( { - "Api-Key": api_key, + **get_azure_ai_auth_headers( + api_key=AzureFoundryModelInfo.get_api_key(api_key), + litellm_params=litellm_params, + api_key_header="Api-Key", + ), "Content-Type": "application/json", } ) diff --git a/litellm/llms/azure_ai/image_edit/mai_transformation.py b/litellm/llms/azure_ai/image_edit/mai_transformation.py index aa1092b0a53..fdac9912193 100644 --- a/litellm/llms/azure_ai/image_edit/mai_transformation.py +++ b/litellm/llms/azure_ai/image_edit/mai_transformation.py @@ -3,7 +3,10 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast import httpx from httpx._types import RequestFiles -from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo +from litellm.llms.azure_ai.common_utils import ( + AzureFoundryModelInfo, + get_azure_ai_auth_headers, +) from litellm.llms.azure_ai.image_generation.mai_transformation import ( AzureFoundryMAIImageGenerationConfig, ) @@ -91,15 +94,13 @@ class AzureFoundryMAIImageEditConfig(OpenAIImageEditConfig): litellm_params: Optional[dict] = None, api_base: Optional[str] = None, ) -> dict: - api_key = AzureFoundryModelInfo.get_api_key(api_key) - - if not api_key: - raise ValueError( - f"Azure AI API key is required for model {model}. " - "Set AZURE_AI_API_KEY environment variable or pass api_key parameter." + headers.update( + get_azure_ai_auth_headers( + api_key=AzureFoundryModelInfo.get_api_key(api_key), + litellm_params=litellm_params, + api_key_header="api-key", ) - - headers.update({"api-key": api_key}) + ) return headers def get_complete_url( diff --git a/litellm/llms/azure_ai/image_edit/transformation.py b/litellm/llms/azure_ai/image_edit/transformation.py index 5393a0ba55f..22b0b169faf 100644 --- a/litellm/llms/azure_ai/image_edit/transformation.py +++ b/litellm/llms/azure_ai/image_edit/transformation.py @@ -3,7 +3,10 @@ from typing import Optional import httpx import litellm -from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo +from litellm.llms.azure_ai.common_utils import ( + AzureFoundryModelInfo, + get_azure_ai_auth_headers, +) from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig from litellm.secret_managers.main import get_secret_str from litellm.utils import _add_path_to_api_base @@ -30,19 +33,14 @@ class AzureFoundryFluxImageEditConfig(OpenAIImageEditConfig): ) -> dict: """ Validate Azure AI Foundry environment and set up authentication - Uses Api-Key header format + Uses the Api-Key header format, or an Entra ID / OAuth bearer token when no key is set """ - api_key = AzureFoundryModelInfo.get_api_key(api_key) - - if not api_key: - raise ValueError( - f"Azure AI API key is required for model {model}. Set AZURE_AI_API_KEY environment variable or pass api_key parameter." - ) - headers.update( - { - "Api-Key": api_key, # Azure AI Foundry uses Api-Key header format - } + get_azure_ai_auth_headers( + api_key=AzureFoundryModelInfo.get_api_key(api_key), + litellm_params=litellm_params, + api_key_header="Api-Key", + ) ) return headers diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index 7d915892a28..4db4472dfc2 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -25,6 +25,7 @@ from litellm.constants import ( AZURE_OPERATION_POLLING_TIMEOUT, ) from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.llms.azure_ai.common_utils import get_azure_ai_auth_headers from litellm.llms.base_llm.ocr.transformation import ( BaseOCRConfig, DocumentType, @@ -215,17 +216,13 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): """ Validate environment and return headers for Azure Document Intelligence. - Authentication uses Ocp-Apim-Subscription-Key header. + Authentication uses the Ocp-Apim-Subscription-Key header, or an Entra ID / OAuth bearer + token when no subscription key is set. """ # Get API key from environment if not provided if api_key is None: api_key = get_secret_str(AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR) - if api_key is None: - raise ValueError( - "Missing Azure Document Intelligence API Key - Set AZURE_DOCUMENT_INTELLIGENCE_API_KEY environment variable or pass api_key parameter" - ) - # Validate API base/endpoint is provided if api_base is None: api_base = get_secret_str("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") @@ -236,7 +233,12 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): ) headers = { - "Ocp-Apim-Subscription-Key": api_key, + **get_azure_ai_auth_headers( + api_key=api_key, + litellm_params=litellm_params, + api_key_header="Ocp-Apim-Subscription-Key", + api_key_env_var=AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR, + ), "Content-Type": "application/json", **headers, } diff --git a/litellm/llms/azure_ai/ocr/transformation.py b/litellm/llms/azure_ai/ocr/transformation.py index a57e3e869cf..abc23008f6a 100644 --- a/litellm/llms/azure_ai/ocr/transformation.py +++ b/litellm/llms/azure_ai/ocr/transformation.py @@ -9,6 +9,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_convert_url_to_base64, convert_url_to_base64, ) +from litellm.llms.azure_ai.common_utils import get_azure_ai_auth_headers from litellm.llms.base_llm.ocr.transformation import DocumentType, OCRRequestData from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.secret_managers.main import get_secret_str @@ -47,17 +48,12 @@ class AzureAIOCRConfig(MistralOCRConfig): """ Validate environment and return headers for Azure AI OCR. - Azure AI uses Bearer token authentication with AZURE_AI_API_KEY. + Authenticates with AZURE_AI_API_KEY, or with an Entra ID / OAuth token when no key is set. """ # Get API key from environment if not provided if api_key is None: api_key = get_secret_str(AZURE_AI_OCR_API_KEY_ENV_VAR) - if api_key is None: - raise ValueError( - "Missing Azure AI API Key - A call is being made to Azure AI but no key is set either in the environment variables or via params" - ) - # Validate API base is provided if api_base is None: api_base = get_secret_str("AZURE_AI_API_BASE") @@ -68,7 +64,7 @@ class AzureAIOCRConfig(MistralOCRConfig): ) headers = { - "Authorization": f"Bearer {api_key}", + **get_azure_ai_auth_headers(api_key=api_key, litellm_params=litellm_params), "Content-Type": "application/json", **headers, } diff --git a/litellm/llms/azure_ai/rerank/transformation.py b/litellm/llms/azure_ai/rerank/transformation.py index 928f53bd485..24cdc67a23b 100644 --- a/litellm/llms/azure_ai/rerank/transformation.py +++ b/litellm/llms/azure_ai/rerank/transformation.py @@ -2,12 +2,14 @@ Translate between Cohere's `/rerank` format and Azure AI's `/rerank` format. """ +from collections.abc import Mapping from typing import Optional import httpx import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.azure_ai.common_utils import get_azure_ai_auth_headers from litellm.llms.cohere.rerank.transformation import CohereRerankConfig from litellm.secret_managers.main import get_secret_str from litellm.types.utils import RerankResponse @@ -64,15 +66,13 @@ class AzureAIRerankConfig(CohereRerankConfig): model: str, api_key: Optional[str] = None, optional_params: Optional[dict] = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: if api_key is None: api_key = get_secret_str("AZURE_AI_API_KEY") or litellm.azure_key - if api_key is None: - raise ValueError("Azure AI API key is required. Please set 'AZURE_AI_API_KEY' or 'litellm.azure_key'") - default_headers = { - "Authorization": f"Bearer {api_key}", + **get_azure_ai_auth_headers(api_key=api_key, litellm_params=litellm_params), "accept": "application/json", "content-type": "application/json", } diff --git a/litellm/llms/base_llm/rerank/transformation.py b/litellm/llms/base_llm/rerank/transformation.py index eac44ba85c5..e9f210fb31c 100644 --- a/litellm/llms/base_llm/rerank/transformation.py +++ b/litellm/llms/base_llm/rerank/transformation.py @@ -1,4 +1,5 @@ from abc import ABC, abstractmethod +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Union import httpx @@ -24,6 +25,7 @@ class BaseRerankConfig(ABC): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: pass diff --git a/litellm/llms/cohere/rerank/transformation.py b/litellm/llms/cohere/rerank/transformation.py index e494e89fbf2..86d9a3d224d 100644 --- a/litellm/llms/cohere/rerank/transformation.py +++ b/litellm/llms/cohere/rerank/transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Any, Dict, List, Union import httpx @@ -81,6 +82,7 @@ class CohereRerankConfig(BaseRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: if api_key is None: api_key = get_secret_str("COHERE_API_KEY") or get_secret_str("CO_API_KEY") or litellm.cohere_key diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ec1301e5923..e41fbea94d6 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1063,6 +1063,7 @@ class BaseLLMHTTPHandler: headers=headers or {}, model=model, optional_params=optional_rerank_params, + litellm_params=litellm_params, ) api_base = provider_config.get_complete_url( diff --git a/litellm/llms/dashscope/rerank/transformation.py b/litellm/llms/dashscope/rerank/transformation.py index 365e15fdd7a..b8c369b892c 100644 --- a/litellm/llms/dashscope/rerank/transformation.py +++ b/litellm/llms/dashscope/rerank/transformation.py @@ -22,6 +22,7 @@ as supported only for gte-rerank-v2 / qwen3-vl-rerank. Docs - https://help.aliyun.com/zh/model-studio/text-rerank-api """ +from collections.abc import Mapping from typing import Any, Dict, List, Union import httpx @@ -85,6 +86,7 @@ class DashScopeRerankConfig(BaseRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: if api_key is None: api_key = get_secret_str("DASHSCOPE_API_KEY") diff --git a/litellm/llms/deepinfra/rerank/transformation.py b/litellm/llms/deepinfra/rerank/transformation.py index 82069e4e195..87a6ecc7120 100644 --- a/litellm/llms/deepinfra/rerank/transformation.py +++ b/litellm/llms/deepinfra/rerank/transformation.py @@ -2,6 +2,7 @@ Translate between Cohere's `/rerank` format and Deepinfra's `/rerank` format. """ +from collections.abc import Mapping from typing import Any, Dict, List, Union import httpx @@ -67,6 +68,7 @@ class DeepinfraRerankConfig(BaseRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: if api_key is None: api_key = get_secret_str("DEEPINFRA_API_KEY") diff --git a/litellm/llms/fireworks_ai/rerank/transformation.py b/litellm/llms/fireworks_ai/rerank/transformation.py index 393a6c5a8e5..e727f5c1d2b 100644 --- a/litellm/llms/fireworks_ai/rerank/transformation.py +++ b/litellm/llms/fireworks_ai/rerank/transformation.py @@ -4,6 +4,7 @@ Fireworks AI Rerank API transformation Reference: https://docs.fireworks.ai/inference-api-reference/rerank """ +from collections.abc import Mapping from typing import Any, Dict, List, Union import httpx @@ -104,6 +105,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: api_key = self._get_api_key(api_key) if api_key is None: diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 77504eba04a..cd35cc72492 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -2,6 +2,7 @@ Transformation logic for Hosted VLLM rerank """ +from collections.abc import Mapping from typing import Any, Dict, List, Union import httpx @@ -107,6 +108,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: if api_key is None: api_key = get_secret_str("HOSTED_VLLM_API_KEY") or "fake-api-key" diff --git a/litellm/llms/huggingface/rerank/transformation.py b/litellm/llms/huggingface/rerank/transformation.py index cdad77a9815..245551cf4f2 100644 --- a/litellm/llms/huggingface/rerank/transformation.py +++ b/litellm/llms/huggingface/rerank/transformation.py @@ -1,4 +1,5 @@ import os +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Union import httpx @@ -125,6 +126,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig): api_key: str | None = None, optional_params: dict | None = None, api_base: str | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: # Get API credentials api_key, api_base = self.get_api_credentials(api_key=api_key, api_base=api_base) diff --git a/litellm/llms/infinity/rerank/transformation.py b/litellm/llms/infinity/rerank/transformation.py index 94746da4609..7451b06c01a 100644 --- a/litellm/llms/infinity/rerank/transformation.py +++ b/litellm/llms/infinity/rerank/transformation.py @@ -4,12 +4,13 @@ Transformation logic from Cohere's /v1/rerank format to Infinity's `/v1/rerank` Why separate file? Make it easy to see how transformation works """ -from litellm._uuid import uuid +from collections.abc import Mapping from typing import List, Optional import httpx import litellm +from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.cohere.rerank.transformation import CohereRerankConfig from litellm.secret_managers.main import get_secret_str @@ -46,6 +47,7 @@ class InfinityRerankConfig(CohereRerankConfig): model: str, api_key: Optional[str] = None, optional_params: Optional[dict] = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: if api_key is None: api_key = get_secret_str("INFINITY_API_KEY") or get_secret_str("INFINITY_API_KEY") or litellm.infinity_key diff --git a/litellm/llms/jina_ai/rerank/transformation.py b/litellm/llms/jina_ai/rerank/transformation.py index 7f4c0709bdd..903e629803b 100644 --- a/litellm/llms/jina_ai/rerank/transformation.py +++ b/litellm/llms/jina_ai/rerank/transformation.py @@ -6,6 +6,7 @@ Why separate file? Make it easy to see how transformation works Docs - https://jina.ai/reranker """ +from collections.abc import Mapping from typing import Any, Dict, List, Tuple, Union from httpx import URL, Response @@ -139,6 +140,7 @@ class JinaAIRerankConfig(BaseRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> Dict: if api_key is None: raise ValueError("api_key is required. Set via `api_key` parameter or `JINA_API_KEY` environment variable.") diff --git a/litellm/llms/nvidia_nim/rerank/transformation.py b/litellm/llms/nvidia_nim/rerank/transformation.py index 2d72d52f991..07b792468c9 100644 --- a/litellm/llms/nvidia_nim/rerank/transformation.py +++ b/litellm/llms/nvidia_nim/rerank/transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Any, Dict, List, Literal, Union import httpx @@ -148,6 +149,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: """ Validate that the Nvidia NIM API key is present. diff --git a/litellm/llms/vertex_ai/rerank/transformation.py b/litellm/llms/vertex_ai/rerank/transformation.py index b9680af20cc..055a02aa40d 100644 --- a/litellm/llms/vertex_ai/rerank/transformation.py +++ b/litellm/llms/vertex_ai/rerank/transformation.py @@ -4,6 +4,7 @@ Translates from Cohere's `/v1/rerank` input format to Vertex AI Discovery Engine Why separate file? Make it easy to see how transformation works """ +from collections.abc import Mapping from typing import Any, Dict, List, Union import httpx @@ -74,14 +75,15 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): model: str, api_key: str | None = None, optional_params: Dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: """ Validate and set up authentication for Vertex AI Discovery Engine API """ # Get credentials and project info from optional_params (which contains vertex_credentials, etc.) - litellm_params = optional_params.copy() if optional_params else {} - vertex_credentials = self.safe_get_vertex_ai_credentials(litellm_params) - vertex_project = self.safe_get_vertex_ai_project(litellm_params) + vertex_params = optional_params.copy() if optional_params else {} + vertex_credentials = self.safe_get_vertex_ai_credentials(vertex_params) + vertex_project = self.safe_get_vertex_ai_project(vertex_params) # Get access token using the base class method access_token, project_id = self._ensure_access_token( diff --git a/litellm/llms/voyage/rerank/transformation.py b/litellm/llms/voyage/rerank/transformation.py index e426e39962b..df9f32dd96d 100644 --- a/litellm/llms/voyage/rerank/transformation.py +++ b/litellm/llms/voyage/rerank/transformation.py @@ -4,6 +4,7 @@ Transformation logic for Voyage AI's /v1/rerank endpoint. Docs - https://docs.voyageai.com/docs/reranker """ +from collections.abc import Mapping from typing import Any, Dict, List, Tuple, Union import httpx @@ -137,6 +138,7 @@ class VoyageRerankConfig(BaseRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> Dict: if api_key is None: api_key = get_secret_str("VOYAGE_API_KEY") or get_secret_str("VOYAGE_AI_API_KEY") diff --git a/litellm/llms/watsonx/rerank/transformation.py b/litellm/llms/watsonx/rerank/transformation.py index 25b593f1c0a..549ccca4748 100644 --- a/litellm/llms/watsonx/rerank/transformation.py +++ b/litellm/llms/watsonx/rerank/transformation.py @@ -5,6 +5,7 @@ Docs - https://cloud.ibm.com/apidocs/watsonx-ai#text-rerank """ import uuid +from collections.abc import Mapping from typing import Any, Dict, List, Union, cast import httpx @@ -60,6 +61,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> Dict: optional_params = optional_params or {} diff --git a/litellm/main.py b/litellm/main.py index acdec7385da..b167a257d18 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6739,6 +6739,8 @@ def embedding( aembedding=aembedding, ) elif custom_llm_provider == "azure_ai": + from litellm.llms.azure_ai.common_utils import get_azure_ai_entra_token + api_base = ( api_base # for deepinfra/perplexity/anyscale/groq/friendliai we check in get_llm_provider and pass in the api base from there or litellm.api_base @@ -6748,8 +6750,8 @@ def embedding( api_key = ( api_key or litellm.api_key # for deepinfra/perplexity/anyscale/friendliai we check in get_llm_provider and pass in the api key from there - or litellm.openai_key or get_secret_str("AZURE_AI_API_KEY") + or get_azure_ai_entra_token(litellm_params=litellm_params_dict) ) ## EMBEDDING CALL diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index a3280b90fe3..450920f1f44 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -11,7 +11,12 @@ sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path import litellm -from litellm.llms.azure.common_utils import BaseAzureLLM, get_azure_ad_token +from litellm.llms.azure.common_utils import ( + BaseAzureLLM, + _cached_entra_id_token_provider, + get_azure_ad_token, + get_azure_ad_token_from_entra_id, +) from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, ) @@ -2034,3 +2039,59 @@ def test_azure_traditional_api_uses_azure_openai_client(): assert isinstance( async_client, AsyncAzureOpenAI ), f"Expected AsyncAzureOpenAI client for api_version={api_version}" + + +class TestEntraIdTokenProviderCache: + def setup_method(self): + _cached_entra_id_token_provider.cache_clear() + + def teardown_method(self): + _cached_entra_id_token_provider.cache_clear() + + def test_reuses_credential_for_the_same_service_principal(self): + with ( + patch("azure.identity.ClientSecretCredential") as mock_credential, + patch("azure.identity.get_bearer_token_provider", side_effect=lambda credential, scope: lambda: "token"), + ): + first = get_azure_ad_token_from_entra_id( + tenant_id="tenant", + client_id="client", + client_secret="secret", + scope="https://cognitiveservices.azure.com/.default", + ) + second = get_azure_ad_token_from_entra_id( + tenant_id="tenant", + client_id="client", + client_secret="secret", + scope="https://cognitiveservices.azure.com/.default", + ) + + assert first is second + assert mock_credential.call_count == 1 + + @pytest.mark.parametrize( + "second_call_kwargs", + [ + {"tenant_id": "other-tenant"}, + {"client_id": "other-client"}, + {"client_secret": "other-secret"}, + {"scope": "https://ai.azure.com/.default"}, + ], + ) + def test_does_not_share_a_provider_across_credentials_or_scopes(self, second_call_kwargs): + base_kwargs = { + "tenant_id": "tenant", + "client_id": "client", + "client_secret": "secret", + "scope": "https://cognitiveservices.azure.com/.default", + } + + with ( + patch("azure.identity.ClientSecretCredential") as mock_credential, + patch("azure.identity.get_bearer_token_provider", side_effect=lambda credential, scope: lambda: "token"), + ): + first = get_azure_ad_token_from_entra_id(**base_kwargs) + second = get_azure_ad_token_from_entra_id(**{**base_kwargs, **second_call_kwargs}) + + assert first is not second + assert mock_credential.call_count == 2 diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py b/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py index da1041f3d60..9c9401fa8e9 100644 --- a/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py +++ b/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py @@ -5,6 +5,10 @@ sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path +import litellm +from litellm.llms.azure_ai.image_edit.flux2_transformation import ( + AzureFoundryFlux2ImageEditConfig, +) from litellm.llms.azure_ai.image_edit.transformation import ( AzureFoundryFluxImageEditConfig, ) @@ -32,3 +36,32 @@ def test_azure_ai_url_generation(): ) expected_url = f"{api_base}/openai/deployments/FLUX.1-Kontext-pro/images/edits?api-version=2025-04-01-preview" assert complete_url == expected_url + + +def test_azure_ai_validate_environment_with_entra_token(monkeypatch): + monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) + monkeypatch.setattr(litellm, "api_key", None) + config = AzureFoundryFluxImageEditConfig() + + headers = config.validate_environment( + {}, + "FLUX.1-Kontext-pro", + litellm_params={"azure_ad_token": "entra-token"}, + ) + + assert headers == {"Authorization": "Bearer entra-token"} + + +def test_flux2_validate_environment_with_entra_token(monkeypatch): + monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) + monkeypatch.setattr(litellm, "api_key", None) + config = AzureFoundryFlux2ImageEditConfig() + + headers = config.validate_environment( + {}, + "flux.2-pro", + litellm_params={"azure_ad_token": "entra-token"}, + ) + + assert headers["Authorization"] == "Bearer entra-token" + assert headers["Content-Type"] == "application/json" diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py b/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py index d5256be02d7..c4e39a26aeb 100644 --- a/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py +++ b/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py @@ -8,6 +8,7 @@ import pytest sys.path.insert(0, os.path.abspath("../../../../../..")) +import litellm from litellm.llms.azure_ai.image_edit import ( AzureFoundryMAIImageEditConfig, get_azure_ai_image_edit_config, @@ -169,3 +170,16 @@ class TestAzureMAIImageEdit: assert image_response.data[0].b64_json == "abc123" assert image_response.usage.output_tokens == 1024 assert image_response.usage.total_tokens == 1024 + + +def test_mai_validate_environment_with_entra_token(monkeypatch): + monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) + monkeypatch.setattr(litellm, "api_key", None) + + headers = AzureFoundryMAIImageEditConfig().validate_environment( + headers={}, + model="MAI-Image-2.5", + litellm_params={"azure_ad_token": "entra-token"}, + ) + + assert headers == {"Authorization": "Bearer entra-token"} diff --git a/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py b/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py index ffabce6e00c..150f5794ab1 100644 --- a/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py +++ b/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py @@ -7,6 +7,7 @@ sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path +import litellm from litellm.llms.azure_ai.rerank.transformation import AzureAIRerankConfig @@ -97,3 +98,26 @@ class TestAzureAIRerankConfigGetCompleteUrl: model=self.model, ) assert url == "https://my-resource.services.ai.azure.com/v1/rerank?r=1" + + +class TestAzureAIRerankConfigValidateEnvironment: + def test_uses_api_key_when_set(self): + headers = AzureAIRerankConfig().validate_environment( + headers={}, + model="azure_ai/cohere-rerank-v3-english", + api_key="my-key", + ) + + assert headers["Authorization"] == "Bearer my-key" + + def test_falls_back_to_entra_token(self, monkeypatch): + monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) + monkeypatch.setattr(litellm, "azure_key", None) + + headers = AzureAIRerankConfig().validate_environment( + headers={}, + model="azure_ai/cohere-rerank-v3-english", + litellm_params={"azure_ad_token": "entra-token"}, + ) + + assert headers["Authorization"] == "Bearer entra-token" diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_entra_auth.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_entra_auth.py new file mode 100644 index 00000000000..8ac37feee4b --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_entra_auth.py @@ -0,0 +1,153 @@ +""" +Entra ID / OAuth auth for Azure AI Foundry routes. + +Every azure_ai route must authenticate with an Entra ID token when no API key is configured, +instead of requiring an API key. +""" + +from unittest.mock import patch + +import pytest + +import litellm +from litellm.llms.azure_ai.common_utils import get_azure_ai_auth_headers +from litellm.llms.azure_ai.ocr.transformation import AzureAIOCRConfig + +ENTRA_PARAMS = {"azure_ad_token": "entra-token"} + + +@pytest.fixture(autouse=True) +def clear_azure_env(monkeypatch): + for env_var in ( + "AZURE_AI_API_KEY", + "AZURE_API_KEY", + "AZURE_AD_TOKEN", + "AZURE_TENANT_ID", + "AZURE_CLIENT_ID", + "AZURE_CLIENT_SECRET", + "AZURE_SCOPE", + "OPENAI_API_KEY", + "AZURE_DOCUMENT_INTELLIGENCE_API_KEY", + ): + monkeypatch.delenv(env_var, raising=False) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + + +def test_api_key_wins_over_entra_credentials(): + headers = get_azure_ai_auth_headers(api_key="my-key", litellm_params=ENTRA_PARAMS, api_key_header="Api-Key") + + assert headers == {"Api-Key": "my-key"} + + +def test_entra_token_used_when_no_api_key(): + headers = get_azure_ai_auth_headers(api_key=None, litellm_params=ENTRA_PARAMS, api_key_header="Api-Key") + + assert headers == {"Authorization": "Bearer entra-token"} + + +def test_service_principal_token_is_requested_with_the_configured_scope(): + with patch("litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id") as mock_entra_id: + mock_entra_id.return_value = lambda: "sp-token" + + headers = get_azure_ai_auth_headers( + api_key=None, + litellm_params={ + "tenant_id": "tenant", + "client_id": "client", + "client_secret": "secret", + "azure_scope": "https://ai.azure.com/.default", + }, + ) + + mock_entra_id.assert_called_once_with( + tenant_id="tenant", + client_id="client", + client_secret="secret", + scope="https://ai.azure.com/.default", + ) + assert headers == {"Authorization": "Bearer sp-token"} + + +def test_error_mentions_both_credential_types_when_nothing_is_configured(): + with pytest.raises(ValueError) as exc_info: + get_azure_ai_auth_headers(api_key=None, litellm_params={}) + + message = str(exc_info.value) + assert "AZURE_AI_API_KEY" in message + assert "client_secret" in message + + +def test_ocr_authenticates_with_entra_token(): + headers = AzureAIOCRConfig().validate_environment( + headers={}, + model="azure_ai/mistral-ocr", + api_base="https://my-resource.services.ai.azure.com", + litellm_params=ENTRA_PARAMS, + ) + + assert headers["Authorization"] == "Bearer entra-token" + + +def test_embedding_falls_back_to_entra_token_instead_of_openai_key(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "sk-openai-key") + + with patch.object(litellm.main.azure_ai_embedding, "embedding") as mock_embedding: + mock_embedding.return_value = litellm.EmbeddingResponse() + + litellm.embedding( + model="azure_ai/cohere-embed-v3-english", + input=["hello"], + api_base="https://my-resource.services.ai.azure.com", + azure_ad_token="entra-token", + ) + + assert mock_embedding.call_args.kwargs["api_key"] == "entra-token" + + +def test_image_generation_authenticates_with_entra_token(): + with patch.object(litellm.images.main.azure_chat_completions, "image_generation") as mock_image_generation: + mock_image_generation.return_value = litellm.ImageResponse() + + litellm.image_generation( + model="azure_ai/FLUX-1.1-pro", + prompt="a red circle", + api_base="https://my-resource.services.ai.azure.com", + azure_ad_token="entra-token", + ) + + headers = mock_image_generation.call_args.kwargs["headers"] + assert headers["Authorization"] == "Bearer entra-token" + assert "api-key" not in headers + + +def test_image_generation_keeps_caller_supplied_authorization_header(): + with patch.object(litellm.images.main.azure_chat_completions, "image_generation") as mock_image_generation: + mock_image_generation.return_value = litellm.ImageResponse() + + litellm.image_generation( + model="azure_ai/FLUX-1.1-pro", + prompt="a red circle", + api_base="https://my-resource.services.ai.azure.com", + headers={"Authorization": "Bearer caller-token"}, + ) + + headers = mock_image_generation.call_args.kwargs["headers"] + assert headers["Authorization"] == "Bearer caller-token" + assert "api-key" not in headers + + +def test_image_generation_still_uses_api_key_header(): + with patch.object(litellm.images.main.azure_chat_completions, "image_generation") as mock_image_generation: + mock_image_generation.return_value = litellm.ImageResponse() + + litellm.image_generation( + model="azure_ai/FLUX-1.1-pro", + prompt="a red circle", + api_base="https://my-resource.services.ai.azure.com", + api_key="my-key", + ) + + headers = mock_image_generation.call_args.kwargs["headers"] + assert headers["api-key"] == "my-key" + assert "Authorization" not in headers diff --git a/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py b/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py index 39d6f1dc355..b8e11a0bfcb 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py @@ -248,3 +248,30 @@ def test_get_complete_url_combines_pages_and_features(): assert "&pages=1,2,3" in url assert "&features=keyValuePairs,languages" in url + + +def test_validate_environment_uses_subscription_key(monkeypatch): + monkeypatch.delenv("AZURE_DOCUMENT_INTELLIGENCE_API_KEY", raising=False) + + headers = AzureDocumentIntelligenceOCRConfig().validate_environment( + headers={}, + model="prebuilt-layout", + api_key="my-key", + api_base="https://example.cognitiveservices.azure.com", + ) + + assert headers["Ocp-Apim-Subscription-Key"] == "my-key" + + +def test_validate_environment_falls_back_to_entra_token(monkeypatch): + monkeypatch.delenv("AZURE_DOCUMENT_INTELLIGENCE_API_KEY", raising=False) + + headers = AzureDocumentIntelligenceOCRConfig().validate_environment( + headers={}, + model="prebuilt-layout", + api_base="https://example.cognitiveservices.azure.com", + litellm_params={"azure_ad_token": "entra-token"}, + ) + + assert headers["Authorization"] == "Bearer entra-token" + assert "Ocp-Apim-Subscription-Key" not in headers From c5d50817a70373f3443423fa8a6a97980ad156a9 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 31 Jul 2026 22:01:28 +0000 Subject: [PATCH 018/542] fix(azure_ai): detect caller auth headers case-insensitively in image generation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/images/main.py | 3 ++- .../llms/azure_ai/test_azure_ai_entra_auth.py | 9 +++++---- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/litellm/images/main.py b/litellm/images/main.py index 4c88eb52cd8..3bee6000d3f 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -440,7 +440,8 @@ def image_generation( if extra_headers is not None: optional_params["extra_headers"] = extra_headers - caller_set_auth = "api-key" in headers or "Authorization" in headers + caller_header_names = frozenset(name.lower() for name in headers) + caller_set_auth = "api-key" in caller_header_names or "authorization" in caller_header_names auth_headers = ( headers if caller_set_auth diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_entra_auth.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_entra_auth.py index 8ac37feee4b..1145439a7b4 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_ai_entra_auth.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_entra_auth.py @@ -121,7 +121,8 @@ def test_image_generation_authenticates_with_entra_token(): assert "api-key" not in headers -def test_image_generation_keeps_caller_supplied_authorization_header(): +@pytest.mark.parametrize("header_name", ["Authorization", "authorization", "api-key", "API-KEY"]) +def test_image_generation_keeps_caller_supplied_auth_header(header_name): with patch.object(litellm.images.main.azure_chat_completions, "image_generation") as mock_image_generation: mock_image_generation.return_value = litellm.ImageResponse() @@ -129,12 +130,12 @@ def test_image_generation_keeps_caller_supplied_authorization_header(): model="azure_ai/FLUX-1.1-pro", prompt="a red circle", api_base="https://my-resource.services.ai.azure.com", - headers={"Authorization": "Bearer caller-token"}, + headers={header_name: "caller-credential"}, ) headers = mock_image_generation.call_args.kwargs["headers"] - assert headers["Authorization"] == "Bearer caller-token" - assert "api-key" not in headers + assert headers[header_name] == "caller-credential" + assert len(headers) == 2 def test_image_generation_still_uses_api_key_header(): From 8aa9d3dfe58bffa4efee69db579bfd7be5e9fc02 Mon Sep 17 00:00:00 2001 From: bhuvan2134686 Date: Fri, 7 Aug 2026 11:22:01 +1000 Subject: [PATCH 019/542] feat(models): swap SCX.ai catalog to GLM-5.2 and Qwen3.8 Max Replaces the five launch models with the two that SCX.ai now leads on. Both are live on api.scx.ai and both were verified against it for tool calling, json_object and json_schema output, reasoning, prompt caching, and, for Qwen3.8 Max, image input Pricing follows SCX's published USD rates. GLM-5.2 lands at $0.55/M input and $1.9255/M output, tracking the recent GLM-5.2 market repricing; Qwen3.8 Max at $1.815/M and $5.4461/M sits under the only other seller of that model, and is the first Qwen3.8 Max entry in the catalog Also corrects a metadata bug the removed entries carried: they set max_tokens equal to max_input_tokens, conflating the context window with the output cap. Both new entries declare a max_output_tokens of 131072, which is what the endpoint's own validator enforces The Add Model placeholder moves to scx-ai/GLM-5.2 now that MiniMax-M2.7 is no longer in the catalog --- ...odel_prices_and_context_window_backup.json | 84 ++++++------------- .../provider_create_fields.json | 2 +- model_prices_and_context_window.json | 84 ++++++------------- .../llms/openai_like/test_scx_ai_provider.py | 34 ++++---- .../components/provider_info_helpers.test.tsx | 2 +- .../src/components/provider_info_helpers.tsx | 2 +- 6 files changed, 75 insertions(+), 133 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2f47643da50..dee3cbdd054 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -33546,72 +33546,40 @@ "supports_vision": true, "source": "https://cloud.sambanova.ai/plans/pricing" }, - "scx-ai/gemma-4-31B-it": { - "input_cost_per_token": 3e-07, + "scx-ai/GLM-5.2": { + "cache_read_input_token_cost": 1.375e-07, + "input_cost_per_token": 5.5e-07, "litellm_provider": "scx-ai", - "max_input_tokens": 131000, - "max_tokens": 131000, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 9.1e-07, - "source": "https://scx.ai/pricing", + "output_cost_per_token": 1.9255e-06, + "source": "https://llmgateway.io/models/glm-5.2/scx-ai-gp", "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "scx-ai/Qwen3.8-Max": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 1.815e-06, + "litellm_provider": "scx-ai", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 5.4461e-06, + "source": "https://llmgateway.io/models/qwen3.8-max/scx-ai-gp", + "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true }, - "scx-ai/gpt-oss-120b": { - "input_cost_per_token": 1.7e-07, - "litellm_provider": "scx-ai", - "max_input_tokens": 131000, - "max_tokens": 131000, - "mode": "chat", - "output_cost_per_token": 5.5e-07, - "source": "https://scx.ai/pricing", - "supports_function_calling": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_tool_choice": true - }, - "scx-ai/Llama-4-Maverick-17B-128E-Instruct": { - "input_cost_per_token": 5.3e-07, - "litellm_provider": "scx-ai", - "max_input_tokens": 131000, - "max_tokens": 131000, - "mode": "chat", - "output_cost_per_token": 1.62e-06, - "source": "https://scx.ai/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "scx-ai/MiniMax-M2.7": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 4.8e-07, - "litellm_provider": "scx-ai", - "max_input_tokens": 192000, - "max_tokens": 192000, - "mode": "chat", - "output_cost_per_token": 1.79e-06, - "source": "https://scx.ai/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_tool_choice": true - }, - "scx-ai/Qwen3-32B": { - "input_cost_per_token": 3.6e-07, - "litellm_provider": "scx-ai", - "max_input_tokens": 33000, - "max_tokens": 33000, - "mode": "chat", - "output_cost_per_token": 8.7e-07, - "source": "https://scx.ai/pricing", - "supports_function_calling": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_tool_choice": true - }, "snowflake/claude-3-5-sonnet": { "litellm_provider": "snowflake", "max_input_tokens": 200000, diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 4cfe9e4ef5e..99d0de262a6 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -2686,7 +2686,7 @@ "default_value": null } ], - "default_model_placeholder": "scx-ai/MiniMax-M2.7" + "default_model_placeholder": "scx-ai/GLM-5.2" }, { "provider": "Snowflake", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8464aec769f..30624d375f7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -33637,72 +33637,40 @@ "supports_vision": true, "source": "https://cloud.sambanova.ai/plans/pricing" }, - "scx-ai/gemma-4-31B-it": { - "input_cost_per_token": 3e-07, + "scx-ai/GLM-5.2": { + "cache_read_input_token_cost": 1.375e-07, + "input_cost_per_token": 5.5e-07, "litellm_provider": "scx-ai", - "max_input_tokens": 131000, - "max_tokens": 131000, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 9.1e-07, - "source": "https://scx.ai/pricing", + "output_cost_per_token": 1.9255e-06, + "source": "https://llmgateway.io/models/glm-5.2/scx-ai-gp", "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "scx-ai/Qwen3.8-Max": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 1.815e-06, + "litellm_provider": "scx-ai", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 5.4461e-06, + "source": "https://llmgateway.io/models/qwen3.8-max/scx-ai-gp", + "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true }, - "scx-ai/gpt-oss-120b": { - "input_cost_per_token": 1.7e-07, - "litellm_provider": "scx-ai", - "max_input_tokens": 131000, - "max_tokens": 131000, - "mode": "chat", - "output_cost_per_token": 5.5e-07, - "source": "https://scx.ai/pricing", - "supports_function_calling": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_tool_choice": true - }, - "scx-ai/Llama-4-Maverick-17B-128E-Instruct": { - "input_cost_per_token": 5.3e-07, - "litellm_provider": "scx-ai", - "max_input_tokens": 131000, - "max_tokens": 131000, - "mode": "chat", - "output_cost_per_token": 1.62e-06, - "source": "https://scx.ai/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "scx-ai/MiniMax-M2.7": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_token": 4.8e-07, - "litellm_provider": "scx-ai", - "max_input_tokens": 192000, - "max_tokens": 192000, - "mode": "chat", - "output_cost_per_token": 1.79e-06, - "source": "https://scx.ai/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_tool_choice": true - }, - "scx-ai/Qwen3-32B": { - "input_cost_per_token": 3.6e-07, - "litellm_provider": "scx-ai", - "max_input_tokens": 33000, - "max_tokens": 33000, - "mode": "chat", - "output_cost_per_token": 8.7e-07, - "source": "https://scx.ai/pricing", - "supports_function_calling": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_tool_choice": true - }, "snowflake/claude-3-5-sonnet": { "litellm_provider": "snowflake", "max_input_tokens": 200000, diff --git a/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py b/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py index 92921067728..6b293cae303 100644 --- a/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py +++ b/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py @@ -34,13 +34,13 @@ class TestSCXAIProviderConfig: from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider model, provider, api_key, api_base = get_llm_provider( - model="scx-ai/gpt-oss-120b", + model="scx-ai/GLM-5.2", custom_llm_provider=None, api_base=None, api_key=None, ) - assert model == "gpt-oss-120b" + assert model == "GLM-5.2" assert provider == "scx-ai" assert api_base == "https://api.scx.ai/v1" @@ -48,7 +48,7 @@ class TestSCXAIProviderConfig: from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider model, provider, api_key, api_base = get_llm_provider( - model="scx-ai/gpt-oss-120b", + model="scx-ai/GLM-5.2", custom_llm_provider=None, api_base="https://custom.scx.ai/v1", api_key="sk-test", @@ -62,7 +62,7 @@ class TestSCXAIProviderConfig: from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider model, provider, api_key, api_base = get_llm_provider( - model="gpt-oss-120b", + model="GLM-5.2", custom_llm_provider=None, api_base="https://api.scx.ai/v1", api_key=None, @@ -81,7 +81,7 @@ class TestSCXAIProviderConfig: optional_params = config.map_openai_params( non_default_params={"temperature": 1.7}, optional_params={}, - model="gpt-oss-120b", + model="GLM-5.2", drop_params=False, ) assert optional_params["temperature"] == 1.0 @@ -89,7 +89,7 @@ class TestSCXAIProviderConfig: optional_params = config.map_openai_params( non_default_params={"temperature": 0.4}, optional_params={}, - model="gpt-oss-120b", + model="GLM-5.2", drop_params=False, ) assert optional_params["temperature"] == 0.4 @@ -105,7 +105,7 @@ class TestSCXAIProviderConfig: optional_params = config.map_openai_params( non_default_params={"max_completion_tokens": 256}, optional_params={}, - model="gpt-oss-120b", + model="GLM-5.2", drop_params=False, ) assert optional_params["max_tokens"] == 256 @@ -119,7 +119,7 @@ class TestSCXAIProviderConfig: { "model_name": "scx-chat", "litellm_params": { - "model": "scx-ai/gpt-oss-120b", + "model": "scx-ai/GLM-5.2", "api_key": "test-key", }, } @@ -132,13 +132,10 @@ class TestSCXAIProviderConfig: class TestSCXAIModelMetadata: SCX_MODELS = ( - "scx-ai/Llama-4-Maverick-17B-128E-Instruct", - "scx-ai/gemma-4-31B-it", - "scx-ai/Qwen3-32B", - "scx-ai/MiniMax-M2.7", - "scx-ai/gpt-oss-120b", + "scx-ai/GLM-5.2", + "scx-ai/Qwen3.8-Max", ) - VISION_MODELS = ("scx-ai/Llama-4-Maverick-17B-128E-Instruct", "scx-ai/gemma-4-31B-it") + VISION_MODELS = ("scx-ai/Qwen3.8-Max",) @staticmethod def _load(path_parts): @@ -160,8 +157,17 @@ class TestSCXAIModelMetadata: assert info["output_cost_per_token"] > 0 assert info["supports_function_calling"] is True assert info["supports_tool_choice"] is True + assert info["supports_reasoning"] is True + assert info["supports_response_schema"] is True assert info.get("supports_vision", False) is (model in self.VISION_MODELS) + assert info["supports_prompt_caching"] is True + assert 0 < info["cache_read_input_token_cost"] < info["input_cost_per_token"] + + assert info["max_output_tokens"] == 131072 + assert info["max_tokens"] == info["max_output_tokens"] + assert info["max_input_tokens"] >= 1_000_000 + def test_scx_ai_models_synced_to_backup(self): model_cost = self._load(("model_prices_and_context_window.json",)) backup = self._load(("litellm", "model_prices_and_context_window_backup.json")) diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index 8a4c9dfd24d..75af4cffa06 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -173,7 +173,7 @@ describe("provider_info_helpers", () => { }); it("should return an scx-ai model placeholder for SCX_AI provider", () => { - expect(getPlaceholder(Providers.SCX_AI)).toBe("scx-ai/MiniMax-M2.7"); + expect(getPlaceholder(Providers.SCX_AI)).toBe("scx-ai/GLM-5.2"); }); it("should return claude-3-opus placeholder for Anthropic provider", () => { diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index fa1f9c11eb3..41898ce27ed 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -444,7 +444,7 @@ export const getPlaceholder = (selectedProvider: string): string => { } else if (selectedProvider === Providers.ZAI) { return "zai/glm-4.5"; } else if (selectedProvider === Providers.SCX_AI) { - return "scx-ai/MiniMax-M2.7"; + return "scx-ai/GLM-5.2"; } else { return "gpt-3.5-turbo"; } From a028c8857e2d9ff23308ff256c2723220868b200 Mon Sep 17 00:00:00 2001 From: bhuvan2134686 Date: Fri, 7 Aug 2026 11:37:38 +1000 Subject: [PATCH 020/542] fix(scx-ai): correct the temperature ceiling to match the endpoint The constraint was 1.0, so anything above that was silently clamped down. SCX accepts [0.0, 2.0), verified live against both GLM-5.2 and Qwen3.8 Max: 1.5, 1.99 and 1.999 all return 200, while 2.0 returns 400 with "Temperature should be in [0.0, 2.0)" Since the clamp is an inclusive min(), 2.0 cannot be the ceiling or it would pass through a value the endpoint rejects. 1.99 is the practical maximum The clamp test now pins both ends: 2.5 comes back as 1.99, and 1.7 rides through untouched where it used to be flattened to 1.0 --- litellm/llms/openai_like/providers.json | 2 +- .../llms/openai_like/test_scx_ai_provider.py | 12 ++++++++++-- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index d796d140878..b43a44c2d3e 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -192,7 +192,7 @@ "max_completion_tokens": "max_tokens" }, "constraints": { - "temperature_max": 1.0 + "temperature_max": 1.99 }, "supported_endpoints": ["/v1/chat/completions"] } diff --git a/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py b/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py index 6b293cae303..1ce2da65fef 100644 --- a/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py +++ b/tests/test_litellm/llms/openai_like/test_scx_ai_provider.py @@ -23,7 +23,7 @@ class TestSCXAIProviderConfig: assert scx.base_url == "https://api.scx.ai/v1" assert scx.api_key_env == "SCX_API_KEY" assert scx.param_mappings.get("max_completion_tokens") == "max_tokens" - assert scx.constraints.get("temperature_max") == 1.0 + assert scx.constraints.get("temperature_max") == 1.99 def test_scx_ai_in_openai_compatible_providers(self): from litellm.constants import openai_compatible_providers @@ -78,13 +78,21 @@ class TestSCXAIProviderConfig: assert provider is not None config = create_config_class(provider)() + optional_params = config.map_openai_params( + non_default_params={"temperature": 2.5}, + optional_params={}, + model="GLM-5.2", + drop_params=False, + ) + assert optional_params["temperature"] == 1.99 + optional_params = config.map_openai_params( non_default_params={"temperature": 1.7}, optional_params={}, model="GLM-5.2", drop_params=False, ) - assert optional_params["temperature"] == 1.0 + assert optional_params["temperature"] == 1.7 optional_params = config.map_openai_params( non_default_params={"temperature": 0.4}, From e3da917e679ff68aaa86f5ed67dcb6528117f072 Mon Sep 17 00:00:00 2001 From: Souravrajvi0 <144546710+Souravrajvi0@users.noreply.github.com> Date: Tue, 11 Aug 2026 06:29:07 +0000 Subject: [PATCH 021/542] fix(proxy): parse form-encoded video edit/extension bodies after auth Fixes #36487 video_edit, video_extension, and video_remix called request.body() after user_api_key_auth had already parsed multipart/form bodies via _read_request_body(), causing RuntimeError Stream consumed and 500s for OpenAI SDK clients. Use _read_request_body consistently and normalize bare-string or JSON-string video references from form fields into video_id. --- litellm/proxy/video_endpoints/endpoints.py | 28 ++++------- litellm/proxy/video_endpoints/utils.py | 21 +++++++++ .../proxy/video_endpoints/test_endpoints.py | 24 ++++++++++ tests/test_litellm/test_video_generation.py | 47 +++++++++++++++++++ 4 files changed, 100 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/video_endpoints/endpoints.py b/litellm/proxy/video_endpoints/endpoints.py index 6c6b004fd17..cb014bcceee 100644 --- a/litellm/proxy/video_endpoints/endpoints.py +++ b/litellm/proxy/video_endpoints/endpoints.py @@ -2,7 +2,6 @@ from typing import Any, Final -import orjson from fastapi import APIRouter, Depends, File, Form, Request, Response, UploadFile from fastapi.responses import ORJSONResponse @@ -20,6 +19,7 @@ from litellm.proxy.video_endpoints.utils import ( encode_character_id_in_response, extract_model_from_target_model_names, get_custom_provider_from_data, + pop_video_reference_to_video_id, ) from litellm.types.videos.utils import ( decode_character_id_with_provider, @@ -451,9 +451,7 @@ async def video_remix( version, ) - # Read request body - body: Final = await request.body() - data: Final = orjson.loads(body) + data: Final = await _read_request_body(request=request) data["video_id"] = video_id decoded: Final = decode_video_id_with_provider(video_id) @@ -760,15 +758,10 @@ async def video_edit( version, ) - body: Final = await request.body() - data: Final = orjson.loads(body) + data: Final = await _read_request_body(request=request) + pop_video_reference_to_video_id(data) - # Extract video_id from nested video object - video_ref: Final = data.pop("video", {}) - video_id: Final = video_ref.get("id", "") if isinstance(video_ref, dict) else "" - data["video_id"] = video_id - - decoded: Final = decode_video_id_with_provider(video_id) + decoded: Final = decode_video_id_with_provider(data["video_id"]) provider_from_id: Final = decoded.get("custom_llm_provider") model_id_from_decoded: Final = decoded.get("model_id") @@ -860,15 +853,10 @@ async def video_extension( version, ) - body: Final = await request.body() - data: Final = orjson.loads(body) + data: Final = await _read_request_body(request=request) + pop_video_reference_to_video_id(data) - # Extract video_id from nested video object - video_ref: Final = data.pop("video", {}) - video_id: Final = video_ref.get("id", "") if isinstance(video_ref, dict) else "" - data["video_id"] = video_id - - decoded: Final = decode_video_id_with_provider(video_id) + decoded: Final = decode_video_id_with_provider(data["video_id"]) provider_from_id: Final = decoded.get("custom_llm_provider") model_id_from_decoded: Final = decoded.get("model_id") diff --git a/litellm/proxy/video_endpoints/utils.py b/litellm/proxy/video_endpoints/utils.py index d6b398e3476..5f508cd02ca 100644 --- a/litellm/proxy/video_endpoints/utils.py +++ b/litellm/proxy/video_endpoints/utils.py @@ -13,6 +13,27 @@ def extract_model_from_target_model_names(target_model_names: Any) -> str | None return target_model_names[0] if target_model_names else None +def pop_video_reference_to_video_id(data: dict[str, Any]) -> None: + """ + Normalize OpenAI video edit/extension payloads into ``video_id``. + + JSON bodies use ``video: {"id": ...}``. Multipart and form-urlencoded bodies + may send a bare id string or a JSON-encoded reference object as a string field. + """ + video_ref: Final = data.pop("video", {}) + if isinstance(video_ref, dict): + video_id: Final = video_ref.get("id", "") + elif isinstance(video_ref, str): + try: + parsed_ref: Final = orjson.loads(video_ref) + except orjson.JSONDecodeError: + parsed_ref = None + video_id = parsed_ref.get("id", "") if isinstance(parsed_ref, dict) else video_ref + else: + video_id = "" + data["video_id"] = video_id + + def get_custom_provider_from_data(data: dict[str, Any]) -> str | None: custom_llm_provider: Final = data.get("custom_llm_provider") if custom_llm_provider: diff --git a/tests/test_litellm/proxy/video_endpoints/test_endpoints.py b/tests/test_litellm/proxy/video_endpoints/test_endpoints.py index 40a26fad3c3..78f32600f73 100644 --- a/tests/test_litellm/proxy/video_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/video_endpoints/test_endpoints.py @@ -375,6 +375,7 @@ async def test_content__model_encoded_id(harness): async def call_edit( harness: Harness, *, body: Dict[str, Any], headers=None, query=None ): + harness.read_body.return_value = dict(body) return await endpoints.video_edit( request=FakeRequest(headers=headers, query=query, raw_body=orjson.dumps(body)), fastapi_response=Response(), @@ -431,6 +432,27 @@ async def test_edit__missing_video_object_defaults_to_openai(harness): assert "video" not in data +@pytest.mark.asyncio +async def test_edit__bare_string_video_id_from_form_field(harness): + await call_edit(harness, body={"prompt": "brighter", "video": "video_plain"}) + + assert harness.processor_data() == { + "prompt": "brighter", + "video_id": "video_plain", + "custom_llm_provider": "openai", + } + + +@pytest.mark.asyncio +async def test_edit__json_string_video_reference_from_form_field(harness): + await call_edit( + harness, + body={"prompt": "brighter", "video": orjson.dumps({"id": "video_plain"}).decode()}, + ) + + assert harness.processor_data()["video_id"] == "video_plain" + + # =========================================================================== # # GET /v1/videos - video_list # # =========================================================================== # @@ -474,6 +496,7 @@ async def test_list__provider_from_header(harness): async def call_remix( harness: Harness, video_id: str, *, body, headers=None, query=None ): + harness.read_body.return_value = dict(body) return await endpoints.video_remix( video_id=video_id, request=FakeRequest(headers=headers, query=query, raw_body=orjson.dumps(body)), @@ -632,6 +655,7 @@ async def test_get_character__plain_id_defaults_openai_no_encode(harness): async def call_extension(harness: Harness, *, body, headers=None, query=None): + harness.read_body.return_value = dict(body) return await endpoints.video_extension( request=FakeRequest(headers=headers, query=query, raw_body=orjson.dumps(body)), fastapi_response=Response(), diff --git a/tests/test_litellm/test_video_generation.py b/tests/test_litellm/test_video_generation.py index 3d0472ef96e..fc7ba773c38 100644 --- a/tests/test_litellm/test_video_generation.py +++ b/tests/test_litellm/test_video_generation.py @@ -2321,6 +2321,53 @@ def test_edit_and_extension_support_custom_provider_from_extra_body( assert captured_data["custom_llm_provider"] == "vertex_ai" +@pytest.mark.parametrize("endpoint", ["/v1/videos/edits", "/v1/videos/extensions"]) +def test_edit_and_extension_accept_form_encoded_after_auth_reads_body( + video_proxy_test_client, endpoint +): + from fastapi import Request + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.http_parsing_utils import _read_request_body + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + captured_data = {} + + async def _mock_base_process(self, **kwargs): + captured_data.update(self.data) + return { + "id": "video_resp_123", + "object": "video", + "status": "queued", + "created_at": 1712697600, + } + + async def auth_that_reads_body_first(request: Request): + await _read_request_body(request=request) + return MagicMock() + + app = video_proxy_test_client.app + app.dependency_overrides[user_api_key_auth] = auth_that_reads_body_first + + with patch.object( + ProxyBaseLLMRequestProcessing, + "base_process_llm_request", + new=_mock_base_process, + ): + response = video_proxy_test_client.post( + endpoint, + headers={"Authorization": "Bearer sk-1234"}, + data={ + "model": "my-video-model", + "prompt": "brighter", + "video": "video_123", + }, + ) + + assert response.status_code == 200, response.text + assert captured_data["video_id"] == "video_123" + assert captured_data["prompt"] == "brighter" + + @pytest.mark.parametrize("endpoint", ["/v1/videos/edits", "/v1/videos/extensions"]) def test_edit_and_extension_route_with_encoded_video_ids( video_proxy_test_client, endpoint From 05c91aa5f23a8f354a0076ab85ae9701a433bd3a Mon Sep 17 00:00:00 2001 From: ozolam Date: Thu, 16 Jul 2026 14:12:28 +0300 Subject: [PATCH 022/542] fix(claude-code): correct skill install command and marketplace setup UX - formatInstallCommand now produces /plugin install {name}@litellm instead of /plugin marketplace add {source} - extraKnownMarketplaces snippet fixed: source must be a nested object not a flat string; the flat string caused Claude Code to reject the settings file - marketplace key renamed from my-org to litellm to match the name the proxy returns in marketplace.json - setup tab now shows /plugin marketplace add command as primary option with settings.json as secondary - usage tab now shows a hint to run /plugin marketplace update litellm when a plugin is not found --- .../claude_code_plugins/helpers.test.ts | 26 +----- .../components/claude_code_plugins/helpers.ts | 22 ++--- .../claude_code_plugins/skill_detail.tsx | 93 ++++++++++++++++++- 3 files changed, 102 insertions(+), 39 deletions(-) diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts index c16eba24f7b..a04134a63d8 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts @@ -25,7 +25,7 @@ describe("buildMarketplaceSettingsSnippet", () => { it("nests the url under a source object so Claude Code accepts the marketplace", () => { expect(JSON.parse(buildMarketplaceSettingsSnippet("https://proxy.example.com"))).toEqual({ extraKnownMarketplaces: { - "my-org": { + litellm: { source: { source: "url", url: "https://proxy.example.com/claude-code/marketplace.json", @@ -37,28 +37,12 @@ describe("buildMarketplaceSettingsSnippet", () => { }); describe("formatInstallCommand", () => { - it("formats github source with repo", () => { - const source: PluginSource = { source: "github", repo: "org/repo" }; - expect(formatInstallCommand({ name: "my-plugin", source })).toBe("/plugin marketplace add org/repo"); + it("produces a /plugin install command scoped to the litellm marketplace", () => { + expect(formatInstallCommand({ name: "my-plugin" })).toBe("/plugin install my-plugin@litellm"); }); - it("formats url source", () => { - const source: PluginSource = { source: "url", url: "https://example.com/plugin" }; - expect(formatInstallCommand({ name: "my-plugin", source })).toBe( - "/plugin marketplace add https://example.com/plugin", - ); - }); - - it("formats git-subdir source using its url", () => { - const source: PluginSource = { source: "git-subdir", url: "https://github.com/org/repo", path: "plugins/x" }; - expect(formatInstallCommand({ name: "my-plugin", source })).toBe( - "/plugin marketplace add https://github.com/org/repo", - ); - }); - - it("falls back to plugin name when no repo or url", () => { - const source: PluginSource = { source: "github" }; - expect(formatInstallCommand({ name: "my-plugin", source })).toBe("/plugin marketplace add my-plugin"); + it("uses the plugin name as the identifier", () => { + expect(formatInstallCommand({ name: "code-review" })).toBe("/plugin install code-review@litellm"); }); }); diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts index a4e70f78af1..85b4ef357f7 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts @@ -179,13 +179,14 @@ export const parseSkillSource = (rawUrl: string, subPath?: string): SkillSourceP /** * Build the `~/.claude/settings.json` snippet that registers the proxy as a marketplace. * Claude Code expects `extraKnownMarketplaces..source` to be a source object, not a - * bare `"url"` string, so the url/source pair is nested one level deeper. + * bare `"url"` string, so the url/source pair is nested one level deeper. The key must be + * "litellm" to match the name the proxy returns in marketplace.json. */ export const buildMarketplaceSettingsSnippet = (proxyOrigin: string): string => JSON.stringify( { extraKnownMarketplaces: { - "my-org": { + litellm: { source: { source: "url", url: `${proxyOrigin}/claude-code/marketplace.json`, @@ -198,20 +199,11 @@ export const buildMarketplaceSettingsSnippet = (proxyOrigin: string): string => ); /** - * Generate install command for Claude Code CLI - * Format: /plugin marketplace add org/repo OR /plugin marketplace add url + * Generate install command for Claude Code CLI. + * Installs the named plugin from the "litellm" marketplace registered in settings.json. */ -export const formatInstallCommand = (plugin: { name: string; source: PluginSource }): string => { - const { source } = plugin; - if (source.source === "github" && source.repo) { - return `/plugin marketplace add ${source.repo}`; - } - if ((source.source === "url" || source.source === "git-subdir") && source.url) { - return `/plugin marketplace add ${source.url}`; - } - // Fallback to plugin name - return `/plugin marketplace add ${plugin.name}`; -}; +export const formatInstallCommand = (plugin: { name: string }): string => + `/plugin install ${plugin.name}@litellm`; /** * Extract unique categories from plugins list diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx index fe001641135..8c99e9a3301 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx @@ -261,6 +261,32 @@ const SkillDetail: React.FC = ({ skill, onBack }) => {
+ {/* Shown when the marketplace catalog is stale and the plugin isn't found yet */} +
+

+ If you see "Plugin {skill.name} not found in marketplace", update the catalog first: +

+ 
+              /plugin marketplace update litellm
+
+
+

Don't have the marketplace configured yet?{" "} setActiveTab("setup")} style={{ color: "#1a73e8", cursor: "pointer" }}> @@ -276,12 +302,73 @@ const SkillDetail: React.FC = ({ skill, onBack }) => {

One-time marketplace setup

-

- Add this to{" "} + + {/* Option 1: single command — fastest path for most users */} +

+ Run this command in Claude Code to register the marketplace: +

+
+
+ Run in Claude Code + +
+ 
+              {`/plugin marketplace add ${typeof window !== "undefined" ? window.location.origin : ""}/claude-code/marketplace.json`}
+
+
+ + {/* Option 2: settings.json — for persistent config or managed deployments. + extraKnownMarketplaces requires source to be a nested object, not a flat string. */} +

+ Or add this to{" "} ~/.claude/settings.json {" "} - to point Claude Code at your proxy: + for a persistent configuration:

Date: Thu, 16 Jul 2026 14:12:29 +0300 Subject: [PATCH 023/542] fix(claude-code): fix prettier formatting and remove unused import --- .../src/components/claude_code_plugins/helpers.test.ts | 2 +- .../src/components/claude_code_plugins/helpers.ts | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts index a04134a63d8..1d07ca09df5 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts @@ -19,7 +19,7 @@ import { isValidSubPath, buildMarketplaceSettingsSnippet, } from "./helpers"; -import { MarketplacePluginEntry, PluginSource } from "./types"; +import { MarketplacePluginEntry } from "./types"; describe("buildMarketplaceSettingsSnippet", () => { it("nests the url under a source object so Claude Code accepts the marketplace", () => { diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts index 85b4ef357f7..5b3d4f20c1e 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts @@ -202,8 +202,7 @@ export const buildMarketplaceSettingsSnippet = (proxyOrigin: string): string => * Generate install command for Claude Code CLI. * Installs the named plugin from the "litellm" marketplace registered in settings.json. */ -export const formatInstallCommand = (plugin: { name: string }): string => - `/plugin install ${plugin.name}@litellm`; +export const formatInstallCommand = (plugin: { name: string }): string => `/plugin install ${plugin.name}@litellm`; /** * Extract unique categories from plugins list From 141281b1d658e192a2b188ddf03d19f6eab58814 Mon Sep 17 00:00:00 2001 From: ozolam Date: Thu, 16 Jul 2026 14:15:43 +0300 Subject: [PATCH 024/542] ci: retrigger CI From 346c065fe0c7516b563996d2cee7061ce276fc3a Mon Sep 17 00:00:00 2001 From: ozolam Date: Tue, 11 Aug 2026 21:34:41 +0300 Subject: [PATCH 025/542] fix(claude-code): restrict marketplace catalog mutations to proxy admins --- .../claude_code_marketplace.py | 24 +++++++ .../test_claude_code_marketplace.py | 71 +++++++++++++++++++ 2 files changed, 95 insertions(+) diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index 46ee9b0911d..4730372511a 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -28,6 +28,7 @@ from fastapi.responses import JSONResponse from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.resource_ownership import is_proxy_admin from litellm.repositories.table_repositories import ClaudeCodePluginRepository from litellm.types.proxy.claude_code_endpoints import ( ListPluginsResponse, @@ -221,6 +222,18 @@ def _name_conflict_error(name: str) -> HTTPException: ) +def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None: + """Catalog mutations are restricted to proxy admins: marketplace.json is served + unauthenticated and any registered/updated entry is immediately installable by + every user, so a non-admin key must never be able to add or overwrite one. + """ + if not is_proxy_admin(user_api_key_dict): + raise HTTPException( + status_code=403, + detail={"error": "Only proxy admins may modify the Claude Code plugin marketplace."}, + ) + + @router.post( "/claude-code/plugins", tags=["Claude Code Marketplace"], @@ -271,6 +284,8 @@ async def register_plugin( from prisma.errors import UniqueViolationError try: + _require_proxy_admin(user_api_key_dict) + prisma_client: Final = await _get_prisma_client() if not re.match(r"^[a-z0-9-]+$", request.name): @@ -468,6 +483,7 @@ async def get_plugin( async def update_plugin( plugin_name: str, request: UpdatePluginRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ Update an existing plugin in the LiteLLM marketplace. @@ -509,6 +525,8 @@ async def update_plugin( from prisma.errors import PrismaError try: + _require_proxy_admin(user_api_key_dict) + prisma_client: Final = await _get_prisma_client() _validate_plugin_source(request.source) @@ -570,6 +588,8 @@ async def enable_plugin( - plugin_name: The name of the plugin to enable """ try: + _require_proxy_admin(user_api_key_dict) + prisma_client: Final = await _get_prisma_client() plugin: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique( @@ -615,6 +635,8 @@ async def disable_plugin( - plugin_name: The name of the plugin to disable """ try: + _require_proxy_admin(user_api_key_dict) + prisma_client: Final = await _get_prisma_client() plugin: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique( @@ -660,6 +682,8 @@ async def delete_plugin( - plugin_name: The name of the plugin to delete """ try: + _require_proxy_admin(user_api_key_dict) + prisma_client: Final = await _get_prisma_client() plugin: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique( diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py b/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py index 18e0f2cb559..d1ccd86044c 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py @@ -18,6 +18,9 @@ from litellm.types.proxy.claude_code_endpoints import ( UpdatePluginRequest, ) from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketplace import ( + delete_plugin, + disable_plugin, + enable_plugin, get_marketplace, register_plugin, update_plugin, @@ -72,6 +75,12 @@ _USER = UserAPIKeyAuth( user_id="test-user", ) +_NON_ADMIN_USER = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-5678", + user_id="regular-user", +) + _GIT_SUBDIR_SOURCE = { "source": "git-subdir", "url": "https://github.com/org/monorepo.git", @@ -151,6 +160,7 @@ async def test_update_plugin_replaces_existing_source(): response = await update_plugin( plugin_name=name, request=UpdatePluginRequest(source=new_source, version="2.0.0", description="updated"), + user_api_key_dict=_USER, ) assert response.status == "success" @@ -170,6 +180,7 @@ async def test_update_plugin_not_found(): await update_plugin( plugin_name="does-not-exist", request=UpdatePluginRequest(source=_GIT_SUBDIR_SOURCE), + user_api_key_dict=_USER, ) assert exc_info.value.status_code == 404 @@ -213,6 +224,7 @@ async def test_update_plugin_db_error_maps_to_structured_500(): await update_plugin( plugin_name=name, request=UpdatePluginRequest(source={"source": "github", "repo": "org/replacement"}), + user_api_key_dict=_USER, ) assert exc_info.value.status_code == 500 @@ -341,3 +353,62 @@ async def test_register_plugin_unknown_source_type(): assert exc_info.value.status_code == 400 assert "git-subdir" in exc_info.value.detail["error"] + + +@pytest.mark.asyncio +async def test_register_plugin_rejects_non_admin(): + """A non-admin key cannot add an entry to the marketplace catalog.""" + request = RegisterPluginRequest(name="attacker-plugin", source=_GIT_SUBDIR_SOURCE) + + with pytest.raises(HTTPException) as exc_info: + await register_plugin(request=request, user_api_key_dict=_NON_ADMIN_USER) + + assert exc_info.value.status_code == 403 + + table = litellm.proxy.proxy_server.prisma_client.db.litellm_claudecodeplugintable + assert await table.find_unique(where={"name": "attacker-plugin"}) is None + + +@pytest.mark.asyncio +async def test_update_plugin_rejects_non_admin_overwrite(): + """A non-admin key cannot overwrite an existing plugin's source.""" + name = "trusted-plugin" + await register_plugin( + request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"), + user_api_key_dict=_USER, + ) + + malicious_source = {"source": "github", "repo": "attacker/malicious-repo"} + with pytest.raises(HTTPException) as exc_info: + await update_plugin( + plugin_name=name, + request=UpdatePluginRequest(source=malicious_source), + user_api_key_dict=_NON_ADMIN_USER, + ) + + assert exc_info.value.status_code == 403 + + stored = await _read_stored_manifest(name) + assert stored["source"] == _GIT_SUBDIR_SOURCE + + +@pytest.mark.asyncio +async def test_enable_disable_delete_plugin_reject_non_admin(): + """Non-admin keys cannot enable, disable, or delete catalog entries.""" + name = "trusted-plugin-2" + await register_plugin( + request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"), + user_api_key_dict=_USER, + ) + + for coro in ( + enable_plugin(plugin_name=name, user_api_key_dict=_NON_ADMIN_USER), + disable_plugin(plugin_name=name, user_api_key_dict=_NON_ADMIN_USER), + delete_plugin(plugin_name=name, user_api_key_dict=_NON_ADMIN_USER), + ): + with pytest.raises(HTTPException) as exc_info: + await coro + assert exc_info.value.status_code == 403 + + table = litellm.proxy.proxy_server.prisma_client.db.litellm_claudecodeplugintable + assert (await table.find_unique(where={"name": name})).enabled is True From 168ee2d9fdbd63cb3f2f4fd6de665595cbcbd8b8 Mon Sep 17 00:00:00 2001 From: ozolam Date: Tue, 11 Aug 2026 22:38:25 +0300 Subject: [PATCH 026/542] fix(claude-code): use Annotated dependency style to stay within B008 lint budget --- .../claude_code_marketplace.py | 30 ++++++++++++------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index 4730372511a..09ba2c93ea0 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -6,21 +6,21 @@ Plugins are stored as metadata + git source references in LiteLLM database. Actual plugin files are hosted on GitHub/GitLab/Bitbucket. Endpoints: -/claude-code/marketplace.json - GET - List plugins for Claude Code discovery -/claude-code/plugins - POST - Register a new plugin (create-only) -/claude-code/plugins - GET - List plugins (admin) -/claude-code/plugins/{name} - GET - Get plugin details -/claude-code/plugins/{name} - PUT - Update an existing plugin -/claude-code/plugins/{name}/enable - POST - Enable a plugin -/claude-code/plugins/{name}/disable - POST - Disable a plugin -/claude-code/plugins/{name} - DELETE - Delete a plugin +/claude-code/marketplace.json - GET - List plugins for Claude Code discovery (unauthenticated) +/claude-code/plugins - POST - Register a new plugin (create-only, proxy admin only) +/claude-code/plugins - GET - List plugins (any authenticated key) +/claude-code/plugins/{name} - GET - Get plugin details (any authenticated key) +/claude-code/plugins/{name} - PUT - Update an existing plugin (proxy admin only) +/claude-code/plugins/{name}/enable - POST - Enable a plugin (proxy admin only) +/claude-code/plugins/{name}/disable - POST - Disable a plugin (proxy admin only) +/claude-code/plugins/{name} - DELETE - Delete a plugin (proxy admin only) """ import json import re from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Final, Protocol, TypedDict +from typing import Annotated, Final, Protocol, TypedDict from fastapi import APIRouter, Depends, HTTPException from fastapi.responses import JSONResponse @@ -255,6 +255,8 @@ async def register_plugin( the same name already exists it returns 409 Conflict; use PUT /claude-code/plugins/{plugin_name} to update an existing plugin. + Requires a proxy admin API key. + Parameters: - name: Plugin name (kebab-case) - source: Git source reference (github, url, or git-subdir format) @@ -483,7 +485,7 @@ async def get_plugin( async def update_plugin( plugin_name: str, request: UpdatePluginRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """ Update an existing plugin in the LiteLLM marketplace. @@ -497,6 +499,8 @@ async def update_plugin( Returns 404 if no plugin with the given name exists; use POST /claude-code/plugins to create a new plugin. + Requires a proxy admin API key. + Parameters: - plugin_name: Name of the plugin to update (path parameter) - source: Git source reference (github, url, or git-subdir format) @@ -584,6 +588,8 @@ async def enable_plugin( """ Enable a disabled plugin. + Requires a proxy admin API key. + Parameters: - plugin_name: The name of the plugin to enable """ @@ -631,6 +637,8 @@ async def disable_plugin( """ Disable a plugin without deleting it. + Requires a proxy admin API key. + Parameters: - plugin_name: The name of the plugin to disable """ @@ -678,6 +686,8 @@ async def delete_plugin( """ Delete a plugin from the marketplace. + Requires a proxy admin API key. + Parameters: - plugin_name: The name of the plugin to delete """ From 3db1759d04ba8d888a08dbe8308cc75a348141d8 Mon Sep 17 00:00:00 2001 From: mateo Date: Thu, 13 Aug 2026 19:45:32 +0000 Subject: [PATCH 027/542] fix(bedrock): stop emitting an empty assistant delta after the finish_reason chunk Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/chat/invoke_handler.py | 8 +++- .../llms/bedrock/chat/test_invoke_handler.py | 46 +++++++++++++++++++ 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 8d2b3dae71b..595884ae630 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -561,6 +561,10 @@ class AWSEventStreamDecoder: elif "usage" in chunk_data: usage = converse_config._transform_usage(chunk_data.get("usage", {})) + carries_message_content: Final = any( + key in chunk_data for key in ("start", "delta", "contentBlockIndex", "stopReason") + ) + model_response_provider_specific_fields: Final = {} if "trace" in chunk_data: trace: Final = chunk_data.get("trace") @@ -571,8 +575,8 @@ class AWSEventStreamDecoder: finish_reason=finish_reason, index=0, # Always 0 - Bedrock never returns multiple choices delta=Delta( - content=text, - role="assistant", + content=text if carries_message_content else None, + role="assistant" if carries_message_content else None, tool_calls=[tool_use] if tool_use else None, provider_specific_fields=(provider_specific_fields if provider_specific_fields else None), thinking_blocks=thinking_blocks, 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..9783976db4a 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -1,3 +1,4 @@ +import datetime import os import sys from unittest.mock import AsyncMock, MagicMock @@ -8,6 +9,8 @@ sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.bedrock.chat.invoke_handler import ( AWSEventStreamDecoder, make_call, @@ -293,3 +296,46 @@ def test_make_sync_call_honors_explicit_stream_chunk_size(): response.iter_bytes.assert_called_once_with(chunk_size=2048) + +@pytest.mark.asyncio +async def test_converse_stream_ends_on_finish_reason_chunk(): + """The usage-only metadata event Bedrock sends after messageStop must not reach the caller as an extra + assistant delta following the finish_reason chunk.""" + model = "anthropic.claude-sonnet-4-6" + events = ( + {"role": "assistant"}, + {"contentBlockIndex": 0, "delta": {"text": "Hello"}}, + {"contentBlockIndex": 0, "delta": {"text": " world"}}, + {"contentBlockIndex": 0}, + {"stopReason": "end_turn"}, + {"usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15}, "metrics": {"latencyMs": 100}}, + ) + + async def bedrock_stream(): + decoder = AWSEventStreamDecoder(model=model) + for event in events: + yield decoder._chunk_parser(chunk_data=event) + + wrapper = CustomStreamWrapper( + completion_stream=bedrock_stream(), + model=model, + custom_llm_provider="bedrock", + logging_obj=LiteLLMLoggingObj( + model=model, + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=datetime.datetime.now(), + litellm_call_id="1234", + function_id="1234", + ), + ) + + chunks = [chunk async for chunk in wrapper] + + assert [choice.finish_reason for chunk in chunks for choice in chunk.choices].count("stop") == 1 + assert chunks[-1].choices[0].finish_reason == "stop", ( + f"stream must end on the finish_reason chunk, got trailing {chunks[-1].model_dump(exclude_none=True)}" + ) + assert any(getattr(chunk, "usage", None) is not None for chunk in wrapper.chunks) + From 0e3f52a4c0b071e8f297decce690a6ba2b6615ff Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Thu, 13 Aug 2026 19:18:22 -0500 Subject: [PATCH 028/542] feat(model_prices): add gemini-3.1-flash-lite-image Register Nano Banana 2 Lite on the unprefixed, gemini/, and vertex_ai/ keys so completion_cost and pass-through spend tracking no longer treat the model as unmapped --- ...odel_prices_and_context_window_backup.json | 95 +++++++ model_prices_and_context_window.json | 95 +++++++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 1 + ...ini_3_1_flash_lite_image_model_metadata.py | 242 ++++++++++++++++++ tests/test_litellm/test_utils.py | 2 + 5 files changed, 435 insertions(+) create mode 100644 tests/test_litellm/test_gemini_3_1_flash_lite_image_model_metadata.py diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1eb72c887b5..ca67d5d6844 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -18141,6 +18141,44 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-3.1-flash-lite-image": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_video_input": true, + "supports_vision": true + }, "gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 5e-07, @@ -19949,6 +19987,42 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/gemini-3.1-flash-lite-image": { + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "gemini", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "rpm": 1000, + "tpm": 4000000, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_vision": true + }, "gemini/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -38760,6 +38834,27 @@ "supports_reasoning": false, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" }, + "vertex_ai/gemini-3.1-flash-lite-image": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_vision": true, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, "vertex_ai/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 5e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1eb72c887b5..ca67d5d6844 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -18141,6 +18141,44 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-3.1-flash-lite-image": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_video_input": true, + "supports_vision": true + }, "gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 5e-07, @@ -19949,6 +19987,42 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/gemini-3.1-flash-lite-image": { + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "gemini", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "rpm": 1000, + "tpm": 4000000, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_vision": true + }, "gemini/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -38760,6 +38834,27 @@ "supports_reasoning": false, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" }, + "vertex_ai/gemini-3.1-flash-lite-image": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_vision": true, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, "vertex_ai/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 5e-07, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 3aa41e18f1e..36fc98a1f09 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -1546,6 +1546,7 @@ def test_service_tier_fallback_pricing(): [ "gemini-3-pro-image-preview", "gemini-3.1-flash-image-preview", + "gemini-3.1-flash-lite-image", ], ) def test_gemini_image_generation_cost_with_zero_text_tokens(model: str): diff --git a/tests/test_litellm/test_gemini_3_1_flash_lite_image_model_metadata.py b/tests/test_litellm/test_gemini_3_1_flash_lite_image_model_metadata.py new file mode 100644 index 00000000000..aa6f03a47ff --- /dev/null +++ b/tests/test_litellm/test_gemini_3_1_flash_lite_image_model_metadata.py @@ -0,0 +1,242 @@ +import json +from pathlib import Path + +import pytest + +import litellm +from litellm import completion_cost +from litellm.cost_calculator import cost_per_token +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +from litellm.llms.gemini.image_generation.cost_calculator import ( + cost_calculator as gemini_image_generation_cost_calculator, +) +from litellm.llms.vertex_ai.image_generation.cost_calculator import ( + cost_calculator as vertex_image_generation_cost_calculator, +) +from litellm.types.utils import ( + CompletionTokensDetailsWrapper, + ImageObject, + ImageResponse, + ImageUsage, + ImageUsageInputTokensDetails, + ModelResponse, + PromptTokensDetailsWrapper, + Usage, +) + +REPO_ROOT = Path(__file__).parents[2] +MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" + +UNPREFIXED = "gemini-3.1-flash-lite-image" +GEMINI = "gemini/gemini-3.1-flash-lite-image" +VERTEX = "vertex_ai/gemini-3.1-flash-lite-image" +ALL_KEYS = (UNPREFIXED, GEMINI, VERTEX) + +INPUT_COST = 2.5e-07 +INPUT_COST_BATCHES = 1.25e-07 +OUTPUT_TEXT_COST = 1.5e-06 +OUTPUT_TEXT_COST_BATCHES = 7.5e-07 +OUTPUT_IMAGE_TOKEN_COST = 3e-05 +OUTPUT_COST_PER_1K_IMAGE = 0.0336 +INPUT_COST_PER_IMAGE = 0.00028 +CACHE_READ_COST = 2.5e-08 +MAX_INPUT_TOKENS = 65536 +MAX_OUTPUT_TOKENS = 4096 +TOKENS_PER_1K_IMAGE = 1120 + + +def _load(path: Path) -> dict: + with open(path) as f: + return json.load(f) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + +@pytest.mark.parametrize("model", ALL_KEYS) +def test_gemini_3_1_flash_lite_image_is_registered(model: str): + info = _load(MAIN_PATH).get(model) + assert info is not None, f"{model} not found in model_prices_and_context_window.json" + + assert info["mode"] == "image_generation" + assert info["input_cost_per_token"] == INPUT_COST + assert info["input_cost_per_token_batches"] == INPUT_COST_BATCHES + assert info["output_cost_per_token"] == OUTPUT_TEXT_COST + assert info["output_cost_per_token_batches"] == OUTPUT_TEXT_COST_BATCHES + assert info["output_cost_per_image"] == OUTPUT_COST_PER_1K_IMAGE + assert info["output_cost_per_image_token"] == OUTPUT_IMAGE_TOKEN_COST + assert info["max_input_tokens"] == MAX_INPUT_TOKENS + assert info["max_output_tokens"] == MAX_OUTPUT_TOKENS + assert info["max_tokens"] == MAX_OUTPUT_TOKENS + assert info["supports_reasoning"] is False + assert info["supports_response_schema"] is False + assert info["supports_vision"] is True + for field in ("supports_web_search", "search_context_cost_per_query", "web_search_billing_unit"): + assert field not in info + + +def test_gemini_3_1_flash_lite_image_provider_specific_fields(): + cost_map = _load(MAIN_PATH) + + unprefixed = cost_map[UNPREFIXED] + assert unprefixed["litellm_provider"] == "vertex_ai-language-models" + assert unprefixed["cache_read_input_token_cost"] == CACHE_READ_COST + assert unprefixed["input_cost_per_image"] == INPUT_COST_PER_IMAGE + assert unprefixed["supports_function_calling"] is False + assert unprefixed["supports_prompt_caching"] is True + assert unprefixed["supports_pdf_input"] is True + assert unprefixed["supports_video_input"] is True + assert unprefixed["supported_modalities"] == ["text", "image", "video"] + + gemini = cost_map[GEMINI] + assert gemini["litellm_provider"] == "gemini" + assert gemini["supports_function_calling"] is True + assert gemini["supports_prompt_caching"] is False + assert "cache_read_input_token_cost" not in gemini + assert gemini["supported_modalities"] == ["text", "image"] + assert gemini["supported_output_modalities"] == ["text", "image"] + assert gemini["rpm"] == 1000 + assert gemini["tpm"] == 4000000 + assert gemini["input_cost_per_image"] == INPUT_COST_PER_IMAGE + + vertex = cost_map[VERTEX] + assert vertex["litellm_provider"] == "vertex_ai-language-models" + assert vertex["cache_read_input_token_cost"] == CACHE_READ_COST + assert vertex["input_cost_per_image"] == INPUT_COST_PER_IMAGE + assert vertex["supports_function_calling"] is False + assert vertex["supports_prompt_caching"] is True + + +def test_one_k_image_price_matches_official_token_math(): + assert TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST == OUTPUT_COST_PER_1K_IMAGE + assert TOKENS_PER_1K_IMAGE * INPUT_COST == INPUT_COST_PER_IMAGE + + +@pytest.mark.parametrize("model", ALL_KEYS) +def test_backup_matches_main(model: str): + main_cost = _load(MAIN_PATH) + backup_cost = _load(BACKUP_PATH) + assert backup_cost.get(model) == main_cost.get(model), f"{model} differs between main and backup model cost maps" + + +def test_gemini_prefix_routes_to_gemini(): + routed_model, provider, _, _ = get_llm_provider(model=GEMINI) + assert routed_model == UNPREFIXED + assert provider == "gemini" + + +def test_vertex_prefix_routes_to_vertex(): + routed_model, provider, _, _ = get_llm_provider(model=VERTEX) + assert routed_model == UNPREFIXED + assert provider == "vertex_ai" + + +def test_text_token_cost(local_model_cost_map): + prompt_cost, text_completion_cost = cost_per_token(model=GEMINI, prompt_tokens=1000, completion_tokens=500) + assert prompt_cost == pytest.approx(1000 * INPUT_COST) + assert text_completion_cost == pytest.approx(500 * OUTPUT_TEXT_COST) + + +def test_completion_cost_bills_one_k_image(local_model_cost_map): + response = ModelResponse() + response.model = UNPREFIXED + response.usage = Usage( + prompt_tokens=7, + completion_tokens=TOKENS_PER_1K_IMAGE, + total_tokens=7 + TOKENS_PER_1K_IMAGE, + completion_tokens_details=CompletionTokensDetailsWrapper(image_tokens=TOKENS_PER_1K_IMAGE, text_tokens=0), + ) + billed = completion_cost( + completion_response=response, + model=UNPREFIXED, + custom_llm_provider="vertex_ai", + ) + expected = TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + 7 * INPUT_COST + assert billed == pytest.approx(expected) + + +def test_image_tokens_are_not_billed_as_text(local_model_cost_map): + usage = Usage( + completion_tokens=1345, + prompt_tokens=10, + total_tokens=1355, + completion_tokens_details=CompletionTokensDetailsWrapper( + accepted_prediction_tokens=None, + audio_tokens=None, + reasoning_tokens=225, + rejected_prediction_tokens=None, + text_tokens=0, + image_tokens=TOKENS_PER_1K_IMAGE, + ), + prompt_tokens_details=PromptTokensDetailsWrapper( + audio_tokens=None, cached_tokens=None, text_tokens=10, image_tokens=None + ), + ) + + _, image_completion_cost = generic_cost_per_token( + model=UNPREFIXED, + usage=usage, + custom_llm_provider="vertex_ai", + ) + + expected_completion_cost = TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + 225 * OUTPUT_TEXT_COST + bugged_text_only_cost = 1345 * OUTPUT_TEXT_COST + assert image_completion_cost > bugged_text_only_cost * 2 + assert image_completion_cost == pytest.approx(expected_completion_cost) + + +def test_gemini_image_generation_uses_token_pricing(local_model_cost_map): + image_response = ImageResponse( + data=[ImageObject(b64_json="img1")], + usage=ImageUsage( + input_tokens=50 + TOKENS_PER_1K_IMAGE, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=50, + image_tokens=TOKENS_PER_1K_IMAGE, + ), + output_tokens=TOKENS_PER_1K_IMAGE, + total_tokens=50 + TOKENS_PER_1K_IMAGE + TOKENS_PER_1K_IMAGE, + ), + ) + + cost = gemini_image_generation_cost_calculator(model=GEMINI, image_response=image_response) + expected = (50 + TOKENS_PER_1K_IMAGE) * INPUT_COST + TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + assert cost == pytest.approx(expected) + assert cost != OUTPUT_COST_PER_1K_IMAGE + + +def test_vertex_image_generation_uses_token_pricing(local_model_cost_map): + image_response = ImageResponse( + data=[ImageObject(b64_json="img1")], + usage=ImageUsage( + input_tokens=50 + TOKENS_PER_1K_IMAGE, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=50, + image_tokens=TOKENS_PER_1K_IMAGE, + ), + output_tokens=TOKENS_PER_1K_IMAGE, + total_tokens=50 + TOKENS_PER_1K_IMAGE + TOKENS_PER_1K_IMAGE, + ), + ) + + cost = vertex_image_generation_cost_calculator(model=UNPREFIXED, image_response=image_response) + expected = (50 + TOKENS_PER_1K_IMAGE) * INPUT_COST + TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + assert cost == pytest.approx(expected) + + +def test_vertex_image_generation_falls_back_to_flat_image_price(local_model_cost_map): + image_response = ImageResponse(data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")]) + cost = vertex_image_generation_cost_calculator(model=UNPREFIXED, image_response=image_response) + assert cost == pytest.approx(2 * OUTPUT_COST_PER_1K_IMAGE) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 8e9e6167fb9..ada053ee38c 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4312,11 +4312,13 @@ class TestVertexEmbeddingEncodingFormat: "vertex_ai/gemini-3-pro-image-preview", "vertex_ai/gemini-3.1-flash-image", "vertex_ai/gemini-3.1-flash-image-preview", + "vertex_ai/gemini-3.1-flash-lite-image", "gemini/gemini-2.5-flash-image", "gemini/gemini-3-pro-image", "gemini/gemini-3-pro-image-preview", "gemini/gemini-3.1-flash-image", "gemini/gemini-3.1-flash-image-preview", + "gemini/gemini-3.1-flash-lite-image", ], ) def test_gemini_image_models_do_not_support_reasoning( 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 029/542] 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 030/542] 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 031/542] 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 032/542] 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 033/542] 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 034/542] 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 035/542] 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 b6ee13803d1c21bc0b07dac6e6b9ad2ede7ad7aa Mon Sep 17 00:00:00 2001 From: HarryZhou <2373256746@qq.com> Date: Sun, 9 Aug 2026 22:59:25 +0800 Subject: [PATCH 036/542] fix(responses-bridge): preserve reasoning input items as reasoning_content --- .../transformation.py | 164 +++++++++++++++++- .../test_reasoning_input_item_preservation.py | 147 ++++++++++++++++ 2 files changed, 310 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 4892e3b348c..5d3ed0477e3 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -557,7 +557,108 @@ class LiteLLMCompletionResponsesConfig: continue messages.extend(chat_completion_messages) - return messages + return LiteLLMCompletionResponsesConfig._merge_reasoning_only_assistant_messages(messages) + + @staticmethod + def _merge_reasoning_only_assistant_messages( + messages: list[ + AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionMessageToolCall + | ChatCompletionResponseMessage + ], + ) -> list[ + AllMessageValues | GenericChatCompletionMessage | ChatCompletionMessageToolCall | ChatCompletionResponseMessage + ]: + """ + Responses API emits prior-turn reasoning as its own ``reasoning`` input + item, which becomes a standalone assistant message with + ``content=None`` + ``reasoning_content``. Chat-completions providers + (e.g. DeepSeek V4, Kimi K2.6) expect the chain-of-thought on the + assistant message that carries the answer or tool calls. This pass + merges standalone reasoning-only assistant messages into the + immediately following assistant message. + + If the reasoning item is not followed by an assistant message (e.g. a + stateless chain replays ``reasoning`` + ``user``), the standalone + reasoning message is preserved so the reasoning is still passed back. + """ + + def _role(msg: Any) -> str: + if isinstance(msg, dict): + return str(msg.get("role") or "") + return str(getattr(msg, "role", "") or "") + + def _reasoning_text(msg: Any) -> str | None: + if isinstance(msg, dict): + value = msg.get("reasoning_content") + else: + value = getattr(msg, "reasoning_content", None) + return value if isinstance(value, str) and value else None + + def _content(msg: Any) -> Any: + if isinstance(msg, dict): + return msg.get("content") + return getattr(msg, "content", None) + + def _tool_calls(msg: Any) -> Any: + if isinstance(msg, dict): + return msg.get("tool_calls") + return getattr(msg, "tool_calls", None) + + merged: list[ + AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionMessageToolCall + | ChatCompletionResponseMessage + ] = [] + pending_reasoning: list[str] = [] + + for msg in messages: + if ( + _role(msg) == "assistant" + and _content(msg) is None + and not _tool_calls(msg) + and _reasoning_text(msg) is not None + ): + pending_reasoning.append(_reasoning_text(msg) or "") + continue + + if pending_reasoning and _role(msg) == "assistant": + combined = "\n".join(pending_reasoning) + existing = _reasoning_text(msg) + if existing: + combined = existing + "\n" + combined + if isinstance(msg, dict): + msg["reasoning_content"] = combined + else: + setattr(msg, "reasoning_content", combined) + pending_reasoning = [] + elif pending_reasoning: + # Not followed by an assistant message — keep the reasoning + # standalone instead of dropping it. + for text in pending_reasoning: + merged.append( + ChatCompletionResponseMessage( + role="assistant", + content=None, + reasoning_content=text, + ) + ) + pending_reasoning = [] + + merged.append(msg) + + for text in pending_reasoning: + merged.append( + ChatCompletionResponseMessage( + role="assistant", + content=None, + reasoning_content=text, + ) + ) + + return merged @staticmethod def _merged_trailing_assistant_message( @@ -1026,6 +1127,25 @@ class LiteLLMCompletionResponsesConfig: return LiteLLMCompletionResponsesConfig._transform_responses_api_function_call_to_chat_completion_message( function_call=input_item ) + elif input_item.get("type") == "reasoning": + # A ResponseReasoningItemParam carries the prior-turn chain-of-thought. + # Chat-completions providers (DeepSeek V4, Kimi K2.6, ...) expect this + # to be replayed as `reasoning_content` on an assistant message, not as + # visible `content` (prompt pollution) and not dropped (DeepSeek V4 + # rejects multi-turn requests with a missing `reasoning_content`). + reasoning_text = LiteLLMCompletionResponsesConfig._extract_reasoning_text_from_input_item(input_item) + if not reasoning_text: + # No plaintext reasoning is available (e.g. encrypted_content only). + # Chat-completions providers cannot consume opaque encrypted blobs, + # so skip the item instead of polluting the prompt. + return [] + return [ + ChatCompletionResponseMessage( + role="assistant", + content=None, + reasoning_content=reasoning_text, + ) + ] else: content: Final[object] = input_item.get("content") # Handle None content: Responses API allows None content, but GenericChatCompletionMessage requires content @@ -1041,6 +1161,48 @@ class LiteLLMCompletionResponsesConfig: ) ] + @staticmethod + def _extract_reasoning_text_from_input_item(input_item: Mapping[str, object]) -> str | None: + """ + Extract plaintext reasoning from a ResponseReasoningItemParam. + + Handles: + - content as a string + - content as a list of blocks (output_text / summary_text / text) + - summary as a list of summary_text blocks (fallback) + + Returns None when only opaque forms (e.g. encrypted_content) are present. + """ + content: Final[object] = input_item.get("content") + if isinstance(content, str) and content.strip(): + return content + if isinstance(content, list): + text_parts: list[str] = [] + for block in content: + if not isinstance(block, Mapping): + continue + block_type = block.get("type") + if block_type in ("encrypted_content", "redacted_thinking"): + continue + text = block.get("text") + if isinstance(text, str) and text.strip(): + text_parts.append(text.strip()) + if text_parts: + return "\n".join(text_parts) + + summary: Final[object] = input_item.get("summary") + if isinstance(summary, list): + text_parts = [] + for block in summary: + if not isinstance(block, Mapping): + continue + text = block.get("text") + if isinstance(text, str) and text.strip(): + text_parts.append(text.strip()) + if text_parts: + return "\n".join(text_parts) + return None + @staticmethod def _is_input_item_tool_call_output(input_item: Mapping[str, object]) -> bool: """ diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py new file mode 100644 index 00000000000..5fcd4df3ff8 --- /dev/null +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py @@ -0,0 +1,147 @@ +""" +Unit tests for preserving prior-turn ``reasoning`` input items when the +Responses API is bridged to chat completions. + +Without this handling, a ``ResponseReasoningItemParam`` falls through to the +generic message branch, polluting the prompt as visible assistant ``content`` +or being silently dropped. Chat-completions providers such as DeepSeek V4 and +Kimi K2.6 require the chain-of-thought to be replayed as ``reasoning_content`` +on an assistant message. +""" + +from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, +) + + +def _transform_item(item): + return LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( + input_item=item + ) + + +def _transform_input(input_items): + return LiteLLMCompletionResponsesConfig._transform_response_input_param_to_chat_completion_message( + input=input_items + ) + + +class TestReasoningInputItemHandler: + """Reasoning input items map to assistant ``reasoning_content``.""" + + def test_reasoning_item_with_output_text_content(self): + """Standard Responses-API reasoning item with output_text blocks.""" + item = { + "type": "reasoning", + "id": "rs_abc", + "summary": [], + "content": [{"type": "output_text", "text": "step 1: think about X"}], + } + messages = _transform_item(item) + assert len(messages) == 1 + assert messages[0]["role"] == "assistant" + assert messages[0]["content"] is None + assert messages[0]["reasoning_content"] == "step 1: think about X" + + def test_reasoning_item_with_string_content(self): + """Variant: reasoning content as a plain string.""" + item = {"type": "reasoning", "id": "rs_1", "content": "step 1: ..."} + messages = _transform_item(item) + assert messages[0]["reasoning_content"] == "step 1: ..." + + def test_reasoning_item_with_summary_only(self): + """SDK form: reasoning carried in summary list, no content.""" + item = { + "type": "reasoning", + "id": "rs_2", + "summary": [{"type": "summary_text", "text": "..."}], + } + messages = _transform_item(item) + assert messages[0]["reasoning_content"] == "..." + + def test_reasoning_item_with_encrypted_content_only_dropped(self): + """Opaque encrypted reasoning cannot be forwarded to chat completions.""" + item = {"type": "reasoning", "id": "rs_3", "encrypted_content": "opaque-blob"} + assert _transform_item(item) == [] + + def test_reasoning_item_empty_dropped(self): + """Reasoning item with neither content nor summary drops cleanly.""" + assert _transform_item({"type": "reasoning", "id": "rs_4"}) == [] + + +class TestReasoningInputItemMerging: + """Standalone reasoning messages merge into the following assistant turn.""" + + def test_reasoning_merged_into_following_assistant_message(self): + """Reasoning + assistant answer become one assistant message.""" + messages = _transform_input( + [ + { + "type": "reasoning", + "id": "rs_1", + "content": [{"type": "output_text", "text": "secret reasoning"}], + }, + {"type": "message", "role": "assistant", "content": "The answer."}, + ] + ) + assert len(messages) == 1 + assert messages[0]["role"] == "assistant" + assert messages[0]["content"] == "The answer." + assert messages[0]["reasoning_content"] == "secret reasoning" + + def test_reasoning_preserved_when_followed_by_user_message(self): + """Stateless chain: reasoning + user prompt keeps the reasoning turn.""" + messages = _transform_input( + [ + { + "type": "reasoning", + "id": "rs_1", + "content": [{"type": "output_text", "text": "secret BLUEBERRY"}], + }, + {"role": "user", "content": "What is the secret word?"}, + ] + ) + assert len(messages) == 2 + assert messages[0]["role"] == "assistant" + assert messages[0]["content"] is None + assert messages[0]["reasoning_content"] == "secret BLUEBERRY" + assert messages[1]["role"] == "user" + + def test_reasoning_merged_into_function_call_assistant(self): + """Reasoning + function_call becomes one assistant tool-call message.""" + messages = _transform_input( + [ + { + "type": "reasoning", + "id": "rs_1", + "content": [{"type": "output_text", "text": "I should look this up"}], + }, + { + "type": "function_call", + "call_id": "call_1", + "name": "lookup", + "arguments": '{"cwe": "79"}', + }, + ] + ) + assert len(messages) == 1 + assert messages[0]["role"] == "assistant" + assert messages[0]["reasoning_content"] == "I should look this up" + assert len(messages[0]["tool_calls"]) == 1 + + +class TestNonReasoningInputItemUnchanged: + """Non-reasoning items still flow through the existing branches.""" + + def test_user_message_unchanged(self): + item = {"role": "user", "content": "hello"} + out = _transform_item(item) + assert len(out) == 1 + assert out[0]["role"] == "user" + + def test_assistant_message_unchanged(self): + item = {"role": "assistant", "content": "hi"} + out = _transform_item(item) + assert len(out) == 1 + assert out[0]["role"] == "assistant" + assert out[0]["content"] == "hi" From 3a77556dc14660e88a7d20f54e9762c39f24b749 Mon Sep 17 00:00:00 2001 From: HarryZhou <2373256746@qq.com> Date: Mon, 10 Aug 2026 00:45:25 +0800 Subject: [PATCH 037/542] fix(responses-bridge): preserve reasoning merge order when assistant already has reasoning_content --- .../transformation.py | 2 +- .../test_reasoning_input_item_preservation.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 5d3ed0477e3..0604c3636ff 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -628,7 +628,7 @@ class LiteLLMCompletionResponsesConfig: combined = "\n".join(pending_reasoning) existing = _reasoning_text(msg) if existing: - combined = existing + "\n" + combined + combined = combined + "\n" + existing if isinstance(msg, dict): msg["reasoning_content"] = combined else: diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py index 5fcd4df3ff8..ecc024b7d04 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py @@ -129,6 +129,18 @@ class TestReasoningInputItemMerging: assert messages[0]["reasoning_content"] == "I should look this up" assert len(messages[0]["tool_calls"]) == 1 + def test_reasoning_merged_into_assistant_with_existing_reasoning_content(self): + """Old reasoning precedes existing reasoning on the target assistant turn.""" + messages = LiteLLMCompletionResponsesConfig._merge_reasoning_only_assistant_messages( + [ + {"role": "assistant", "content": None, "reasoning_content": "old reasoning"}, + {"role": "assistant", "content": "The answer.", "reasoning_content": "new reasoning"}, + ] + ) + assert len(messages) == 1 + assert messages[0]["content"] == "The answer." + assert messages[0]["reasoning_content"] == "old reasoning\nnew reasoning" + class TestNonReasoningInputItemUnchanged: """Non-reasoning items still flow through the existing branches.""" From 5911124f1dbba1e9c58f3b53619c3f875752a20f Mon Sep 17 00:00:00 2001 From: HarryZhou <2373256746@qq.com> Date: Mon, 10 Aug 2026 00:58:47 +0800 Subject: [PATCH 038/542] fix(responses-bridge): satisfy ruff strict-rule budget in reasoning merge --- .../transformation.py | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 0604c3636ff..2c506d4a4c7 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -584,24 +584,24 @@ class LiteLLMCompletionResponsesConfig: reasoning message is preserved so the reasoning is still passed back. """ - def _role(msg: Any) -> str: + def _role(msg: object) -> str: if isinstance(msg, dict): return str(msg.get("role") or "") return str(getattr(msg, "role", "") or "") - def _reasoning_text(msg: Any) -> str | None: + def _reasoning_text(msg: object) -> str | None: if isinstance(msg, dict): value = msg.get("reasoning_content") else: value = getattr(msg, "reasoning_content", None) return value if isinstance(value, str) and value else None - def _content(msg: Any) -> Any: + def _content(msg: object) -> object | None: if isinstance(msg, dict): return msg.get("content") return getattr(msg, "content", None) - def _tool_calls(msg: Any) -> Any: + def _tool_calls(msg: object) -> object | None: if isinstance(msg, dict): return msg.get("tool_calls") return getattr(msg, "tool_calls", None) @@ -632,31 +632,35 @@ class LiteLLMCompletionResponsesConfig: if isinstance(msg, dict): msg["reasoning_content"] = combined else: - setattr(msg, "reasoning_content", combined) + setattr(msg, "reasoning_content", combined) # noqa: B010 pending_reasoning = [] elif pending_reasoning: # Not followed by an assistant message — keep the reasoning # standalone instead of dropping it. - for text in pending_reasoning: - merged.append( + merged.extend( + [ ChatCompletionResponseMessage( role="assistant", content=None, reasoning_content=text, ) - ) + for text in pending_reasoning + ] + ) pending_reasoning = [] merged.append(msg) - for text in pending_reasoning: - merged.append( + merged.extend( + [ ChatCompletionResponseMessage( role="assistant", content=None, reasoning_content=text, ) - ) + for text in pending_reasoning + ] + ) return merged From 438c1850fe1223feec1e2e6e5b48f0a6c15a1328 Mon Sep 17 00:00:00 2001 From: HarryZhou <2373256746@qq.com> Date: Mon, 10 Aug 2026 01:09:26 +0800 Subject: [PATCH 039/542] fix(responses-bridge): satisfy type-discipline budget in reasoning merge --- .../transformation.py | 34 ++++++++++--------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 2c506d4a4c7..e3e62ab3c55 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -561,13 +561,13 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _merge_reasoning_only_assistant_messages( - messages: list[ + messages: list[ # mutable-ok: input sequence AllMessageValues | GenericChatCompletionMessage | ChatCompletionMessageToolCall | ChatCompletionResponseMessage ], - ) -> list[ + ) -> list[ # mutable-ok: fresh merged list AllMessageValues | GenericChatCompletionMessage | ChatCompletionMessageToolCall | ChatCompletionResponseMessage ]: """ @@ -591,9 +591,9 @@ class LiteLLMCompletionResponsesConfig: def _reasoning_text(msg: object) -> str | None: if isinstance(msg, dict): - value = msg.get("reasoning_content") + value = msg.get("reasoning_content") # rebind-ok: branch lookup else: - value = getattr(msg, "reasoning_content", None) + value = getattr(msg, "reasoning_content", None) # rebind-ok: branch lookup return value if isinstance(value, str) and value else None def _content(msg: object) -> object | None: @@ -606,13 +606,13 @@ class LiteLLMCompletionResponsesConfig: return msg.get("tool_calls") return getattr(msg, "tool_calls", None) - merged: list[ + merged: list[ # mutable-ok: accumulator # rebind-ok: accumulator AllMessageValues | GenericChatCompletionMessage | ChatCompletionMessageToolCall | ChatCompletionResponseMessage - ] = [] - pending_reasoning: list[str] = [] + ] = [] # mutable-ok: accumulator + pending_reasoning: list[str] = [] # mutable-ok: accumulator # rebind-ok: accumulator for msg in messages: if ( @@ -633,11 +633,11 @@ class LiteLLMCompletionResponsesConfig: msg["reasoning_content"] = combined else: setattr(msg, "reasoning_content", combined) # noqa: B010 - pending_reasoning = [] + pending_reasoning = [] # mutable-ok: reset accumulator elif pending_reasoning: # Not followed by an assistant message — keep the reasoning # standalone instead of dropping it. - merged.extend( + merged.extend( # mutable-ok: append reasoning messages [ ChatCompletionResponseMessage( role="assistant", @@ -647,11 +647,11 @@ class LiteLLMCompletionResponsesConfig: for text in pending_reasoning ] ) - pending_reasoning = [] + pending_reasoning = [] # mutable-ok: reset accumulator merged.append(msg) - merged.extend( + merged.extend( # mutable-ok: append trailing reasoning [ ChatCompletionResponseMessage( role="assistant", @@ -1137,13 +1137,15 @@ class LiteLLMCompletionResponsesConfig: # to be replayed as `reasoning_content` on an assistant message, not as # visible `content` (prompt pollution) and not dropped (DeepSeek V4 # rejects multi-turn requests with a missing `reasoning_content`). - reasoning_text = LiteLLMCompletionResponsesConfig._extract_reasoning_text_from_input_item(input_item) + reasoning_text = LiteLLMCompletionResponsesConfig._extract_reasoning_text_from_input_item( # rebind-ok: extraction result + input_item + ) if not reasoning_text: # No plaintext reasoning is available (e.g. encrypted_content only). # Chat-completions providers cannot consume opaque encrypted blobs, # so skip the item instead of polluting the prompt. - return [] - return [ + return [] # mutable-ok: empty drop result + return [ # mutable-ok: single message result ChatCompletionResponseMessage( role="assistant", content=None, @@ -1181,7 +1183,7 @@ class LiteLLMCompletionResponsesConfig: if isinstance(content, str) and content.strip(): return content if isinstance(content, list): - text_parts: list[str] = [] + text_parts: list[str] = [] # mutable-ok: text accumulator # rebind-ok: text accumulator for block in content: if not isinstance(block, Mapping): continue @@ -1196,7 +1198,7 @@ class LiteLLMCompletionResponsesConfig: summary: Final[object] = input_item.get("summary") if isinstance(summary, list): - text_parts = [] + text_parts = [] # mutable-ok: text accumulator # rebind-ok: text accumulator for block in summary: if not isinstance(block, Mapping): continue From de95372dfbd7bbba8c478815340dd49c1b21da11 Mon Sep 17 00:00:00 2001 From: HarryZhou <2373256746@qq.com> Date: Mon, 10 Aug 2026 01:24:24 +0800 Subject: [PATCH 040/542] fix(responses-bridge): type-safe reasoning_content assignment in merge pass --- .../litellm_completion_transformation/transformation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index e3e62ab3c55..d7b6b8c7b8f 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -630,7 +630,7 @@ class LiteLLMCompletionResponsesConfig: if existing: combined = combined + "\n" + existing if isinstance(msg, dict): - msg["reasoning_content"] = combined + cast(dict[str, Any], msg)["reasoning_content"] = combined # cast-ok: mutable reasoning carrier else: setattr(msg, "reasoning_content", combined) # noqa: B010 pending_reasoning = [] # mutable-ok: reset accumulator From 2d4e6afe1c7d6d3233a18668d53fafe4cffa50b3 Mon Sep 17 00:00:00 2001 From: HarryZhou <2373256746@qq.com> Date: Tue, 18 Aug 2026 21:11:29 +0800 Subject: [PATCH 041/542] fix(guardrails): inspect responses reasoning content and summary text --- litellm/proxy/guardrails/_content_utils.py | 59 +++++++++++++------ .../transformation.py | 6 +- .../proxy/guardrails/test_content_utils.py | 54 +++++++++++++++++ 3 files changed, 100 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 6ed6f0013df..ae92adcb1ee 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -8,7 +8,7 @@ skip the other shapes — these helpers normalise that so every hook sees every text fragment. """ -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Mapping from typing import Any, Final # Call types whose body carries free-form chat / prompt text that @@ -33,7 +33,9 @@ def is_text_content_call_type(call_type: str) -> bool: return call_type in TEXT_CONTENT_CALL_TYPES -TEXT_PART_TYPES: Final[frozenset[str]] = frozenset({"text", "input_text", "output_text"}) +TEXT_PART_TYPES: Final[frozenset[str]] = frozenset( + {"text", "input_text", "output_text", "summary_text", "reasoning_text"} +) # Responses-API item types whose ``output`` field carries user/tool text # that guardrails should inspect. ``function_call_output`` is the @@ -42,6 +44,16 @@ TEXT_PART_TYPES: Final[frozenset[str]] = frozenset({"text", "input_text", "outpu _OUTPUT_ITEM_TYPES: Final[frozenset[str]] = frozenset({"function_call_output", "custom_tool_call_output"}) +def _part_text(part: Mapping[str, object]) -> str | None: + """Return non-empty plaintext from any content part that carries ``text``.""" + if not isinstance(part, dict): + return None + text = part.get("text") + if isinstance(text, str) and text: + return text + return None + + def _iter_text_parts_in_content(content: Any) -> Iterator[str]: """Yield text fragments from a ``message.content`` value (string or multimodal list). Non-text parts (images, audio, …) are skipped.""" @@ -58,10 +70,9 @@ def _iter_text_parts_in_content(content: Any) -> Iterator[str]: continue if not isinstance(part, dict): continue - if part.get("type") in TEXT_PART_TYPES: - text = part.get("text") - if isinstance(text, str) and text: - yield text + text = _part_text(part) + if text is not None: + yield text def _coerce_input_to_messages(input_value: Any) -> list[dict[str, Any]]: @@ -75,8 +86,23 @@ def _coerce_input_to_messages(input_value: Any) -> list[dict[str, Any]]: if isinstance(item, str): messages.append({"role": "user", "content": item}) elif isinstance(item, dict): - if item.get("type") in TEXT_PART_TYPES: + if _part_text(item) is not None: messages.append({"role": item.get("role") or "user", "content": [item]}) + elif item.get("type") == "reasoning": + if "content" in item: + messages.append( + { # mutable-ok: append reasoning content + "role": item.get("role") or "assistant", + "content": item["content"], + } + ) + if isinstance(item.get("summary"), list): + messages.append( + { # mutable-ok: append reasoning summary + "role": item.get("role") or "assistant", + "content": item["summary"], + } + ) elif "content" in item: messages.append({"role": item.get("role") or "user", "content": item["content"]}) elif item.get("type") in _OUTPUT_ITEM_TYPES and "output" in item: @@ -126,12 +152,7 @@ def walk_user_text(data: dict[str, Any], visit: Callable[[str], str]) -> int: if isinstance(part, str) and part: visited += 1 new_parts.append(visit(part)) - elif ( - isinstance(part, dict) - and part.get("type") in TEXT_PART_TYPES - and isinstance(part.get("text"), str) - and part["text"] - ): + elif isinstance(part, dict) and _part_text(part) is not None: visited += 1 new_parts.append({**part, "text": visit(part["text"])}) else: @@ -158,10 +179,14 @@ def walk_user_text(data: dict[str, Any], visit: Callable[[str], str]) -> int: visited += 1 input_value[idx] = visit(item) elif isinstance(item, dict): - if item.get("type") in TEXT_PART_TYPES: - if isinstance(item.get("text"), str) and item["text"]: - visited += 1 - input_value[idx] = {**item, "text": visit(item["text"])} + if _part_text(item) is not None: + visited += 1 + input_value[idx] = {**item, "text": visit(item["text"])} # mutable-ok: rewrite text part in place + elif item.get("type") == "reasoning": + if "content" in item: + item["content"] = _rewrite_content(item["content"]) + if isinstance(item.get("summary"), list): + item["summary"] = _rewrite_content(item["summary"]) elif "content" in item: item["content"] = _rewrite_content(item["content"]) elif item.get("type") in _OUTPUT_ITEM_TYPES and "output" in item: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index d7b6b8c7b8f..c5f40242bfd 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -638,7 +638,7 @@ class LiteLLMCompletionResponsesConfig: # Not followed by an assistant message — keep the reasoning # standalone instead of dropping it. merged.extend( # mutable-ok: append reasoning messages - [ + [ # mutable-ok: append reasoning messages ChatCompletionResponseMessage( role="assistant", content=None, @@ -652,7 +652,7 @@ class LiteLLMCompletionResponsesConfig: merged.append(msg) merged.extend( # mutable-ok: append trailing reasoning - [ + [ # mutable-ok: append trailing reasoning ChatCompletionResponseMessage( role="assistant", content=None, @@ -1196,6 +1196,8 @@ class LiteLLMCompletionResponsesConfig: if text_parts: return "\n".join(text_parts) + # Guardrail traversal in litellm/proxy/guardrails/_content_utils.py + # inspects and rewrites these summary blocks before they are forwarded. summary: Final[object] = input_item.get("summary") if isinstance(summary, list): text_parts = [] # mutable-ok: text accumulator # rebind-ok: text accumulator diff --git a/tests/test_litellm/proxy/guardrails/test_content_utils.py b/tests/test_litellm/proxy/guardrails/test_content_utils.py index 3dfb98c12ea..d9e079c6d92 100644 --- a/tests/test_litellm/proxy/guardrails/test_content_utils.py +++ b/tests/test_litellm/proxy/guardrails/test_content_utils.py @@ -149,6 +149,22 @@ def test_iter_message_text_responses_api_tool_call_taxonomy(): assert list(iter_message_text(data)) == ["hello", "sunny"] +def test_iter_message_text_inspects_reasoning_content_and_summary(): + """VERIA: reasoning items forwarded as ``reasoning_content`` must be + inspected, including ``summary`` blocks the bridge reads as a fallback.""" + data = { + "input": [ + { + "type": "reasoning", + "id": "rs_1", + "content": [{"type": "summary_text", "text": "content secret"}], + "summary": [{"type": "summary_text", "text": "summary secret"}], + } + ] + } + assert list(iter_message_text(data)) == ["content secret", "summary secret"] + + # ── walk_user_text ──────────────────────────────────────────────────────────── @@ -308,6 +324,27 @@ def test_walk_user_text_redacts_mixed_list_input(): assert data["input"][2] == {"type": "image_url", "image_url": {"url": "..."}} +def test_walk_user_text_redacts_reasoning_content_and_summary(): + """VERIA: in-place redaction must cover both plaintext shapes the bridge + forwards from a reasoning item.""" + data = { + "input": [ + { + "type": "reasoning", + "id": "rs_1", + "content": [{"type": "summary_text", "text": "AKIAEXAMPLE content"}], + "summary": [{"type": "summary_text", "text": "AKIAEXAMPLE summary"}], + } + ] + } + visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]")) + assert visited == 2 + item = data["input"][0] + assert item["content"][0]["text"] == "[REDACTED] content" + assert item["summary"][0]["text"] == "[REDACTED] summary" + assert item["id"] == "rs_1" + + # ── build_inspection_messages ───────────────────────────────────────────────── @@ -462,6 +499,23 @@ def test_build_inspection_messages_empty_data(): assert build_inspection_messages({"input": ""}) == [] +def test_build_inspection_messages_includes_reasoning_summary(): + """VERIA: remote guardrail APIs must see reasoning summaries even when + the reasoning item has no ``content`` field.""" + data = { + "input": [ + { + "type": "reasoning", + "id": "rs_1", + "summary": [{"type": "summary_text", "text": "secret summary"}], + } + ] + } + assert build_inspection_messages(data) == [ + {"role": "assistant", "content": "secret summary"} + ] + + # ── has_non_string_content ──────────────────────────────────────────────────── 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 042/542] 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 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 043/542] 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 044/542] 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 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 045/542] 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 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 046/542] 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 047/542] 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 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 048/542] 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 049/542] 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 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 050/542] 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 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 051/542] 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 34e692c903c9d75b52065b4093c7d80b7eb2e00b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 20:16:22 -0700 Subject: [PATCH 052/542] fix(batches): decode model-encoded output file id so completed batches book spend Adds e2e coverage for batches terminal state and cost write-back, failure paths, per-backend file content downloads, and two-gateway routing (LIT-5730). --- litellm/batches/batch_utils.py | 40 +- tests/e2e/batches/COVERAGE.md | 87 +++- tests/e2e/batches/batch_client.py | 29 +- tests/e2e/batches/capabilities.py | 10 + tests/e2e/batches/test_batches_e2e.py | 471 +++++++++++++++++- .../llm_nonconversational.yaml | 10 + .../test_litellm/batches/test_batch_utils.py | 37 ++ 7 files changed, 646 insertions(+), 38 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 0cf22d82ca6..6eb13d2cba7 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -296,6 +296,32 @@ def calculate_vertex_ai_batch_cost_and_usage( ) +def _provider_output_file_id(output_file_id: str) -> str: + """ + Resolve the file id the provider actually knows: unified ids yield their embedded + llm_output_file_id, model-encoded ids decode to the raw provider id, raw ids pass through. + """ + from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, + get_original_file_id, + ) + + unified_file_id: Final = _is_base64_encoded_unified_file_id(output_file_id) + if not unified_file_id: + return get_original_file_id(output_file_id) + try: + extracted: Final = unified_file_id.split("llm_output_file_id,")[1].split(";")[0] + except (IndexError, AttributeError) as e: + verbose_logger.error( + "Failed to extract LLM output file ID from unified file ID: %s, error: %s", + output_file_id, + e, + ) + return output_file_id + verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", extracted) + return extracted + + async def _fetch_batch_output_file_content( batch: Batch, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", @@ -311,23 +337,11 @@ async def _fetch_batch_output_file_content( Required for Azure and other providers that need authentication """ from litellm.files.main import afile_content - from litellm.proxy.openai_files_endpoints.common_utils import ( - _is_base64_encoded_unified_file_id, - ) if batch.output_file_id is None: raise ValueError("Output file id is None cannot retrieve file content") - file_id = batch.output_file_id - is_base64_unified_file_id: Final = _is_base64_encoded_unified_file_id(file_id) - if is_base64_unified_file_id: - try: - file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(";")[0] - verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", file_id) - except (IndexError, AttributeError) as e: - verbose_logger.error( - "Failed to extract LLM output file ID from unified file ID: %s, error: %s", batch.output_file_id, e - ) + file_id: Final = _provider_output_file_id(batch.output_file_id) # Build kwargs for afile_content with credentials from litellm_params file_content_kwargs: Final = { diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index ca48204962a..f02d4eb4fe4 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -1,9 +1,11 @@ # Batches Test Coverage Matrix Live e2e coverage of the Batches API over a real proxy, real provider keys, and -real cost. Synchronous tier only: a batch's completion window is 24h, so these -tests never wait for `completed`. They assert the proxy accepts, routes, retrieves, -cancels, and lists a batch; everything created is deleted on teardown. +real cost. Mostly synchronous tier: a batch's completion window is 24h, so the +lifecycle matrix never waits for `completed`. It asserts the proxy accepts, routes, +retrieves, cancels, and lists a batch; everything created is deleted on teardown. +The exception is `TestBatchTerminalState`, which covers the completed state and +cost write-back via a cross-run marker baton (design below). ## Provider x operation @@ -12,19 +14,26 @@ row per supported (provider, scenario) pair, so there are no skipped cells in th parametrized run. The batches suite never skips: missing provider creds or upstream failures are hard test failures (see `tests/e2e/CLAUDE.md`). -| Provider | create | retrieve | cancel | list | file backing | -|-----------|--------|----------|--------|------|--------------| -| OpenAI | yes | yes | yes | yes | OpenAI Files | -| Azure | yes | yes | yes | yes | Azure Files | -| Vertex AI | yes | yes | yes | yes | GCS (`gcs_bucket_name` / `GCS_BUCKET_NAME` on model) | -| Bedrock | yes (unified only) | yes | no (limited upstream) | no | S3 (`s3_bucket_name` + `aws_*` + `AWS_BATCH_ROLE_ARN` on model) | +| Provider | create | retrieve | cancel | list | content download | file backing | +|-----------|--------|----------|--------|------|------------------|--------------| +| OpenAI | yes | yes | yes | yes | yes (lifecycle + terminal output) | OpenAI Files | +| Azure | yes | yes | yes | yes | yes (byte-verbatim) | Azure Files | +| Vertex AI | yes | yes | yes | yes | yes (provider-transformed) | GCS (`gcs_bucket_name` / `GCS_BUCKET_NAME` on model) | +| Bedrock | yes (unified only) | yes | no (limited upstream) | no | yes (provider-transformed) | S3 (`s3_bucket_name` + `aws_*` + `AWS_BATCH_ROLE_ARN` on model) | Bedrock cancel is unreliable upstream and list is unsupported, so both are gated off -(`can_cancel=False`, `can_list=False`) when that provider is enabled in the matrix. +(`can_cancel=False`, `can_list=False`) when that provider is enabled in the matrix; +flipping those gates is tracked in LIT-4774 and deliberately not part of this suite. Bedrock file upload requires a model on the request (`encoded` / `unified` scenarios only); `model_param` and `provider_fallback` are omitted because `POST /bedrock/v1/files` has no model-less passthrough path. +`GET /v1/files/{id}/content` is exercised for the unified upload path per backend in +`test_unified_file_content_downloads`. Azure stores the JSONL verbatim, so its download +is asserted byte-equal to the upload. Vertex (GCS) and Bedrock (S3) transform lines at +upload time, so those assert a 200 with non-empty parseable JSON lines instead. Gemini +(non-Vertex) raises `NotImplementedError` for file content and has no cell here. + ## Routing scenarios (per `litellm/proxy/batches_endpoints/endpoints.py`) Each create-capable provider runs all four. The test asserts the returned file id @@ -71,11 +80,59 @@ File delete asserts `object=="file"` and `deleted==True`. | `batch_client.py` | typed file upload/download + batch create/retrieve/cancel/list/delete over the shared ProxyClient; runtime batch model registration via /model/new; denial helpers | | `capabilities.py` | the provider x scenario matrix + per-provider /model/new params + id-shape classifiers + per-provider raw-id assertion | | `conftest.py` | session-scoped batch deployment registration and teardown | -| `test_batches_e2e.py` | parametrized lifecycle with per-endpoint output assertions, file upload/delete outputs, key-model-access denial | +| `test_batches_e2e.py` | parametrized lifecycle with per-endpoint output assertions, file upload/delete outputs, key-model-access denial, per-backend content download, failure paths, second-hop routing, terminal state + cost | + +## Failure paths + +`TestBatchFailurePaths` pins the customer-facing error contracts. A malformed input +file is a 400 at upload naming the bad content. A JSONL line whose url contradicts +the batch endpoint passes create (providers validate asynchronously) and drives the +batch to `failed` with structured `errors.data` (code/line/message), a null +`output_file_id`, and a $0 spend row keyed `{batch_id}_batch_cost` (LIT-4852: a +failed batch books $0 instead of crashing cost tracking). Cancelling that failed +batch is a 409 naming the terminal status. A file id encoded for one deployment wins +over a conflicting `model` param on create: the batch routes and re-encodes by the +file's embedded model (foreign-id precedence). + +## Second hop (two chained gateways) + +`TestBatchSecondHop` registers a `litellm_proxy/` deployment pointing at +the proxy's own base URL with a freshly minted virtual key, so unified upload and +create traverse gateway -> gateway -> OpenAI (LIT-5347, PR #36240). The pin: +`target_model_names` is rewritten to the inner deployment on the second hop and the +nested managed ids round-trip retrieve. This self-chaining only needs the proxy to +reach its own `PROXY_BASE_URL`, which holds both locally and on the e2e stage. + +## Terminal state + cost write-back (cross-run marker baton) + +The 24h completion window rules out submit-and-wait inside one run, so +`TestBatchTerminalState` amortizes across runs. Each run submits a 1-line marker +batch (stable metadata key/value plus a per-run field) and deliberately never +cancels or deletes it or its input file: the marker is the baton the next run picks +up (OpenAI files expire on their own after ~30 days). Polling is list-only, up to 5 +minutes, because retrieving a non-terminal batch books a $0 spend row whose +request_id then blocks the later real-cost row (`skip_duplicates`); the single +retrieve happens only once a completed marker exists. The assertion target is the +newest completed marker from ANY run: run-scoped deployment names mean the list +re-encodes prior-run batches under new encoded ids, so their spend keys are fresh +and a prior-run marker is billable by this run. On the 6h stage cadence the full +assertions are therefore deterministic from run 2 onward. On a cold start (no +completed marker within the poll budget) the test passes on the submission +assertions alone: a documented vacuous pass, not a skip. Markers aged past the 24h +window (25h-73h band, within the newest 100-item list page) must be terminal. + +The cost assertion is the LIT-5730 headline: retrieving a completed model-encoded +batch must write a positive spend row with call_type `aretrieve_batch` and token +usage. Before the fix in `litellm/batches/batch_utils.py`, the retrieve endpoint +re-encoded the response's `output_file_id` in place before the queued logging +worker ran, the worker sent that encoded id to OpenAI, got a 404, and the spend row +never landed. ## Out of scope (intentionally) -Driving a batch to `completed`, cost tracking on completion, and the DB write-back -are not covered here; the 24h window makes them unfit for a synchronous gate. That -logic belongs in a DI-stubbed proxy integration test under `tests/test_litellm/proxy/` -where the provider client is injected to return `completed` deterministically. +Unified (managed) batch cost is owned by the hourly `CheckBatchCost` poller, and a +terminal DB status short-circuits retrieve for those ids, so the terminal-state cell +uses the encoded path; poller timing does not fit an e2e gate and belongs in a +DI-stubbed proxy integration test under `tests/test_litellm/proxy/`. Bedrock +cancel/list stay gated pending LIT-4774. Gemini (non-Vertex) file content raises +`NotImplementedError` upstream and is not a coverage cell. diff --git a/tests/e2e/batches/batch_client.py b/tests/e2e/batches/batch_client.py index 5cc5d1dae3b..21a56f3398f 100644 --- a/tests/e2e/batches/batch_client.py +++ b/tests/e2e/batches/batch_client.py @@ -44,6 +44,17 @@ class FileList(BaseModel): data: list[FileObject] = [] +class BatchErrorItem(BaseModel): + code: str | None = None + line: int | None = None + message: str | None = None + + +class BatchErrorList(BaseModel): + object: str | None = None + data: list[BatchErrorItem] = [] + + class BatchObject(BaseModel): id: str object: str | None = None @@ -51,6 +62,9 @@ class BatchObject(BaseModel): endpoint: str | None = None input_file_id: str | None = None output_file_id: str | None = None + error_file_id: str | None = None + errors: BatchErrorList | None = None + metadata: dict[str, str] | None = None completion_window: str | None = None created_at: int | None = None model: str | None = None @@ -72,12 +86,18 @@ class BatchCreateBody(BaseModel): endpoint: str = "/v1/chat/completions" completion_window: str = "24h" model: str | None = None + metadata: dict[str, str] | None = None class ModelQuery(BaseModel): model: str | None = None +class BatchListQuery(BaseModel): + model: str | None = None + limit: int | None = None + + def is_model_access_denied(resp: StreamingResponse) -> bool: """True if the proxy rejected the call because the key may not access the model.""" return resp.status_code == 403 and "key_model_access_denied" in resp.body @@ -168,12 +188,17 @@ class BatchClient: ) def list_batches( - self, *, key: str, provider: str | None = None + self, + *, + key: str, + provider: str | None = None, + model: str | None = None, + limit: int | None = None, ) -> Result[BatchList]: return self.proxy.transport.get( _batches_path(provider), headers=self.proxy.transport.bearer(key), - params=NoBody(), + params=BatchListQuery(model=model, limit=limit), response_type=BatchList, ) diff --git a/tests/e2e/batches/capabilities.py b/tests/e2e/batches/capabilities.py index 3988fb5e7e1..ce1f68184a7 100644 --- a/tests/e2e/batches/capabilities.py +++ b/tests/e2e/batches/capabilities.py @@ -210,6 +210,16 @@ def is_model_encoded_id(id_str: str) -> bool: return False +def decoded_model_from_id(id_str: str) -> str | None: + """Deployment name embedded in a model-encoded file/batch id, or None.""" + for prefix in ("file-", "batch_"): + if id_str.startswith(prefix): + decoded = _b64_decode(id_str[len(prefix) :]) + if decoded.startswith("litellm:") and ";model," in decoded: + return decoded.split(";model,", 1)[1].split(";")[0] + return None + + def matches_id_shape(shape: IdShape, id_str: str) -> bool: if shape == "managed": return is_managed_id(id_str) diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 1376bdbed38..75b3a6cd758 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -1,11 +1,12 @@ """Live e2e for the Batches API across every provider LiteLLM supports. -Synchronous tier only: a batch's completion window is 24h, so these never wait for -"completed". Each case uploads a tiny JSONL, creates the batch through one of the -four routing scenarios, asserts it was accepted (non-terminal status) and routed to -the right provider, then retrieves / cancels / lists where the provider supports it. -Everything created is deleted on teardown. Completion + cost tracking are out of -scope here (see COVERAGE.md). +Mostly synchronous tier: a batch's completion window is 24h, so the lifecycle +matrix never waits for "completed". Each case uploads a tiny JSONL, creates the +batch through one of the four routing scenarios, asserts it was accepted +(non-terminal status) and routed to the right provider, then retrieves / cancels / +lists where the provider supports it. Everything created is deleted on teardown. +The exception is TestBatchTerminalState, which carries completed-state + cost +write-back coverage via a cross-run marker baton (design in COVERAGE.md). Routing signal: for provider_fallback the raw batch id discriminates the provider; for the encoded/unified/model_param scenarios the proxy re-encodes the id, so the @@ -22,8 +23,9 @@ from datetime import datetime, timedelta, timezone from typing import Callable import pytest +from pydantic import BaseModel -from e2e_config import unique_marker +from e2e_config import PROXY_BASE_URL, unique_marker from batch_client import ( UPLOAD_FILENAME, @@ -40,9 +42,12 @@ from capabilities import ( CAPABILITIES, FILE_ID_SHAPE, OPENAI_BATCH_MODEL, + PROVIDERS, Capability, + Provider, batch_model_name, coverage_cells_for_lifecycle, + decoded_model_from_id, is_managed_id, matches_id_shape, raw_id_matches_provider, @@ -475,9 +480,22 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row( OPENAI_FILE_CONTENT_BACKEND = "gpt-4o-mini" +FILE_CONTENT_CELLS = { + "azure": "llm.files.azure_openai.content.nonstream.works", + "vertex_ai": "llm.files.vertex.content.nonstream.works", + "bedrock": "llm.files.bedrock.content.nonstream.works", +} +BYTE_FIDELITY_CONTENT_PROVIDERS = frozenset({"azure"}) + class TestBatchFileContent: - """GET /v1/files/{id}/content returns the uploaded batch JSONL bytes.""" + """GET /v1/files/{id}/content returns the uploaded batch JSONL bytes. + + Azure stores the upload verbatim, so its download is asserted byte-equal. + Vertex (GCS) and Bedrock (S3) transform each JSONL line into the provider's + request format at upload time, so their downloads assert 200 plus non-empty + parseable JSON lines instead of byte equality. + """ @pytest.mark.covers( "llm.files.openai.content.nonstream.works", @@ -521,6 +539,62 @@ class TestBatchFileContent: "downloaded file content must match the uploaded JSONL bytes" ) + @pytest.mark.parametrize( + "provider", + [ + pytest.param( + p, + id=p.name, + marks=pytest.mark.covers( + FILE_CONTENT_CELLS[p.name], exercised_on=["files"] + ), + ) + for p in PROVIDERS + if p.name in FILE_CONTENT_CELLS + ], + ) + def test_unified_file_content_downloads( + self, + provider: Provider, + client: BatchClient, + resources: ResourceManager, + batch_deployments: None, + ) -> None: + key = resources.key() + payload = render_jsonl(provider.raw_model) + file = unwrap( + client.upload_file( + content=payload, + form=FileUploadForm(purpose="batch", target_model_names=provider.model), + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + assert_file_object(file, provider=provider.name) + assert is_managed_id(file.id), ( + f"{provider.name}: unified upload must return a managed file id, got {file.id!r}" + ) + + downloaded = client.proxy.transport.download( + f"/v1/files/{file.id}/content", + headers=client.proxy.transport.bearer(key), + ) + assert downloaded.status_code == 200, ( + f"{provider.name}: file content must be 200, " + f"got {downloaded.status_code}: {downloaded.body[:300]}" + ) + body = downloaded.body.strip() + assert body, f"{provider.name}: file content download returned an empty body" + if provider.name in BYTE_FIDELITY_CONTENT_PROVIDERS: + assert body == payload.decode().strip(), ( + f"{provider.name}: downloaded content must match the uploaded JSONL bytes" + ) + else: + for line in body.splitlines(): + assert json.loads(line), ( + f"{provider.name}: content line is not JSON: {line[:200]}" + ) + class TestOpenAIFiles: """GET /v1/files (list) and GET /v1/files/{id} (retrieve) over the OpenAI route. @@ -866,3 +940,384 @@ class TestHostedVllmBatch: f"hosted_vllm batch has non-transitional status {batch.status!r}" ) assert_batch_object(batch) + + +BATCH_TERMINAL_STATUSES = frozenset({"completed", "failed", "expired", "cancelled"}) +FAILED_BATCH_POLL_SECONDS = 120.0 +FAILED_BATCH_POLL_INTERVAL_SECONDS = 5.0 + +AZURE_BATCH_RAW_MODEL = next(p.raw_model for p in PROVIDERS if p.name == "azure") + + +def _mismatched_endpoint_jsonl(model: str) -> bytes: + line = { + "custom_id": "req-1", + "method": "POST", + "url": "/v1/embeddings", + "body": {"model": model, "input": "ping"}, + } + return (json.dumps(line) + "\n").encode() + + +def _poll_until_terminal(client: BatchClient, batch_id: str, key: str) -> BatchObject: + deadline = time.monotonic() + FAILED_BATCH_POLL_SECONDS + fetched = retrieve_batch(client, batch_id, key=key, provider=None) + while fetched.status not in BATCH_TERMINAL_STATUSES and time.monotonic() < deadline: + time.sleep(FAILED_BATCH_POLL_INTERVAL_SECONDS) + fetched = retrieve_batch(client, batch_id, key=key, provider=None) + return fetched + + +class TestBatchFailurePaths: + """Customer-facing failure contracts for /v1/batches. + + A malformed input file is rejected at upload with a 400 naming the bad + content. A JSONL line whose url contradicts the batch endpoint is accepted + at create (providers validate asynchronously) and drives the batch to + "failed" with structured per-line errors, a null output_file_id, and a + zero-cost spend row (LIT-4852: a failed batch must book $0, not crash cost + tracking). Cancelling that already-failed batch returns a 409 naming the + terminal status. A file id encoded for one deployment wins over a + conflicting model param on create: the batch routes (and re-encodes) by the + file's embedded model, pinning that precedence. + """ + + @pytest.mark.covers( + "llm.batches.openai.malformed_jsonl.nonstream.works", + exercised_on=["files"], + ) + def test_malformed_jsonl_upload_rejected( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + result = client.upload_file( + content=b"this is not json\n", + form=FileUploadForm(purpose="batch"), + model=OPENAI_BATCH_MODEL, + key=resources.key(), + ) + match result: + case UnknownApiError(status_code=400, body=body): + assert "json" in body.lower(), ( + f"400 must name the malformed JSONL so users can fix the file, got: {body[:300]}" + ) + case _: + pytest.fail(f"malformed JSONL upload must be rejected with a 400, got: {result}") + + @pytest.mark.covers( + "llm.batches.openai.jsonl_endpoint_mismatch.nonstream.works", + "llm.batches.openai.cancel_terminal.nonstream.works", + exercised_on=["batches", "files"], + ) + def test_endpoint_mismatch_fails_batch_and_cancel_conflicts( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + file = unwrap( + client.upload_file( + content=_mismatched_endpoint_jsonl("gpt-4o-mini"), + form=FileUploadForm(purpose="batch"), + model=OPENAI_BATCH_MODEL, + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + + created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) + require_successful_call(created) + batch = BatchObject.model_validate_json(created.body) + + fetched = _poll_until_terminal(client, batch.id, key) + assert fetched.status == "failed", ( + f"endpoint-mismatched batch must fail, got {fetched.status!r}" + ) + assert fetched.output_file_id is None, ( + f"failed batch must have no output file, got {fetched.output_file_id!r}" + ) + assert fetched.errors is not None and fetched.errors.data, ( + "failed batch must surface structured errors so users can fix the JSONL" + ) + first_error = fetched.errors.data[0] + assert first_error.message, "batch error item has no message" + assert first_error.code, "batch error item has no code" + + rows = client.proxy.poll_logs_for_request_id(f"{fetched.id}_batch_cost") + assert rows, ( + f"failed batch {fetched.id} wrote no spend row; retrieve must book $0 (LIT-4852)" + ) + assert all((row.spend or 0) == 0 for row in rows), ( + f"failed batch must cost $0, got {[(r.request_id, r.spend) for r in rows]}" + ) + assert rows[0].call_type == "aretrieve_batch", ( + f"batch cost row call_type={rows[0].call_type!r}" + ) + + conflict = client.cancel_batch(batch.id, key=key) + match conflict: + case UnknownApiError(status_code=409, body=body): + assert "failed" in body.lower(), ( + f"409 must name the terminal status blocking the cancel, got: {body[:300]}" + ) + case _: + pytest.fail(f"cancel of a failed batch must return a 409 conflict, got: {conflict}") + + @pytest.mark.covers( + "llm.batches.openai.foreign_file_id.nonstream.works", + exercised_on=["batches", "files"], + ) + def test_foreign_encoded_file_id_routes_by_file_model( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + file = unwrap( + client.upload_file( + content=render_jsonl(AZURE_BATCH_RAW_MODEL), + form=FileUploadForm(purpose="batch"), + model=AZURE_BATCH_MODEL, + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + assert decoded_model_from_id(file.id) == AZURE_BATCH_MODEL, ( + f"upload did not encode the azure deployment into the file id: {file.id!r}" + ) + + created = client.create_batch( + body=BatchCreateBody(input_file_id=file.id, model=OPENAI_BATCH_MODEL), key=key + ) + require_successful_call(created) + batch = BatchObject.model_validate_json(created.body) + resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + + assert decoded_model_from_id(batch.id) == AZURE_BATCH_MODEL, ( + "create with a foreign encoded file id must route by the file's embedded model, " + f"but the batch id encodes {decoded_model_from_id(batch.id)!r} " + f"(model param was {OPENAI_BATCH_MODEL!r})" + ) + fetched = retrieve_batch(client, batch.id, key=key, provider=None) + assert fetched.id == batch.id + assert fetched.status, "retrieved foreign-file batch has no status" + + +class TestBatchSecondHop: + """Two-proxy batch routing: a litellm_proxy deployment chained to the gateway + itself (LIT-5347, PR #36240). + + The hop deployment's litellm_params point litellm_proxy/ at this + gateway's own base URL with a freshly minted virtual key, so the unified + upload and batch create traverse gateway -> gateway -> OpenAI. The regression + this pins: target_model_names must be rewritten to the inner deployment on + the second hop and the nested managed ids must round-trip retrieve. + """ + + @pytest.mark.covers( + "llm.batches.openai.second_hop.nonstream.works", + exercised_on=["batches", "files"], + ) + def test_unified_create_and_retrieve_via_chained_gateway( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + hop_name = batch_model_name("openai-batch-hop") + model_id = client.create_model( + hop_name, + LiteLLMParamsBody( + model=f"litellm_proxy/{OPENAI_BATCH_MODEL}", + api_base=PROXY_BASE_URL, + api_key=key, + ), + ) + resources.defer(lambda: client.delete_model(model_id)) + + file = unwrap( + client.upload_file( + content=render_jsonl("gpt-4o-mini"), + form=FileUploadForm(purpose="batch", target_model_names=hop_name), + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + assert is_managed_id(file.id), ( + f"second-hop unified upload must return a managed file id, got {file.id!r}" + ) + + created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) + require_successful_call(created) + batch = BatchObject.model_validate_json(created.body) + resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + + assert is_managed_id(batch.id), ( + f"second-hop create must return a managed batch id, got {batch.id!r}" + ) + assert batch.status in CREATED_BATCH_STATUSES, ( + f"second-hop batch has non-transitional status {batch.status!r}" + ) + assert_batch_object(batch) + + fetched = retrieve_batch(client, batch.id, key=key, provider=None) + assert fetched.id == batch.id + assert fetched.status, "second-hop retrieve returned no status" + + +class BatchOutputBody(BaseModel): + choices: list[object] = [] + + +class BatchOutputResponse(BaseModel): + status_code: int | None = None + body: BatchOutputBody | None = None + + +class BatchOutputLine(BaseModel): + response: BatchOutputResponse + + +TERMINAL_MARKER_KEY = "litellm_e2e_suite" +TERMINAL_MARKER_VALUE = "batches-terminal-baton" +TERMINAL_POLL_SECONDS = 300.0 +TERMINAL_POLL_INTERVAL_SECONDS = 10.0 +TERMINAL_LIST_LIMIT = 100 +TERMINAL_BAND_MIN_AGE_SECONDS = 25 * 3600 +TERMINAL_BAND_MAX_AGE_SECONDS = 73 * 3600 + + +def _marker_batches(client: BatchClient, key: str) -> list[BatchObject]: + listed = unwrap( + client.list_batches(key=key, model=OPENAI_BATCH_MODEL, limit=TERMINAL_LIST_LIMIT) + ) + return [ + b + for b in listed.data + if (b.metadata or {}).get(TERMINAL_MARKER_KEY) == TERMINAL_MARKER_VALUE + ] + + +def _await_completed_marker( + client: BatchClient, key: str +) -> tuple[BatchObject | None, list[BatchObject]]: + deadline = time.monotonic() + TERMINAL_POLL_SECONDS + while True: + markers = _marker_batches(client, key) + completed = max( + (b for b in markers if b.status == "completed"), + key=lambda b: b.created_at or 0, + default=None, + ) + if completed is not None or time.monotonic() >= deadline: + return completed, markers + time.sleep(TERMINAL_POLL_INTERVAL_SECONDS) + + +def _assert_aged_markers_terminal(markers: list[BatchObject]) -> None: + now = time.time() + stuck = [ + b + for b in markers + if b.created_at is not None + and TERMINAL_BAND_MIN_AGE_SECONDS <= now - b.created_at <= TERMINAL_BAND_MAX_AGE_SECONDS + and b.status not in BATCH_TERMINAL_STATUSES + ] + assert not stuck, ( + "marker batches past their 24h completion window must be terminal; stuck: " + f"{[(b.id, b.status, b.created_at) for b in stuck]}" + ) + + +class TestBatchTerminalState: + """Terminal state + cost write-back via a cross-run marker baton. + + Each run submits a 1-line marker batch (stable metadata key/value plus a + per-run field) and never cancels or deletes it: the marker is the baton the + next run picks up. Polling is list-only for up to 5 minutes because a + retrieve of a non-terminal batch books a $0 spend row whose request_id then + blocks the real-cost row (skip_duplicates); the single retrieve happens only + once a completed marker exists. The assertion target is the newest completed + marker from ANY run, so on the 6h stage cadence the full assertions are + deterministic from run 2 onward. On a cold start (no marker has ever + completed within the poll budget) the test passes on the submission + assertions alone: that is a documented vacuous pass, not a skip, and this + run's marker becomes the next run's target. Markers aged past OpenAI's 24h + completion window (25h-73h band, within the newest list page) must be + terminal. The cost assertion is the LIT-5730 headline: retrieving a + completed model-encoded batch must write a positive spend row keyed + {batch_id}_batch_cost; before the fix the logging worker fetched the + re-encoded output_file_id, 404d, and the row never landed. + """ + + @pytest.mark.covers( + "llm.batches.openai.terminal_state.nonstream.works", + "llm.batches.openai.terminal_state.nonstream.cost_logged", + exercised_on=["batches", "files"], + ) + def test_completed_batch_downloads_output_and_books_cost( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + file = unwrap( + client.upload_file( + content=render_jsonl("gpt-4o-mini"), + form=FileUploadForm(purpose="batch"), + model=OPENAI_BATCH_MODEL, + key=key, + ) + ) + created = client.create_batch( + body=BatchCreateBody( + input_file_id=file.id, + metadata={ + TERMINAL_MARKER_KEY: TERMINAL_MARKER_VALUE, + "run": unique_marker(), + }, + ), + key=key, + ) + require_successful_call(created) + submitted = BatchObject.model_validate_json(created.body) + assert submitted.status in CREATED_BATCH_STATUSES, ( + f"marker batch has non-transitional status {submitted.status!r}" + ) + assert (submitted.metadata or {}).get(TERMINAL_MARKER_KEY) == TERMINAL_MARKER_VALUE, ( + f"create dropped the marker metadata: {submitted.metadata!r}" + ) + + completed, markers = _await_completed_marker(client, key) + _assert_aged_markers_terminal(markers) + if completed is None: + return + + fetched = retrieve_batch(client, completed.id, key=key, provider=None) + assert fetched.status == "completed", ( + f"listed-completed marker retrieved as {fetched.status!r}" + ) + assert fetched.output_file_id, "completed batch has no output_file_id" + + downloaded = client.proxy.transport.download( + f"/v1/files/{fetched.output_file_id}/content", + headers=client.proxy.transport.bearer(key), + ) + assert downloaded.status_code == 200, ( + f"output content must be 200, got {downloaded.status_code}: {downloaded.body[:300]}" + ) + first_line = BatchOutputLine.model_validate_json(downloaded.body.strip().splitlines()[0]) + assert first_line.response.status_code == 200, ( + f"batch output line reports failure: {downloaded.body[:400]}" + ) + assert first_line.response.body is not None and first_line.response.body.choices, ( + "batch output line has no choices" + ) + + rows = client.proxy.poll_logs_for_request_id( + f"{fetched.id}_batch_cost", + predicate=lambda found: any((row.spend or 0) > 0 for row in found), + ) + priced = [row for row in rows if (row.spend or 0) > 0] + assert priced, ( + f"completed batch {fetched.id} wrote no positive-cost spend row under " + f"request_id {fetched.id}_batch_cost; cost write-back is broken (LIT-5730)" + ) + cost_row = priced[0] + assert cost_row.call_type == "aretrieve_batch", ( + f"batch cost row call_type={cost_row.call_type!r}" + ) + assert (cost_row.total_tokens or 0) > 0, ( + f"batch cost row has no token usage: {cost_row.total_tokens!r}" + ) diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index bb7169509eb..3de462683d2 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -24,6 +24,13 @@ - {id: llm.batches.hosted_vllm.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: hosted_vllm, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "hosted_vllm OpenAI-compatible batch create"} - {id: llm.batches.openai.key_model_access_denied.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Key model restriction 403 on upload/create"} - {id: llm.batches.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.18 / LIT-4778", rationale: "Missing input_file_id and invalid batch id rejected"} +- {id: llm.batches.openai.terminal_state.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "A batch actually reaches completed and its output file downloads through GET /v1/files/{id}/content with per-line provider responses"} +- {id: llm.batches.openai.terminal_state.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "test_batches_e2e.py / LIT-5730", fail_before_fix: proven, rationale: "Retrieving a completed model-encoded batch writes a positive spend row keyed {batch_id}_batch_cost (pins LIT-4852/LIT-5666; before the fix the logging worker 404d fetching the re-encoded output_file_id and the row was never written)"} +- {id: llm.batches.openai.malformed_jsonl.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "Uploading a non-JSON batch file is rejected with a 400 naming the bad line"} +- {id: llm.batches.openai.jsonl_endpoint_mismatch.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "JSONL line url that contradicts the batch endpoint drives the batch to failed with structured errors, retrieve stays clean, and the terminal retrieve books a zero-cost spend row (LIT-4852)"} +- {id: llm.batches.openai.cancel_terminal.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "Cancelling an already-terminal batch returns a 409 conflict naming the terminal status"} +- {id: llm.batches.openai.foreign_file_id.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "Create with one deployment's encoded file id and a conflicting model param routes by the file's embedded model; the returned batch id pins that precedence"} +- {id: llm.batches.openai.second_hop.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5347", rationale: "A litellm_proxy deployment chained to the gateway itself preserves target_model_names through nested unified ids; upload, create, and retrieve work over the two-hop chain (PR #36240)"} - {id: llm.files.openai.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "openai_files_endpoints/files_endpoints.py:46", rationale: "File upload returns OpenAIFileObject"} - {id: llm.files.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.16 / LIT-4778", rationale: "File upload without purpose rejected"} - {id: llm.files.openai.retrieve.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "files_endpoints.py", rationale: "File retrieve by id"} @@ -36,6 +43,9 @@ - {id: llm.files.hosted_vllm.upload.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: hosted_vllm, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "hosted_vllm OpenAI-compatible file upload"} - {id: llm.rerank.cohere.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: rerank, route: cohere, capability: basic, streaming: nonstream, assertions: [works], source: "test_rerank_e2e.py:29", rationale: "Cohere rerank, top_n + relevance_score"} - {id: llm.files.openai.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "GET /v1/files/{id}/content returns uploaded batch JSONL bytes"} +- {id: llm.files.azure_openai.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "GET /v1/files/{id}/content on an Azure unified file returns the uploaded JSONL bytes verbatim"} +- {id: llm.files.vertex.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "GET /v1/files/{id}/content on a Vertex unified file streams the GCS object back (provider-transformed JSONL, so asserts non-empty JSON lines rather than byte equality)"} +- {id: llm.files.bedrock.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "GET /v1/files/{id}/content on a Bedrock unified file streams the S3 object back (provider-transformed JSONL, so asserts non-empty JSON lines rather than byte equality)"} - {id: llm.realtime.bedrock_converse.basic.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: bedrock_converse, capability: basic, streaming: stream, assertions: [works], source: "test_realtime_bedrock_e2e.py", rationale: "Nova Sonic realtime session emits response.done (LIT-2239)"} - {id: llm.google_native.gemini.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: google_native, route: gemini, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "LIT-4076 / proxy/google_endpoints/endpoints.py", fail_before_fix: proven, rationale: "google-native generateContent must stamp x-litellm-response-cost so SDK traffic reconciles against spend"} - {id: llm.google_native.gemini.basic.stream.works, module: llm, tier: P0, subject_endpoint: google_native, route: gemini, capability: basic, streaming: stream, assertions: [works], source: "PR #28213 / proxy/proxy_server.py async_data_generator", fail_before_fix: proven, rationale: "streamGenerateContent must relay single-prefixed SSE frames with no [DONE] sentinel; doubled data: prefixes and the OpenAI terminator both break the Vertex Java SDK"} diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index ebe093c591c..08cdf945b80 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -800,6 +800,43 @@ async def test_output_file_content_vertex_unified_file_id_extracts_gcs_uri(monke assert captured["custom_llm_provider"] == "vertex_ai" +@pytest.mark.asyncio +async def test_output_file_content_model_encoded_file_id_decoded_to_provider_id(monkeypatch): + import litellm.files.main as files_main + from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model + + captured: dict = {} + + async def fake_afile_content(**kw): + captured.update(kw) + return type("R", (), {"content": b'{"a": 1}'})() + + monkeypatch.setattr(files_main, "afile_content", fake_afile_content) + encoded_id = encode_file_id_with_model("file-Y3FHrMpi7uCkDpY6fgWGeR", "my-batch-model") + + await bu._fetch_batch_output_file_content(_batch(encoded_id), custom_llm_provider="openai") + + assert captured["file_id"] == "file-Y3FHrMpi7uCkDpY6fgWGeR" + assert captured["custom_llm_provider"] == "openai" + + +@pytest.mark.asyncio +async def test_output_file_content_raw_openai_file_id_passes_through(monkeypatch): + import litellm.files.main as files_main + + captured: dict = {} + + async def fake_afile_content(**kw): + captured.update(kw) + return type("R", (), {"content": b'{"a": 1}'})() + + monkeypatch.setattr(files_main, "afile_content", fake_afile_content) + + await bu._fetch_batch_output_file_content(_batch("file-abc123"), custom_llm_provider="openai") + + assert captured["file_id"] == "file-abc123" + + def _vertex_predictions_row(custom_id, prompt_tokens, completion_tokens): return { "request": { 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 053/542] 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 054/542] 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 42cffe93c821ce60a8a6f24d95cd01d995e52f32 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 19 Aug 2026 16:45:52 -0700 Subject: [PATCH 055/542] Add moonshot/kimi-k3 to model prices and context window map Pricing per https://platform.kimi.ai/docs/pricing/chat-k3: - $3.00/M input (cache miss), $0.30/M cache read, $15.00/M output - 1,048,576 context window; max_completion_tokens settable up to 1,048,576 - Supports reasoning (reasoning_effort low/high/max), tool calling, structured output, vision and video input Co-Authored-By: Claude Fable 5 --- .../model_prices_and_context_window_backup.json | 17 +++++++++++++++++ model_prices_and_context_window.json | 17 +++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d0eca17272d..a7c9825ee7a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -30213,6 +30213,23 @@ "supports_video_input": true, "supports_vision": true }, + "moonshot/kimi-k3": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "moonshot", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://platform.kimi.ai/docs/pricing/chat-k3", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, "moonshot/kimi-latest": { "cache_read_input_token_cost": 1.5e-07, "deprecation_date": "2026-01-28", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d0eca17272d..a7c9825ee7a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -30213,6 +30213,23 @@ "supports_video_input": true, "supports_vision": true }, + "moonshot/kimi-k3": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "moonshot", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://platform.kimi.ai/docs/pricing/chat-k3", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, "moonshot/kimi-latest": { "cache_read_input_token_cost": 1.5e-07, "deprecation_date": "2026-01-28", From 3c73a39877fa97db3dff4e22dc685bcb77fa4041 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 21:10:59 -0700 Subject: [PATCH 056/542] fix(cli): verify every credential store transition before reporting it done A keyring backend can accept a write and keep nothing. That is exactly what `keyring --disable` and PYTHON_KEYRING_BACKEND=keyring.backends.null.Keyring select, and it raises nothing to distinguish itself, so `lite login` was handing the credential to a black hole, scrubbing its own copy from token.json, and printing a success message over a login that no longer worked. Reading the value back is the only way to tell that backend apart from a keychain that really stored the secret. The same rule closes the rest of the gaps. A credential the token file will not record is taken back out of the keychain instead of being left live on a machine with no record of it, and is reported rather than raised. The migration stages its scrubbed file before the keychain is handed anything, so a directory that will not accept the rewrite stops the move rather than leaving the secret in two places. Logout no longer reads a key in the file as proof that the keychain is clear, which was never sound across two separate runs, and only draws that conclusion when the `keyring` package is missing outright, where nothing could have reached a keychain at all. --- litellm/litellm_core_utils/cli_keyring.py | 26 +++- litellm/litellm_core_utils/cli_token_utils.py | 104 +++++++++----- litellm/litellm_core_utils/private_json.py | 32 ++++- litellm/proxy/client/cli/commands/auth.py | 63 +++++++-- .../test_cli_token_utils.py | 131 ++++++++++++++++-- .../proxy/client/cli/test_auth_commands.py | 62 ++++++++- 6 files changed, 356 insertions(+), 62 deletions(-) diff --git a/litellm/litellm_core_utils/cli_keyring.py b/litellm/litellm_core_utils/cli_keyring.py index 0497991c2a0..15282fc522c 100644 --- a/litellm/litellm_core_utils/cli_keyring.py +++ b/litellm/litellm_core_utils/cli_keyring.py @@ -6,8 +6,12 @@ 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, 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. +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. + +A write is only reported as stored once it has been read back, because keyring's +null backend, which `keyring --disable` and headless CI images both select, +accepts every write and keeps nothing. """ import os @@ -61,7 +65,12 @@ class KeyringUnreachable: pass -KeyringUnusable: TypeAlias = KeyringNotInstalled | KeyringDisabled | KeyringUnreachable +@dataclass(frozen=True, slots=True) +class KeyringDiscardsWrites: + pass + + +KeyringUnusable: TypeAlias = KeyringNotInstalled | KeyringDisabled | KeyringUnreachable | KeyringDiscardsWrites SecretRead: TypeAlias = SecretFound | SecretMissing | KeyringUnusable SecretWrite: TypeAlias = SecretStored | KeyringUnusable SecretErase: TypeAlias = SecretErased | SecretStranded | KeyringUnusable @@ -119,6 +128,13 @@ class KeyringVault: return SecretMissing() if blob is None else SecretFound(blob) def write(self, blob: str) -> SecretWrite: + """Store the secret, reporting stored only once the keychain hands the same bytes back. + + A backend that accepts writes and keeps nothing, which is exactly what `keyring --disable` + and `PYTHON_KEYRING_BACKEND=keyring.backends.null.Keyring` select, raises nothing to + distinguish itself. Reading the value back is the only way to tell it apart from a keychain + that really stored the credential, and the caller is about to drop its own copy on our word. + """ api: Final = _keyring_api() if isinstance(api, (KeyringNotInstalled, KeyringDisabled)): return api @@ -126,7 +142,7 @@ class KeyringVault: 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 KeyringUnreachable() - return SecretStored() + return SecretStored() if self.read() == SecretFound(blob) else KeyringDiscardsWrites() def erase(self) -> SecretErase: """Remove our entry, reporting whether the keychain is guaranteed to be free of it. @@ -137,7 +153,7 @@ class KeyringVault: the caller knows whether this machine ever put a secret in a keychain. """ match self.read(): - case KeyringNotInstalled() | KeyringDisabled() | KeyringUnreachable() as unusable: + case KeyringNotInstalled() | KeyringDisabled() | KeyringUnreachable() | KeyringDiscardsWrites() as unusable: return unusable case SecretMissing(): return SecretErased() diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index 53289770c1a..af1b918fc2a 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -13,15 +13,17 @@ This module has no dependencies on proxy code and can be safely imported at the """ import time +from dataclasses import dataclass from pathlib import Path from types import MappingProxyType -from typing import Final +from typing import Final, TypeAlias from pydantic import BaseModel, ConfigDict, ValidationError from litellm.litellm_core_utils.cli_keyring import ( SYSTEM_KEYRING, KeyringDisabled, + KeyringDiscardsWrites, KeyringNotInstalled, KeyringUnreachable, SecretErase, @@ -33,7 +35,23 @@ from litellm.litellm_core_utils.cli_keyring import ( SecretVault, SecretWrite, ) -from litellm.litellm_core_utils.private_json import ensure_private_dir, write_private_json +from litellm.litellm_core_utils.private_json import ( + commit_staged_json, + discard_staged_json, + ensure_private_dir, + stage_private_json, + write_private_json, +) + + +@dataclass(frozen=True, slots=True) +class CredentialNotSaved: + """The credential was minted but no store would keep it, so this machine has none.""" + + detail: str + + +SecretSave: TypeAlias = SecretWrite | CredentialNotSaved class CliTokenRecord(BaseModel): @@ -85,14 +103,24 @@ 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) -> SecretWrite: - """Store a freshly minted credential. Reports where its secret material ended up, and why""" +def save_cli_token(record: CliTokenRecord, *, vault: SecretVault = SYSTEM_KEYRING) -> SecretSave: + """Store a freshly minted credential. Reports where its secret material ended up, and why. + + The token file is what makes a keychain-backed credential findable again, so a file that will + not be written takes the keychain copy down with it rather than leaving a live credential + stored under a machine that has no record of it. + """ 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) + try: + _write_token_file(_without_secret(record) if isinstance(outcome, SecretStored) else record) + except OSError as error: + if record.key is not None and isinstance(outcome, SecretStored): + vault.erase() + return CredentialNotSaved(str(error)) return outcome @@ -105,24 +133,30 @@ def clear_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> SecretErase: def _nothing_left_behind(outcome: SecretErase) -> bool: - """Whether the keychain can be trusted to hold no credential of ours once the file is gone""" + """Whether the keychain can be trusted to hold no credential of ours once the file is gone. + + A keychain that exists but is out of reach right now is never trusted, whatever the token file + looks like: the login that stored a secret there and the logout that cannot remove it are + separate runs, free to differ in whether the keychain was usable at the time. + """ match outcome: case SecretErased(): return True - case SecretStranded(): + case SecretStranded() | KeyringDisabled() | KeyringUnreachable() | KeyringDiscardsWrites(): return False - case KeyringNotInstalled() | KeyringDisabled() | KeyringUnreachable(): - return not _secret_lives_in_keychain() + case KeyringNotInstalled(): + return _file_holds_its_own_secret() -def _secret_lives_in_keychain() -> bool: - """Whether the token file is the metadata half of a pair whose secret half went to a keychain. +def _file_holds_its_own_secret() -> bool: + """Whether the stored login keeps its secret in the token file, ruling out a keychain entry. - 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. + Sound only against a missing `keyring` package, the one way to lose the keychain that had to + hold at storage time too, since nothing here can reach a keychain without it. A file whose + secret half is absent went to a keychain by definition, and so rules nothing out. """ record: Final = _read_token_file() - return record is not None and record.key is None and not record.jwt_token + return record is not None and record.key is not None def get_litellm_gateway_api_key( @@ -179,7 +213,7 @@ def is_cli_token_fresh(token_data: CliTokenRecord, buffer_hours: float = 0.1) -> def _read_token_file() -> CliTokenRecord | None: try: raw: Final = Path(get_cli_token_file_path()).read_text() - except OSError: + except (OSError, ValueError): return None try: return CliTokenRecord.model_validate_json(raw) @@ -193,7 +227,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 KeyringNotInstalled() | KeyringDisabled() | KeyringUnreachable(): + case KeyringNotInstalled() | KeyringDisabled() | KeyringUnreachable() | KeyringDiscardsWrites(): return record @@ -216,38 +250,46 @@ 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. + """Move a file-held secret into the vault, but only once 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. + The scrubbed file is staged first so a directory that will not accept it stops the migration + before the keychain is handed anything. Copying the credential into a second store and only + then discovering the first one cannot be cleaned would widen exposure instead of narrowing it, + which is the opposite of what moving it into the keychain is for. """ if record.key is None: return None - if not isinstance(vault.write(_encode_secret(record.base_url, record.key, record.jwt_token)), SecretStored): + staged: Final = _stage_scrubbed_file(record) + if staged is None: return record - if not _scrub_file_secret(record): + if not isinstance(vault.write(_encode_secret(record.base_url, record.key, record.jwt_token)), SecretStored): + discard_staged_json(staged) + return record + if not _commit_scrubbed_file(staged): vault.erase() return record 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. - """ + """Leave no secret material in the token file once the vault holds it""" if record.key is None and not record.jwt_token: return True + staged: Final = _stage_scrubbed_file(record) + return staged is not None and _commit_scrubbed_file(staged) + + +def _stage_scrubbed_file(record: CliTokenRecord) -> str | None: + path: Final = Path(get_cli_token_file_path()) try: - _write_token_file(_without_secret(record)) + ensure_private_dir(path.parent) + return stage_private_json(str(path), _without_secret(record).model_dump(exclude_none=True)) except OSError: - return _discard_token_file() - return True + return None -def _discard_token_file() -> bool: +def _commit_scrubbed_file(staged: str) -> bool: try: - Path(get_cli_token_file_path()).unlink(missing_ok=True) + commit_staged_json(staged, get_cli_token_file_path()) except OSError: return False return True diff --git a/litellm/litellm_core_utils/private_json.py b/litellm/litellm_core_utils/private_json.py index 32bc2e169e2..fbeb74aab5a 100644 --- a/litellm/litellm_core_utils/private_json.py +++ b/litellm/litellm_core_utils/private_json.py @@ -16,8 +16,12 @@ def ensure_private_dir(directory: Path) -> None: 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)""" +def stage_private_json(path: str, data: Mapping[str, object]) -> str: + """Write JSON to a private temp file beside `path`, ready for `commit_staged_json`. + + Staging is the half that can fail on a read-only or full directory, so callers with something + to lose can find that out before they act on the assumption that the rewrite will land. + """ parent: Final = Path(path).parent parent.mkdir(parents=True, exist_ok=True) fd, tmp_path = tempfile.mkstemp(dir=str(parent), prefix=".tmp-", suffix=".json") @@ -26,6 +30,26 @@ def write_private_json(path: str, data: Mapping[str, object]) -> None: json.dump(data, f, indent=2) f.flush() os.fsync(f.fileno()) - os.replace(tmp_path, path) - finally: + except BaseException: Path(tmp_path).unlink(missing_ok=True) + raise + return tmp_path + + +def commit_staged_json(staged: str, path: str) -> None: + """Move a staged file into place, replacing whatever is there in one step""" + try: + os.replace(staged, path) + except OSError: + Path(staged).unlink(missing_ok=True) + raise + + +def discard_staged_json(staged: str) -> None: + """Throw a staged file away when the change it was part of is abandoned""" + Path(staged).unlink(missing_ok=True) + + +def write_private_json(path: str, data: Mapping[str, object]) -> None: + """Atomically write JSON to path with owner-only permissions (0600)""" + commit_staged_json(stage_private_json(path, data), path) diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 4cf18435473..eba9994f7ec 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -15,16 +15,20 @@ from litellm.litellm_core_utils.cli_keyring import ( DISABLE_KEYRING_ENV_VAR, SYSTEM_KEYRING, KeyringDisabled, + KeyringDiscardsWrites, KeyringNotInstalled, KeyringUnreachable, SecretErased, + SecretFound, + SecretMissing, SecretStored, SecretStranded, SecretVault, - SecretWrite, ) from litellm.litellm_core_utils.cli_token_utils import ( CliTokenRecord, + CredentialNotSaved, + SecretSave, clear_cli_token, get_cli_token_file_path, get_litellm_gateway_api_key, @@ -84,17 +88,19 @@ class CliAuthResult(TypedDict): KEYRING_INSTALL_HINT: Final = "pip install 'litellm[cli]'" +KEYRING_ENABLE_HINT: Final = "keyring --enable (or unset PYTHON_KEYRING_BACKEND)" + STRANDED_CREDENTIAL_MESSAGE: Final = ( "Logged out locally, but your credential is still in the OS keychain and could not be removed." ) -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." +UNCHECKED_KEYCHAIN_MESSAGE: Final = ( + "Logged out locally, but your OS keychain could not be checked, so a credential stored there by " + "an earlier login may still be usable." ) -def storage_notice(outcome: SecretWrite) -> str: +def storage_notice(outcome: SecretSave) -> 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: @@ -109,6 +115,38 @@ def storage_notice(outcome: SecretWrite) -> str: 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)." + case KeyringDiscardsWrites(): + return ( + f"Your keyring backend keeps nothing it is given, so the credential was stored in {path} " + f"(owner-only) instead. For OS keychain storage, run: {KEYRING_ENABLE_HINT}" + ) + case CredentialNotSaved(detail=detail): + return ( + f"Signed in, but the credential could not be saved to {path}: {detail}. " + "Nothing was kept, so run 'lite login' again once that path is writable." + ) + + +def keychain_unreadable_notice(vault: SecretVault) -> str: + """Explain why the secret half of a stored login cannot be produced, and what fixes it""" + match vault.read(): + case KeyringNotInstalled(): + return ( + "Your credential is in your OS keychain, which this install cannot read without the " + f"keyring package. Install it with: {KEYRING_INSTALL_HINT}, or run 'lite login' to start over." + ) + case KeyringDisabled(): + return ( + f"Your credential is in your OS keychain, which {DISABLE_KEYRING_ENV_VAR} is blocking. " + "Unset it, or run 'lite login' to start over." + ) + case KeyringUnreachable() | KeyringDiscardsWrites(): + return ( + "Your credential is in your OS keychain, which could not be read. Unlock it, or run " + "'lite login' to start over." + ) + case SecretFound() | SecretMissing(): + return "Your credential could not be read from your OS keychain. Run 'lite login' to start over." def context_secret_vault(ctx: click.Context) -> SecretVault: @@ -715,6 +753,8 @@ def login(ctx: click.Context, config_claude: bool): click.echo("\nLogin successful!") click.echo(f"JWT Token: {api_key[:20]}...") click.echo(storage_notice(stored)) + if isinstance(stored, CredentialNotSaved): + return click.echo("You can now use the CLI without specifying --api-key") if config_claude: @@ -751,14 +791,17 @@ def logout(ctx: click.Context): match clear_cli_token(vault=context_secret_vault(ctx)): case SecretErased(): click.echo("Logged out successfully. Authentication token cleared.") + case SecretStranded(): + click.echo(STRANDED_CREDENTIAL_MESSAGE) + click.echo("Unlock your keychain and run 'lite logout' again to clear it.") 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(UNCHECKED_KEYCHAIN_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) + case KeyringUnreachable() | KeyringDiscardsWrites(): + click.echo(UNCHECKED_KEYCHAIN_MESSAGE) click.echo("Unlock your keychain and run 'lite logout' again to clear it.") @@ -795,7 +838,7 @@ def print_token(ctx: click.Context): api_key: Final = token_data.key if not api_key: - click.echo(KEYCHAIN_UNREACHABLE_MESSAGE, err=True) + click.echo(keychain_unreadable_notice(context_secret_vault(ctx)), err=True) sys.exit(1) click.echo(api_key) @@ -821,7 +864,7 @@ def whoami(ctx: click.Context): click.echo(f"Token age: {age_hours:.1f} hours") if token_data.key is None: - click.echo(KEYCHAIN_UNREACHABLE_MESSAGE) + click.echo(keychain_unreadable_notice(context_secret_vault(ctx))) 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/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py b/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py index cbc93bbde71..e0f5f99dc1b 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 @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.cli_keyring import ( SecretFound, SecretMissing, KeyringDisabled, + KeyringDiscardsWrites, KeyringNotInstalled, KeyringUnreachable, SecretErased, @@ -23,6 +24,7 @@ from litellm.litellm_core_utils.cli_keyring import ( ) from litellm.litellm_core_utils.cli_token_utils import ( CliTokenRecord, + CredentialNotSaved, clear_cli_token, get_cli_token_file_path, get_litellm_gateway_api_key, @@ -225,6 +227,15 @@ class TestLoadCliToken: assert record.key == "sk-legacy" + def test_a_token_file_that_is_not_text_is_not_a_login(self, isolated_home, secret_vault_factory): + """A truncated write or a half-synced backup can leave bytes that are not UTF-8 at all. + Reading them must fail the way an absent file does, not crash every `lite` command.""" + path = _token_file(isolated_home) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"\xff\xfe not utf-8 at all") + + assert load_cli_token(vault=secret_vault_factory()) is None + 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 {{{") @@ -301,6 +312,40 @@ class TestSaveCliToken: assert stat.S_IMODE(config_dir.stat().st_mode) == 0o700 + def test_a_credential_no_store_would_keep_is_reported_rather_than_raised( + self, isolated_home, secret_vault_factory, monkeypatch + ): + """`lite login` catches whatever escapes here and calls it an authentication failure, which + is the one thing that did not happen: the proxy minted a real credential. Saying so lets the + user act on the actual problem instead of retrying a sign-in that already worked.""" + + def _explode(*args, **kwargs): + raise OSError("read-only file system") + + monkeypatch.setattr("litellm.litellm_core_utils.private_json.json.dump", _explode) + + outcome = save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-new"), vault=secret_vault_factory()) + + assert isinstance(outcome, CredentialNotSaved) + assert "read-only file system" in outcome.detail + + def test_a_credential_the_file_will_not_record_is_taken_back_out_of_the_keychain( + self, isolated_home, secret_vault_factory, monkeypatch + ): + """The token file is what makes a keychain entry findable again. Leaving the secret in the + keychain with nothing pointing at it strands a live credential under a machine that has no + idea it is there, and no `lite logout` would ever go looking for it.""" + vault = secret_vault_factory() + + def _explode(*args, **kwargs): + raise OSError("read-only file system") + + monkeypatch.setattr("litellm.litellm_core_utils.private_json.json.dump", _explode) + + save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-new"), vault=vault) + + assert vault.blob is None + 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() @@ -340,9 +385,11 @@ class TestScrubFailure: 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( + def test_a_full_disk_stops_the_migration_before_the_keychain_is_handed_anything( self, isolated_home, secret_vault_factory, monkeypatch ): + """The scrubbed file is staged first precisely so this is knowable in advance. A disk that + cannot take the rewrite leaves the credential where it already was, in one store.""" path = _write_legacy_file(isolated_home) vault = secret_vault_factory() @@ -354,8 +401,8 @@ class TestScrubFailure: record = load_cli_token(vault=vault) assert record.key == "sk-legacy" - assert json.loads(vault.blob)["key"] == "sk-legacy" - assert not path.exists() + assert vault.blob is None + assert json.loads(path.read_text())["key"] == "sk-legacy" assert list(path.parent.glob(".tmp-*")) == [] @@ -376,12 +423,39 @@ class TestClearCliToken: assert clear_cli_token(vault=vault) == SecretStranded() assert not _token_file(isolated_home).exists() + @pytest.mark.parametrize( + "failure", [KeyringDisabled(), KeyringUnreachable(), KeyringDiscardsWrites()] + ) + def test_a_secret_in_the_file_is_no_evidence_about_a_keychain_that_exists( + self, isolated_home, secret_vault_factory, failure + ): + """Store a secret in the keychain, sign in again while the keychain is unusable so the new + secret lands in the file, then log out while it is still unusable. The file now carries its + own secret and the first login's entry is still there, so reading the file as proof of a + clean keychain reports a logout that did not happen.""" + _write_legacy_file(isolated_home) + vault = secret_vault_factory(available=False, failure=failure) + + assert clear_cli_token(vault=vault) == failure + assert not _token_file(isolated_home).exists() + + def test_a_second_logout_still_reports_the_keychain_it_could_not_clear( + self, isolated_home, secret_vault_factory + ): + """The first logout deletes the file and tells the user to run it again once the keychain is + reachable. If the second run reads that missing file as proof of a clean keychain, the advice + turns into the very false all-clear it was issued to prevent.""" + _write_metadata_only_file(isolated_home) + vault = secret_vault_factory(available=False, failure=KeyringUnreachable()) + + assert clear_cli_token(vault=vault) == KeyringUnreachable() + assert clear_cli_token(vault=vault) == KeyringUnreachable() + 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.""" + """A file holding only metadata put its secret in a keychain by definition. Losing the + package that reaches it does not take the entry with it, so this cannot report success.""" _write_metadata_only_file(isolated_home) vault = secret_vault_factory(available=False, failure=KeyringNotInstalled()) @@ -389,8 +463,9 @@ class TestClearCliToken: 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.""" + """The complement, and the one inference the file does support: nothing here can reach a + keychain without the package, so an install that lacks it and a file that still holds its + own secret between them account for the whole credential.""" _write_legacy_file(isolated_home) vault = secret_vault_factory(available=False, failure=KeyringNotInstalled()) @@ -417,11 +492,12 @@ class TestIsCliTokenFresh: class _FakeKeyringModule: - def __init__(self, stored=None, *, get_error=None, set_error=None, delete_error=None): + def __init__(self, stored=None, *, get_error=None, set_error=None, delete_error=None, discard=False): self.stored = stored self.get_error = get_error self.set_error = set_error self.delete_error = delete_error + self.discard = discard self.calls = [] def get_password(self, service_name, username): @@ -434,6 +510,8 @@ class _FakeKeyringModule: self.calls.append(("set", service_name, username)) if self.set_error is not None: raise self.set_error + if self.discard: + return self.stored = password def delete_password(self, service_name, username): @@ -503,6 +581,41 @@ class TestKeyringVault: assert KeyringVault().erase() == SecretStranded() + def test_a_backend_that_keeps_nothing_is_not_a_successful_write(self, install_fake_keyring): + """keyring's null backend accepts every write, stores nothing, and raises nothing to say so. + Taking its silence for success is how a credential gets deleted: the caller drops its own + copy on our word. Only reading the value back tells the two apart.""" + fake = install_fake_keyring(_FakeKeyringModule(discard=True)) + + assert KeyringVault().write("blob-1") == KeyringDiscardsWrites() + assert fake.stored is None + + def test_the_real_null_backend_is_rejected(self, monkeypatch): + """Pinned against the actual library rather than the double above, because the whole risk is + that upstream's no-op write looks exactly like a successful one.""" + keyring = pytest.importorskip("keyring") + null_backend = pytest.importorskip("keyring.backends.null") + monkeypatch.delenv(DISABLE_KEYRING_ENV_VAR, raising=False) + previous = keyring.get_keyring() + keyring.set_keyring(null_backend.Keyring()) + try: + assert KeyringVault().write("blob-1") == KeyringDiscardsWrites() + finally: + keyring.set_keyring(previous) + + def test_a_credential_survives_a_backend_that_keeps_nothing( + self, isolated_home, install_fake_keyring + ): + """The end of the same story: the credential must still be usable afterwards. Reporting the + discard is only worth anything if the token file then keeps the copy the keychain refused.""" + install_fake_keyring(_FakeKeyringModule(discard=True)) + + outcome = save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-only-copy")) + + assert outcome == KeyringDiscardsWrites() + assert json.loads(_token_file(isolated_home).read_text())["key"] == "sk-only-copy" + assert load_cli_token().key == "sk-only-copy" + def test_erasing_a_locked_keychain_is_a_failure(self, install_fake_keyring): install_fake_keyring(_FakeKeyringModule(get_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 33bd8307c21..8e1551c0720 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -16,12 +16,13 @@ from litellm.constants import CLI_JWT_EXPIRATION_HOURS from litellm.litellm_core_utils.cli_keyring import ( DISABLE_KEYRING_ENV_VAR, KeyringDisabled, + KeyringDiscardsWrites, 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 ( - KEYCHAIN_UNREACHABLE_MESSAGE, + DISABLE_KEYRING_ENV_VAR, get_stored_api_key, login, logout, @@ -461,6 +462,20 @@ class TestLogoutCommand: assert "still in the OS keychain" in result.output assert "pip install 'litellm[cli]'" in result.output + def test_logout_does_not_call_an_unusable_keychain_clean(self, isolated_home, secret_vault_factory): + """A keychain-backed login, then a login that fell back to the file because the keychain had + become unusable, leaves the first entry live. The file's own secret says nothing about it, + so a clean bill of health here is the one answer that cannot be justified.""" + _write_token_file(isolated_home, key="sk-in-file") + vault = secret_vault_factory(available=False, failure=KeyringDisabled()) + + result = self.runner.invoke(logout, obj={"secret_vault": vault}) + + assert result.exit_code == 0 + assert "Logged out successfully" not in result.output + assert "could not be checked" in result.output + assert DISABLE_KEYRING_ENV_VAR in result.output + def test_logout_warns_when_the_keychain_refuses_to_release_the_entry( self, isolated_home, secret_vault_factory ): @@ -1003,6 +1018,19 @@ class TestKeychainBackedCommands: assert "No OS keychain available" not in result.output assert json.loads(token_file.read_text())["key"] == "sk-minted" + def test_login_keeps_the_credential_when_the_backend_keeps_nothing( + self, isolated_home, secret_vault_factory + ): + """A backend that accepts writes and stores nothing must not be reported as keychain + storage, because the file is then told to drop the only remaining copy.""" + result = self._login(secret_vault_factory(available=False, failure=KeyringDiscardsWrites())) + + token_file = isolated_home / ".litellm" / "token.json" + assert result.exit_code == 0 + assert "Credential stored in your OS keychain." not in result.output + assert "keyring --enable" 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 ): @@ -1061,7 +1089,8 @@ class TestKeychainBackedCommands: result = self.runner.invoke(print_token, obj=obj) assert result.exit_code == 1 - assert KEYCHAIN_UNREACHABLE_MESSAGE in result.output + assert "could not be read" in result.output + assert "lite login" in result.output def test_whoami_flags_a_locked_keychain(self, isolated_home, secret_vault_factory): _write_home_json( @@ -1074,7 +1103,34 @@ class TestKeychainBackedCommands: result = self.runner.invoke(whoami, obj=obj) assert "Authenticated" in result.output - assert KEYCHAIN_UNREACHABLE_MESSAGE in result.output + assert "could not be read" in result.output + + def test_whoami_names_the_kill_switch_rather_than_a_missing_package( + self, isolated_home, secret_vault_factory + ): + """Every unreachable keychain used to be described as a locked one needing the keyring + package installed. Someone who set the kill switch has the package and an unlocked keychain, + so that advice sends them to fix two things that were never wrong.""" + _write_token_file(isolated_home, key=None) + vault = secret_vault_factory(available=False, failure=KeyringDisabled()) + + result = self.runner.invoke(whoami, obj={"base_url": "https://test.example.com", "secret_vault": vault}) + + assert DISABLE_KEYRING_ENV_VAR in result.output + assert "pip install" not in result.output + + def test_print_token_points_an_install_without_keyring_at_the_package( + self, isolated_home, secret_vault_factory + ): + _write_token_file(isolated_home, key=None) + vault = secret_vault_factory(available=False, failure=KeyringNotInstalled()) + obj = {"base_url": "https://test.example.com", "secret_vault": vault} + + result = self.runner.invoke(print_token, obj=obj) + + assert result.exit_code == 1 + assert "pip install 'litellm[cli]'" in result.output + assert DISABLE_KEYRING_ENV_VAR not in result.output class TestApiKeyPrecedence: From 26f62377451369bd1b96d2fbab11205c1018575c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 01:39:43 -0700 Subject: [PATCH 057/542] build: skip deleted files in the changed-file ruff format check `make lint` hands every path in the diff against the base branch to `ruff format --check`, including the ones the branch deleted, so any branch that moves or removes a file under `litellm/` fails the gate with "No such file or directory" instead of a formatting complaint. test-linting.yml already filters those out with `--diff-filter=ACMR`, so the Makefile was the half that drifted. Match it. --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 5e5f7c80027..ab6eba880bf 100644 --- a/Makefile +++ b/Makefile @@ -146,7 +146,7 @@ lint-install: # only the litellm Python files changed vs the base are checked, so a pre-existing # format issue elsewhere doesn't block an unrelated commit. lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE) - @files=$$(git diff --name-only origin/litellm_internal_staging...HEAD -- 'litellm/**/*.py' | grep -v '^litellm/enterprise/' || true); \ + @files=$$(git diff --name-only --diff-filter=ACMR origin/litellm_internal_staging...HEAD -- 'litellm/**/*.py' | grep -v '^litellm/enterprise/' || true); \ if [ -z "$$files" ]; then \ echo "No changed litellm Python files to format-check."; \ else \ From ba637553f8cded70ddab429c429c9c030f29dcc5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 01:39:43 -0700 Subject: [PATCH 058/542] fix(cli): keep lite login and logout honest when the keychain will not answer Three ways the credential commands could mislead or hang. `lite logout` on a machine that never logged in warned that a credential may be stranded in a keychain it could not check, and told the user to install keyring to go clear it. There was nothing there. A missing token file is now read as the evidence it is, because logout keeps a secret-free file behind whenever the keychain is left unconfirmed, so a later run can tell a machine with a credential it cannot reach apart from one that never had a login. That holds on the LITELLM_CLI_DISABLE_KEYRING path too. `KeyringDiscardsWrites` was handled on the read and erase paths, which cannot produce it: the null backend returns None from `get_password` rather than raising, so only a write ever detects it. It now lives on `SecretWrite` alone and the unreachable arms are gone. `keyring.set_password` blocks forever under a HOME with no usable login keychain, which is what containers, CI images, `sudo -H`, and service accounts run with, and reads answer normally there so nothing cheaper tells them apart. `lite login` never touched a keychain before this, so a sign-in that simply never returns would be a new way for it to fail. Writes are pre-flighted with a throwaway value on a bounded wait, and a keychain that stays silent falls back to the token file. The real credential is never the thing handed to a call that might land long after we stopped waiting. Saving also stages the token file before the keychain is given anything, since the file is the half a read-only or full directory refuses. A save that cannot land now leaves both stores as it found them, which matters most when the login it failed to replace still works. --- litellm/litellm_core_utils/cli_keyring.py | 53 +++++++- litellm/litellm_core_utils/cli_token_utils.py | 88 +++++++----- litellm/proxy/client/cli/commands/auth.py | 7 +- pyproject.toml | 2 +- .../test_cli_token_utils.py | 125 +++++++++++++++++- .../proxy/client/cli/test_auth_commands.py | 2 +- 6 files changed, 227 insertions(+), 50 deletions(-) diff --git a/litellm/litellm_core_utils/cli_keyring.py b/litellm/litellm_core_utils/cli_keyring.py index 15282fc522c..8da3e5226d4 100644 --- a/litellm/litellm_core_utils/cli_keyring.py +++ b/litellm/litellm_core_utils/cli_keyring.py @@ -11,18 +11,24 @@ token file and tell the user what to do about it. A write is only reported as stored once it has been read back, because keyring's null backend, which `keyring --disable` and headless CI images both select, -accepts every write and keeps nothing. +accepts every write and keeps nothing. Writes are also pre-flighted with a +throwaway value, because a keychain can answer neither way and block forever. """ import os +import threading +from contextlib import suppress from dataclasses import dataclass from typing import Final, Protocol, TypeAlias KEYRING_SERVICE: Final = "litellm-cli" KEYRING_ACCOUNT: Final = "credential" +KEYRING_PREFLIGHT_ACCOUNT: Final = "credential-preflight" DISABLE_KEYRING_ENV_VAR: Final = "LITELLM_CLI_DISABLE_KEYRING" _DISABLED_VALUES: Final = frozenset(("1", "true", "yes", "on")) +_PREFLIGHT_VALUE: Final = "preflight" +_PREFLIGHT_TIMEOUT_SECONDS: Final = 5.0 @dataclass(frozen=True, slots=True) @@ -70,9 +76,9 @@ class KeyringDiscardsWrites: pass -KeyringUnusable: TypeAlias = KeyringNotInstalled | KeyringDisabled | KeyringUnreachable | KeyringDiscardsWrites +KeyringUnusable: TypeAlias = KeyringNotInstalled | KeyringDisabled | KeyringUnreachable SecretRead: TypeAlias = SecretFound | SecretMissing | KeyringUnusable -SecretWrite: TypeAlias = SecretStored | KeyringUnusable +SecretWrite: TypeAlias = SecretStored | KeyringUnusable | KeyringDiscardsWrites SecretErase: TypeAlias = SecretErased | SecretStranded | KeyringUnusable @@ -113,10 +119,43 @@ def _keyring_api() -> KeyringApi | KeyringNotInstalled | KeyringDisabled: return KeyringNotInstalled() if api is None else api +def _answers_a_write(api: KeyringApi, timeout_seconds: float) -> bool: + """Whether the keychain answers a write at all, asked with a value worth nothing. + + macOS derives the login keychain from `$HOME`, and `set_password` against a HOME with no usable + one blocks forever with no timeout of its own. Containers, CI images, `sudo -H`, and service + accounts all run there, and reads answer normally, so nothing cheaper tells them apart. Asking + with a throwaway value keeps a keychain that never answers from taking `lite login` down with + it, and keeps the real credential out of a store that might accept it long after we gave up. + A keychain that refuses the probe outright still answered it, so only silence counts against it. + """ + answered: Final = threading.Event() + + def ask() -> None: + with suppress(Exception): + api.set_password(KEYRING_SERVICE, KEYRING_PREFLIGHT_ACCOUNT, _PREFLIGHT_VALUE) + answered.set() + + threading.Thread(target=ask, daemon=True, name="litellm-cli-keyring-preflight").start() + return answered.wait(timeout_seconds) + + +def _forget_the_preflight(api: KeyringApi) -> None: + """Take the throwaway probe back out. + + A backend that kept nothing has nothing to remove, and the probe is worth nothing either way, + so a keychain that refuses to give it up costs the caller nothing. + """ + with suppress(Exception): + api.delete_password(KEYRING_SERVICE, KEYRING_PREFLIGHT_ACCOUNT) + + @dataclass(frozen=True, slots=True) class KeyringVault: """The OS keychain, reached through the optional `keyring` package.""" + preflight_timeout_seconds: float = _PREFLIGHT_TIMEOUT_SECONDS + def read(self) -> SecretRead: api: Final = _keyring_api() if isinstance(api, (KeyringNotInstalled, KeyringDisabled)): @@ -134,10 +173,16 @@ class KeyringVault: and `PYTHON_KEYRING_BACKEND=keyring.backends.null.Keyring` select, raises nothing to distinguish itself. Reading the value back is the only way to tell it apart from a keychain that really stored the credential, and the caller is about to drop its own copy on our word. + + The keychain is pre-flighted first, because one that blocks rather than answering would + otherwise hang `lite login` outright. """ api: Final = _keyring_api() if isinstance(api, (KeyringNotInstalled, KeyringDisabled)): return api + if not _answers_a_write(api, self.preflight_timeout_seconds): + return KeyringUnreachable() + _forget_the_preflight(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 @@ -153,7 +198,7 @@ class KeyringVault: the caller knows whether this machine ever put a secret in a keychain. """ match self.read(): - case KeyringNotInstalled() | KeyringDisabled() | KeyringUnreachable() | KeyringDiscardsWrites() as unusable: + case KeyringNotInstalled() | KeyringDisabled() | KeyringUnreachable() as unusable: return unusable case SecretMissing(): return SecretErased() diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index af1b918fc2a..825967e6866 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -23,7 +23,6 @@ from pydantic import BaseModel, ConfigDict, ValidationError from litellm.litellm_core_utils.cli_keyring import ( SYSTEM_KEYRING, KeyringDisabled, - KeyringDiscardsWrites, KeyringNotInstalled, KeyringUnreachable, SecretErase, @@ -53,6 +52,8 @@ class CredentialNotSaved: SecretSave: TypeAlias = SecretWrite | CredentialNotSaved +_UNREPLACEABLE_FILE: Final = "the staged file could not replace the one already there" + class CliTokenRecord(BaseModel): """A stored CLI credential. @@ -106,57 +107,71 @@ def load_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> CliTokenRecord | N def save_cli_token(record: CliTokenRecord, *, vault: SecretVault = SYSTEM_KEYRING) -> SecretSave: """Store a freshly minted credential. Reports where its secret material ended up, and why. - The token file is what makes a keychain-backed credential findable again, so a file that will - not be written takes the keychain copy down with it rather than leaving a live credential - stored under a machine that has no record of it. + The token file is what makes a keychain-backed credential findable again, and it is also the + half that a read-only or full directory refuses, so it is staged before the keychain is handed + anything. A save that cannot land then leaves both stores exactly as it found them, which + matters most when the login it failed to replace is still perfectly good. """ + staged: Final = _stage_token_file(_without_secret(record)) + if isinstance(staged, CredentialNotSaved): + return staged outcome: Final = ( SecretStored() if record.key is None else vault.write(_encode_secret(record.base_url, record.key, record.jwt_token)) ) + if isinstance(outcome, SecretStored): + return outcome if _commit_token_file(staged) else CredentialNotSaved(_UNREPLACEABLE_FILE) + discard_staged_json(staged) + return _keep_the_secret_in_the_file(record, outcome) + + +def _keep_the_secret_in_the_file(record: CliTokenRecord, outcome: SecretWrite) -> SecretSave: + """Fall back to the owner-only file, which is all that is left when no keychain took the secret""" try: - _write_token_file(_without_secret(record) if isinstance(outcome, SecretStored) else record) + _write_token_file(record) except OSError as error: - if record.key is not None and isinstance(outcome, SecretStored): - vault.erase() return CredentialNotSaved(str(error)) return outcome def clear_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> SecretErase: - """Remove the credential from both stores. Reports whether the keychain is now free of it""" + """Remove the credential from both stores. Reports whether the keychain is now free of it. + + A file that holds no secret of its own is kept when the keychain will not confirm the entry is + gone, because it is the only remaining record that something is still in there to remove. That + is what lets a later run tell a machine with a credential it cannot reach apart from one that + never had a login at all. Anything still holding a secret is removed either way. + """ outcome: Final = vault.erase() - settled: Final = _nothing_left_behind(outcome) - Path(get_cli_token_file_path()).unlink(missing_ok=True) + record: Final = _read_token_file() + settled: Final = _nothing_left_behind(outcome, record) + if settled or record is None or record.key is not None: + Path(get_cli_token_file_path()).unlink(missing_ok=True) return SecretErased() if settled else outcome -def _nothing_left_behind(outcome: SecretErase) -> bool: +def _nothing_left_behind(outcome: SecretErase, record: CliTokenRecord | None) -> bool: """Whether the keychain can be trusted to hold no credential of ours once the file is gone. - A keychain that exists but is out of reach right now is never trusted, whatever the token file - looks like: the login that stored a secret there and the logout that cannot remove it are - separate runs, free to differ in whether the keychain was usable at the time. + A machine with no token file has no stored login to end, and `clear_cli_token` keeps one behind + whenever the keychain is left unconfirmed, so a missing file is real evidence rather than the + absence of it. Past that, a keychain that exists but is out of reach right now is + never trusted, whatever the file looks like: the login that stored a secret there and the + logout that cannot remove it are separate runs, free to differ in whether the keychain was + usable at the time. The exception is a missing `keyring` package, which had to be missing when + the credential was stored too, so a file still holding its own secret proves no keychain was + ever involved. `SecretStranded` is the keychain answering for itself and outranks the file. """ match outcome: case SecretErased(): return True - case SecretStranded() | KeyringDisabled() | KeyringUnreachable() | KeyringDiscardsWrites(): + case SecretStranded(): return False + case KeyringDisabled() | KeyringUnreachable(): + return record is None case KeyringNotInstalled(): - return _file_holds_its_own_secret() - - -def _file_holds_its_own_secret() -> bool: - """Whether the stored login keeps its secret in the token file, ruling out a keychain entry. - - Sound only against a missing `keyring` package, the one way to lose the keychain that had to - hold at storage time too, since nothing here can reach a keychain without it. A file whose - secret half is absent went to a keychain by definition, and so rules nothing out. - """ - record: Final = _read_token_file() - return record is not None and record.key is not None + return record is None or record.key is not None def get_litellm_gateway_api_key( @@ -227,7 +242,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 KeyringNotInstalled() | KeyringDisabled() | KeyringUnreachable() | KeyringDiscardsWrites(): + case KeyringNotInstalled() | KeyringDisabled() | KeyringUnreachable(): return record @@ -265,7 +280,7 @@ def _migrate_file_secret(record: CliTokenRecord, vault: SecretVault) -> CliToken if not isinstance(vault.write(_encode_secret(record.base_url, record.key, record.jwt_token)), SecretStored): discard_staged_json(staged) return record - if not _commit_scrubbed_file(staged): + if not _commit_token_file(staged): vault.erase() return record @@ -275,19 +290,24 @@ def _scrub_file_secret(record: CliTokenRecord) -> bool: if record.key is None and not record.jwt_token: return True staged: Final = _stage_scrubbed_file(record) - return staged is not None and _commit_scrubbed_file(staged) + return staged is not None and _commit_token_file(staged) def _stage_scrubbed_file(record: CliTokenRecord) -> str | None: + staged: Final = _stage_token_file(_without_secret(record)) + return None if isinstance(staged, CredentialNotSaved) else staged + + +def _stage_token_file(record: CliTokenRecord) -> str | CredentialNotSaved: path: Final = Path(get_cli_token_file_path()) try: ensure_private_dir(path.parent) - return stage_private_json(str(path), _without_secret(record).model_dump(exclude_none=True)) - except OSError: - return None + return stage_private_json(str(path), record.model_dump(exclude_none=True)) + except OSError as error: + return CredentialNotSaved(str(error)) -def _commit_scrubbed_file(staged: str) -> bool: +def _commit_token_file(staged: str) -> bool: try: commit_staged_json(staged, get_cli_token_file_path()) except OSError: diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index eba9994f7ec..d89641d2366 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -123,7 +123,8 @@ def storage_notice(outcome: SecretSave) -> str: case CredentialNotSaved(detail=detail): return ( f"Signed in, but the credential could not be saved to {path}: {detail}. " - "Nothing was kept, so run 'lite login' again once that path is writable." + "Any login you already had is untouched. Run 'lite login' again once that path is " + "writable, or 'lite logout' to clear whatever is stored now." ) @@ -140,7 +141,7 @@ def keychain_unreadable_notice(vault: SecretVault) -> str: f"Your credential is in your OS keychain, which {DISABLE_KEYRING_ENV_VAR} is blocking. " "Unset it, or run 'lite login' to start over." ) - case KeyringUnreachable() | KeyringDiscardsWrites(): + case KeyringUnreachable(): return ( "Your credential is in your OS keychain, which could not be read. Unlock it, or run " "'lite login' to start over." @@ -800,7 +801,7 @@ def logout(ctx: click.Context): case KeyringDisabled(): click.echo(UNCHECKED_KEYCHAIN_MESSAGE) click.echo(f"Unset {DISABLE_KEYRING_ENV_VAR} and run 'lite logout' again to clear it.") - case KeyringUnreachable() | KeyringDiscardsWrites(): + case KeyringUnreachable(): click.echo(UNCHECKED_KEYCHAIN_MESSAGE) click.echo("Unlock your keychain and run 'lite logout' again to clear it.") diff --git a/pyproject.toml b/pyproject.toml index 32921e14d31..64adb9cd595 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -79,7 +79,7 @@ proxy = [ ] # Thin client install for the `lite` CLI on developer laptops. The CLI's heavy # imports (fastapi, cryptography, ...) are all guarded, so it runs on the base -# SDK plus just these four; none of the server runtime in `proxy` is pulled in. +# SDK plus just these five; none of the server runtime in `proxy` is pulled in. cli = [ "rich>=13.9.4,<14.0", "pyyaml>=6.0.3,<7.0", 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 e0f5f99dc1b..69ce47b25b3 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 @@ -2,6 +2,7 @@ import json import os import stat import sys +import threading import time import pytest @@ -10,6 +11,7 @@ from litellm.constants import CLI_JWT_EXPIRATION_HOURS from litellm.litellm_core_utils.cli_keyring import ( DISABLE_KEYRING_ENV_VAR, KEYRING_ACCOUNT, + KEYRING_PREFLIGHT_ACCOUNT, KEYRING_SERVICE, KeyringVault, SecretFound, @@ -329,12 +331,12 @@ class TestSaveCliToken: assert isinstance(outcome, CredentialNotSaved) assert "read-only file system" in outcome.detail - def test_a_credential_the_file_will_not_record_is_taken_back_out_of_the_keychain( + def test_a_file_that_will_not_be_written_stops_the_save_before_the_keychain_is_touched( self, isolated_home, secret_vault_factory, monkeypatch ): - """The token file is what makes a keychain entry findable again. Leaving the secret in the - keychain with nothing pointing at it strands a live credential under a machine that has no - idea it is there, and no `lite logout` would ever go looking for it.""" + """The token file is what makes a keychain entry findable again, so it is staged first. + Handing the keychain a secret and only then finding out that nothing will point at it + would strand a live credential under a machine with no idea it is there.""" vault = secret_vault_factory() def _explode(*args, **kwargs): @@ -346,6 +348,26 @@ class TestSaveCliToken: assert vault.blob is None + @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores directory permissions") + def test_a_login_that_cannot_be_saved_leaves_the_working_one_alone( + self, isolated_home, secret_vault_factory + ): + """Signing in again on a machine whose ~/.litellm has gone read-only must not cost the user + the credential they already had. Overwriting the keychain and then failing to record it, or + undoing that write afterwards, would take a login that still works out from under them.""" + _write_legacy_file(isolated_home, key=None) + vault = secret_vault_factory(blob=_blob(key="sk-in-use")) + path = _token_file(isolated_home) + path.parent.chmod(0o500) + try: + outcome = save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-new"), vault=vault) + finally: + path.parent.chmod(0o700) + + assert isinstance(outcome, CredentialNotSaved) + assert json.loads(vault.blob)["key"] == "sk-in-use" + assert load_cli_token(vault=vault).key == "sk-in-use" + 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() @@ -460,8 +482,44 @@ class TestClearCliToken: vault = secret_vault_factory(available=False, failure=KeyringNotInstalled()) assert clear_cli_token(vault=vault) == KeyringNotInstalled() + assert json.loads(_token_file(isolated_home).read_text()).get("key") is None + + def test_a_logout_that_cannot_clear_the_keychain_keeps_the_record_that_it_has_to( + self, isolated_home, secret_vault_factory + ): + """The file left behind holds no secret. It is what a later run reads to tell a machine with + a credential it cannot reach apart from one that never had a login, which is the difference + between warning the user and inventing a credential for them to worry about.""" + _write_metadata_only_file(isolated_home) + vault = secret_vault_factory(available=False, failure=KeyringUnreachable()) + + clear_cli_token(vault=vault) + + assert json.loads(_token_file(isolated_home).read_text()).get("key") is None + + def test_a_logout_that_cannot_clear_the_keychain_still_takes_the_file_secret_away( + self, isolated_home, secret_vault_factory + ): + """Keeping a record of the unreachable keychain must never mean keeping the cleartext copy + the user just asked to be rid of.""" + _write_legacy_file(isolated_home) + vault = secret_vault_factory(available=False, failure=KeyringUnreachable()) + + clear_cli_token(vault=vault) + assert not _token_file(isolated_home).exists() + @pytest.mark.parametrize("failure", [KeyringNotInstalled(), KeyringDisabled(), KeyringUnreachable()]) + def test_logging_out_of_a_machine_that_never_logged_in_invents_nothing_to_warn_about( + self, isolated_home, secret_vault_factory, failure + ): + """`lite logout` with no token file has nothing to end. Warning that a credential may be + stranded in a keychain it cannot check sends the user after something that was never there, + and `pip install keyring` will not make it appear.""" + vault = secret_vault_factory(available=False, failure=failure) + + assert clear_cli_token(vault=vault) == SecretErased() + def test_a_file_backed_login_logs_out_quietly_without_keyring(self, isolated_home, secret_vault_factory): """The complement, and the one inference the file does support: nothing here can reach a keychain without the package, so an install that lacks it and a file that still holds its @@ -510,7 +568,7 @@ class _FakeKeyringModule: self.calls.append(("set", service_name, username)) if self.set_error is not None: raise self.set_error - if self.discard: + if self.discard or username != KEYRING_ACCOUNT: return self.stored = password @@ -518,7 +576,22 @@ class _FakeKeyringModule: self.calls.append(("delete", service_name, username)) if self.delete_error is not None: raise self.delete_error - self.stored = None + if username == KEYRING_ACCOUNT: + self.stored = None + + +class _NeverAnsweringKeyringModule(_FakeKeyringModule): + """A keychain whose writes block instead of returning, the way macOS does under a HOME that + has no usable login keychain.""" + + def __init__(self): + super().__init__() + self.blocked = threading.Event() + + def set_password(self, service_name, username, password): + self.calls.append(("set", service_name, username)) + self.blocked.set() + threading.Event().wait() @pytest.fixture @@ -540,7 +613,8 @@ class TestKeyringVault: assert vault.read() == SecretFound("blob-1") assert vault.erase() == SecretErased() assert vault.read() == SecretMissing() - assert {call[1:] for call in fake.calls} == {(KEYRING_SERVICE, KEYRING_ACCOUNT)} + assert {call[1] for call in fake.calls} == {KEYRING_SERVICE} + assert {call[2] for call in fake.calls} == {KEYRING_ACCOUNT, KEYRING_PREFLIGHT_ACCOUNT} def test_the_kill_switch_reports_no_keychain(self, monkeypatch): """`LITELLM_CLI_DISABLE_KEYRING` has to work without importing keyring, because keyring @@ -590,6 +664,43 @@ class TestKeyringVault: assert KeyringVault().write("blob-1") == KeyringDiscardsWrites() assert fake.stored is None + def test_a_keychain_that_never_answers_does_not_hang_the_login(self, install_fake_keyring): + """macOS derives the login keychain from `$HOME`, and `set_password` under a HOME with no + usable one blocks forever with no timeout of its own. Containers, CI images, `sudo -H`, and + service accounts all run there, and `lite login` never touched a keychain before this, so a + sign-in that simply never returns would be a new way for it to fail.""" + fake = install_fake_keyring(_NeverAnsweringKeyringModule()) + vault = KeyringVault(preflight_timeout_seconds=0.2) + + started = time.monotonic() + outcome = vault.write("blob-1") + + assert outcome == KeyringUnreachable() + assert time.monotonic() - started < 5 + assert fake.blocked.is_set() + + def test_a_keychain_that_never_answers_is_never_handed_the_credential(self, install_fake_keyring): + """Giving up on the write is only safe if the secret was never the thing being written. A + blocked call can still land later, and a keychain copy nobody waited for would sit beside + the file copy the user was told about.""" + fake = install_fake_keyring(_NeverAnsweringKeyringModule()) + + KeyringVault(preflight_timeout_seconds=0.2).write("blob-1") + + assert [call[2] for call in fake.calls] == [KEYRING_PREFLIGHT_ACCOUNT] + + def test_a_login_survives_a_keychain_that_never_answers(self, isolated_home, install_fake_keyring): + """The end of the same story: the credential still has to be usable afterwards.""" + install_fake_keyring(_NeverAnsweringKeyringModule()) + + outcome = save_cli_token( + CliTokenRecord(base_url=SERVER, key="sk-only-copy"), + vault=KeyringVault(preflight_timeout_seconds=0.2), + ) + + assert outcome == KeyringUnreachable() + assert json.loads(_token_file(isolated_home).read_text())["key"] == "sk-only-copy" + def test_the_real_null_backend_is_rejected(self, monkeypatch): """Pinned against the actual library rather than the double above, because the whole risk is that upstream's no-op write looks exactly like a successful one.""" 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 8e1551c0720..e491dbd5aca 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -1074,7 +1074,7 @@ class TestKeychainBackedCommands: assert result.exit_code == 0 assert "could not be removed" in result.output - assert not (isolated_home / ".litellm" / "token.json").exists() + assert json.loads((isolated_home / ".litellm" / "token.json").read_text()).get("key") is None def test_print_token_explains_a_locked_keychain_instead_of_printing_nothing( self, isolated_home, secret_vault_factory From c7da91d47f95aa9aa992264c1979748b916de796 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 01:47:50 -0700 Subject: [PATCH 059/542] fix(cli): say so when the keychain took a credential the file cannot name Staging the token file can succeed and the replacement still fail afterwards, and that is the one save path where the keychain has already taken the new secret. It was reported as a save that kept nothing, which sends the user looking for a credential that is sitting in their keychain, and it claimed the previous login was untouched when the one keychain slot had just been written over. Give that path its own outcome and its own notice. The new secret stays where it is: the entry it replaced went the moment it landed, so no rollback brings that back, and removing the new one too would turn a login this machine may still be able to use into no login at all. The remaining `CredentialNotSaved` paths all leave both stores untouched, so the reassurance they carry is now true wherever it is printed. --- litellm/litellm_core_utils/cli_token_utils.py | 23 ++++++++++--- litellm/proxy/client/cli/commands/auth.py | 9 +++++- .../test_cli_token_utils.py | 32 +++++++++++++++++++ 3 files changed, 59 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index 825967e6866..15b1390f337 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -45,14 +45,25 @@ from litellm.litellm_core_utils.private_json import ( @dataclass(frozen=True, slots=True) class CredentialNotSaved: - """The credential was minted but no store would keep it, so this machine has none.""" + """The credential was minted but no store would keep it, so this machine has none. + + Nothing was touched on the way to this, so a login that already worked still does. + """ detail: str -SecretSave: TypeAlias = SecretWrite | CredentialNotSaved +@dataclass(frozen=True, slots=True) +class CredentialNotRecorded: + """The keychain took the credential, but the file that names it could not be replaced. -_UNREPLACEABLE_FILE: Final = "the staged file could not replace the one already there" + The keychain holds one entry, so the secret that was there is already gone and no rollback + brings it back. Removing the new one as well would only turn a login this machine may still + be able to use into no login at all, so it stays, and the user is told what is where. + """ + + +SecretSave: TypeAlias = SecretWrite | CredentialNotSaved | CredentialNotRecorded class CliTokenRecord(BaseModel): @@ -111,6 +122,10 @@ def save_cli_token(record: CliTokenRecord, *, vault: SecretVault = SYSTEM_KEYRIN half that a read-only or full directory refuses, so it is staged before the keychain is handed anything. A save that cannot land then leaves both stores exactly as it found them, which matters most when the login it failed to replace is still perfectly good. + + Staging can still succeed and the replacement fail afterwards. That is the one case where the + keychain has already taken the new secret, and it reports itself as such rather than claiming + the previous login survived. """ staged: Final = _stage_token_file(_without_secret(record)) if isinstance(staged, CredentialNotSaved): @@ -121,7 +136,7 @@ def save_cli_token(record: CliTokenRecord, *, vault: SecretVault = SYSTEM_KEYRIN else vault.write(_encode_secret(record.base_url, record.key, record.jwt_token)) ) if isinstance(outcome, SecretStored): - return outcome if _commit_token_file(staged) else CredentialNotSaved(_UNREPLACEABLE_FILE) + return outcome if _commit_token_file(staged) else CredentialNotRecorded() discard_staged_json(staged) return _keep_the_secret_in_the_file(record, outcome) diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index d89641d2366..eb956cd536c 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -27,6 +27,7 @@ from litellm.litellm_core_utils.cli_keyring import ( ) from litellm.litellm_core_utils.cli_token_utils import ( CliTokenRecord, + CredentialNotRecorded, CredentialNotSaved, SecretSave, clear_cli_token, @@ -126,6 +127,12 @@ def storage_notice(outcome: SecretSave) -> str: "Any login you already had is untouched. Run 'lite login' again once that path is " "writable, or 'lite logout' to clear whatever is stored now." ) + case CredentialNotRecorded(): + return ( + f"Signed in, and the credential is in your OS keychain, but {path} could not be " + "replaced, so this machine may still be using your previous login. Run 'lite login' " + "again once that path is writable, or 'lite logout' to clear both." + ) def keychain_unreadable_notice(vault: SecretVault) -> str: @@ -754,7 +761,7 @@ def login(ctx: click.Context, config_claude: bool): click.echo("\nLogin successful!") click.echo(f"JWT Token: {api_key[:20]}...") click.echo(storage_notice(stored)) - if isinstance(stored, CredentialNotSaved): + if isinstance(stored, (CredentialNotSaved, CredentialNotRecorded)): return click.echo("You can now use the CLI without specifying --api-key") 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 69ce47b25b3..0cf1e55d363 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 @@ -26,6 +26,7 @@ from litellm.litellm_core_utils.cli_keyring import ( ) from litellm.litellm_core_utils.cli_token_utils import ( CliTokenRecord, + CredentialNotRecorded, CredentialNotSaved, clear_cli_token, get_cli_token_file_path, @@ -368,6 +369,37 @@ class TestSaveCliToken: assert json.loads(vault.blob)["key"] == "sk-in-use" assert load_cli_token(vault=vault).key == "sk-in-use" + def test_a_keychain_write_the_file_cannot_be_pointed_at_is_reported_as_that( + self, isolated_home, secret_vault_factory + ): + """Staging the file can succeed and the replacement still fail, and that is the one path + where the keychain already took the new secret. Reporting it as a save that kept nothing + would send the user looking for a credential that is sitting in their keychain.""" + vault = secret_vault_factory() + path = _token_file(isolated_home) + path.parent.mkdir(parents=True, exist_ok=True) + path.mkdir() + + outcome = save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-new"), vault=vault) + + assert isinstance(outcome, CredentialNotRecorded) + assert json.loads(vault.blob)["key"] == "sk-new" + + def test_the_credential_the_file_cannot_name_is_left_in_the_keychain( + self, isolated_home, secret_vault_factory + ): + """The keychain holds one entry, so the secret that was there went the moment this one + landed. Taking the new one back out would turn a login this machine may still be able to + use into no login at all, and it cannot restore the old one either way.""" + vault = secret_vault_factory(blob=_blob(key="sk-in-use")) + path = _token_file(isolated_home) + path.parent.mkdir(parents=True, exist_ok=True) + path.mkdir() + + save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-new"), vault=vault) + + assert vault.blob is not None + 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() From bcb6a6eaab2f5080fb905014fc60f413df84fad3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 01:50:01 -0700 Subject: [PATCH 060/542] test(e2e): pin prompt-cache, service-tier, and cost-header billing Seven live e2e tests covering cost-tracking regressions that currently ship unnoticed: cache-write tokens billed at the cache-creation rate (#34046), per-component cost_breakdown on the spend row (#31686), cache reads billed at the cache-read discount on streamed calls (#34812), cache tokens surviving the anthropic-messages to Responses bridge (#34957), priority-tier rates applied to input, output and reasoning (#35923, #35925), the per-component response cost headers summing to the total (#36965), and cost injected into the final usage frame of an /openai passthrough stream (#36503). Every test registers its own deployment with a distinct custom rate per component, so a component billed at the wrong rate cannot pass. The shared helpers in cost_rows.py encode the one thing the two surfaces disagree on: the spend row's input_cost is gross of cache while the response's cost-input header is net of it. --- .../coverage_registry/quota_management.yaml | 7 + tests/e2e/models.py | 25 +- .../spend_tracking/cost_rows.py | 204 ++++++++++++++ .../test_cache_cost_accounting_e2e.py | 263 ++++++++++++++++++ .../spend_tracking/test_cost_headers_e2e.py | 136 +++++++++ .../test_passthrough_stream_cost_e2e.py | 68 +++++ .../test_service_tier_pricing_e2e.py | 116 ++++++++ 7 files changed, 817 insertions(+), 2 deletions(-) create mode 100644 tests/e2e/quota_management/spend_tracking/cost_rows.py create mode 100644 tests/e2e/quota_management/spend_tracking/test_cache_cost_accounting_e2e.py create mode 100644 tests/e2e/quota_management/spend_tracking/test_cost_headers_e2e.py create mode 100644 tests/e2e/quota_management/spend_tracking/test_passthrough_stream_cost_e2e.py create mode 100644 tests/e2e/quota_management/spend_tracking/test_service_tier_pricing_e2e.py diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index 2dfa7adddea..5438ed8534a 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -44,3 +44,10 @@ - {id: quota_management.spend_tracking.failure.writes_failure_row, module: quota_management, tier: P1, behavior: spend_tracking, variant: failure, assertions: [writes_failure_row], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_log_error_logger.py", rationale: "A failed call writes a failure-status spend row"} - {id: quota_management.spend_tracking.spend_calculate.returns_cost, module: quota_management, tier: P2, behavior: spend_tracking, variant: spend_calculate, assertions: [returns_cost], exercised_on: [spend_calculate], source: "proxy/spend_tracking/spend_management_endpoints.py", rationale: "/spend/calculate prices a hypothetical request at nonzero cost"} - {id: quota_management.spend_tracking.pagination.keeps_total, module: quota_management, tier: P2, behavior: spend_tracking, variant: pagination, assertions: [keeps_total], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_management_endpoints.py", rationale: "Spend-logs v2 pagination caps page size without losing the total"} +- {id: quota_management.spend_tracking.cache_write.bills_cache_creation_rate, module: quota_management, tier: P1, behavior: spend_tracking, variant: cache_write, assertions: [bills_cache_creation_rate], exercised_on: [chat_completions], source: "litellm_core_utils/llm_cost_calc/utils.py", rationale: "OpenAI cache-write tokens land on the spend row as cache-creation tokens billed at the cache-creation rate, not silently at the input rate (#34046)"} +- {id: quota_management.spend_tracking.cost_breakdown.reports_component_costs, module: quota_management, tier: P1, behavior: spend_tracking, variant: cost_breakdown, assertions: [reports_component_costs], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "The spend row's metadata.cost_breakdown itemizes cache-read, cache-creation, output, and reasoning costs at the deployment's own rates and they sum to the row's spend (#31686)"} +- {id: quota_management.spend_tracking.stream_cache_read.bills_cache_read_rate, module: quota_management, tier: P1, behavior: spend_tracking, variant: stream_cache_read, assertions: [bills_cache_read_rate], exercised_on: [chat_completions], source: "litellm_core_utils/streaming_chunk_builder_utils.py", rationale: "A streamed call's reassembled usage keeps the cached-token detail so cache reads bill at the cache-read discount, not full input price (#34812)"} +- {id: quota_management.spend_tracking.messages_bridge.keeps_cache_tokens, module: quota_management, tier: P1, behavior: spend_tracking, variant: messages_bridge, assertions: [keeps_cache_tokens], exercised_on: [messages], source: "llms/anthropic/experimental_pass_through/responses_adapters/handler.py", rationale: "A /v1/messages request served by a Responses-only OpenAI model keeps its cache-read tokens and their discounted billing across the bridge (#34957)"} +- {id: quota_management.spend_tracking.service_tier.bills_tier_rates, module: quota_management, tier: P1, behavior: spend_tracking, variant: service_tier, assertions: [bills_tier_rates], exercised_on: [chat_completions], source: "cost_calculator.py", rationale: "A priority service_tier call bills input, output, and reasoning at the deployment's *_priority rates and records the tier on the row (#35923, #35925)"} +- {id: quota_management.spend_tracking.cost_headers.additive_components, module: quota_management, tier: P1, behavior: spend_tracking, variant: cost_headers, assertions: [additive_components], exercised_on: [chat_completions], source: "proxy/common_request_processing.py", rationale: "The x-litellm-response-cost-* component headers sum to the total, input covers only fresh tokens, and reasoning stays a subset of output (#36965)"} +- {id: quota_management.spend_tracking.passthrough_stream.injects_usage_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: passthrough_stream, assertions: [injects_usage_cost], exercised_on: [openai_passthrough], source: "proxy/pass_through_endpoints/streaming_handler.py", rationale: "With include_cost_in_streaming_usage on, the /openai passthrough's final streaming usage frame carries the proxy-computed cost (#36503)"} diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 619f4dcacfe..fe93b13e0a4 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -216,10 +216,19 @@ class McpChatTool(BaseModel): allowed_tools: list[str] | None = None +class StreamOptions(BaseModel): + """OpenAI `stream_options`: `include_usage` asks for a final usage-only SSE + frame, which is where the proxy's `include_cost_in_streaming_usage` setting + injects `usage.cost`.""" + + include_usage: bool = True + + class ChatBody(BaseModel): model: str messages: list[ChatMessage] stream: bool = False + stream_options: StreamOptions | None = None max_tokens: int | None = None max_completion_tokens: int | None = None temperature: float | None = None @@ -322,6 +331,9 @@ class CompletionTokensDetails(BaseModel): class Usage(BaseModel): + """`cost` exists only on streaming usage frames from a proxy running with + `include_cost_in_streaming_usage: true`; providers never send it.""" + prompt_tokens: int | None = None completion_tokens: int | None = None total_tokens: int | None = None @@ -329,6 +341,7 @@ class Usage(BaseModel): cache_creation_input_tokens: int | None = None prompt_tokens_details: PromptTokensDetails | None = None completion_tokens_details: CompletionTokensDetails | None = None + cost: float | None = None class ChatResponse(BaseModel): @@ -449,9 +462,11 @@ class AnthropicMessagesResponse(BaseModel): for triage.""" model_config = ConfigDict(extra="allow") + id: str | None = None model: str | None = None content: list[AnthropicContentBlock] | None = None choices: list[ChatChoice] | None = None + usage: Usage | None = None class CountTokensResponse(BaseModel): @@ -716,8 +731,10 @@ class FineTuningJobsResponse(BaseModel): class LiteLLMParamsBody(BaseModel): """POST /model/new litellm_params: `model` is the only required field; `api_key` et al may be an `os.environ/FOO` reference the proxy resolves at call time. - `input_cost_per_token`/`output_cost_per_token` register a per-deployment custom - pricing override; left None (and dropped from the body) the deployment keeps the + The `*_cost_per_token` / `*_token_cost` fields register a per-deployment custom + pricing override (the cache and `_priority` rates only apply when both base + rates are set, which is what makes the proxy register the deployment's full + pricing entry); left None (and dropped from the body) the deployment keeps the backend's canonical rate.""" model: str @@ -744,6 +761,10 @@ class LiteLLMParamsBody(BaseModel): aws_external_id: str | None = None input_cost_per_token: float | None = None output_cost_per_token: float | None = None + cache_read_input_token_cost: float | None = None + cache_creation_input_token_cost: float | None = None + input_cost_per_token_priority: float | None = None + output_cost_per_token_priority: float | None = None extra_headers: dict[str, str] | None = None use_in_pass_through: bool | None = None complexity_router_config: dict[str, object] | None = None diff --git a/tests/e2e/quota_management/spend_tracking/cost_rows.py b/tests/e2e/quota_management/spend_tracking/cost_rows.py new file mode 100644 index 00000000000..87af54fe83f --- /dev/null +++ b/tests/e2e/quota_management/spend_tracking/cost_rows.py @@ -0,0 +1,204 @@ +"""Cost-accounting helpers for the spend-tracking suite: the /spend/logs row shape +that carries the per-component cost breakdown, a poll that waits for it, and the +builders the cache-pricing tests share. + +The shared SpendLogRow deliberately stays thin (most tests only read totals), so +the component-cost tests model the metadata they assert on here instead: +`metadata.cost_breakdown` (input/output/cache-read/cache-creation/reasoning costs +plus the service-tier pricing basis) and `metadata.additional_usage_values` (the +cache token counts the biller derived from the provider's usage). + +Determinism strategy: every test registers its own deployment with explicit custom +rates for each component it asserts on (`register_priced_model`), so expected cost +is exactly tokens-on-the-row times configured rate, immune to provider price +changes. The rates are chosen ~100x above canonical and distinct from one another, +so a component billed at the wrong rate can never accidentally match. + +OpenAI prompt caching is implicit and keyed on the exact token prefix, with a +1024-token minimum. `cacheable_prefix` builds a prefix whose first word is the +run's unique marker: unique marker = the whole prefix is novel (a fresh cache +write), same marker + different question = a cache read that still misses the +proxy's own response cache. How long the prefix has to be before the provider +actually reports a read varies by model, so callers pass `words` to suit theirs. + +Two facts about the recorded bill that the assertions here encode, because the +two surfaces disagree on purpose. On the spend row, `input_cost` is gross: it +already contains the cache-read and cache-creation costs, so the row's total is +input + output + tool-usage and the fresh-token cost is input minus the two cache +components. In the response headers, `x-litellm-response-cost-input` is net of +cache, which is what makes the component headers sum to the total. +""" + +import time +from collections.abc import Callable + +from pydantic import BaseModel, RootModel + +from e2e_config import unique_marker +from e2e_http import Success +from lifecycle import ResourceManager +from models import LiteLLMParamsBody, SpendLogsParams +from proxy_client import ProxyClient + + +class CostBreakdownRow(BaseModel): + input_cost: float | None = None + output_cost: float | None = None + cache_read_cost: float | None = None + cache_creation_cost: float | None = None + reasoning_cost: float | None = None + tool_usage_cost: float | None = None + total_cost: float | None = None + service_tier: str | None = None + + +class AdditionalUsageValues(BaseModel): + cache_read_input_tokens: int | None = None + cache_creation_input_tokens: int | None = None + + +class CostRowMetadata(BaseModel): + cost_breakdown: CostBreakdownRow | None = None + additional_usage_values: AdditionalUsageValues | None = None + + +class CostRow(BaseModel): + request_id: str | None = None + spend: float | None = None + prompt_tokens: int | None = None + completion_tokens: int | None = None + metadata: CostRowMetadata | None = None + + @property + def breakdown(self) -> CostBreakdownRow: + assert self.metadata and self.metadata.cost_breakdown, ( + f"spend row {self.request_id} landed without a cost breakdown" + ) + return self.metadata.cost_breakdown + + @property + def cache_read_tokens(self) -> int: + if self.metadata and self.metadata.additional_usage_values: + return self.metadata.additional_usage_values.cache_read_input_tokens or 0 + return 0 + + @property + def cache_creation_tokens(self) -> int: + if self.metadata and self.metadata.additional_usage_values: + return self.metadata.additional_usage_values.cache_creation_input_tokens or 0 + return 0 + + +class CostRows(RootModel[list[CostRow]]): + pass + + +def approx_equal(actual: float, expected: float) -> bool: + """Within 1% or 1e-9 absolute - spend math, not exact float identity.""" + return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2) + + +def assert_total_is_sum_of_components(row: CostRow) -> None: + """The row's total is input + output + tool usage. The cache components are + already inside the gross input cost, so adding them again would double-bill.""" + breakdown = row.breakdown + components = sum( + cost or 0.0 + for cost in (breakdown.input_cost, breakdown.output_cost, breakdown.tool_usage_cost) + ) + assert breakdown.total_cost is not None and approx_equal(breakdown.total_cost, components), ( + f"total_cost {breakdown.total_cost} != input + output + tool usage ({components}): {breakdown}" + ) + assert row.spend is not None and approx_equal(row.spend, breakdown.total_cost), ( + f"row spend {row.spend} != breakdown total {breakdown.total_cost}" + ) + + +def assert_fresh_tokens_billed_at(row: CostRow, input_rate: float) -> None: + """Strip the cache components out of the gross input cost and what is left must + be the freshly-read tokens at the deployment's input rate.""" + breakdown = row.breakdown + fresh_tokens = (row.prompt_tokens or 0) - row.cache_read_tokens - row.cache_creation_tokens + fresh_cost = ( + (breakdown.input_cost or 0.0) + - (breakdown.cache_read_cost or 0.0) + - (breakdown.cache_creation_cost or 0.0) + ) + assert breakdown.input_cost is not None and approx_equal(fresh_cost, fresh_tokens * input_rate), ( + f"input_cost {breakdown.input_cost} less cache read {breakdown.cache_read_cost} and " + f"cache creation {breakdown.cache_creation_cost} leaves {fresh_cost}, not " + f"{fresh_tokens} fresh tokens * {input_rate} (prompt {row.prompt_tokens}, " + f"cache read {row.cache_read_tokens}, cache creation {row.cache_creation_tokens}); " + "cached tokens are being billed at the input rate" + ) + + +def poll_cost_row(proxy: ProxyClient, request_id: str) -> CostRow | None: + """Poll /spend/logs for the call's row until it lands with a cost breakdown + (rows flush ~60s behind the call via proxy_batch_write_at); None on timeout.""" + deadline = time.monotonic() + proxy.poll_timeout + while time.monotonic() < deadline: + result = proxy.transport.get( + "/spend/logs", + headers=proxy.transport.master, + params=SpendLogsParams(request_id=request_id), + response_type=CostRows, + ) + match result: + case Success(data=data): + rows = data.root + case _: + rows = [] + for row in rows: + if row.metadata and row.metadata.cost_breakdown: + return row + time.sleep(proxy.poll_interval) + return None + + +def poll_cost_row_where( + proxy: ProxyClient, api_key: str, predicate: Callable[[CostRow], bool] +) -> CostRow | None: + """Poll the key's own /spend/logs until one of its rows carries a cost breakdown + the predicate accepts; None on timeout. For calls whose response id is not the + id the bill is filed under, which is how a user finds the row in the UI anyway.""" + deadline = time.monotonic() + proxy.poll_timeout + while time.monotonic() < deadline: + result = proxy.transport.get( + "/spend/logs", + headers=proxy.transport.master, + params=SpendLogsParams(api_key=api_key), + response_type=CostRows, + ) + match result: + case Success(data=data): + rows = data.root + case _: + rows = [] + for row in rows: + if row.metadata and row.metadata.cost_breakdown and predicate(row): + return row + time.sleep(proxy.poll_interval) + return None + + +def register_priced_model( + proxy: ProxyClient, + resources: ResourceManager, + name_prefix: str, + litellm_params: LiteLLMParamsBody, +) -> str: + """Register a deployment with explicit custom rates (deleted on teardown) and + return its unique model name.""" + model_name = f"{name_prefix}-{unique_marker()}" + model_id = proxy.create_model(model_name, litellm_params) + resources.defer(lambda: proxy.delete_model(model_id)) + return model_name + + +def cacheable_prefix(marker: str, *, words: int = 1200) -> str: + """A prompt prefix above OpenAI's 1024-token caching minimum whose identity is + fully determined by `marker` (it is the first word, and prefix caching matches + from token zero). Raise `words` for models that only report a cache read on a + substantially longer prefix.""" + return " ".join(marker if i == 0 else f"token{i:04d}" for i in range(words)) diff --git a/tests/e2e/quota_management/spend_tracking/test_cache_cost_accounting_e2e.py b/tests/e2e/quota_management/spend_tracking/test_cache_cost_accounting_e2e.py new file mode 100644 index 00000000000..ca5985fcda6 --- /dev/null +++ b/tests/e2e/quota_management/spend_tracking/test_cache_cost_accounting_e2e.py @@ -0,0 +1,263 @@ +"""Live e2e: prompt-cache token accounting bills each cache component at its own rate. + +Four regressions the gateway has shipped fixes for, pinned against real OpenAI +prompt caching (implicit, keyed on the token prefix). Every test registers its own +deployment with distinct custom rates for input / output / cache-read / +cache-creation, so the expected bill is exactly the row's token counts times the +configured rates and a component billed at the wrong rate can never pass: + +- cache writes: gpt-5.6's cache-write tokens must land on the spend row as + cache-creation tokens billed at the cache-creation rate, not silently at the + input rate (#34046) +- breakdown components: the row's metadata.cost_breakdown must itemize cache-read, + cache-creation, and reasoning costs, with reasoning a subset of output (#31686) +- streaming: a streamed call's reassembled usage must keep the cached-token detail + so cache reads bill at the cache-read discount, not full input price (#34812) +- /v1/messages bridge: a request served by a Responses-only OpenAI model crosses + the anthropic-messages -> Responses adapter and must keep its cache-read tokens + and their discounted billing (#34957) + +Each test drives the model that actually reports the component it bills, which is +not the same model throughout. gpt-5.6-luna reports cache-write tokens on every +call over the caching minimum and never reports a cache read, so it is the one +model that can prove cache-write billing and the one model that can never prove +cache-read billing. gpt-5.5 is the reverse: it reports cached tokens on the second +call and no cache writes at all. gpt-5.3-codex is Responses-only, which is what +forces the /v1/messages bridge, and it starts reporting cache reads once the +prefix is a few thousand tokens rather than one. + +OpenAI caching is best-effort, so each test retries with a fresh prefix (new +marker = brand-new cache identity) up to three times before failing; the prime and +measured calls share the prefix but differ in the trailing question, which defeats +the proxy's own response cache without touching the provider's prefix cache. +""" + +import pytest + +from cost_rows import ( + CostRow, + approx_equal, + assert_fresh_tokens_billed_at, + assert_total_is_sum_of_components, + cacheable_prefix, + poll_cost_row, + poll_cost_row_where, + register_priced_model, +) +from e2e_config import unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from models import AnthropicMessagesBody, ChatBody, ChatMessage, LiteLLMParamsBody +from pydantic import BaseModel +from spend_e2e_client import SpendClient + +pytestmark = pytest.mark.e2e + +CACHE_WRITE_BACKEND = "openai/gpt-5.6-luna" +CACHE_READ_BACKEND = "openai/gpt-5.5" +BRIDGE_BACKEND = "openai/gpt-5.3-codex" +BRIDGE_PREFIX_WORDS = 3000 +OPENAI_API_KEY = "os.environ/OPENAI_API_KEY" +CACHE_ATTEMPTS = 3 + +INPUT_RATE = 4e-05 +OUTPUT_RATE = 8e-05 +CACHE_READ_RATE = 1e-05 +CACHE_WRITE_RATE = 5e-05 + +PRIME_QUESTION = "Reply with the single word ready." +REASONING_QUESTION = "Compute 47*83 - 19*7 step by step, then reply with just the final number." + + +class _StreamChunk(BaseModel): + id: str | None = None + + +def _cache_priced_params(backend: str) -> LiteLLMParamsBody: + return LiteLLMParamsBody( + model=backend, + api_key=OPENAI_API_KEY, + input_cost_per_token=INPUT_RATE, + output_cost_per_token=OUTPUT_RATE, + cache_read_input_token_cost=CACHE_READ_RATE, + cache_creation_input_token_cost=CACHE_WRITE_RATE, + ) + + +def _chat_body(model: str, content: str, *, stream: bool = False) -> ChatBody: + return ChatBody( + model=model, + messages=[ChatMessage(role="user", content=content)], + stream=stream, + max_completion_tokens=4000, + ) + + +def _require_row(client: SpendClient, request_id: str) -> CostRow: + row = poll_cost_row(client.proxy, request_id) + assert row is not None, f"no spend row with a cost breakdown landed for {request_id}" + return row + + +def _assert_cache_read_billed(row: CostRow) -> None: + assert row.breakdown.cache_read_cost is not None and approx_equal( + row.breakdown.cache_read_cost, row.cache_read_tokens * CACHE_READ_RATE + ), ( + f"cache_read_cost {row.breakdown.cache_read_cost} != " + f"{row.cache_read_tokens} cached tokens * {CACHE_READ_RATE}" + ) + assert_fresh_tokens_billed_at(row, INPUT_RATE) + assert_total_is_sum_of_components(row) + + +class TestCacheCostAccounting: + @pytest.mark.covers("quota_management.spend_tracking.cache_write.bills_cache_creation_rate") + def test_cache_write_tokens_billed_at_cache_creation_rate( + self, client: SpendClient, resources: ResourceManager, scoped_key: str + ) -> None: + model = register_priced_model( + client.proxy, resources, "cache-write-priced", _cache_priced_params(CACHE_WRITE_BACKEND) + ) + + for _ in range(CACHE_ATTEMPTS): + prompt = f"{cacheable_prefix(unique_marker())}\n{PRIME_QUESTION}" + chat = unwrap(client.proxy.chat(scoped_key, _chat_body(model, prompt))) + assert chat.id, f"chat response carried no id: {chat}" + row = _require_row(client, chat.id) + if row.cache_creation_tokens > 0: + break + else: + pytest.fail( + f"OpenAI reported no cache-write tokens across {CACHE_ATTEMPTS} fresh " + "~2k-token prompts; the cache-write billing path was never exercised" + ) + + assert row.breakdown.cache_creation_cost is not None and approx_equal( + row.breakdown.cache_creation_cost, row.cache_creation_tokens * CACHE_WRITE_RATE + ), ( + f"cache_creation_cost {row.breakdown.cache_creation_cost} != " + f"{row.cache_creation_tokens} cache-write tokens * {CACHE_WRITE_RATE}" + ) + assert_fresh_tokens_billed_at(row, INPUT_RATE) + assert_total_is_sum_of_components(row) + + @pytest.mark.covers("quota_management.spend_tracking.cost_breakdown.reports_component_costs") + def test_cost_breakdown_reports_component_costs( + self, client: SpendClient, resources: ResourceManager, scoped_key: str + ) -> None: + model = register_priced_model( + client.proxy, resources, "breakdown-priced", _cache_priced_params(CACHE_READ_BACKEND) + ) + + for _ in range(CACHE_ATTEMPTS): + prefix = cacheable_prefix(unique_marker()) + unwrap(client.proxy.chat(scoped_key, _chat_body(model, f"{prefix}\n{PRIME_QUESTION}"))) + chat = unwrap( + client.proxy.chat(scoped_key, _chat_body(model, f"{prefix}\n{REASONING_QUESTION}")) + ) + assert chat.id, f"chat response carried no id: {chat}" + row = _require_row(client, chat.id) + if row.cache_read_tokens > 0: + break + else: + pytest.fail( + f"no cache read landed across {CACHE_ATTEMPTS} prime+read rounds; " + "the component-cost breakdown was never exercised with cached input" + ) + + usage = chat.usage + assert usage is not None and usage.completion_tokens_details is not None, ( + f"no completion token details on the measured call: {chat}" + ) + reasoning_tokens = usage.completion_tokens_details.reasoning_tokens or 0 + assert reasoning_tokens > 0, f"the reasoning question produced no reasoning tokens: {usage}" + + breakdown = row.breakdown + assert breakdown.output_cost is not None and approx_equal( + breakdown.output_cost, (row.completion_tokens or 0) * OUTPUT_RATE + ), ( + f"output_cost {breakdown.output_cost} != " + f"{row.completion_tokens} completion tokens * {OUTPUT_RATE}" + ) + assert breakdown.reasoning_cost is not None and approx_equal( + breakdown.reasoning_cost, reasoning_tokens * OUTPUT_RATE + ), ( + f"reasoning_cost {breakdown.reasoning_cost} != " + f"{reasoning_tokens} reasoning tokens * {OUTPUT_RATE}" + ) + assert breakdown.reasoning_cost <= (breakdown.output_cost or 0.0) * 1.01, ( + f"reasoning_cost {breakdown.reasoning_cost} exceeds output_cost " + f"{breakdown.output_cost}; reasoning must be a subset of output" + ) + _assert_cache_read_billed(row) + + @pytest.mark.covers("quota_management.spend_tracking.stream_cache_read.bills_cache_read_rate") + def test_streaming_cache_read_billed_at_cache_read_rate( + self, client: SpendClient, resources: ResourceManager, scoped_key: str + ) -> None: + model = register_priced_model( + client.proxy, resources, "stream-cache-priced", _cache_priced_params(CACHE_READ_BACKEND) + ) + + for _ in range(CACHE_ATTEMPTS): + prefix = cacheable_prefix(unique_marker()) + unwrap(client.proxy.chat(scoped_key, _chat_body(model, f"{prefix}\n{PRIME_QUESTION}"))) + result = client.proxy.chat_stream( + scoped_key, + _chat_body(model, f"{prefix}\nReply with the single word cached.", stream=True), + ) + assert result.ok and result.stream_events, ( + f"streamed chat failed (status {result.status_code}): {result.body[:300]}" + ) + stream_id = _StreamChunk.model_validate_json(result.stream_events[0]).id + assert stream_id, f"first stream chunk carried no id: {result.stream_events[0][:200]}" + row = _require_row(client, stream_id) + if row.cache_read_tokens > 0: + break + else: + pytest.fail( + f"no cache read landed across {CACHE_ATTEMPTS} prime+stream rounds; " + "streaming cache-read billing was never exercised" + ) + + _assert_cache_read_billed(row) + + @pytest.mark.covers("quota_management.spend_tracking.messages_bridge.keeps_cache_tokens") + def test_messages_bridge_keeps_cache_tokens( + self, client: SpendClient, resources: ResourceManager, scoped_key: str + ) -> None: + model = register_priced_model( + client.proxy, resources, "bridge-cache-priced", _cache_priced_params(BRIDGE_BACKEND) + ) + + def bridge_call(content: str) -> int: + response = unwrap( + client.proxy.messages( + scoped_key, + AnthropicMessagesBody( + model=model, + messages=[ChatMessage(role="user", content=content)], + max_tokens=4000, + ), + ) + ) + assert response.usage is not None, f"bridged response carried no usage: {response}" + return response.usage.cache_read_input_tokens or 0 + + for _ in range(CACHE_ATTEMPTS): + prefix = cacheable_prefix(unique_marker(), words=BRIDGE_PREFIX_WORDS) + bridge_call(f"{prefix}\n{PRIME_QUESTION}") + if bridge_call(f"{prefix}\nReply with the single word bridged.") > 0: + break + else: + pytest.fail( + f"no cache read survived {CACHE_ATTEMPTS} bridged prime+read rounds; " + "cache tokens are not surviving the anthropic-messages -> Responses bridge" + ) + + row = poll_cost_row_where(client.proxy, scoped_key, lambda r: r.cache_read_tokens > 0) + assert row is not None, ( + "the bridged call reported cached tokens but no spend row for the key " + "recorded any; the cache tokens were dropped on the way to the bill" + ) + _assert_cache_read_billed(row) diff --git a/tests/e2e/quota_management/spend_tracking/test_cost_headers_e2e.py b/tests/e2e/quota_management/spend_tracking/test_cost_headers_e2e.py new file mode 100644 index 00000000000..203be611905 --- /dev/null +++ b/tests/e2e/quota_management/spend_tracking/test_cost_headers_e2e.py @@ -0,0 +1,136 @@ +"""Live e2e: the per-component x-litellm-response-cost-* headers keep their contract. + +Pins the header contract shipped in #36965: alongside the x-litellm-response-cost +total, every response carries the component costs (input, output, cache-read, +cache-creation, reasoning, tool-usage), where input covers only fresh tokens (the +cache components are subtracted out) so the components sum to the total, and +reasoning stays a subset of output. + +The deployment carries distinct custom rates per component, a prime call fills the +provider's prefix cache, and the measured call re-reads it, so the cache-read +header is exercised with a real nonzero value instead of passing vacuously. The +backend is gpt-5.5 because it reports cached tokens on the second call; the +gpt-5.6 line reports cache writes and never a read, which would leave the +cache-read header at zero forever. The raw-transport send is used because the +typed chat client validates bodies and drops headers. OpenAI caching is +best-effort, so the prime+measure round retries with a fresh prefix before +failing. +""" + +import pytest + +from cost_rows import approx_equal, cacheable_prefix, register_priced_model +from e2e_config import unique_marker +from e2e_http import StreamingResponse +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody +from spend_e2e_client import SpendClient + +pytestmark = pytest.mark.e2e + +BACKEND = "openai/gpt-5.5" +OPENAI_API_KEY = "os.environ/OPENAI_API_KEY" +CACHE_ATTEMPTS = 3 + +INPUT_RATE = 4e-05 +OUTPUT_RATE = 8e-05 +CACHE_READ_RATE = 1e-05 +CACHE_WRITE_RATE = 5e-05 + +COMPONENT_HEADERS = ( + "x-litellm-response-cost-input", + "x-litellm-response-cost-cache-read", + "x-litellm-response-cost-cache-creation", + "x-litellm-response-cost-output", + "x-litellm-response-cost-tool-usage", +) + + +def _header_cost(response: StreamingResponse, name: str) -> float: + value = response.headers.get(name) + return float(value) if value not in (None, "", "None") else 0.0 + + +class TestCostHeaders: + @pytest.mark.covers("quota_management.spend_tracking.cost_headers.additive_components") + def test_component_cost_headers_sum_to_total( + self, client: SpendClient, resources: ResourceManager, scoped_key: str + ) -> None: + model = register_priced_model( + client.proxy, + resources, + "header-priced", + LiteLLMParamsBody( + model=BACKEND, + api_key=OPENAI_API_KEY, + input_cost_per_token=INPUT_RATE, + output_cost_per_token=OUTPUT_RATE, + cache_read_input_token_cost=CACHE_READ_RATE, + cache_creation_input_token_cost=CACHE_WRITE_RATE, + ), + ) + + def priced_call(content: str) -> StreamingResponse: + response = client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(scoped_key), + json=ChatBody( + model=model, + messages=[ChatMessage(role="user", content=content)], + max_completion_tokens=4000, + ), + ) + assert response.ok, f"chat failed (status {response.status_code}): {response.body[:300]}" + return response + + for _ in range(CACHE_ATTEMPTS): + prefix = cacheable_prefix(unique_marker()) + priced_call(f"{prefix}\nReply with the single word ready.") + measured = priced_call(f"{prefix}\nReply with the single word measured.") + if _header_cost(measured, "x-litellm-response-cost-cache-read") > 0: + break + else: + pytest.fail( + f"no cache read landed across {CACHE_ATTEMPTS} prime+measure rounds; " + "the cache-read cost header was never exercised with a nonzero value" + ) + + total = measured.response_cost + assert total is not None and total > 0, ( + f"x-litellm-response-cost missing or zero: {measured.headers}" + ) + component_sum = sum(_header_cost(measured, name) for name in COMPONENT_HEADERS) + assert approx_equal(component_sum, total), ( + f"component headers sum to {component_sum}, not the total {total}: " + f"{ {name: measured.headers.get(name) for name in COMPONENT_HEADERS} }" + ) + + reasoning = _header_cost(measured, "x-litellm-response-cost-reasoning") + output = _header_cost(measured, "x-litellm-response-cost-output") + assert reasoning <= output * 1.01, ( + f"reasoning header {reasoning} exceeds output header {output}; " + "reasoning must be a subset of output" + ) + + usage = ChatResponse.model_validate_json(measured.body).usage + assert usage is not None, f"measured response carried no usage: {measured.body[:300]}" + cached_tokens = ( + usage.prompt_tokens_details.cached_tokens or 0 if usage.prompt_tokens_details else 0 + ) + cache_creation_tokens = usage.cache_creation_input_tokens or 0 + assert cached_tokens > 0, f"cache-read header nonzero but usage shows no cached tokens: {usage}" + assert approx_equal( + _header_cost(measured, "x-litellm-response-cost-cache-read"), + cached_tokens * CACHE_READ_RATE, + ), ( + f"cache-read header {measured.headers.get('x-litellm-response-cost-cache-read')} != " + f"{cached_tokens} cached tokens * {CACHE_READ_RATE}" + ) + fresh_tokens = (usage.prompt_tokens or 0) - cached_tokens - cache_creation_tokens + assert approx_equal( + _header_cost(measured, "x-litellm-response-cost-input"), fresh_tokens * INPUT_RATE + ), ( + f"input header {measured.headers.get('x-litellm-response-cost-input')} != " + f"{fresh_tokens} fresh tokens * {INPUT_RATE}; the input component is not " + "subtracting the cache components" + ) diff --git a/tests/e2e/quota_management/spend_tracking/test_passthrough_stream_cost_e2e.py b/tests/e2e/quota_management/spend_tracking/test_passthrough_stream_cost_e2e.py new file mode 100644 index 00000000000..4c3a2a4509f --- /dev/null +++ b/tests/e2e/quota_management/spend_tracking/test_passthrough_stream_cost_e2e.py @@ -0,0 +1,68 @@ +"""Live e2e: the /openai passthrough injects usage.cost into streaming usage frames. + +Pins #36503: with the proxy running `include_cost_in_streaming_usage: true`, a +streamed call through the provider passthrough surface must carry the computed +cost inside the final usage-only SSE frame, the same contract the native +/chat/completions stream has. Providers never send `cost` themselves, so a +nonzero value proves the proxy computed and injected it on the passthrough path. + +The row-side spend accounting for passthrough calls is covered elsewhere; this +test pins only the in-stream cost surface, which clients read without ever +touching /spend/logs. +""" + +import pytest +from pydantic import BaseModel + +from e2e_config import unique_marker +from models import ChatBody, ChatMessage, StreamOptions, Usage +from spend_e2e_client import SpendClient + +pytestmark = pytest.mark.e2e + +OPENAI_MODEL = "gpt-5.6-luna" + + +class _StreamFrame(BaseModel): + usage: Usage | None = None + + +class TestPassthroughStreamCost: + @pytest.mark.covers("quota_management.spend_tracking.passthrough_stream.injects_usage_cost") + def test_passthrough_stream_final_usage_frame_carries_cost( + self, client: SpendClient, scoped_key: str + ) -> None: + result = client.proxy.transport.send( + "/openai/v1/chat/completions", + headers=client.proxy.transport.bearer(scoped_key), + json=ChatBody( + model=OPENAI_MODEL, + messages=[ + ChatMessage( + role="user", + content=f"{unique_marker()} Reply with the single word passthrough.", + ) + ], + stream=True, + stream_options=StreamOptions(), + ), + stream=True, + ) + assert result.ok and result.stream_events, ( + f"passthrough stream failed (status {result.status_code}): {result.body[:300]}" + ) + + usage_frames = [ + frame.usage + for frame in (_StreamFrame.model_validate_json(event) for event in result.stream_events) + if frame.usage is not None + ] + assert usage_frames, ( + f"no usage frame in the passthrough stream despite stream_options.include_usage; " + f"last event: {result.stream_events[-1][:300]}" + ) + + final_usage = usage_frames[-1] + assert final_usage.cost is not None and final_usage.cost > 0, ( + f"final passthrough usage frame carries no injected cost: {final_usage}" + ) diff --git a/tests/e2e/quota_management/spend_tracking/test_service_tier_pricing_e2e.py b/tests/e2e/quota_management/spend_tracking/test_service_tier_pricing_e2e.py new file mode 100644 index 00000000000..171c849fb4c --- /dev/null +++ b/tests/e2e/quota_management/spend_tracking/test_service_tier_pricing_e2e.py @@ -0,0 +1,116 @@ +"""Live e2e: a service_tier request bills every component at the tier's own rates. + +Pins the tier-billing fixes (#35923, #35925): a priority-tier call must price +input and output at the deployment's `*_priority` rates, including the reasoning +tokens inside output (the shipped bug billed reasoning at the default-tier rate), +and the spend row must record the tier the bill was computed on. + +The deployment carries custom base AND priority rates, each distinct, so a bill +computed from the wrong tier (or a mix) cannot match the expected numbers. The +prompt is a fresh unique marker per run, keeping cached tokens out of the math. +The response's own `service_tier` echo is asserted first: if OpenAI ever declined +priority processing and served the default tier, the test fails there instead of +producing a vacuous rate comparison. +""" + +import pytest + +from cost_rows import ( + approx_equal, + assert_fresh_tokens_billed_at, + assert_total_is_sum_of_components, + poll_cost_row, + register_priced_model, +) +from e2e_config import unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, LiteLLMParamsBody +from spend_e2e_client import SpendClient + +pytestmark = pytest.mark.e2e + +BACKEND = "openai/gpt-5.6-luna" +OPENAI_API_KEY = "os.environ/OPENAI_API_KEY" + +INPUT_RATE = 4e-05 +OUTPUT_RATE = 8e-05 +PRIORITY_INPUT_RATE = 6e-05 +PRIORITY_OUTPUT_RATE = 1.6e-04 + + +class TestServiceTierPricing: + @pytest.mark.covers("quota_management.spend_tracking.service_tier.bills_tier_rates") + def test_priority_tier_bills_priority_rates( + self, client: SpendClient, resources: ResourceManager, scoped_key: str + ) -> None: + model = register_priced_model( + client.proxy, + resources, + "tier-priced", + LiteLLMParamsBody( + model=BACKEND, + api_key=OPENAI_API_KEY, + input_cost_per_token=INPUT_RATE, + output_cost_per_token=OUTPUT_RATE, + input_cost_per_token_priority=PRIORITY_INPUT_RATE, + output_cost_per_token_priority=PRIORITY_OUTPUT_RATE, + ), + ) + + chat = unwrap( + client.proxy.chat( + scoped_key, + ChatBody( + model=model, + messages=[ + ChatMessage( + role="user", + content=( + f"{unique_marker()} Compute 47*83 - 19*7 step by step, " + "then reply with just the final number." + ), + ) + ], + max_completion_tokens=4000, + service_tier="priority", + ), + ) + ) + assert chat.service_tier == "priority", ( + f"OpenAI served tier {chat.service_tier!r} instead of priority; " + "tier billing was never exercised" + ) + assert chat.id, f"chat response carried no id: {chat}" + + row = poll_cost_row(client.proxy, chat.id) + assert row is not None, f"no spend row with a cost breakdown landed for {chat.id}" + breakdown = row.breakdown + + assert breakdown.service_tier == "priority", ( + f"the bill records pricing basis {breakdown.service_tier!r}, not priority" + ) + + assert_fresh_tokens_billed_at(row, PRIORITY_INPUT_RATE) + assert breakdown.output_cost is not None and approx_equal( + breakdown.output_cost, (row.completion_tokens or 0) * PRIORITY_OUTPUT_RATE + ), ( + f"output_cost {breakdown.output_cost} != {row.completion_tokens} tokens * priority rate " + f"{PRIORITY_OUTPUT_RATE} (base rate would give {(row.completion_tokens or 0) * OUTPUT_RATE})" + ) + + usage = chat.usage + assert usage is not None and usage.completion_tokens_details is not None, ( + f"no completion token details on the priority call: {chat}" + ) + reasoning_tokens = usage.completion_tokens_details.reasoning_tokens or 0 + assert reasoning_tokens > 0, f"the reasoning question produced no reasoning tokens: {usage}" + assert breakdown.reasoning_cost is not None and approx_equal( + breakdown.reasoning_cost, reasoning_tokens * PRIORITY_OUTPUT_RATE + ), ( + f"reasoning_cost {breakdown.reasoning_cost} != {reasoning_tokens} reasoning tokens * " + f"priority rate {PRIORITY_OUTPUT_RATE} (the default-tier rate would give " + f"{reasoning_tokens * OUTPUT_RATE})" + ) + + assert_total_is_sum_of_components(row) From ef6af5c615815592da85e26315b6212046ba82ba Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 02:03:27 -0700 Subject: [PATCH 061/542] test(e2e): accept both model-not-found phrasings on a shared proxy /audio/transcriptions answers a model-less request with one of two 400s depending on whether any wildcard deployment is registered at the time, and every suite shares one proxy, so run order decided which message came back. The assertion pinned only the no-wildcard wording, so it went red whenever the model-access-group suite had registered its wildcards first. It now accepts either message and still holds the error to naming the model Verified against a live proxy in both states: with a wildcard registered (the message CI was seeing) and with none (the message the assertion expected), the suite passes 3/3 either way --- .../llm_translation/test_audio_transcriptions_e2e.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py b/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py index 92b33fef85f..735f1a4a703 100644 --- a/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py +++ b/tests/e2e/llm_translation/test_audio_transcriptions_e2e.py @@ -3,7 +3,10 @@ Registers an OpenAI speech-to-text deployment at runtime and uploads a spoken weather question (the realtime suite's 24kHz WAV fixture) as multipart, asserting the returned transcript is non-empty and mentions the word it was asked about. -Also pins missing file/model negatives. +Also pins missing file/model negatives. A model-less request comes back as one of +two 400s depending on whether any wildcard deployment happens to be registered on +the shared proxy, so the assertion accepts either phrasing and holds both to naming +the model as the problem. """ from __future__ import annotations @@ -25,6 +28,8 @@ WEATHER_WAV = ( Path(__file__).resolve().parent / "realtime" / "fixtures" / "weather_question_24k.wav" ) +MISSING_MODEL_PHRASES: Final = ("model=none", "invalid model", "model is required") + class _OptionalTranscriptionForm(BaseModel): model: str | None = None @@ -105,8 +110,8 @@ class TestAudioTranscriptions: match result: case UnknownApiError(status_code=400, body=body): lowered: Final = body.lower() - assert "model" in lowered and ("required" in lowered or "invalid model" in lowered), ( - f"missing model error must identify the required model: {body[:300]}" + assert any(phrase in lowered for phrase in MISSING_MODEL_PHRASES), ( + f"missing model error must name the model as the problem: {body[:300]}" ) case other: pytest.fail(f"missing model expected a model-specific 400, got {other!r}") From b6fef179ff151e2cb88990ed24fe2613a1824704 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 02:08:04 -0700 Subject: [PATCH 062/542] fix(cli): stop a repeat logout from retracting its own keychain warning A logout that could not reach the keychain deleted the token file whenever it still held its own secret, and the next logout read that missing file as proof the keychain was clean. It answered the warning the first run had just issued with "Logged out successfully" while the entry an earlier login left behind was still live. The file is the only record that something may still be in there, which is what `_nothing_left_behind` already says it relies on, so keep it and take only the secret out. A keychain that did answer is a different case. `SecretStranded` means the entry is confirmed there and would not delete, and that needs no note in the file, while keeping one lets every later command read the credential straight back out of the keychain, which makes "Logged out locally" untrue. That one drops the file, as it did before. The secret still goes first either way: a copy that cannot be replaced with a secret-free one is removed rather than kept. --- litellm/litellm_core_utils/cli_token_utils.py | 24 ++++++++++--- .../test_cli_token_utils.py | 34 +++++++++++++++++-- .../proxy/client/cli/test_auth_commands.py | 2 +- 3 files changed, 52 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index 15b1390f337..4e01dd723ce 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -153,19 +153,33 @@ def _keep_the_secret_in_the_file(record: CliTokenRecord, outcome: SecretWrite) - def clear_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> SecretErase: """Remove the credential from both stores. Reports whether the keychain is now free of it. - A file that holds no secret of its own is kept when the keychain will not confirm the entry is - gone, because it is the only remaining record that something is still in there to remove. That - is what lets a later run tell a machine with a credential it cannot reach apart from one that - never had a login at all. Anything still holding a secret is removed either way. + A logout the keychain never answered keeps the token file, with its secret taken out, because + that file is the only remaining record that something may still be in there to remove. It is + what lets a later run tell a machine with a credential it cannot reach apart from one that never + had a login at all, and taking it away would leave the next logout answering the warning this + one just issued with a false all-clear. The secret goes either way. """ outcome: Final = vault.erase() record: Final = _read_token_file() settled: Final = _nothing_left_behind(outcome, record) - if settled or record is None or record.key is not None: + if settled or not _keep_the_unchecked_keychain_on_record(outcome, record): Path(get_cli_token_file_path()).unlink(missing_ok=True) return SecretErased() if settled else outcome +def _keep_the_unchecked_keychain_on_record(outcome: SecretErase, record: CliTokenRecord | None) -> bool: + """Whether the token file, stripped of its secret, is worth keeping as the note that says so. + + Only a keychain that could not be reached leaves the question open. One that answered for itself + is remembered without any help from the file, and a file it can still pair a live entry with + would leave the machine signed in to the login that was just ended. A copy that cannot be + replaced with a secret-free one is not kept either, because the secret goes first. + """ + if record is None or isinstance(outcome, SecretStranded): + return False + return _scrub_file_secret(record) + + def _nothing_left_behind(outcome: SecretErase, record: CliTokenRecord | None) -> bool: """Whether the keychain can be trusted to hold no credential of ours once the file is gone. 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 0cf1e55d363..162e5dd4b67 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 @@ -477,6 +477,18 @@ class TestClearCliToken: assert clear_cli_token(vault=vault) == SecretStranded() assert not _token_file(isolated_home).exists() + def test_a_keychain_that_will_not_release_the_secret_still_ends_the_local_login( + self, isolated_home, secret_vault_factory + ): + """The warning this returns says the machine is logged out locally and the keychain entry is + what is left over. Keeping the file that names that entry makes the first half untrue: every + later command reads the credential straight back out of the keychain and keeps working.""" + _write_metadata_only_file(isolated_home) + vault = secret_vault_factory(blob=_blob(), erasable=False) + + assert clear_cli_token(vault=vault) == SecretStranded() + assert load_cli_token(vault=vault) is None + @pytest.mark.parametrize( "failure", [KeyringDisabled(), KeyringUnreachable(), KeyringDiscardsWrites()] ) @@ -491,7 +503,7 @@ class TestClearCliToken: vault = secret_vault_factory(available=False, failure=failure) assert clear_cli_token(vault=vault) == failure - assert not _token_file(isolated_home).exists() + assert json.loads(_token_file(isolated_home).read_text()).get("key") is None def test_a_second_logout_still_reports_the_keychain_it_could_not_clear( self, isolated_home, secret_vault_factory @@ -539,7 +551,25 @@ class TestClearCliToken: clear_cli_token(vault=vault) - assert not _token_file(isolated_home).exists() + assert "sk-legacy" not in _token_file(isolated_home).read_text() + + def test_a_repeat_logout_never_answers_its_own_warning_with_an_all_clear( + self, isolated_home, secret_vault_factory + ): + """Sign in while the keychain works, sign in again once it has gone out of reach so the + second secret lands in the file, then log out twice. The first logout cannot say the first + login's entry is gone, and says so. If the second one reads the file the first one took + away as proof of a clean keychain, it retracts that warning while the credential behind it + is still live.""" + vault = secret_vault_factory() + save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-first"), vault=vault) + vault.available = False + save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-second"), vault=vault) + + assert clear_cli_token(vault=vault) == KeyringUnreachable() + assert clear_cli_token(vault=vault) == KeyringUnreachable() + assert vault.blob is not None + assert "sk-second" not in _token_file(isolated_home).read_text() @pytest.mark.parametrize("failure", [KeyringNotInstalled(), KeyringDisabled(), KeyringUnreachable()]) def test_logging_out_of_a_machine_that_never_logged_in_invents_nothing_to_warn_about( 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 e491dbd5aca..8e1551c0720 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -1074,7 +1074,7 @@ class TestKeychainBackedCommands: assert result.exit_code == 0 assert "could not be removed" in result.output - assert json.loads((isolated_home / ".litellm" / "token.json").read_text()).get("key") is None + 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 From f86aeba1e7b1ed1395a5b6d5b5c47c2e6c94d7ca Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 02:21:21 -0700 Subject: [PATCH 063/542] fix(cli): stop a file-held secret from vouching for an unreadable keychain A logout run from an install without the keyring package treated a token file holding its own secret as proof that no keychain entry could exist. That only holds for the login which wrote the file. A login before it may have had the package and put its credential in the keychain, where it outlives both the uninstall and the file that replaced it, so logout reported a clean sweep over a live credential. Every keychain that cannot be reached is now treated the same way, and the message says the keychain went unchecked rather than asserting what is in it. --- litellm/litellm_core_utils/cli_token_utils.py | 15 +++++------- litellm/proxy/client/cli/commands/auth.py | 2 +- .../test_cli_token_utils.py | 23 ++++++++++++------- .../proxy/client/cli/test_auth_commands.py | 16 ++++++++----- 4 files changed, 32 insertions(+), 24 deletions(-) diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index 4e01dd723ce..78b52f33e2d 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -185,22 +185,19 @@ def _nothing_left_behind(outcome: SecretErase, record: CliTokenRecord | None) -> A machine with no token file has no stored login to end, and `clear_cli_token` keeps one behind whenever the keychain is left unconfirmed, so a missing file is real evidence rather than the - absence of it. Past that, a keychain that exists but is out of reach right now is - never trusted, whatever the file looks like: the login that stored a secret there and the - logout that cannot remove it are separate runs, free to differ in whether the keychain was - usable at the time. The exception is a missing `keyring` package, which had to be missing when - the credential was stored too, so a file still holding its own secret proves no keychain was - ever involved. `SecretStranded` is the keychain answering for itself and outranks the file. + absence of it. Past that, a keychain that could not be reached is never trusted, whatever the + file looks like. Even a file holding its own secret says only that the login which wrote it had + no keychain to write to, and the login before it may well have had one: the entry that login + left outlives both the uninstalled package and the file that replaced it. `SecretStranded` is + the keychain answering for itself and outranks the file. """ match outcome: case SecretErased(): return True case SecretStranded(): return False - case KeyringDisabled() | KeyringUnreachable(): + case KeyringDisabled() | KeyringNotInstalled() | KeyringUnreachable(): return record is None - case KeyringNotInstalled(): - return record is None or record.key is not None def get_litellm_gateway_api_key( diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index eb956cd536c..5fcc53e69eb 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -803,7 +803,7 @@ def logout(ctx: click.Context): click.echo(STRANDED_CREDENTIAL_MESSAGE) click.echo("Unlock your keychain and run 'lite logout' again to clear it.") case KeyringNotInstalled(): - click.echo(STRANDED_CREDENTIAL_MESSAGE) + click.echo(UNCHECKED_KEYCHAIN_MESSAGE) click.echo(f"Install the keyring package with: {KEYRING_INSTALL_HINT}, then run 'lite logout' again.") case KeyringDisabled(): click.echo(UNCHECKED_KEYCHAIN_MESSAGE) 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 162e5dd4b67..143c9e738a9 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 @@ -582,15 +582,22 @@ class TestClearCliToken: assert clear_cli_token(vault=vault) == SecretErased() - def test_a_file_backed_login_logs_out_quietly_without_keyring(self, isolated_home, secret_vault_factory): - """The complement, and the one inference the file does support: nothing here can reach a - keychain without the package, so an install that lacks it and a file that still holds its - own secret between them account for the whole credential.""" - _write_legacy_file(isolated_home) - vault = secret_vault_factory(available=False, failure=KeyringNotInstalled()) + def test_a_file_backed_login_cannot_vouch_for_a_keychain_no_package_can_reach( + self, isolated_home, secret_vault_factory + ): + """Sign in with the keyring package installed, lose the package, then sign in again so the + second secret lands in the file. The first login's entry outlives both, and the file that + replaced it holds a secret of its own, which is the shape a logout must not read as proof + that no keychain was ever involved.""" + vault = secret_vault_factory() + save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-keychain"), vault=vault) + vault.available = False + vault.failure = KeyringNotInstalled() + save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-in-file"), vault=vault) - assert clear_cli_token(vault=vault) == SecretErased() - assert not _token_file(isolated_home).exists() + assert clear_cli_token(vault=vault) == KeyringNotInstalled() + assert vault.blob is not None + assert "sk-in-file" not in _token_file(isolated_home).read_text() def test_is_safe_when_nothing_was_ever_stored(self, isolated_home, secret_vault_factory): assert clear_cli_token(vault=secret_vault_factory()) == SecretErased() 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 8e1551c0720..9e4ede88337 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -459,7 +459,7 @@ class TestLogoutCommand: assert result.exit_code == 0 assert "Logged out successfully" not in result.output - assert "still in the OS keychain" in result.output + assert "could not be checked" in result.output assert "pip install 'litellm[cli]'" in result.output def test_logout_does_not_call_an_unusable_keychain_clean(self, isolated_home, secret_vault_factory): @@ -492,9 +492,12 @@ class TestLogoutCommand: 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.""" + def test_logout_without_the_keyring_package_still_warns_about_a_file_held_secret( + self, isolated_home, secret_vault_factory + ): + """A file holding its own secret only says the login that wrote it had no keychain to write + to. An earlier login on this machine may have had one, and no install without the package + can look, so the honest answer is that the keychain went unchecked.""" _write_token_file(isolated_home, key="sk-in-file") result = self.runner.invoke( @@ -502,8 +505,9 @@ class TestLogoutCommand: ) assert result.exit_code == 0 - assert "Logged out successfully" in result.output - assert "still in the OS keychain" not in result.output + assert "Logged out successfully" not in result.output + assert "could not be checked" in result.output + assert "pip install 'litellm[cli]'" in result.output class TestWhoamiCommand: From ef104acdaf9ec791f1b1674fd7d6f7eabadeb6fc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 02:28:48 -0700 Subject: [PATCH 064/542] fix(cli): report a token file logout cannot remove instead of crashing A ~/.litellm that has gone read-only, or one left root-owned by a sudo login, refuses both the scrubbed rewrite and the removal. The removal was unguarded, so 'lite logout' ended in a PermissionError traceback with the credential still readable in the file. It now comes back as an outcome the command reports, naming the file and what to do about it, and a file that holds no secret is still not worth alarming anyone over. --- litellm/litellm_core_utils/cli_token_utils.py | 34 ++++++- litellm/proxy/client/cli/commands/auth.py | 5 + .../test_cli_token_utils.py | 91 +++++++++++++++++++ .../proxy/client/cli/test_auth_commands.py | 17 ++++ 4 files changed, 143 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index 78b52f33e2d..f2c0f8ed001 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -63,8 +63,22 @@ class CredentialNotRecorded: """ +@dataclass(frozen=True, slots=True) +class CredentialNotCleared: + """The token file still holds the secret, because it could not be removed or rewritten. + + Logging out of the keychain is only half of it. A `~/.litellm` that refuses both the scrubbed + rewrite and the removal leaves the credential readable on disk, which is the one thing a logout + is for, so it is reported instead of being counted as a clean sweep. + """ + + detail: str + + SecretSave: TypeAlias = SecretWrite | CredentialNotSaved | CredentialNotRecorded +SecretClear: TypeAlias = SecretErase | CredentialNotCleared + class CliTokenRecord(BaseModel): """A stored CLI credential. @@ -150,23 +164,35 @@ def _keep_the_secret_in_the_file(record: CliTokenRecord, outcome: SecretWrite) - return outcome -def clear_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> SecretErase: +def clear_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> SecretClear: """Remove the credential from both stores. Reports whether the keychain is now free of it. A logout the keychain never answered keeps the token file, with its secret taken out, because that file is the only remaining record that something may still be in there to remove. It is what lets a later run tell a machine with a credential it cannot reach apart from one that never had a login at all, and taking it away would leave the next logout answering the warning this - one just issued with a false all-clear. The secret goes either way. + one just issued with a false all-clear. The secret goes either way, and a file that will give up + neither its copy nor itself outranks whatever the keychain had to say. """ outcome: Final = vault.erase() record: Final = _read_token_file() settled: Final = _nothing_left_behind(outcome, record) - if settled or not _keep_the_unchecked_keychain_on_record(outcome, record): - Path(get_cli_token_file_path()).unlink(missing_ok=True) + if not settled and _keep_the_unchecked_keychain_on_record(outcome, record): + return outcome + removal: Final = _remove_token_file() + if removal is not None and record is not None and record.key is not None: + return removal return SecretErased() if settled else outcome +def _remove_token_file() -> CredentialNotCleared | None: + try: + Path(get_cli_token_file_path()).unlink(missing_ok=True) + except OSError as error: + return CredentialNotCleared(str(error)) + return None + + def _keep_the_unchecked_keychain_on_record(outcome: SecretErase, record: CliTokenRecord | None) -> bool: """Whether the token file, stripped of its secret, is worth keeping as the note that says so. diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 5fcc53e69eb..b3a7db4bed4 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -27,6 +27,7 @@ from litellm.litellm_core_utils.cli_keyring import ( ) from litellm.litellm_core_utils.cli_token_utils import ( CliTokenRecord, + CredentialNotCleared, CredentialNotRecorded, CredentialNotSaved, SecretSave, @@ -796,9 +797,13 @@ def login(ctx: click.Context, config_claude: bool): @click.pass_context def logout(ctx: click.Context): """Logout and clear stored authentication""" + path: Final = get_cli_token_file_path() match clear_cli_token(vault=context_secret_vault(ctx)): case SecretErased(): click.echo("Logged out successfully. Authentication token cleared.") + case CredentialNotCleared(detail=detail): + click.echo(f"Your credential is still in {path}, which could not be removed: {detail}.") + click.echo("Delete that file, or make the directory writable and run 'lite logout' again.") case SecretStranded(): click.echo(STRANDED_CREDENTIAL_MESSAGE) click.echo("Unlock your keychain and run 'lite logout' again to clear it.") 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 143c9e738a9..2830fad58b4 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 @@ -27,6 +27,7 @@ from litellm.litellm_core_utils.cli_keyring import ( from litellm.litellm_core_utils.cli_token_utils import ( CliTokenRecord, CredentialNotRecorded, + CredentialNotCleared, CredentialNotSaved, clear_cli_token, get_cli_token_file_path, @@ -81,6 +82,26 @@ def _blob(base_url=SERVER, key="sk-vault", jwt_token=""): return json.dumps({"base_url": base_url, "key": key, "jwt_token": jwt_token}) +_REAL_REPLACE = os.replace + + +def _refuse_replace(*args, **kwargs): + raise OSError("device or resource busy") + + +class _ReplaceThatStartsRefusing: + """`os.replace` standing in for a path that cannot be replaced yet: a file another process holds + open on Windows, a directory that went read-only between staging and the rewrite.""" + + def __init__(self): + self.allowed = False + + def __call__(self, src, dst): + if not self.allowed: + raise OSError("device or resource busy") + _REAL_REPLACE(src, dst) + + 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") @@ -459,6 +480,43 @@ class TestScrubFailure: assert json.loads(path.read_text())["key"] == "sk-legacy" assert list(path.parent.glob(".tmp-*")) == [] + def test_a_rewrite_that_fails_after_the_keychain_took_the_secret_hands_it_back( + self, isolated_home, secret_vault_factory, monkeypatch + ): + """Staging can succeed and the rewrite still fail afterwards, which is the one window where + both stores hold the credential. The keychain copy goes back, so the file is left exactly as + it was found and the move can be tried again.""" + path = _write_legacy_file(isolated_home) + vault = secret_vault_factory() + monkeypatch.setattr("litellm.litellm_core_utils.private_json.os.replace", _refuse_replace) + + record = load_cli_token(vault=vault) + + assert record.key == "sk-legacy" + assert vault.blob is None + assert json.loads(path.read_text())["key"] == "sk-legacy" + assert list(path.parent.glob(".tmp-*")) == [] + + def test_a_rollback_the_keychain_refuses_is_finished_by_the_next_read( + self, isolated_home, secret_vault_factory, monkeypatch + ): + """A keychain that will not give back what it just took leaves the credential in both stores. + Nothing is lost by that, and nothing is abandoned either: the next read carries the move the + rest of the way, so the duplicate outlives only the condition that caused it.""" + path = _write_legacy_file(isolated_home) + vault = secret_vault_factory(erasable=False) + replace = _ReplaceThatStartsRefusing() + monkeypatch.setattr("litellm.litellm_core_utils.private_json.os.replace", replace) + + assert load_cli_token(vault=vault).key == "sk-legacy" + assert vault.blob is not None + assert json.loads(path.read_text())["key"] == "sk-legacy" + + replace.allowed = True + + assert load_cli_token(vault=vault).key == "sk-legacy" + assert json.loads(path.read_text()).get("key") is None + class TestClearCliToken: def test_removes_the_credential_from_both_stores(self, isolated_home, secret_vault_factory): @@ -599,6 +657,39 @@ class TestClearCliToken: assert vault.blob is not None assert "sk-in-file" not in _token_file(isolated_home).read_text() + @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores directory permissions") + def test_a_file_that_can_be_neither_scrubbed_nor_removed_is_reported_not_raised( + self, isolated_home, secret_vault_factory + ): + """A `~/.litellm` gone read-only, or one left root-owned by a `sudo lite login`, refuses the + scrubbed rewrite and the removal alike. The credential is still readable on disk, which is + the one thing logging out is for, so it has to come back as an answer rather than as a + traceback the user has to read the code to understand.""" + path = _write_legacy_file(isolated_home) + path.parent.chmod(0o500) + try: + outcome = clear_cli_token(vault=secret_vault_factory()) + finally: + path.parent.chmod(0o700) + + assert isinstance(outcome, CredentialNotCleared) + assert json.loads(path.read_text())["key"] == "sk-legacy" + + @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores directory permissions") + def test_a_metadata_file_that_will_not_go_is_not_worth_alarming_the_user_over( + self, isolated_home, secret_vault_factory + ): + """The secret was in the keychain and the keychain gave it up. What is stuck on disk names a + credential that no longer exists, so the logout it describes really did happen.""" + path = _write_metadata_only_file(isolated_home) + path.parent.chmod(0o500) + try: + outcome = clear_cli_token(vault=secret_vault_factory(blob=_blob())) + finally: + path.parent.chmod(0o700) + + assert outcome == SecretErased() + def test_is_safe_when_nothing_was_ever_stored(self, isolated_home, secret_vault_factory): assert clear_cli_token(vault=secret_vault_factory()) == SecretErased() 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 9e4ede88337..71fc1cd3e6a 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -492,6 +492,23 @@ class TestLogoutCommand: assert "still in the OS keychain" in result.output assert "Unlock your keychain" in result.output + @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores directory permissions") + def test_logout_reports_a_token_file_it_cannot_remove(self, isolated_home, secret_vault_factory): + """`lite logout` on a read-only ~/.litellm used to end in a PermissionError traceback with + the credential still sitting in the file. The user has to be told what is left and where.""" + _write_token_file(isolated_home, key="sk-in-file") + config_dir = isolated_home / ".litellm" + config_dir.chmod(0o500) + try: + result = self.runner.invoke(logout, obj={"secret_vault": secret_vault_factory()}) + finally: + config_dir.chmod(0o700) + + assert result.exit_code == 0 + assert "Logged out successfully" not in result.output + assert "still in" in result.output + assert str(config_dir / "token.json") in result.output + def test_logout_without_the_keyring_package_still_warns_about_a_file_held_secret( self, isolated_home, secret_vault_factory ): From 4ca1f3148a9df608eda0a4b562badc2c90d25f28 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 02:58:53 -0700 Subject: [PATCH 065/542] fix(cli): finish a refused token file rewrite in place Taking the secret out of ~/.litellm/token.json stages a replacement and moves it into place, which needs room for a second file and a directory that will accept a new entry. A full disk refuses the first and a read-only ~/.litellm the second, and logout gave up there: it removed the file when it could, dropping the record that the keychain had never been confirmed clear, so the logout after it reported a clean keychain it never checked Shortening the file already in place needs neither, so the logout scrub and the legacy migration now fall back to overwriting it where it lies. On a read-only ~/.litellm the logout the user asked for now happens, instead of coming back with instructions to delete the file by hand --- litellm/litellm_core_utils/cli_token_utils.py | 45 ++++++++++---- litellm/litellm_core_utils/private_json.py | 15 +++++ .../test_cli_token_utils.py | 58 ++++++++++++++----- .../litellm_core_utils/test_private_json.py | 37 ++++++++++++ .../proxy/client/cli/test_auth_commands.py | 38 +++++++++--- 5 files changed, 162 insertions(+), 31 deletions(-) create mode 100644 tests/test_litellm/litellm_core_utils/test_private_json.py diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index f2c0f8ed001..2f45742ce03 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -38,6 +38,7 @@ from litellm.litellm_core_utils.private_json import ( commit_staged_json, discard_staged_json, ensure_private_dir, + overwrite_private_json, stage_private_json, write_private_json, ) @@ -180,7 +181,7 @@ def clear_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> SecretClear: if not settled and _keep_the_unchecked_keychain_on_record(outcome, record): return outcome removal: Final = _remove_token_file() - if removal is not None and record is not None and record.key is not None: + if removal is not None and record is not None and not _scrub_file_secret(record): return removal return SecretErased() if settled else outcome @@ -198,8 +199,9 @@ def _keep_the_unchecked_keychain_on_record(outcome: SecretErase, record: CliToke Only a keychain that could not be reached leaves the question open. One that answered for itself is remembered without any help from the file, and a file it can still pair a live entry with - would leave the machine signed in to the login that was just ended. A copy that cannot be - replaced with a secret-free one is not kept either, because the secret goes first. + would leave the machine signed in to the login that was just ended. A copy that will give up + its secret neither to a staged replacement nor to an overwrite is not kept either, because the + secret goes first. """ if record is None or isinstance(outcome, SecretStranded): return False @@ -210,11 +212,12 @@ def _nothing_left_behind(outcome: SecretErase, record: CliTokenRecord | None) -> """Whether the keychain can be trusted to hold no credential of ours once the file is gone. A machine with no token file has no stored login to end, and `clear_cli_token` keeps one behind - whenever the keychain is left unconfirmed, so a missing file is real evidence rather than the - absence of it. Past that, a keychain that could not be reached is never trusted, whatever the - file looks like. Even a file holding its own secret says only that the login which wrote it had - no keychain to write to, and the login before it may well have had one: the entry that login - left outlives both the uninstalled package and the file that replaced it. `SecretStranded` is + whenever the keychain is left unconfirmed, taking the secret out in place when it cannot stage a + replacement, so a missing file is real evidence rather than the absence of it. Past that, a + keychain that could not be reached is never trusted, whatever the file looks like. Even a file + holding its own secret says only that the login which wrote it had no keychain to write to, and + the login before it may well have had one: the entry that login left outlives both the + uninstalled package and the file that replaced it. `SecretStranded` is the keychain answering for itself and outranks the file. """ match outcome: @@ -323,6 +326,11 @@ def _migrate_file_secret(record: CliTokenRecord, vault: SecretVault) -> CliToken before the keychain is handed anything. Copying the credential into a second store and only then discovering the first one cannot be cleaned would widen exposure instead of narrowing it, which is the opposite of what moving it into the keychain is for. + + A staged file that will not go into place is overwritten where it lies before the keychain is + asked to take the new entry back, so the migration finishes on a directory that would only ever + have refused it. Rolling back is the last resort, and a rollback the keychain also refuses + leaves the secret in both stores until the next read, which retries this same migration. """ if record.key is None: return None @@ -332,7 +340,7 @@ def _migrate_file_secret(record: CliTokenRecord, vault: SecretVault) -> CliToken if not isinstance(vault.write(_encode_secret(record.base_url, record.key, record.jwt_token)), SecretStored): discard_staged_json(staged) return record - if not _commit_token_file(staged): + if not _commit_token_file(staged) and not _overwrite_file_secret(record): vault.erase() return record @@ -342,7 +350,24 @@ def _scrub_file_secret(record: CliTokenRecord) -> bool: if record.key is None and not record.jwt_token: return True staged: Final = _stage_scrubbed_file(record) - return staged is not None and _commit_token_file(staged) + if staged is not None and _commit_token_file(staged): + return True + return _overwrite_file_secret(record) + + +def _overwrite_file_secret(record: CliTokenRecord) -> bool: + """Take the secret out of the token file where it lies, when no replacement can be put in place. + + The atomic rewrite wants room for a second file and a directory that will accept it. A full disk + refuses the first and a read-only `~/.litellm` the second, and neither stands in the way of + shortening the file that is already there. It is worth the loss of atomicity because a partial + write reads as no login at all, which is where the refused rewrite left the next run anyway. + """ + try: + overwrite_private_json(get_cli_token_file_path(), _without_secret(record).model_dump(exclude_none=True)) + except OSError: + return False + return True def _stage_scrubbed_file(record: CliTokenRecord) -> str | None: diff --git a/litellm/litellm_core_utils/private_json.py b/litellm/litellm_core_utils/private_json.py index fbeb74aab5a..30f64c8fc27 100644 --- a/litellm/litellm_core_utils/private_json.py +++ b/litellm/litellm_core_utils/private_json.py @@ -45,6 +45,21 @@ def commit_staged_json(staged: str, path: str) -> None: raise +def overwrite_private_json(path: str, data: Mapping[str, object]) -> None: + """Rewrite a file that is already there, in place, keeping the mode it was created with. + + `write_private_json` needs room for a second file and a directory that will accept it, which is + what a full disk and a read-only `~/.litellm` respectively refuse. Shortening the file already + in place needs neither. It is not atomic, so an interrupted write leaves a partial file, and it + never creates one, so it cannot put a world-readable file where a private one was. + """ + fd: Final = os.open(path, os.O_WRONLY | os.O_TRUNC) + with os.fdopen(fd, "w") as f: + json.dump(data, f, indent=2) + f.flush() + os.fsync(f.fileno()) + + def discard_staged_json(staged: str) -> None: """Throw a staged file away when the change it was part of is abandoned""" Path(staged).unlink(missing_ok=True) 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 2830fad58b4..e3d14c6da46 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 @@ -480,12 +480,13 @@ class TestScrubFailure: assert json.loads(path.read_text())["key"] == "sk-legacy" assert list(path.parent.glob(".tmp-*")) == [] - def test_a_rewrite_that_fails_after_the_keychain_took_the_secret_hands_it_back( + def test_a_rewrite_the_directory_refuses_is_finished_in_place( self, isolated_home, secret_vault_factory, monkeypatch ): """Staging can succeed and the rewrite still fail afterwards, which is the one window where - both stores hold the credential. The keychain copy goes back, so the file is left exactly as - it was found and the move can be tried again.""" + both stores hold the credential. Shortening the file already there needs neither a second + file nor a cooperative directory, so the move finishes rather than handing the keychain copy + back and leaving the cleartext where it was.""" path = _write_legacy_file(isolated_home) vault = secret_vault_factory() monkeypatch.setattr("litellm.litellm_core_utils.private_json.os.replace", _refuse_replace) @@ -493,26 +494,30 @@ class TestScrubFailure: record = load_cli_token(vault=vault) assert record.key == "sk-legacy" - assert vault.blob is None - assert json.loads(path.read_text())["key"] == "sk-legacy" + assert vault.blob is not None + assert json.loads(path.read_text()).get("key") is None assert list(path.parent.glob(".tmp-*")) == [] + @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores file permissions") def test_a_rollback_the_keychain_refuses_is_finished_by_the_next_read( self, isolated_home, secret_vault_factory, monkeypatch ): - """A keychain that will not give back what it just took leaves the credential in both stores. - Nothing is lost by that, and nothing is abandoned either: the next read carries the move the - rest of the way, so the duplicate outlives only the condition that caused it.""" + """A file that will take neither a replacement nor an overwrite, and a keychain that will not + give back what it just took, leave the credential in both stores. Nothing is lost by that, + and nothing is abandoned either: the next read carries the move the rest of the way, so the + duplicate outlives only the conditions that caused it.""" path = _write_legacy_file(isolated_home) vault = secret_vault_factory(erasable=False) replace = _ReplaceThatStartsRefusing() monkeypatch.setattr("litellm.litellm_core_utils.private_json.os.replace", replace) + path.chmod(0o400) assert load_cli_token(vault=vault).key == "sk-legacy" assert vault.blob is not None assert json.loads(path.read_text())["key"] == "sk-legacy" replace.allowed = True + path.chmod(0o600) assert load_cli_token(vault=vault).key == "sk-legacy" assert json.loads(path.read_text()).get("key") is None @@ -657,24 +662,49 @@ class TestClearCliToken: assert vault.blob is not None assert "sk-in-file" not in _token_file(isolated_home).read_text() - @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores directory permissions") - def test_a_file_that_can_be_neither_scrubbed_nor_removed_is_reported_not_raised( + @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores file permissions") + def test_a_file_that_gives_up_neither_its_secret_nor_itself_is_reported_not_raised( self, isolated_home, secret_vault_factory ): - """A `~/.litellm` gone read-only, or one left root-owned by a `sudo lite login`, refuses the - scrubbed rewrite and the removal alike. The credential is still readable on disk, which is - the one thing logging out is for, so it has to come back as an answer rather than as a - traceback the user has to read the code to understand.""" + """A `~/.litellm` gone read-only refuses the staged rewrite and the removal, and a token file + left read-only with it, as a `sudo lite login` leaves both, refuses the overwrite too. The + credential is still readable on disk, which is the one thing logging out is for, so it has to + come back as an answer rather than as a traceback the user has to read the code to + understand.""" path = _write_legacy_file(isolated_home) + path.chmod(0o400) path.parent.chmod(0o500) try: outcome = clear_cli_token(vault=secret_vault_factory()) finally: path.parent.chmod(0o700) + path.chmod(0o600) assert isinstance(outcome, CredentialNotCleared) assert json.loads(path.read_text())["key"] == "sk-legacy" + @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores directory permissions") + def test_a_directory_that_takes_no_new_file_still_gives_up_the_secret_in_the_old_one( + self, isolated_home, secret_vault_factory + ): + """A read-only `~/.litellm` accepts no replacement token file and no removal of the one it + has, and still lets that one be shortened. The secret goes, the file stays as the note that + the keychain went unchecked, and the logout after it warns again instead of reading the gap + the removal would have left as a clean keychain. + + The key is a realistic length so the file genuinely shrinks: a rewrite in place that leaves + the tail of the old contents behind hands the next run a file it cannot parse.""" + path = _write_legacy_file(isolated_home, key="sk-" + "a" * 700) + vault = secret_vault_factory(available=False, failure=KeyringUnreachable()) + path.parent.chmod(0o500) + try: + assert clear_cli_token(vault=vault) == KeyringUnreachable() + assert clear_cli_token(vault=vault) == KeyringUnreachable() + finally: + path.parent.chmod(0o700) + + assert json.loads(path.read_text()).get("key") is None + @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores directory permissions") def test_a_metadata_file_that_will_not_go_is_not_worth_alarming_the_user_over( self, isolated_home, secret_vault_factory diff --git a/tests/test_litellm/litellm_core_utils/test_private_json.py b/tests/test_litellm/litellm_core_utils/test_private_json.py new file mode 100644 index 00000000000..cedff61959f --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_private_json.py @@ -0,0 +1,37 @@ +import json +import os +import stat + +import pytest + +from litellm.litellm_core_utils.private_json import overwrite_private_json, write_private_json + + +class TestOverwritePrivateJson: + def test_replaces_the_contents_of_the_file_already_there(self, tmp_path): + path = tmp_path / "token.json" + write_private_json(str(path), {"key": "sk-" + "a" * 700}) + + overwrite_private_json(str(path), {"user_id": "u-1"}) + + assert json.loads(path.read_text()) == {"user_id": "u-1"} + + def test_refuses_to_create_the_file_it_was_asked_to_rewrite(self, tmp_path): + """This is the one writer that does not go through a private temp file, so a path it creates + would land with whatever the umask allows. Refusing keeps it unable to put a world-readable + file where the caller believed a private one already was.""" + path = tmp_path / "token.json" + + with pytest.raises(FileNotFoundError): + overwrite_private_json(str(path), {"user_id": "u-1"}) + + assert not path.exists() + + @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores file permissions") + def test_keeps_the_owner_only_mode_the_file_was_created_with(self, tmp_path): + path = tmp_path / "token.json" + write_private_json(str(path), {"key": "sk-live"}) + + overwrite_private_json(str(path), {"user_id": "u-1"}) + + assert stat.S_IMODE(path.stat().st_mode) == 0o600 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 71fc1cd3e6a..8191a8edca7 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -492,12 +492,37 @@ class TestLogoutCommand: assert "still in the OS keychain" in result.output assert "Unlock your keychain" in result.output - @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores directory permissions") - def test_logout_reports_a_token_file_it_cannot_remove(self, isolated_home, secret_vault_factory): - """`lite logout` on a read-only ~/.litellm used to end in a PermissionError traceback with - the credential still sitting in the file. The user has to be told what is left and where.""" + @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores file permissions") + def test_logout_reports_a_token_file_it_cannot_clear(self, isolated_home, secret_vault_factory): + """`lite logout` on a read-only ~/.litellm holding a read-only token file used to end in a + PermissionError traceback with the credential still sitting in the file. The user has to be + told what is left and where.""" _write_token_file(isolated_home, key="sk-in-file") config_dir = isolated_home / ".litellm" + path = config_dir / "token.json" + path.chmod(0o400) + config_dir.chmod(0o500) + try: + result = self.runner.invoke(logout, obj={"secret_vault": secret_vault_factory()}) + finally: + config_dir.chmod(0o700) + path.chmod(0o600) + + assert result.exit_code == 0 + assert "Logged out successfully" not in result.output + assert "still in" in result.output + assert str(config_dir / "token.json") in result.output + + @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores directory permissions") + def test_logout_on_a_read_only_directory_still_takes_the_secret_out_of_the_file( + self, isolated_home, secret_vault_factory + ): + """A ~/.litellm that will accept no replacement file and no removal still lets the file it + has be shortened, so the logout the user asked for happens rather than being handed back to + them with instructions.""" + _write_token_file(isolated_home, key="sk-in-file") + config_dir = isolated_home / ".litellm" + path = config_dir / "token.json" config_dir.chmod(0o500) try: result = self.runner.invoke(logout, obj={"secret_vault": secret_vault_factory()}) @@ -505,9 +530,8 @@ class TestLogoutCommand: config_dir.chmod(0o700) assert result.exit_code == 0 - assert "Logged out successfully" not in result.output - assert "still in" in result.output - assert str(config_dir / "token.json") in result.output + assert "Logged out successfully" in result.output + assert "sk-in-file" not in path.read_text() def test_logout_without_the_keyring_package_still_warns_about_a_file_held_secret( self, isolated_home, secret_vault_factory From b142d1d76576221e394893c8e7abaf3254e167d3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 02:59:20 -0700 Subject: [PATCH 066/542] fix(cli): stop whoami calling an unreadable credential authenticated `lite whoami` led with "Authenticated" whenever a token file was on disk, even when the keychain holding the credential would not give it up. The notice about that sat below the account lines, so the session read as a working one and sent the user looking for the problem anywhere but the keychain --- litellm/proxy/client/cli/commands/auth.py | 2 +- .../proxy/client/cli/test_auth_commands.py | 10 ++++++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index b3a7db4bed4..ac330ad6892 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -867,7 +867,7 @@ def whoami(ctx: click.Context): click.echo("Not authenticated. Run 'lite login' to authenticate.") return - click.echo("Authenticated") + click.echo("Authenticated" if token_data.key is not None else "Signed in, but the credential cannot be read") 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'}") 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 8191a8edca7..1dd8a7ead92 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -1137,7 +1137,12 @@ class TestKeychainBackedCommands: assert "could not be read" in result.output assert "lite login" in result.output - def test_whoami_flags_a_locked_keychain(self, isolated_home, secret_vault_factory): + def test_whoami_does_not_call_a_credential_it_cannot_read_authenticated( + self, isolated_home, secret_vault_factory + ): + """A login whose secret is stuck in an unreachable keychain authenticates nothing. Leading + with "Authenticated" and a token age reads as a working session, and sends the user looking + for the problem somewhere other than the keychain the notice underneath names.""" _write_home_json( isolated_home, "token.json", @@ -1147,7 +1152,8 @@ class TestKeychainBackedCommands: result = self.runner.invoke(whoami, obj=obj) - assert "Authenticated" in result.output + assert "Authenticated" not in result.output + assert "the credential cannot be read" in result.output assert "could not be read" in result.output def test_whoami_names_the_kill_switch_rather_than_a_missing_package( From 8b7c801d61be5e4d02127ff6e86d743b157678f9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 03:03:35 -0700 Subject: [PATCH 067/542] test(e2e): pin openai_passthrough routing, cost logging, and file list isolation Five e2e tests over routes a customer drives through the gateway, each one pinning a fix that currently has no live coverage. The dedicated /openai_passthrough prefix used to be swallowed by the provider-scoped /{provider}/v1/files and /{provider}/v1/batches routes, which bound "openai_passthrough" as a provider name and failed inside the gateway before ever reaching OpenAI. Two tests now upload a file and list batches through that prefix and assert OpenAI's own objects come back. Streamed /openai_passthrough/v1/responses and /openai_passthrough/v1/embeddings are relayed to OpenAI but still have to be costed, since the customer budgets against this traffic. Both used to land a row the gateway could not use: the streamed responses call logged a zero-cost row under a random id, and embeddings wrote no row at all. Each test now reconciles the logged spend and token counts against the response the caller was actually served. GET /v1/files narrowed its data to the caller's own rows but left first_id and last_id addressing the shared provider account's page, handing any caller raw provider file ids belonging to other tenants. The new test asserts both cursors address rows in the page the caller can see. ResourceManager.defer now accepts any callable rather than one returning None, so a delete that answers with a response model can be deferred as-is. --- tests/e2e/batches/batch_client.py | 7 + tests/e2e/batches/test_batches_e2e.py | 34 +++++ .../coverage_registry/llm_conversational.yaml | 1 + .../llm_nonconversational.yaml | 4 + tests/e2e/lifecycle.py | 9 +- .../e2e/llm_translation/passthrough_client.py | 132 +++++++++++++++++- .../llm_translation/test_passthrough_e2e.py | 124 +++++++++++++++- 7 files changed, 305 insertions(+), 6 deletions(-) diff --git a/tests/e2e/batches/batch_client.py b/tests/e2e/batches/batch_client.py index 5cc5d1dae3b..968a357e8af 100644 --- a/tests/e2e/batches/batch_client.py +++ b/tests/e2e/batches/batch_client.py @@ -40,8 +40,15 @@ class FileObject(BaseModel): class FileList(BaseModel): + """GET /v1/files page. The cursors are modelled because they are part of the + page's isolation contract: they must address rows in `data`, never rows the + caller was not allowed to see.""" + object: str | None = None data: list[FileObject] = [] + first_id: str | None = None + last_id: str | None = None + has_more: bool | None = None class BatchObject(BaseModel): diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 53bf9739983..536bc113a25 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -572,6 +572,40 @@ class TestOpenAIFiles: f"listed file must round-trip the upload purpose, got {match.purpose!r}" ) + @pytest.mark.covers( + "llm.files.openai.list_isolation.nonstream.works", + exercised_on=["files"], + ) + def test_list_page_cursors_address_only_the_callers_own_files( + self, client: BatchClient, resources: ResourceManager + ) -> None: + """A list page's pagination cursors must address rows in that page. + + The proxy fronts one shared provider account, so the upstream page is the + whole organization's. The gateway narrows `data` to the files the caller + owns, and `first_id` / `last_id` have to be narrowed with it: left as the + upstream org's, they hand any caller raw provider file ids belonging to + other tenants, which is the handle the file routes accept. + """ + key = resources.key(user_id=f"e2e-file-list-{unique_marker()}") + + listed = unwrap(client.list_files(key=key)) + + expected_first = listed.data[0].id if listed.data else None + expected_last = listed.data[-1].id if listed.data else None + assert listed.first_id == expected_first, ( + f"first_id {listed.first_id!r} is not the first row this caller can see " + f"({expected_first!r}); the page leaked another caller's file id" + ) + assert listed.last_id == expected_last, ( + f"last_id {listed.last_id!r} is not the last row this caller can see " + f"({expected_last!r}); the page leaked another caller's file id" + ) + assert listed.has_more is not True, ( + "the page advertises another page, but the proxy never forwards a cursor " + "upstream, so following it re-serves this same page forever" + ) + @pytest.mark.covers( "llm.files.openai.retrieve.nonstream.works", exercised_on=["files"], diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index 82bee39b9b2..1ddc146c12d 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -63,6 +63,7 @@ - {id: llm.responses.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.9 / LIT-4778", rationale: "Responses missing/empty input and missing model are rejected"} - {id: llm.responses.openai.basic.stream.works, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: basic, streaming: stream, assertions: [works], source: "response_api_endpoints/endpoints.py:26", rationale: "Streaming via /v1/responses"} - {id: llm.responses.openai.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: basic, streaming: nonstream, assertions: [works, cost_logged], source: "response_api_endpoints/endpoints.py:26", rationale: "Cost logged on responses"} +- {id: llm.responses.openai.passthrough.stream.cost_logged, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: basic, streaming: stream, assertions: [cost_logged], source: "test_passthrough_e2e.py", rationale: "A streamed POST /openai_passthrough/v1/responses is costed and keyed by the provider response id; it used to log a zero-cost row under a random id (LIT-5870)"} - {id: llm.responses.openai.tool_use.nonstream.works, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: tool_use, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Tool calls via Responses API"} - {id: llm.responses.openai.vision.nonstream.works, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: vision, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Vision via Responses API"} - {id: llm.responses.anthropic.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: anthropic, capability: basic, streaming: nonstream, assertions: [works], source: "response_api_endpoints/endpoints.py:26", rationale: "Responses w/ Anthropic translation (smoke)"} diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index bb7169509eb..674d369b49f 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -3,6 +3,7 @@ - {id: llm.embeddings.openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: embeddings, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_embeddings_endpoint_e2e.py:23", rationale: "Core endpoint, live vector response"} - {id: llm.embeddings.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: embeddings, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.3 / LIT-4778", rationale: "Missing model/input on /embeddings return client errors"} - {id: llm.embeddings.openai.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: embeddings, route: openai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "SPEND_TRACKING_COVERAGE_MATRIX.md:34", rationale: "Cost tracking on embeddings"} +- {id: llm.embeddings.openai.passthrough.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: embeddings, route: openai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "test_passthrough_e2e.py", rationale: "POST /openai_passthrough/v1/embeddings is costed; the route wrote no spend row at all, so budgets never saw traffic OpenAI was billing for (LIT-5870)"} - {id: llm.embeddings.azure_openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: embeddings, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "llms/azure/azure.py", rationale: "Azure embeddings via translation"} - {id: llm.embeddings.bedrock.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: embeddings, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "llms/bedrock/embed/embedding.py", rationale: "Bedrock Titan embeddings"} - {id: llm.embeddings.vertex.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: embeddings, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "vertex_embeddings/embedding_handler.py", rationale: "Vertex embeddings"} @@ -13,6 +14,7 @@ - {id: llm.batches.openai.cancel.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Batch cancel"} - {id: llm.batches.openai.list.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Batch list envelope"} - {id: llm.batches.openai.file_lifecycle.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "File upload/retrieve/delete for batch flow"} +- {id: llm.batches.openai.passthrough.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_passthrough_e2e.py", rationale: "GET /openai_passthrough/v1/batches relays OpenAI's own batch page; the dedicated prefix must not bind as a provider name on the /{provider}/v1/batches route (LIT-5870)"} - {id: llm.batches.openai_encoded.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py", rationale: "Encoded scenario lifecycle"} - {id: llm.batches.openai_unified.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py", rationale: "Unified/managed-id scenario"} - {id: llm.batches.openai_model_param.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py", rationale: "Model-param scenario"} @@ -29,6 +31,8 @@ - {id: llm.files.openai.retrieve.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "files_endpoints.py", rationale: "File retrieve by id"} - {id: llm.files.openai.delete.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "files_endpoints.py", rationale: "File delete returns deleted=true"} - {id: llm.files.openai.list.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "files_endpoints.py", rationale: "File list paginated"} +- {id: llm.files.openai.list_isolation.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "GET /v1/files pagination cursors address only rows the caller owns; on a shared provider account the upstream cursors otherwise hand out other tenants' raw provider file ids (LIT-5870)"} +- {id: llm.files.openai.passthrough.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_passthrough_e2e.py", rationale: "POST/DELETE /openai_passthrough/v1/files relay OpenAI's own file object; the dedicated prefix must not bind as a provider name on the /{provider}/v1/files route (LIT-5870)"} - {id: llm.files.azure_openai.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:45", rationale: "Azure file upload managed backend"} - {id: llm.files.vertex.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:52", rationale: "Vertex file upload to GCS"} - {id: llm.files.bedrock.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:59", rationale: "Bedrock file upload to S3"} diff --git a/tests/e2e/lifecycle.py b/tests/e2e/lifecycle.py index 4ef25509905..c9a67ebdb8c 100644 --- a/tests/e2e/lifecycle.py +++ b/tests/e2e/lifecycle.py @@ -52,7 +52,7 @@ class ResourceManager: """ client: ResourceClient - _cleanups: List[Callable[[], None]] = field( + _cleanups: List[Callable[[], object]] = field( default_factory=list ) # mutable-ok: append-only teardown registry @@ -60,8 +60,11 @@ class ResourceManager: """No global setup needed today; present for lifecycle symmetry.""" return None - def defer(self, cleanup: Callable[[], None]) -> None: - """Register a teardown action for any resource the test just created.""" + def defer(self, cleanup: Callable[[], object]) -> None: + """Register a teardown action for any resource the test just created. + + Whatever the action returns is discarded, so a delete that answers with a + response model can be deferred directly.""" self._cleanups.append(cleanup) def key(self, models: list[str] | None = None, user_id: str | None = "e2e-test-user") -> str: diff --git a/tests/e2e/llm_translation/passthrough_client.py b/tests/e2e/llm_translation/passthrough_client.py index 439594f3624..e0dfae679a9 100644 --- a/tests/e2e/llm_translation/passthrough_client.py +++ b/tests/e2e/llm_translation/passthrough_client.py @@ -15,7 +15,7 @@ from dataclasses import dataclass from pydantic import BaseModel, Field from proxy_client import ProxyClient -from e2e_http import Headers, StreamingResponse +from e2e_http import FileUploadForm, Headers, NoBody, Result, StreamingResponse from models import ChatMessage @@ -113,6 +113,76 @@ class OpenAIChatBody(BaseModel): max_completion_tokens: int = 64 +class PassthroughFileObject(BaseModel): + id: str + object: str | None = None + purpose: str | None = None + filename: str | None = None + bytes: int | None = None + + +class PassthroughFileDeleted(BaseModel): + id: str + deleted: bool + + +class PassthroughListEntry(BaseModel): + id: str + + +class ResponsesUsage(BaseModel): + input_tokens: int + output_tokens: int + + +class ResponsesObject(BaseModel): + id: str + usage: ResponsesUsage | None = None + + +class ResponsesStreamEvent(BaseModel): + """One SSE frame of a native Responses stream. Only the terminal frames carry a + `response`, so it stays optional and the deltas validate as themselves.""" + + type: str + response: ResponsesObject | None = None + + +def completed_responses_object(result: StreamingResponse) -> ResponsesObject | None: + """The `response.completed` frame's response object, or None if the stream never + completed. Its `id` is what the spend row is keyed by on this route, and its + usage is what the row is priced from.""" + events = ( + ResponsesStreamEvent.model_validate_json(payload) + for payload in result.stream_events + ) + completed = tuple( + event.response + for event in events + if event.type == "response.completed" and event.response is not None + ) + return completed[-1] if completed else None + + +class OpenAIResponsesBody(BaseModel): + model: str + input: str + stream: bool = False + + +class OpenAIEmbeddingBody(BaseModel): + model: str + input: str + + +class PassthroughBatchList(BaseModel): + """OpenAI's own batch page, relayed verbatim. `object` is required so a body + that is not an OpenAI list fails validation instead of passing vacuously.""" + + object: str + data: list[PassthroughListEntry] + + def _tags_header(tags: list[str] | None) -> str | None: return ",".join(tags) if tags else None @@ -196,6 +266,66 @@ class PassthroughClient: stream=stream, ) + # ---- OpenAI file/batch routes under /openai_passthrough ------------- + # + # Relayed to OpenAI untouched, which is the whole point of the prefix: the + # customer opts out of the gateway's managed-file handling here. + + def openai_passthrough_upload_file( + self, key: str, *, content: bytes, filename: str + ) -> Result[PassthroughFileObject]: + return self.proxy.transport.upload( + "/openai_passthrough/v1/files", + headers=self.proxy.transport.bearer(key), + form=FileUploadForm(purpose="batch"), + filename=filename, + content=content, + response_type=PassthroughFileObject, + ) + + def openai_passthrough_delete_file( + self, key: str, file_id: str + ) -> Result[PassthroughFileDeleted]: + return self.proxy.transport.delete( + f"/openai_passthrough/v1/files/{file_id}", + headers=self.proxy.transport.bearer(key), + json=NoBody(), + response_type=PassthroughFileDeleted, + ) + + def openai_passthrough_list_batches(self, key: str) -> Result[PassthroughBatchList]: + return self.proxy.transport.get( + "/openai_passthrough/v1/batches", + headers=self.proxy.transport.bearer(key), + params=NoBody(), + response_type=PassthroughBatchList, + ) + + # ---- OpenAI inference routes under /openai_passthrough ------------- + # + # Relayed to OpenAI verbatim, but still costed by the gateway: the customer + # budgets against this traffic, so a 200 that logs no spend is money the + # gateway never sees. + + def openai_passthrough_responses( + self, key: str, model: str, text: str, *, stream: bool = False + ) -> StreamingResponse: + return self.proxy.transport.send( + "/openai_passthrough/v1/responses", + headers=self.proxy.transport.bearer(key), + json=OpenAIResponsesBody(model=model, input=text, stream=stream), + stream=stream, + ) + + def openai_passthrough_embed( + self, key: str, model: str, text: str + ) -> StreamingResponse: + return self.proxy.transport.send( + "/openai_passthrough/v1/embeddings", + headers=self.proxy.transport.bearer(key), + json=OpenAIEmbeddingBody(model=model, input=text), + ) + def openai_chat( self, key: str, model: str, text: str, *, max_completion_tokens: int = 64 ) -> StreamingResponse: diff --git a/tests/e2e/llm_translation/test_passthrough_e2e.py b/tests/e2e/llm_translation/test_passthrough_e2e.py index b57164df9bb..b084c711a88 100644 --- a/tests/e2e/llm_translation/test_passthrough_e2e.py +++ b/tests/e2e/llm_translation/test_passthrough_e2e.py @@ -13,8 +13,8 @@ A passthrough call returning non-2xx fails hard (never a skip); once it returns import pytest -from e2e_config import unique_marker -from e2e_http import StreamingResponse, require_successful_call +from e2e_config import CHEAP_OPENAI_MODEL, unique_marker +from e2e_http import StreamingResponse, require_successful_call, unwrap from lifecycle import ResourceManager from models import KeyGenerateBody, SpendLogRow from passthrough_client import ( @@ -24,8 +24,11 @@ from passthrough_client import ( JsonSchema, JsonSchemaProperty, PassthroughClient, + completed_responses_object, ) +EMBEDDING_MODEL = "text-embedding-3-small" + pytestmark = pytest.mark.e2e @@ -210,3 +213,120 @@ class TestPassthroughModelAllowlist: "a key restricted to gemini-2.5-flash must be denied a claude passthrough call, " f"got {result.status_code}: {result.body[:300]}" ) + + +class TestOpenAIPassthroughPrefix: + """The dedicated `/openai_passthrough` prefix must reach OpenAI, not be + swallowed by the provider-scoped `/{provider}/v1/...` routes. + + The customer fronts OpenAI's own file and batch APIs through this prefix + precisely to opt out of the gateway's managed-file handling. `/v1/files` and + `/v1/batches` also answer `/{provider}/v1/files` and `/{provider}/v1/batches`, + so `openai_passthrough` used to bind as a provider name and the request died + inside the gateway with a provider-lookup error, never reaching OpenAI. + """ + + @pytest.mark.covers("llm.files.openai.passthrough.nonstream.works") + def test_passthrough_prefix_uploads_a_file_to_openai( + self, client: PassthroughClient, resources: ResourceManager, scoped_key: str + ) -> None: + content = f'{{"marker":"{unique_marker()}"}}\n'.encode() + uploaded = unwrap( + client.openai_passthrough_upload_file( + scoped_key, content=content, filename="e2e-passthrough-batch.jsonl" + ) + ) + resources.defer( + lambda: client.openai_passthrough_delete_file(scoped_key, uploaded.id) + ) + + assert uploaded.object == "file", ( + f"/openai_passthrough/v1/files did not relay OpenAI's file object: {uploaded}" + ) + assert uploaded.purpose == "batch" + assert uploaded.bytes == len(content) + + @pytest.mark.covers("llm.batches.openai.passthrough.nonstream.works") + def test_passthrough_prefix_lists_batches_from_openai( + self, client: PassthroughClient, scoped_key: str + ) -> None: + listed = unwrap(client.openai_passthrough_list_batches(scoped_key)) + + assert listed.object == "list", ( + f"/openai_passthrough/v1/batches did not relay OpenAI's batch page: {listed}" + ) + + +class TestOpenAIPassthroughSpend: + """A call relayed to OpenAI's own endpoints must still be costed. + + The customer routes native OpenAI traffic through `/openai_passthrough` and + budgets against it, so a call that returns 200 while logging no spend is money + the gateway never sees and a budget that never trips. Streamed Responses calls + and embeddings each used to land exactly that way, on separate code paths. + """ + + @pytest.mark.covers("llm.responses.openai.passthrough.stream.cost_logged") + def test_streamed_responses_call_logs_its_cost( + self, client: PassthroughClient, scoped_key: str + ) -> None: + result = client.openai_passthrough_responses( + scoped_key, + CHEAP_OPENAI_MODEL, + f"Say hi in one word. {unique_marker()}", + stream=True, + ) + require_successful_call(result) + assert result.chunks > 0, "streamed responses passthrough produced no events" + + completed = completed_responses_object(result) + assert completed is not None, ( + f"the stream never delivered a response.completed frame, so there is no " + f"provider id to reconcile against: last events {result.stream_events[-3:]}" + ) + assert completed.usage is not None, ( + f"the completed response carried no usage to price from: {completed}" + ) + + rows = client.proxy.poll_logs_for_request_id( + completed.id, predicate=lambda rows: (rows[0].spend or 0) > 0 + ) + assert rows, ( + f"no spend row for the response the customer was served ({completed.id}); " + "a streamed passthrough call OpenAI bills them for is invisible to the " + "gateway's own spend and budgets" + ) + row = rows[0] + assert (row.spend or 0) > 0, f"streamed responses passthrough was not costed: {row}" + assert row.prompt_tokens == completed.usage.input_tokens, ( + f"logged {row.prompt_tokens} prompt tokens, the response the customer read " + f"reported {completed.usage.input_tokens}" + ) + assert row.completion_tokens == completed.usage.output_tokens, ( + f"logged {row.completion_tokens} completion tokens, the response the customer " + f"read reported {completed.usage.output_tokens}" + ) + + @pytest.mark.covers("llm.embeddings.openai.passthrough.nonstream.cost_logged") + def test_embeddings_call_logs_its_cost( + self, client: PassthroughClient, scoped_key: str + ) -> None: + result = client.openai_passthrough_embed( + scoped_key, EMBEDDING_MODEL, f"cost this sentence {unique_marker()}" + ) + require_successful_call(result) + assert result.call_id, "embeddings passthrough returned no x-litellm-call-id" + + rows = client.proxy.poll_logs_for_request_id( + result.call_id, predicate=lambda rows: (rows[0].spend or 0) > 0 + ) + assert rows, ( + f"no spend row for embeddings call {result.call_id}; the customer is billed " + "by OpenAI for tokens the gateway never counted against their budget" + ) + row = rows[0] + assert (row.spend or 0) > 0, f"embeddings passthrough was not costed: {row}" + assert (row.prompt_tokens or 0) > 0, ( + f"the embeddings row logged no prompt tokens, so whatever cost it carries " + f"was not computed from the real usage: {row}" + ) From aa8e7278e3d27a483ef8039995bce649dc6c0d88 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 03:08:05 -0700 Subject: [PATCH 068/542] test(e2e): drop the passthrough streaming-cost test, it needs a config flag The final streaming usage frame only carries usage.cost when the proxy runs with litellm_settings.include_cost_in_streaming_usage: true, and that flag is readable only off the module-level litellm setting. There is no header, key, or management route that turns it on per request, so a test cannot ask the shared e2e proxy for it, and the proxy's config does not live in this repo. The registry row stays as an uncovered gap with the reason recorded, rather than being deleted, so the behavior is still on the list of things we want covered once the gateway config is reachable. The StreamOptions model, ChatBody.stream_options, Usage.cost, and AnthropicMessagesResponse.id existed only for that test, so they go with it. --- .../coverage_registry/quota_management.yaml | 2 +- tests/e2e/models.py | 14 ---- .../test_passthrough_stream_cost_e2e.py | 68 ------------------- 3 files changed, 1 insertion(+), 83 deletions(-) delete mode 100644 tests/e2e/quota_management/spend_tracking/test_passthrough_stream_cost_e2e.py diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index 5438ed8534a..98a45eefb2d 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -50,4 +50,4 @@ - {id: quota_management.spend_tracking.messages_bridge.keeps_cache_tokens, module: quota_management, tier: P1, behavior: spend_tracking, variant: messages_bridge, assertions: [keeps_cache_tokens], exercised_on: [messages], source: "llms/anthropic/experimental_pass_through/responses_adapters/handler.py", rationale: "A /v1/messages request served by a Responses-only OpenAI model keeps its cache-read tokens and their discounted billing across the bridge (#34957)"} - {id: quota_management.spend_tracking.service_tier.bills_tier_rates, module: quota_management, tier: P1, behavior: spend_tracking, variant: service_tier, assertions: [bills_tier_rates], exercised_on: [chat_completions], source: "cost_calculator.py", rationale: "A priority service_tier call bills input, output, and reasoning at the deployment's *_priority rates and records the tier on the row (#35923, #35925)"} - {id: quota_management.spend_tracking.cost_headers.additive_components, module: quota_management, tier: P1, behavior: spend_tracking, variant: cost_headers, assertions: [additive_components], exercised_on: [chat_completions], source: "proxy/common_request_processing.py", rationale: "The x-litellm-response-cost-* component headers sum to the total, input covers only fresh tokens, and reasoning stays a subset of output (#36965)"} -- {id: quota_management.spend_tracking.passthrough_stream.injects_usage_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: passthrough_stream, assertions: [injects_usage_cost], exercised_on: [openai_passthrough], source: "proxy/pass_through_endpoints/streaming_handler.py", rationale: "With include_cost_in_streaming_usage on, the /openai passthrough's final streaming usage frame carries the proxy-computed cost (#36503)"} +- {id: quota_management.spend_tracking.passthrough_stream.injects_usage_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: passthrough_stream, assertions: [injects_usage_cost], exercised_on: [openai_passthrough], source: "proxy/pass_through_endpoints/streaming_handler.py", rationale: "With include_cost_in_streaming_usage on, the /openai passthrough's final streaming usage frame carries the proxy-computed cost (#36503). Uncovered: the flag is only settable in litellm_settings, and the shared e2e proxy's config is not in this repo"} diff --git a/tests/e2e/models.py b/tests/e2e/models.py index fe93b13e0a4..f1d5733253a 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -216,19 +216,10 @@ class McpChatTool(BaseModel): allowed_tools: list[str] | None = None -class StreamOptions(BaseModel): - """OpenAI `stream_options`: `include_usage` asks for a final usage-only SSE - frame, which is where the proxy's `include_cost_in_streaming_usage` setting - injects `usage.cost`.""" - - include_usage: bool = True - - class ChatBody(BaseModel): model: str messages: list[ChatMessage] stream: bool = False - stream_options: StreamOptions | None = None max_tokens: int | None = None max_completion_tokens: int | None = None temperature: float | None = None @@ -331,9 +322,6 @@ class CompletionTokensDetails(BaseModel): class Usage(BaseModel): - """`cost` exists only on streaming usage frames from a proxy running with - `include_cost_in_streaming_usage: true`; providers never send it.""" - prompt_tokens: int | None = None completion_tokens: int | None = None total_tokens: int | None = None @@ -341,7 +329,6 @@ class Usage(BaseModel): cache_creation_input_tokens: int | None = None prompt_tokens_details: PromptTokensDetails | None = None completion_tokens_details: CompletionTokensDetails | None = None - cost: float | None = None class ChatResponse(BaseModel): @@ -462,7 +449,6 @@ class AnthropicMessagesResponse(BaseModel): for triage.""" model_config = ConfigDict(extra="allow") - id: str | None = None model: str | None = None content: list[AnthropicContentBlock] | None = None choices: list[ChatChoice] | None = None diff --git a/tests/e2e/quota_management/spend_tracking/test_passthrough_stream_cost_e2e.py b/tests/e2e/quota_management/spend_tracking/test_passthrough_stream_cost_e2e.py deleted file mode 100644 index 4c3a2a4509f..00000000000 --- a/tests/e2e/quota_management/spend_tracking/test_passthrough_stream_cost_e2e.py +++ /dev/null @@ -1,68 +0,0 @@ -"""Live e2e: the /openai passthrough injects usage.cost into streaming usage frames. - -Pins #36503: with the proxy running `include_cost_in_streaming_usage: true`, a -streamed call through the provider passthrough surface must carry the computed -cost inside the final usage-only SSE frame, the same contract the native -/chat/completions stream has. Providers never send `cost` themselves, so a -nonzero value proves the proxy computed and injected it on the passthrough path. - -The row-side spend accounting for passthrough calls is covered elsewhere; this -test pins only the in-stream cost surface, which clients read without ever -touching /spend/logs. -""" - -import pytest -from pydantic import BaseModel - -from e2e_config import unique_marker -from models import ChatBody, ChatMessage, StreamOptions, Usage -from spend_e2e_client import SpendClient - -pytestmark = pytest.mark.e2e - -OPENAI_MODEL = "gpt-5.6-luna" - - -class _StreamFrame(BaseModel): - usage: Usage | None = None - - -class TestPassthroughStreamCost: - @pytest.mark.covers("quota_management.spend_tracking.passthrough_stream.injects_usage_cost") - def test_passthrough_stream_final_usage_frame_carries_cost( - self, client: SpendClient, scoped_key: str - ) -> None: - result = client.proxy.transport.send( - "/openai/v1/chat/completions", - headers=client.proxy.transport.bearer(scoped_key), - json=ChatBody( - model=OPENAI_MODEL, - messages=[ - ChatMessage( - role="user", - content=f"{unique_marker()} Reply with the single word passthrough.", - ) - ], - stream=True, - stream_options=StreamOptions(), - ), - stream=True, - ) - assert result.ok and result.stream_events, ( - f"passthrough stream failed (status {result.status_code}): {result.body[:300]}" - ) - - usage_frames = [ - frame.usage - for frame in (_StreamFrame.model_validate_json(event) for event in result.stream_events) - if frame.usage is not None - ] - assert usage_frames, ( - f"no usage frame in the passthrough stream despite stream_options.include_usage; " - f"last event: {result.stream_events[-1][:300]}" - ) - - final_usage = usage_frames[-1] - assert final_usage.cost is not None and final_usage.cost > 0, ( - f"final passthrough usage frame carries no injected cost: {final_usage}" - ) From 69278ae37bba01a5ad19a83273719794c5b47c91 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 03:25:31 -0700 Subject: [PATCH 069/542] docs(e2e): correct the passthrough-stream registry row's uncovered reason --- tests/e2e/coverage_registry/quota_management.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index 98a45eefb2d..e4b7e755bd0 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -50,4 +50,4 @@ - {id: quota_management.spend_tracking.messages_bridge.keeps_cache_tokens, module: quota_management, tier: P1, behavior: spend_tracking, variant: messages_bridge, assertions: [keeps_cache_tokens], exercised_on: [messages], source: "llms/anthropic/experimental_pass_through/responses_adapters/handler.py", rationale: "A /v1/messages request served by a Responses-only OpenAI model keeps its cache-read tokens and their discounted billing across the bridge (#34957)"} - {id: quota_management.spend_tracking.service_tier.bills_tier_rates, module: quota_management, tier: P1, behavior: spend_tracking, variant: service_tier, assertions: [bills_tier_rates], exercised_on: [chat_completions], source: "cost_calculator.py", rationale: "A priority service_tier call bills input, output, and reasoning at the deployment's *_priority rates and records the tier on the row (#35923, #35925)"} - {id: quota_management.spend_tracking.cost_headers.additive_components, module: quota_management, tier: P1, behavior: spend_tracking, variant: cost_headers, assertions: [additive_components], exercised_on: [chat_completions], source: "proxy/common_request_processing.py", rationale: "The x-litellm-response-cost-* component headers sum to the total, input covers only fresh tokens, and reasoning stays a subset of output (#36965)"} -- {id: quota_management.spend_tracking.passthrough_stream.injects_usage_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: passthrough_stream, assertions: [injects_usage_cost], exercised_on: [openai_passthrough], source: "proxy/pass_through_endpoints/streaming_handler.py", rationale: "With include_cost_in_streaming_usage on, the /openai passthrough's final streaming usage frame carries the proxy-computed cost (#36503). Uncovered: the flag is only settable in litellm_settings, and the shared e2e proxy's config is not in this repo"} +- {id: quota_management.spend_tracking.passthrough_stream.injects_usage_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: passthrough_stream, assertions: [injects_usage_cost], exercised_on: [openai_passthrough], source: "proxy/pass_through_endpoints/streaming_handler.py", rationale: "With include_cost_in_streaming_usage on, the /openai passthrough's final streaming usage frame carries the proxy-computed cost (#36503). Uncovered: the flag is only settable in litellm_settings, and the shared e2e stack does not turn it on yet"} From fe11202c2df1a6373ed3230e22847a76e157cb0d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 03:28:28 -0700 Subject: [PATCH 070/542] fix(cli): write the logout note again when the file holding it had to go When a full disk refuses the replacement file and a read-only token file refuses the rewrite in place, the only way left to get the secret off disk is to remove the file carrying it. That file was also the note saying the keychain went unchecked, so its absence made the next logout read a keychain that was never confirmed as one already known to be clean. Removing it is what frees the room the replacement was refused for, so the note is written again on the way out and the logout after this one still warns. --- litellm/litellm_core_utils/cli_token_utils.py | 32 ++++++++++-- .../test_cli_token_utils.py | 52 ++++++++++++++++++- 2 files changed, 79 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index 2f45742ce03..69a77a36882 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -173,7 +173,8 @@ def clear_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> SecretClear: what lets a later run tell a machine with a credential it cannot reach apart from one that never had a login at all, and taking it away would leave the next logout answering the warning this one just issued with a false all-clear. The secret goes either way, and a file that will give up - neither its copy nor itself outranks whatever the keychain had to say. + neither its copy nor itself is removed rather than kept, with the note written again afterwards + so the warning still outlives this run. """ outcome: Final = vault.erase() record: Final = _read_token_file() @@ -183,6 +184,8 @@ def clear_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> SecretClear: removal: Final = _remove_token_file() if removal is not None and record is not None and not _scrub_file_secret(record): return removal + if removal is None and record is not None and _the_keychain_went_unchecked(outcome): + _write_the_note_the_removal_took_with_it(record) return SecretErased() if settled else outcome @@ -194,6 +197,19 @@ def _remove_token_file() -> CredentialNotCleared | None: return None +def _write_the_note_the_removal_took_with_it(record: CliTokenRecord) -> None: + """Put the secret-free note back after the file carrying it had to go to get the secret off disk. + + Reaching here means neither rewrite would take, so the file went instead, and its absence is + what the next logout would read as a keychain already known to be clean. Removing it is also + what frees the room the rewrite was refused for, so the note usually lands on this second try. + When it does not, the warning this logout printed is the only one the user gets. + """ + staged: Final = _stage_scrubbed_file(record) + if staged is not None: + _commit_token_file(staged) + + def _keep_the_unchecked_keychain_on_record(outcome: SecretErase, record: CliTokenRecord | None) -> bool: """Whether the token file, stripped of its secret, is worth keeping as the note that says so. @@ -203,17 +219,27 @@ def _keep_the_unchecked_keychain_on_record(outcome: SecretErase, record: CliToke its secret neither to a staged replacement nor to an overwrite is not kept either, because the secret goes first. """ - if record is None or isinstance(outcome, SecretStranded): + if record is None or not _the_keychain_went_unchecked(outcome): return False return _scrub_file_secret(record) +def _the_keychain_went_unchecked(outcome: SecretErase) -> bool: + """Whether the keychain neither confirmed the erase nor answered that it still holds the secret""" + match outcome: + case SecretErased() | SecretStranded(): + return False + case KeyringDisabled() | KeyringNotInstalled() | KeyringUnreachable(): + return True + + def _nothing_left_behind(outcome: SecretErase, record: CliTokenRecord | None) -> bool: """Whether the keychain can be trusted to hold no credential of ours once the file is gone. A machine with no token file has no stored login to end, and `clear_cli_token` keeps one behind whenever the keychain is left unconfirmed, taking the secret out in place when it cannot stage a - replacement, so a missing file is real evidence rather than the absence of it. Past that, a + replacement and writing the note again when the file holding it had to go, so a missing file is + real evidence rather than the absence of it. Past that, a keychain that could not be reached is never trusted, whatever the file looks like. Even a file holding its own secret says only that the login which wrote it had no keychain to write to, and the login before it may well have had one: the entry that login left outlives both the 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 e3d14c6da46..b8fcf3618cf 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,7 +1,9 @@ +import errno import json import os import stat import sys +import tempfile import threading import time @@ -82,6 +84,25 @@ def _blob(base_url=SERVER, key="sk-vault", jwt_token=""): return json.dumps({"base_url": base_url, "key": key, "jwt_token": jwt_token}) +_REAL_MKSTEMP = tempfile.mkstemp + + +class _MkstempThatNeedsTheOldFileGone: + """A disk with exactly one token file's worth of room left on it. + + Staging a replacement needs room for a second file, which is what a full disk refuses. Removing + the file already there is what gives that room back. + """ + + def __init__(self, path): + self.path = path + + def __call__(self, *args, **kwargs): + if self.path.exists(): + raise OSError(errno.ENOSPC, "No space left on device") + return _REAL_MKSTEMP(*args, **kwargs) + + _REAL_REPLACE = os.replace @@ -553,7 +574,7 @@ class TestClearCliToken: assert load_cli_token(vault=vault) is None @pytest.mark.parametrize( - "failure", [KeyringDisabled(), KeyringUnreachable(), KeyringDiscardsWrites()] + "failure", [KeyringDisabled(), KeyringUnreachable(), KeyringNotInstalled()] ) def test_a_secret_in_the_file_is_no_evidence_about_a_keychain_that_exists( self, isolated_home, secret_vault_factory, failure @@ -561,7 +582,10 @@ class TestClearCliToken: """Store a secret in the keychain, sign in again while the keychain is unusable so the new secret lands in the file, then log out while it is still unusable. The file now carries its own secret and the first login's entry is still there, so reading the file as proof of a - clean keychain reports a logout that did not happen.""" + clean keychain reports a logout that did not happen. + + The three unusable states are the whole of what an erase can answer besides erased and + stranded; a backend that keeps nothing it is given is something only a write finds out.""" _write_legacy_file(isolated_home) vault = secret_vault_factory(available=False, failure=failure) @@ -705,6 +729,30 @@ class TestClearCliToken: assert json.loads(path.read_text()).get("key") is None + @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores file permissions") + def test_a_note_the_logout_had_to_remove_is_written_again_for_the_next_one( + self, isolated_home, secret_vault_factory, monkeypatch + ): + """A full disk refuses the replacement file and a read-only token file refuses the rewrite + in place, so the only way left to get the secret off disk is to remove the file carrying it. + That file was also the note saying the keychain went unchecked, and its absence is what the + next logout would read as a keychain already known to be clean. + + Removing it is what frees the room the replacement was refused for, so the note is written + again on the way out and the logout after this one still warns.""" + path = _write_legacy_file(isolated_home) + path.chmod(0o400) + monkeypatch.setattr( + "litellm.litellm_core_utils.private_json.tempfile.mkstemp", + _MkstempThatNeedsTheOldFileGone(path), + ) + vault = secret_vault_factory(available=False, failure=KeyringUnreachable()) + + assert clear_cli_token(vault=vault) == KeyringUnreachable() + assert clear_cli_token(vault=vault) == KeyringUnreachable() + + assert json.loads(path.read_text()).get("key") is None + @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores directory permissions") def test_a_metadata_file_that_will_not_go_is_not_worth_alarming_the_user_over( self, isolated_home, secret_vault_factory From a2928efc75032f9a20a7685fce1ae3ffd8dc4c9d Mon Sep 17 00:00:00 2001 From: Mateo Date: Thu, 20 Aug 2026 03:45:00 -0700 Subject: [PATCH 071/542] test(cli): cover the keyless token record and keep keyring to the cli extra `lite up` treats a token record whose key the keychain would not hand over as no login at all, and that clause had no test: every existing freshness test passed a record carrying a real key, so deleting the clause left the whole suite green The base install smoke check now also asserts keyring is absent, which is what makes the lazy import in cli_keyring meaningful. keyring ships in the cli extra only, so a plain `pip install litellm` must not be able to reach it --- .../base_sdk_tests/check_base_sdk_install.py | 2 +- .../proxy/client/cli/test_up_commands.py | 22 +++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/tests/base_sdk_tests/check_base_sdk_install.py b/tests/base_sdk_tests/check_base_sdk_install.py index 723f30cad76..6b38de75e2e 100644 --- a/tests/base_sdk_tests/check_base_sdk_install.py +++ b/tests/base_sdk_tests/check_base_sdk_install.py @@ -11,7 +11,7 @@ import sys import traceback from collections.abc import Callable -EXTRAS_ONLY_MODULES = ("fastapi", "uvicorn") +EXTRAS_ONLY_MODULES = ("fastapi", "uvicorn", "keyring") def _require(condition: bool, message: str) -> None: 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 aebf441f777..a8d81f1c4bb 100644 --- a/tests/test_litellm/proxy/client/cli/test_up_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_up_commands.py @@ -262,6 +262,28 @@ class TestEnsureFreshLogin: assert login_calls == ["http://proxy-b:4000"] + def test_forces_a_fresh_login_when_the_cached_token_has_no_readable_key(self, monkeypatch): + monkeypatch.setattr(up_module.sys.stdin, "isatty", lambda: True) + tokens = iter( + [ + _token(None, "http://proxy-a:4000"), + _token("sk-a", "http://proxy-a:4000"), + ] + ) + monkeypatch.setattr(up_module, "load_cli_token", lambda **_: next(tokens)) + monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True) + login_calls = [] + + @click.pass_context + def fake_login(ctx): + login_calls.append(ctx.obj["base_url"]) + + monkeypatch.setattr(up_module, "login", fake_login) + + _ensure_fresh_login(_make_ctx("http://proxy-a:4000")) + + assert login_calls == ["http://proxy-a:4000"] + 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_cli_token", lambda **_: _token("sk-a", "http://proxy-a:4000")) From 975a6806c318976efd1d9adc6ee44aed4d8ff8d4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 04:11:08 -0700 Subject: [PATCH 072/542] test(e2e): request reasoning explicitly on the reasoning-cost assertions The two tests that assert on reasoning cost read reasoning_tokens off the response and required it to be nonzero, without ever asking the model to reason. Both now send reasoning_effort, so the assertion rests on a parameter the test sets rather than on the model's default behavior. The cache-breakdown test sends it on its prime call too: OpenAI's prefix cache keys on the reasoning setting as well as the tokens, so priming at a different effort never produces a read. --- .../test_cache_cost_accounting_e2e.py | 30 +++++++++++++++++-- .../test_service_tier_pricing_e2e.py | 7 ++++- 2 files changed, 33 insertions(+), 4 deletions(-) diff --git a/tests/e2e/quota_management/spend_tracking/test_cache_cost_accounting_e2e.py b/tests/e2e/quota_management/spend_tracking/test_cache_cost_accounting_e2e.py index ca5985fcda6..c50ec3d902f 100644 --- a/tests/e2e/quota_management/spend_tracking/test_cache_cost_accounting_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_cache_cost_accounting_e2e.py @@ -30,6 +30,12 @@ OpenAI caching is best-effort, so each test retries with a fresh prefix (new marker = brand-new cache identity) up to three times before failing; the prime and measured calls share the prefix but differ in the trailing question, which defeats the proxy's own response cache without touching the provider's prefix cache. + +The test that asserts on reasoning cost requests reasoning explicitly with +`reasoning_effort`, so that assertion rests on a parameter the test sets rather +than on whatever the model happens to do by default. Its prime call carries the +same value: OpenAI's prefix cache keys on the reasoning setting as well as the +tokens, so a prime at a different effort never produces a read. """ import pytest @@ -67,6 +73,7 @@ CACHE_WRITE_RATE = 5e-05 PRIME_QUESTION = "Reply with the single word ready." REASONING_QUESTION = "Compute 47*83 - 19*7 step by step, then reply with just the final number." +REASONING_EFFORT = "high" class _StreamChunk(BaseModel): @@ -84,12 +91,15 @@ def _cache_priced_params(backend: str) -> LiteLLMParamsBody: ) -def _chat_body(model: str, content: str, *, stream: bool = False) -> ChatBody: +def _chat_body( + model: str, content: str, *, stream: bool = False, reasoning_effort: str | None = None +) -> ChatBody: return ChatBody( model=model, messages=[ChatMessage(role="user", content=content)], stream=stream, max_completion_tokens=4000, + reasoning_effort=reasoning_effort, ) @@ -151,9 +161,23 @@ class TestCacheCostAccounting: for _ in range(CACHE_ATTEMPTS): prefix = cacheable_prefix(unique_marker()) - unwrap(client.proxy.chat(scoped_key, _chat_body(model, f"{prefix}\n{PRIME_QUESTION}"))) + unwrap( + client.proxy.chat( + scoped_key, + _chat_body( + model, f"{prefix}\n{PRIME_QUESTION}", reasoning_effort=REASONING_EFFORT + ), + ) + ) chat = unwrap( - client.proxy.chat(scoped_key, _chat_body(model, f"{prefix}\n{REASONING_QUESTION}")) + client.proxy.chat( + scoped_key, + _chat_body( + model, + f"{prefix}\n{REASONING_QUESTION}", + reasoning_effort=REASONING_EFFORT, + ), + ) ) assert chat.id, f"chat response carried no id: {chat}" row = _require_row(client, chat.id) diff --git a/tests/e2e/quota_management/spend_tracking/test_service_tier_pricing_e2e.py b/tests/e2e/quota_management/spend_tracking/test_service_tier_pricing_e2e.py index 171c849fb4c..770c5699b4e 100644 --- a/tests/e2e/quota_management/spend_tracking/test_service_tier_pricing_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_service_tier_pricing_e2e.py @@ -10,7 +10,9 @@ computed from the wrong tier (or a mix) cannot match the expected numbers. The prompt is a fresh unique marker per run, keeping cached tokens out of the math. The response's own `service_tier` echo is asserted first: if OpenAI ever declined priority processing and served the default tier, the test fails there instead of -producing a vacuous rate comparison. +producing a vacuous rate comparison. Reasoning is requested explicitly with +`reasoning_effort`, so the reasoning-rate assertion rests on a parameter the test +sets rather than on whatever the model happens to do by default. """ import pytest @@ -38,6 +40,8 @@ OUTPUT_RATE = 8e-05 PRIORITY_INPUT_RATE = 6e-05 PRIORITY_OUTPUT_RATE = 1.6e-04 +REASONING_EFFORT = "high" + class TestServiceTierPricing: @pytest.mark.covers("quota_management.spend_tracking.service_tier.bills_tier_rates") @@ -74,6 +78,7 @@ class TestServiceTierPricing: ], max_completion_tokens=4000, service_tier="priority", + reasoning_effort=REASONING_EFFORT, ), ) ) From 5099379e034533c16793d77ad9548f68b3ed07f1 Mon Sep 17 00:00:00 2001 From: Mateo Date: Thu, 20 Aug 2026 04:12:23 -0700 Subject: [PATCH 073/542] fix(cli): stop asking a keychain that already stopped answering A pre-flight that times out leaves its write parked inside the keychain, holding it against every later call, so the next read blocks on the main thread with no timeout of its own. Anything that resolves the credential more than once in a process hits it: an SDK Client built a second time never returns. The vault now remembers the silence and reports the keychain unreachable for the rest of the process rather than queueing behind the parked call. --- litellm/litellm_core_utils/cli_keyring.py | 16 ++++++- .../test_cli_token_utils.py | 47 +++++++++++++++++++ 2 files changed, 61 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/cli_keyring.py b/litellm/litellm_core_utils/cli_keyring.py index 8da3e5226d4..70b1773739d 100644 --- a/litellm/litellm_core_utils/cli_keyring.py +++ b/litellm/litellm_core_utils/cli_keyring.py @@ -18,7 +18,7 @@ throwaway value, because a keychain can answer neither way and block forever. import os import threading from contextlib import suppress -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Final, Protocol, TypeAlias KEYRING_SERVICE: Final = "litellm-cli" @@ -152,11 +152,20 @@ def _forget_the_preflight(api: KeyringApi) -> None: @dataclass(frozen=True, slots=True) class KeyringVault: - """The OS keychain, reached through the optional `keyring` package.""" + """The OS keychain, reached through the optional `keyring` package. + + A keychain that let the pre-flight time out is not asked anything else for the rest of the + process. The probe that timed out is still sitting in the keychain on a thread of its own, and + it holds the keychain against every later call, so the read after it would block on the main + thread with no timeout to save it. One silence is answer enough. + """ preflight_timeout_seconds: float = _PREFLIGHT_TIMEOUT_SECONDS + stopped_answering: threading.Event = field(default_factory=threading.Event, compare=False, repr=False) def read(self) -> SecretRead: + if self.stopped_answering.is_set(): + return KeyringUnreachable() api: Final = _keyring_api() if isinstance(api, (KeyringNotInstalled, KeyringDisabled)): return api @@ -177,10 +186,13 @@ class KeyringVault: The keychain is pre-flighted first, because one that blocks rather than answering would otherwise hang `lite login` outright. """ + if self.stopped_answering.is_set(): + return KeyringUnreachable() api: Final = _keyring_api() if isinstance(api, (KeyringNotInstalled, KeyringDisabled)): return api if not _answers_a_write(api, self.preflight_timeout_seconds): + self.stopped_answering.set() return KeyringUnreachable() _forget_the_preflight(api) try: 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 b8fcf3618cf..c6fee6d0c25 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 @@ -832,6 +832,26 @@ class _NeverAnsweringKeyringModule(_FakeKeyringModule): threading.Event().wait() +class _KeychainHeldByABlockedWrite(_NeverAnsweringKeyringModule): + """The same keychain, plus what the blocked write does to everything after it: the stuck call + holds the keychain, so every later read blocks behind it too.""" + + def get_password(self, service_name, username): + self.calls.append(("get", service_name, username)) + if self.blocked.is_set(): + threading.Event().wait() + return self.stored + + +def _answered_within(seconds, call): + answers = [] + worker = threading.Thread(target=lambda: answers.append(call()), daemon=True) + worker.start() + worker.join(seconds) + assert not worker.is_alive(), f"{call.__qualname__} never returned" + return answers[0] + + @pytest.fixture def install_fake_keyring(monkeypatch): def _install(fake): @@ -927,6 +947,33 @@ class TestKeyringVault: assert [call[2] for call in fake.calls] == [KEYRING_PREFLIGHT_ACCOUNT] + def test_a_keychain_that_stopped_answering_is_not_asked_again(self, install_fake_keyring): + """The write that timed out is still holding the keychain when we give up on it, so the + call after it is the one that hangs, and read has nothing to time out against. Anything + resolving the credential more than once in a process hits that: an SDK client built twice + pays the pre-flight timeout on the first build and never returns from the second.""" + install_fake_keyring(_KeychainHeldByABlockedWrite()) + vault = KeyringVault(preflight_timeout_seconds=0.05) + + assert vault.write("blob-1") == KeyringUnreachable() + + assert _answered_within(5, vault.read) == KeyringUnreachable() + assert _answered_within(5, vault.erase) == KeyringUnreachable() + assert _answered_within(5, lambda: vault.write("blob-2")) == KeyringUnreachable() + + def test_a_keychain_that_stopped_answering_leaves_the_credential_in_the_file( + self, isolated_home, install_fake_keyring + ): + """The end of the same story: giving up on the keychain has to leave a login that still + works, and loading it back must not go asking the keychain that already stopped answering.""" + install_fake_keyring(_KeychainHeldByABlockedWrite()) + vault = KeyringVault(preflight_timeout_seconds=0.05) + + outcome = save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-only-copy"), vault=vault) + + assert outcome == KeyringUnreachable() + assert _answered_within(5, lambda: load_cli_token(vault=vault)).key == "sk-only-copy" + def test_a_login_survives_a_keychain_that_never_answers(self, isolated_home, install_fake_keyring): """The end of the same story: the credential still has to be usable afterwards.""" install_fake_keyring(_NeverAnsweringKeyringModule()) From 71a583390a4efe0dd171a249b62f895d73a4ad39 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 04:29:34 -0700 Subject: [PATCH 074/542] fix(cli): pick the credential by the sign-in it came from A login the keychain accepted whose token file could not be replaced left the superseded secret on disk, and the next load preferred the file unconditionally, so it served the old credential and erased the new one from the keychain on the way past. The keychain entry now carries the timestamp of the sign-in that minted it, and the two stores are compared on that instead. --- litellm/litellm_core_utils/cli_token_utils.py | 52 ++++++++++++------- litellm/proxy/client/cli/commands/auth.py | 5 +- .../test_cli_token_utils.py | 35 ++++++++++++- 3 files changed, 69 insertions(+), 23 deletions(-) diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index 69a77a36882..134e7e2331c 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -105,7 +105,8 @@ class CliTokenSecret(BaseModel): `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. + metadata file is edited underneath us. `timestamp` is the sign-in this + secret came from, which is what decides it against a secret still on disk. """ model_config = ConfigDict(frozen=True) @@ -113,6 +114,7 @@ class CliTokenSecret(BaseModel): base_url: str key: str jwt_token: str = "" + timestamp: float = 0.0 def get_cli_token_file_path() -> str: @@ -145,11 +147,7 @@ def save_cli_token(record: CliTokenRecord, *, vault: SecretVault = SYSTEM_KEYRIN staged: Final = _stage_token_file(_without_secret(record)) if isinstance(staged, CredentialNotSaved): return staged - outcome: Final = ( - SecretStored() - if record.key is None - else vault.write(_encode_secret(record.base_url, record.key, record.jwt_token)) - ) + outcome: Final = SecretStored() if record.key is None else vault.write(_encode_secret(record, record.key)) if isinstance(outcome, SecretStored): return outcome if _commit_token_file(staged) else CredentialNotRecorded() discard_staged_json(staged) @@ -330,19 +328,24 @@ 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. + The sign-in each secret came from decides it, because either store can be the stale one. A + secret is usually left on disk by a keychain that would not take it, which makes the file the + fresher of the two. It is the older one when a login the keychain did take could not replace + the file afterwards, and serving that one would put a superseded credential back in use. """ - if record.key is not None: - 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: + secret: Final = _decode_secret(blob, record.base_url) + if secret is None or (record.key is not None and secret.timestamp <= record.timestamp): return _migrate_file_secret(record, vault) _scrub_file_secret(record) - return record.model_copy(update=MappingProxyType({"key": secret.key, "jwt_token": secret.jwt_token})) + return record.model_copy( + update=MappingProxyType( + { + "key": secret.key, + "jwt_token": secret.jwt_token, + "timestamp": max(secret.timestamp, record.timestamp), + } + ) + ) def _migrate_file_secret(record: CliTokenRecord, vault: SecretVault) -> CliTokenRecord | None: @@ -363,7 +366,7 @@ def _migrate_file_secret(record: CliTokenRecord, vault: SecretVault) -> CliToken staged: Final = _stage_scrubbed_file(record) if staged is None: return record - if not isinstance(vault.write(_encode_secret(record.base_url, record.key, record.jwt_token)), SecretStored): + if not isinstance(vault.write(_encode_secret(record, record.key)), SecretStored): discard_staged_json(staged) return record if not _commit_token_file(staged) and not _overwrite_file_secret(record): @@ -422,8 +425,19 @@ 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 _encode_secret(record: CliTokenRecord, key: str) -> str: + return CliTokenSecret( + base_url=record.base_url, key=key, jwt_token=record.jwt_token, timestamp=record.timestamp + ).model_dump_json() + + +def _decode_secret(blob: str, base_url: str) -> CliTokenSecret | None: + """The keychain entry, when it is one this metadata file may be paired with""" + try: + secret: Final = CliTokenSecret.model_validate_json(blob) + except ValidationError: + return None + return secret if secret.base_url == base_url else None def _write_token_file(record: CliTokenRecord) -> None: diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index ac330ad6892..e5d90b9b640 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -131,8 +131,9 @@ def storage_notice(outcome: SecretSave) -> str: case CredentialNotRecorded(): return ( f"Signed in, and the credential is in your OS keychain, but {path} could not be " - "replaced, so this machine may still be using your previous login. Run 'lite login' " - "again once that path is writable, or 'lite logout' to clear both." + "replaced, so it still describes your previous login and may still hold its " + "credential. Run 'lite login' again once that path is writable, or 'lite logout' " + "to clear both." ) 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 c6fee6d0c25..4fc835d2db1 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 @@ -80,8 +80,8 @@ def _write_metadata_only_file(home): return path -def _blob(base_url=SERVER, key="sk-vault", jwt_token=""): - return json.dumps({"base_url": base_url, "key": key, "jwt_token": jwt_token}) +def _blob(base_url=SERVER, key="sk-vault", jwt_token="", timestamp=0.0): + return json.dumps({"base_url": base_url, "key": key, "jwt_token": jwt_token, "timestamp": timestamp}) _REAL_MKSTEMP = tempfile.mkstemp @@ -210,6 +210,37 @@ class TestLoadCliToken: assert json.loads(vault.blob)["key"] == "sk-fresh" assert "key" not in json.loads(path.read_text()) + def test_a_login_the_file_could_not_record_is_the_one_that_gets_used( + self, isolated_home, secret_vault_factory + ): + """A login the keychain took and the file could not be pointed at afterwards leaves the + superseded secret sitting on disk in front of the fresh one. Serving the file's copy would + put a credential the user just replaced, and may well have just revoked, back into every + request, and would overwrite the keychain with it on the way past.""" + path = _write_legacy_file(isolated_home, key="sk-superseded", timestamp=1000.0) + vault = secret_vault_factory(blob=_blob(key="sk-fresh", timestamp=2000.0)) + + record = load_cli_token(vault=vault) + + assert record.key == "sk-fresh" + assert record.timestamp == 2000.0 + assert json.loads(vault.blob)["key"] == "sk-fresh" + assert "key" not in json.loads(path.read_text()) + + def test_a_secret_written_to_disk_after_the_keychain_entry_still_wins( + self, isolated_home, secret_vault_factory + ): + """The other direction of the same rule, which is the common one: a login that fell back to + the file because the keychain refused it is newer than whatever the keychain kept.""" + path = _write_legacy_file(isolated_home, key="sk-fresh", timestamp=2000.0) + vault = secret_vault_factory(blob=_blob(key="sk-stale", timestamp=1000.0)) + + 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 ): From b9ce630b0e47480676b943ba2208133509afac3b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 04:36:32 -0700 Subject: [PATCH 075/542] docs(cli): say why a refused scrub does not roll the keychain back The two stores hold different credentials on that path, so the rollback a migration does would hand the superseded one back out. The login that could not replace the file already named the state, and logout reports it too. --- litellm/litellm_core_utils/cli_token_utils.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index 134e7e2331c..085448f8b63 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -332,6 +332,12 @@ def _apply_vault_secret(record: CliTokenRecord, blob: str, vault: SecretVault) - secret is usually left on disk by a keychain that would not take it, which makes the file the fresher of the two. It is the older one when a login the keychain did take could not replace the file afterwards, and serving that one would put a superseded credential back in use. + + A scrub the file refuses leaves that superseded secret where it lies, which is the state the + login already named when it could not replace the file, and which `lite logout` reports rather + than counting as a clean sweep. Rolling the vault back the way a migration does is not the + answer here, because the two stores hold different credentials and the rollback would hand the + superseded one back out. """ secret: Final = _decode_secret(blob, record.base_url) if secret is None or (record.key is not None and secret.timestamp <= record.timestamp): From a329dfbb45ea3ec7ea39bb6621d0093f4862094b Mon Sep 17 00:00:00 2001 From: Mateo Edgeton Date: Thu, 20 Aug 2026 04:53:09 -0700 Subject: [PATCH 076/542] fix(cli): keep each sign-in stamped past the one it replaces The stamp in the keychain entry is what decides that secret against one still sitting in the token file, and it came straight off the wall clock. A clock that stepped backwards between two logins therefore handed the win to the older of them: a login the keychain took but the token file could not be pointed at was resolved back to the credential it replaced, and the fresh one was erased from the keychain on the way past. save_cli_token now reads the stamp already on disk and pins the new sign-in just above it, so the ordering never depends on the clock having moved forwards. On a clock that did, this changes nothing. --- litellm/litellm_core_utils/cli_token_utils.py | 26 ++++++++++-- .../test_cli_token_utils.py | 40 +++++++++++++++++++ 2 files changed, 62 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index 085448f8b63..da8b80ab598 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -12,6 +12,7 @@ first time it reads one. This module has no dependencies on proxy code and can be safely imported at the SDK level. """ +import math import time from dataclasses import dataclass from pathlib import Path @@ -144,14 +145,29 @@ def save_cli_token(record: CliTokenRecord, *, vault: SecretVault = SYSTEM_KEYRIN keychain has already taken the new secret, and it reports itself as such rather than claiming the previous login survived. """ - staged: Final = _stage_token_file(_without_secret(record)) + stamped: Final = _stamped_past_the_login_on_disk(record) + staged: Final = _stage_token_file(_without_secret(stamped)) if isinstance(staged, CredentialNotSaved): return staged - outcome: Final = SecretStored() if record.key is None else vault.write(_encode_secret(record, record.key)) + outcome: Final = SecretStored() if stamped.key is None else vault.write(_encode_secret(stamped, stamped.key)) if isinstance(outcome, SecretStored): return outcome if _commit_token_file(staged) else CredentialNotRecorded() discard_staged_json(staged) - return _keep_the_secret_in_the_file(record, outcome) + return _keep_the_secret_in_the_file(stamped, outcome) + + +def _stamped_past_the_login_on_disk(record: CliTokenRecord) -> CliTokenRecord: + """Keep a sign-in's stamp ahead of the one it replaces, whatever the clock did in between. + + The stamp is what decides a keychain secret against one still on disk, so a clock that stepped + backwards between two logins would hand the older of them the win and put a superseded + credential back in use. The file already names the login being replaced, and pinning the new + stamp just past it costs one read that changes nothing on a clock that only moves forwards. + """ + previous: Final = _read_token_file() + if previous is None or previous.timestamp < record.timestamp: + return record + return record.model_copy(update=MappingProxyType({"timestamp": math.nextafter(previous.timestamp, math.inf)})) def _keep_the_secret_in_the_file(record: CliTokenRecord, outcome: SecretWrite) -> SecretSave: @@ -331,7 +347,9 @@ def _apply_vault_secret(record: CliTokenRecord, blob: str, vault: SecretVault) - The sign-in each secret came from decides it, because either store can be the stale one. A secret is usually left on disk by a keychain that would not take it, which makes the file the fresher of the two. It is the older one when a login the keychain did take could not replace - the file afterwards, and serving that one would put a superseded credential back in use. + the file afterwards, and serving that one would put a superseded credential back in use. Equal + stamps are one login sitting in both stores, left by a migration whose scrub was refused, so + that branch retries the migration rather than trading one credential for another. A scrub the file refuses leaves that superseded secret where it lies, which is the state the login already named when it could not replace the file, and which `lite logout` reports rather 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 4fc835d2db1..d9fc17f908f 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 @@ -488,6 +488,46 @@ class TestSaveCliToken: assert path.read_text() == before assert list(path.parent.glob(".tmp-*")) == [] + def test_a_login_is_stamped_past_the_one_it_replaces_even_on_a_clock_that_went_back( + self, isolated_home, secret_vault_factory + ): + """The stamp is what decides the keychain secret against the one on disk, so a login that + carries an earlier wall clock than the login before it must not be filed as the older of + the two.""" + _write_legacy_file(isolated_home, key="sk-old", timestamp=2000.0) + vault = secret_vault_factory() + + save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-new", timestamp=1000.0), vault=vault) + + assert json.loads(vault.blob)["timestamp"] > 2000.0 + + def test_a_clock_that_went_back_does_not_hand_the_win_to_the_superseded_login( + self, isolated_home, secret_vault_factory + ): + """The disk state a login reports as CredentialNotRecorded: the keychain took the new + secret and the file still holds the previous one. Reading it back has to produce the login + that was just made, and an earlier wall clock is no reason to serve the one it replaced.""" + _write_legacy_file(isolated_home, key="sk-superseded", timestamp=2000.0) + vault = secret_vault_factory() + + save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-fresh", timestamp=1000.0), vault=vault) + _write_legacy_file(isolated_home, key="sk-superseded", timestamp=2000.0) + + assert load_cli_token(vault=vault).key == "sk-fresh" + + def test_a_login_on_a_clock_that_moved_forwards_keeps_its_own_time( + self, isolated_home, secret_vault_factory + ): + """Pinning the stamp above the previous login is only ever a floor. The ordinary case has + to record when the user actually signed in, because that is what decides expiry.""" + _write_legacy_file(isolated_home, key="sk-old", timestamp=1000.0) + vault = secret_vault_factory() + + save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-new", timestamp=2000.0), vault=vault) + + assert json.loads(vault.blob)["timestamp"] == 2000.0 + assert json.loads(_token_file(isolated_home).read_text())["timestamp"] == 2000.0 + class TestScrubFailure: """A keychain that took the secret while the file kept it is the worst of both stores: the From 8a8e8dc8edcf819041a5ff23b6a30f5a4939347e Mon Sep 17 00:00:00 2001 From: Mateo Edgeton Date: Thu, 20 Aug 2026 04:53:09 -0700 Subject: [PATCH 077/542] docs(deps): say that the cli extra pulls cryptography on linux The comment above the extra named cryptography as one of the heavy imports a thin install leaves out. That stopped being true when keyring joined the extra: on Linux it reaches the Secret Service through secretstorage, which depends on cryptography. --- pyproject.toml | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 64adb9cd595..09a69f3771e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,8 +78,10 @@ proxy = [ "expression>=5.6.0,<6.0", ] # Thin client install for the `lite` CLI on developer laptops. The CLI's heavy -# imports (fastapi, cryptography, ...) are all guarded, so it runs on the base -# SDK plus just these five; none of the server runtime in `proxy` is pulled in. +# imports are all guarded, so it runs on the base SDK plus just these five, and +# none of the server runtime in `proxy` is pulled in. On Linux, +# keyring reaches the Secret Service through secretstorage, which brings +# cryptography with it. cli = [ "rich>=13.9.4,<14.0", "pyyaml>=6.0.3,<7.0", From fb69fcf765997610e7f32fc041200cd5b5217e3a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 05:07:21 -0700 Subject: [PATCH 078/542] fix(cli): stamp each sign-in past the keychain as well as the file A login the keychain took but the token file could not record leaves the keychain naming a later sign-in than the file does. Reading only the file then stamps the next login below that keychain entry, and a clock that went back far enough puts the superseded credential back in use. --- litellm/litellm_core_utils/cli_token_utils.py | 40 +++++++++++++++---- .../test_cli_token_utils.py | 13 ++++++ 2 files changed, 45 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index da8b80ab598..c0e57a737e4 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -145,7 +145,7 @@ def save_cli_token(record: CliTokenRecord, *, vault: SecretVault = SYSTEM_KEYRIN keychain has already taken the new secret, and it reports itself as such rather than claiming the previous login survived. """ - stamped: Final = _stamped_past_the_login_on_disk(record) + stamped: Final = _stamped_past_every_stored_login(record, vault) staged: Final = _stage_token_file(_without_secret(stamped)) if isinstance(staged, CredentialNotSaved): return staged @@ -156,18 +156,42 @@ def save_cli_token(record: CliTokenRecord, *, vault: SecretVault = SYSTEM_KEYRIN return _keep_the_secret_in_the_file(stamped, outcome) -def _stamped_past_the_login_on_disk(record: CliTokenRecord) -> CliTokenRecord: - """Keep a sign-in's stamp ahead of the one it replaces, whatever the clock did in between. +def _stamped_past_every_stored_login(record: CliTokenRecord, vault: SecretVault) -> CliTokenRecord: + """Keep a sign-in's stamp ahead of every login already stored, whatever the clock did in between. The stamp is what decides a keychain secret against one still on disk, so a clock that stepped backwards between two logins would hand the older of them the win and put a superseded - credential back in use. The file already names the login being replaced, and pinning the new - stamp just past it costs one read that changes nothing on a clock that only moves forwards. + credential back in use. Pinning the new stamp just past the highest one either store holds costs + one read each and changes nothing on a clock that only moves forwards. + """ + highest: Final = _highest_stamp_already_stored(record.base_url, vault) + if highest < record.timestamp: + return record + return record.model_copy(update=MappingProxyType({"timestamp": math.nextafter(highest, math.inf)})) + + +def _highest_stamp_already_stored(base_url: str, vault: SecretVault) -> float: + """When the latest login either store still holds was made, or minus infinity when neither has one. + + Both are asked because the file names the login being replaced only while the two agree. A login + the keychain took but the file could not record afterwards leaves the keychain holding the later + of the two, and reading only the file would stamp the next sign-in below it. """ previous: Final = _read_token_file() - if previous is None or previous.timestamp < record.timestamp: - return record - return record.model_copy(update=MappingProxyType({"timestamp": math.nextafter(previous.timestamp, math.inf)})) + secret: Final = _stored_secret(base_url, vault) + return max( + -math.inf if previous is None else previous.timestamp, + -math.inf if secret is None else secret.timestamp, + ) + + +def _stored_secret(base_url: str, vault: SecretVault) -> CliTokenSecret | None: + """The keychain's secret for this server, when it holds one this login may be compared against""" + match vault.read(): + case SecretFound(blob=blob): + return _decode_secret(blob, base_url) + case SecretMissing() | KeyringNotInstalled() | KeyringDisabled() | KeyringUnreachable(): + return None def _keep_the_secret_in_the_file(record: CliTokenRecord, outcome: SecretWrite) -> SecretSave: 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 d9fc17f908f..e78add61dd5 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 @@ -528,6 +528,19 @@ class TestSaveCliToken: assert json.loads(vault.blob)["timestamp"] == 2000.0 assert json.loads(_token_file(isolated_home).read_text())["timestamp"] == 2000.0 + def test_a_login_is_stamped_past_the_keychain_the_file_could_not_keep_up_with( + self, isolated_home, secret_vault_factory + ): + """A login reported as CredentialNotRecorded leaves the keychain holding a later sign-in + than the file names, so the file alone is no longer the floor. A later login on a clock + that went back past that keychain entry still has to be the one served.""" + _write_legacy_file(isolated_home, key="sk-superseded", timestamp=1000.0) + vault = secret_vault_factory(blob=_blob(key="sk-recorded", timestamp=2000.0), writable=False) + + save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-fresh", timestamp=1500.0), vault=vault) + + assert load_cli_token(vault=vault).key == "sk-fresh" + class TestScrubFailure: """A keychain that took the secret while the file kept it is the worst of both stores: the From 1f2baf509e1d683e53a11daec6ac673d4dec9d52 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 05:28:19 -0700 Subject: [PATCH 079/542] test(cli): model keyring's null backend in the vault test double FakeSecretVault could only stand in for a discarding backend by passing KeyringDiscardsWrites as its `failure`, which also made read() and erase() hand it back. Neither SecretRead nor SecretErase admits that outcome and the real KeyringVault never produces it there, so the login path's match was falling through on a value it can never see. Give the double a `discards` flag that reports it from write() alone, which is what the null backend does. Also widen lint-format-check-changed's pathspec. Git wildmatch runs without FNM_PATHNAME here, so 'litellm/**/*.py' still requires an intermediate directory and silently skipped all 21 top-level modules, litellm/__init__.py and litellm/main.py among them. All 21 already pass ruff format. --- Makefile | 2 +- tests/test_litellm/conftest.py | 8 +++++++- tests/test_litellm/proxy/client/cli/test_auth_commands.py | 3 +-- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/Makefile b/Makefile index ab6eba880bf..a47b3f0f66a 100644 --- a/Makefile +++ b/Makefile @@ -146,7 +146,7 @@ lint-install: # only the litellm Python files changed vs the base are checked, so a pre-existing # format issue elsewhere doesn't block an unrelated commit. lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE) - @files=$$(git diff --name-only --diff-filter=ACMR origin/litellm_internal_staging...HEAD -- 'litellm/**/*.py' | grep -v '^litellm/enterprise/' || true); \ + @files=$$(git diff --name-only --diff-filter=ACMR origin/litellm_internal_staging...HEAD -- 'litellm/*.py' | grep -v '^litellm/enterprise/' || true); \ if [ -z "$$files" ]; then \ echo "No changed litellm Python files to format-check."; \ else \ diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index b42355fa045..1229642dea0 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -23,6 +23,7 @@ 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 ( + KeyringDiscardsWrites, KeyringUnreachable, KeyringUnusable, SecretErase, @@ -132,7 +133,8 @@ class FakeSecretVault: `available=False` models a keychain that is locked or has no backend, `writable=False` one that refuses to store, `erasable=False` one that will not release what it already holds, and `failure` - picks which unusable state those report. + picks which unusable state those report. `discards=True` is keyring's null backend, which answers + reads and erases like any other yet keeps nothing it is given, so only writes report it. """ def __init__( @@ -142,12 +144,14 @@ class FakeSecretVault: available: bool = True, writable: bool = True, erasable: bool = True, + discards: bool = False, failure: KeyringUnusable = KeyringUnreachable(), ) -> None: self.blob: str | None = blob self.available: bool = available self.writable: bool = writable self.erasable: bool = erasable + self.discards: bool = discards self.failure: KeyringUnusable = failure self.reads: int = 0 self.writes: list[str] = [] @@ -163,6 +167,8 @@ class FakeSecretVault: self.writes.append(blob) if not (self.available and self.writable): return self.failure + if self.discards: + return KeyringDiscardsWrites() self.blob = blob return SecretStored() 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 1dd8a7ead92..4f44514167c 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -16,7 +16,6 @@ from litellm.constants import CLI_JWT_EXPIRATION_HOURS from litellm.litellm_core_utils.cli_keyring import ( DISABLE_KEYRING_ENV_VAR, KeyringDisabled, - KeyringDiscardsWrites, KeyringNotInstalled, ) from litellm.litellm_core_utils.cli_token_utils import CliTokenRecord, save_cli_token @@ -1068,7 +1067,7 @@ class TestKeychainBackedCommands: ): """A backend that accepts writes and stores nothing must not be reported as keychain storage, because the file is then told to drop the only remaining copy.""" - result = self._login(secret_vault_factory(available=False, failure=KeyringDiscardsWrites())) + result = self._login(secret_vault_factory(discards=True)) token_file = isolated_home / ".litellm" / "token.json" assert result.exit_code == 0 From 1fe06a1280f84b8b4e477d1fbaee3555c28489d8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 06:14:40 -0700 Subject: [PATCH 080/542] test(cli): pin the shared stamp's effect on the freshness shortcut The stamp both orders the two stores and drives is_cli_token_fresh, and nothing tied the two together, so a login that inherits a stamp from the future could stop being a deliberate trade without anything failing. Also corrects the lint-format-check-changed comment: git pathspecs match recursively, so the target checks a superset of the CI step rather than an identical set. --- Makefile | 6 ++++-- .../litellm_core_utils/test_cli_token_utils.py | 9 +++++++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/Makefile b/Makefile index a47b3f0f66a..c3aa106cf3e 100644 --- a/Makefile +++ b/Makefile @@ -142,9 +142,11 @@ lint-install: $(UV) sync --inexact --frozen --group proxy-dev --group e2e-dev $(UV_RUN) python scripts/prisma_generate_if_needed.py -# Diff-scoped format check, identical to test-linting.yml's "Check ruff format" step: +# Diff-scoped format check, mirroring test-linting.yml's "Check ruff format" step: # only the litellm Python files changed vs the base are checked, so a pre-existing -# format issue elsewhere doesn't block an unrelated commit. +# format issue elsewhere doesn't block an unrelated commit. Git pathspecs match +# recursively, so 'litellm/*.py' covers nested modules and the top-level files that +# CI's 'litellm/**/*.py' skips, which makes this target a superset of the CI step. lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE) @files=$$(git diff --name-only --diff-filter=ACMR origin/litellm_internal_staging...HEAD -- 'litellm/*.py' | grep -v '^litellm/enterprise/' || true); \ if [ -z "$$files" ]; then \ 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 e78add61dd5..357d3ae1b10 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 @@ -870,6 +870,15 @@ class TestIsCliTokenFresh: assert is_cli_token_fresh(almost, buffer_hours=0.1) is False + def test_a_stamp_left_in_the_future_keeps_reporting_fresh_until_the_clock_catches_up(self): + """The stamp both orders the two stores and drives this shortcut, so a store left stamped + ahead of the clock hands that stamp to the next sign-in and keeps it looking fresh past the + expiry the gateway will actually enforce. Pinning that here so the shared stamp cannot stop + being a deliberate trade without this failing first.""" + ahead = CliTokenRecord(timestamp=time.time() + CLI_JWT_EXPIRATION_HOURS * 3600) + + assert is_cli_token_fresh(ahead) is True + class _FakeKeyringModule: def __init__(self, stored=None, *, get_error=None, set_error=None, delete_error=None, discard=False): From 4fac88790dc07c427ff3b75ca7e63b538cf496a0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 10:45:24 -0700 Subject: [PATCH 081/542] fix(mistral): correct zai-glm-5-2 limits, add cached-input price and glm-5-2 alias Mistral's live /v1/models reports max_context_length 1048576 and capabilities.reasoning true for zai-glm-5-2, and its docs price cached input at $0.14/M. Without cache_read_input_token_cost LiteLLM billed every cached prompt token at $0, so a repeat request against a 21k-token cached prefix logged $0.0000135 instead of its real cost. Mistral also serves the model under the short glm-5-2 name, which had no cost map entry at all and therefore no pricing, so add it alongside. --- ...odel_prices_and_context_window_backup.json | 26 ++++- model_prices_and_context_window.json | 26 ++++- ...test_mistral_zai_glm_5_2_model_metadata.py | 103 ++++++++++++++++++ 3 files changed, 149 insertions(+), 6 deletions(-) create mode 100644 tests/test_litellm/test_mistral_zai_glm_5_2_model_metadata.py diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b73feae90d3..650415279cb 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -29046,16 +29046,36 @@ "supports_tool_choice": true }, "mistral/zai-glm-5-2": { + "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, "litellm_provider": "mistral", - "max_input_tokens": 1000000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, "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_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "mistral/glm-5-2": { + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "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_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b73feae90d3..650415279cb 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -29046,16 +29046,36 @@ "supports_tool_choice": true }, "mistral/zai-glm-5-2": { + "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, "litellm_provider": "mistral", - "max_input_tokens": 1000000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, "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_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "mistral/glm-5-2": { + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "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_prompt_caching": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true }, diff --git a/tests/test_litellm/test_mistral_zai_glm_5_2_model_metadata.py b/tests/test_litellm/test_mistral_zai_glm_5_2_model_metadata.py new file mode 100644 index 00000000000..ad1f3b06e15 --- /dev/null +++ b/tests/test_litellm/test_mistral_zai_glm_5_2_model_metadata.py @@ -0,0 +1,103 @@ +import json +from pathlib import Path + +import pytest + +import litellm +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.types.utils import PromptTokensDetailsWrapper, Usage +from litellm.utils import supports_prompt_caching, supports_reasoning + +REPO_ROOT = Path(__file__).parents[2] +MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" + +GLM_5_2_MODELS = ("mistral/zai-glm-5-2", "mistral/glm-5-2") + +INPUT_COST = 1.4e-06 +CACHED_INPUT_COST = 1.4e-07 +OUTPUT_COST = 4.4e-06 + + +def _load(path): + with open(path) as f: + return json.load(f) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Force get_model_info to resolve against the in-repo cost map instead of the + remote one fetched at import time, which still carries the pre-merge pricing.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + +@pytest.mark.parametrize("model", GLM_5_2_MODELS) +def test_zai_glm_5_2_specs(model): + info = _load(MAIN_PATH).get(model) + assert info is not None, f"{model} missing from model_prices_and_context_window.json" + + assert info["litellm_provider"] == "mistral" + assert info["mode"] == "chat" + + assert info["input_cost_per_token"] == INPUT_COST + assert info["output_cost_per_token"] == OUTPUT_COST + assert info["cache_read_input_token_cost"] == CACHED_INPUT_COST + + assert info["max_input_tokens"] == 1048576 + assert info["max_output_tokens"] == 131072 + assert info["max_tokens"] == 131072 + + assert info["supports_assistant_prefill"] is True + assert info["supports_function_calling"] is True + assert info["supports_prompt_caching"] is True + assert info["supports_reasoning"] is True + assert info["supports_response_schema"] is True + assert info["supports_tool_choice"] is True + + routed_model, provider, _, _ = get_llm_provider(model=model) + assert routed_model == model.split("/", 1)[1] + assert provider == "mistral" + + +@pytest.mark.parametrize("model", GLM_5_2_MODELS) +def test_zai_glm_5_2_capabilities_are_visible_to_callers(local_model_cost_map, model): + """Mistral advertises reasoning and prompt caching on this model, so the helpers + every caller checks before sending a request must say so too.""" + assert supports_reasoning(model=model) is True + assert supports_prompt_caching(model=model) is True + + info = litellm.get_model_info(model=model) + assert info["max_input_tokens"] == 1048576 + assert info["max_output_tokens"] == 131072 + + +@pytest.mark.parametrize("model", GLM_5_2_MODELS) +def test_cached_prompt_tokens_bill_at_the_cached_rate(local_model_cost_map, model): + """A cache hit reports its reused tokens under prompt_tokens_details, and those + tokens cost a tenth of the input rate, not the full rate and not nothing.""" + usage = Usage( + prompt_tokens=21010, + completion_tokens=100, + total_tokens=21110, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=20992), + ) + + prompt_cost, completion_cost = litellm.cost_per_token( + model=model, usage_object=usage, custom_llm_provider="mistral" + ) + + assert prompt_cost == pytest.approx(18 * INPUT_COST + 20992 * CACHED_INPUT_COST) + assert completion_cost == pytest.approx(100 * OUTPUT_COST) + + +@pytest.mark.parametrize("model", GLM_5_2_MODELS) +def test_backup_matches_main(model): + """Ensure the bundled (backup) cost map stays in sync with the canonical file.""" + main_cost = _load(MAIN_PATH) + backup_cost = _load(BACKUP_PATH) + + assert backup_cost.get(model) == main_cost.get(model), f"{model} differs between main and backup model cost maps" From e12833e6b41511ca991f1fe54900ff5c343c6902 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 20 Aug 2026 10:58:27 -0700 Subject: [PATCH 082/542] fix(ui): make dark-mode form controls visible (#37648) * fix(ui): make dark-mode form controls visible Two dark-mode defects left form controls without any visual boundary or fill. `--input` and `--border` share one value in `.dark`, oklch(0.309), which resolves to rgb(48,48,48). Against `--background` (33) that is a 15-step stroke, and against `--popover` (42) it collapses to 6 steps out of 255, so a control inside any dialog is effectively undrawn. The controls also use `bg-transparent`, so there is no fill cue either and only the placeholder text renders. Measured 1.09:1 against the dialog surface where WCAG 1.4.11 asks for 3.0:1 on the boundary of a user interface component. Splitting `--input` off at oklch(0.56) restores 3.07:1 without touching `--border`, which stays where it is because it draws decorative separators rather than control boundaries. 91 controls across 19 routes were measured at the collapsed value, every one with an identical stroke and surface, so a single token covers all of them. Separately, `@tailwindcss/forms` paints a white fill on every bare control. The block above already neutralises that for `combobox-chip-input`, but its audit covered `components/ui` only, and hand-rolled controls elsewhere still render white on a dark page: typed text lands at 1.11:1 and native selects at 2.19:1 on `/model-hub-table`, `/playground`, `/guardrails`, `/mcp-servers` and `/models-and-endpoints`. Tracking `--background` fixes those at 14.51:1 and 7.34:1. Light mode is unchanged by both. The token edit is scoped to `.dark`, and `--background` in `:root` is the same white the plugin was already painting, verified control-by-control on a dev server: backgrounds stay rgb(255,255,255) and ratios stay 20.13:1 and 4.84:1. * fix(ui): keep the combobox chip input transparent under the bare-control fill The new base rule matched at (0,2,1) while the combobox chip-input override sits at (0,1,0), so ComboboxChipsInput lost its transparent background and painted an opaque page-colored rectangle inside the chips container, which carries its own bg-transparent / dark:bg-input/30 fill. Folding the exclusions into one :not() list adds the chip input and drops the selector to (0,1,1). Every @tailwindcss/forms base selector is wrapped in :where(), so it lands at (0,0,1); (0,1,1) still outweighs it and bare inputs, textareas and selects keep the fill this PR gives them. --- ui/litellm-dashboard/src/app/globals.css | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/app/globals.css b/ui/litellm-dashboard/src/app/globals.css index 8fa7cde2d37..04a09693287 100644 --- a/ui/litellm-dashboard/src/app/globals.css +++ b/ui/litellm-dashboard/src/app/globals.css @@ -134,7 +134,7 @@ --warning: oklch(0.828 0.189 84.429); --info: oklch(0.707 0.165 254.624); --border: oklch(0.309 0 0); - --input: oklch(0.309 0 0); + --input: oklch(0.56 0 0); --ring: oklch(0.569 0 0); --chart-1: oklch(0.488 0.243 264.376); --chart-2: oklch(0.696 0.17 162.48); @@ -230,6 +230,13 @@ letter-spacing: inherit; } + /* Same plugin's white fill, on the hand-rolled controls outside components/ui that the audit + above did not cover. Tracks --background so light mode keeps the white it already painted. + Delete along with the plugin. */ + :is(input, textarea, select):not([type="checkbox"], [type="radio"], [data-slot="combobox-chip-input"]) { + background-color: var(--color-background); + } + button:not(:disabled), [role="button"]:not(:disabled) { cursor: pointer; From c794dcb91d449a825106b0f8f3c33ffd93b2a4d5 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 20 Aug 2026 10:58:32 -0700 Subject: [PATCH 083/542] fix(ui): give status colours a readable foreground and drop the muted 70% step (#37649) * fix(ui): give status colours a readable foreground and drop the muted 70% step The four status tokens are lightened for dark mode, which is correct when they are used as text and wrong for the 27 places that use them as a background under `text-white`. Every one of those passes in light and fails in dark: success 1.78:1, warning 1.72:1, info 2.64:1, destructive 2.89:1. The cause is not 27 authoring mistakes, it is that no `--success-foreground` and no sibling ever existed, so `text-white` was the only thing available to write. Adding the four companions and registering them in `@theme` makes the correct pairing expressible, and the call sites then read `text-success-foreground` instead of a hardcoded colour. Dark lands at 9.98, 10.31, 6.72 and 6.15. Light is deliberately pure white rather than the near-white the other `-foreground` tokens use, so the four ratios stay at exactly the 4.95, 5.03, 5.25 and 4.77 they are today instead of drifting down to 4.73, 4.81, 5.02 and 4.56. Separately `text-muted-foreground/70` measures 2.75:1 on a light page and 4.31:1 on a dark one, so the same 183 occurrences fail AA in light and sit under it in dark. Dropping the opacity step takes them to 4.84:1 and 7.34:1. The identical step on the placeholder base rule goes with them, which is what put every input's placeholder at 2.75:1 in light. Residual, not addressed here: `text-muted-foreground` over `bg-muted` reaches 4.39:1 in light, still short of 4.5. Closing that needs `--muted-foreground` itself to move, which changes every secondary label in the product and is a design call rather than a defect fix. * fix(ui): finish the status-foreground swap and repoint no-op muted hovers Four sites still forced text-white on a status fill because the class sat on a child element rather than on the filled container, so the earlier sweep did not reach them. The compliance quick-test bubble was worse: it paired bg-info with text-success-foreground and its paragraph kept text-white on top, so the dark-theme contrast the PR set out to fix was still reachable there Dropping the /70 step also turned 21 existing "text-muted-foreground/70 hover:text-muted-foreground" pairs into hovers that change nothing, which local/no-noop-hover-variant flags as an error. The affordance was "brighten on hover", so these now hover to text-foreground, matching the 74 places that already spell it that way The remaining churn is prettier reflowing the handful of lines whose length changed, since the token names are longer than text-white * fix(ui): let the approve/reject confirm button pick the token its fill uses Both submission review dialogs put text-success-foreground on the shared button class while the fill below it swings between bg-success for Approve and bg-destructive for Reject, so Reject drew a success token over a destructive fill. The two tokens resolve to the same value today, so nothing looks wrong, but the pairing only holds by coincidence and would break the moment either token moves. Moving the token into the branch makes it track the fill * fix(ui): drop the last 70% placeholders, still live on the legacy utility Four inputs spell their placeholder colour with Tailwind's older placeholder- utility rather than placeholder:text-, so the sweep that dropped the 70% step passed over them. Tailwind 4.3 still emits that utility, and utilities sit after base in the layer order, so those four kept overriding the new input::placeholder rule and kept rendering at 70% in dark mode, which is the contrast failure this PR set out to close They now spell it the same way as the three placeholders the PR already converted, which both removes the step and settles on one spelling --- .../_components/CacheLeakageCard.tsx | 4 +- .../_components/TierTurnsChart.tsx | 2 +- .../_components/cost_tracking_settings.tsx | 4 +- .../pricing_calculator/multi_cost_results.tsx | 12 +-- .../_components/provider_margin_table.tsx | 2 +- .../_components/GuardrailsOverview.tsx | 2 +- .../_components/TeamGuardrailsTab.tsx | 12 +-- .../_components/add_guardrail_form.tsx | 2 +- .../_components/EnvVarsSection.tsx | 2 +- .../_components/MCPSubmissionsTab.tsx | 44 ++++++----- .../mcp-servers/_components/mcp_connect.tsx | 2 +- .../components/chat_ui/A2AMetrics.tsx | 8 +- .../chat_ui/AdditionalModelSettings.tsx | 10 +-- .../components/chat_ui/AgentBuilderView.tsx | 2 +- .../components/chat_ui/ChatImageUpload.tsx | 2 +- .../playground/components/chat_ui/ChatUI.tsx | 16 ++-- .../chat_ui/CodeInterpreterOutput.tsx | 4 +- .../chat_ui/CodeInterpreterTool.tsx | 2 +- .../components/chat_ui/FilePreviewCard.tsx | 4 +- .../components/chat_ui/RealtimePlayground.tsx | 10 ++- .../chat_ui/ResponsesImageUpload.tsx | 2 +- .../chat_ui/SearchResultsDisplay.tsx | 6 +- .../components/chat_ui/SessionManagement.tsx | 2 +- .../components/compareUI/CompareUI.tsx | 4 +- .../compareUI/components/ComparisonPanel.tsx | 2 +- .../components/complianceUI/ComplianceUI.tsx | 76 +++++++++---------- .../_components/ai_suggestion_modal.tsx | 12 +-- .../prompts/_components/prompt_info.tsx | 2 +- .../_components/general_settings.tsx | 2 +- .../components/EndpointUsageTable.tsx | 2 +- .../components/EntityUsage/EntityUsage.tsx | 6 +- .../EntityUsage/SpendByProvider.tsx | 4 +- .../components/UsageAIChatPanel.tsx | 10 +-- .../_components/components/UsagePageView.tsx | 14 ++-- ui/litellm-dashboard/src/app/globals.css | 14 +++- .../AIHub/UsefulLinksManagement.tsx | 4 +- .../GuardrailsMonitor/LogViewer.tsx | 6 +- .../GuardrailsMonitor/MetricCard.tsx | 2 +- .../src/components/HelpLink.tsx | 4 +- .../Navbar/UserDropdown/UserDropdown.tsx | 2 +- .../SSOSettings/RoleMappings.tsx | 4 +- .../Fallbacks/FallbackGroupConfig.tsx | 6 +- .../RouterSettings/Fallbacks/Fallbacks.tsx | 2 +- .../components/EntityUsage/TopKeyView.tsx | 10 +-- .../VirtualKeysPage/keyTableColumns.tsx | 4 +- .../src/components/activity_metrics.tsx | 2 +- .../add_model/ClassificationMethodConfig.tsx | 6 +- .../add_model/ComplexityRouterConfig.tsx | 6 +- .../add_model/EscalationKeywords.tsx | 2 +- .../components/add_model/KeywordTierRules.tsx | 2 +- .../add_model/SemanticKeywordMatching.tsx | 2 +- .../components/bulk_create_users_button.tsx | 14 +++- .../src/components/chat/ConversationList.tsx | 2 +- .../components/chat_ui/MCPEventsDisplay.tsx | 4 +- .../common_components/AutoRotationView.tsx | 4 +- .../RouterSettingsSummary.tsx | 4 +- .../BudgetFallbacksEditor.tsx | 4 +- .../src/components/logging_settings_view.tsx | 4 +- .../mcp_tools/ByokCredentialModal.tsx | 8 +- .../HealthChecksTableColumns.tsx | 2 +- .../components/model_group_alias_settings.tsx | 2 +- .../permissions/AgentPermissions.tsx | 2 +- .../permissions/MCPServerPermissions.tsx | 10 +-- .../permissions/VectorStorePermissions.tsx | 2 +- .../src/components/public_model_hub.tsx | 16 ++-- .../shared/advanced_date_picker.tsx | 2 +- .../src/components/shared/chart_loader.tsx | 2 +- .../src/components/team/TeamInfo.tsx | 8 +- .../components/templates/KeyInfoHeader.tsx | 2 +- .../components/templates/key_info_view.tsx | 4 +- .../GuardrailViewer/CompliancePanel.tsx | 4 +- .../GuardrailViewer/GuardrailViewer.tsx | 4 +- .../LogDetailsDrawer/LogDetailsDrawer.tsx | 2 +- .../view_logs/RequestLogsTableColumns.tsx | 4 +- .../src/components/view_logs/TypeBadges.tsx | 2 +- 75 files changed, 248 insertions(+), 232 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx index c0b4150b4f4..ca47b71725d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx @@ -40,7 +40,7 @@ const compareRows = (a: CacheLeakageRow, b: CacheLeakageRow, sort: SortState): n const InfoTooltip = ({ info }: { info: string }) => ( }> - + {info} @@ -72,7 +72,7 @@ const SortableHead = ({ className="inline-flex items-center gap-1 font-medium hover:text-foreground" > {label} - + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/TierTurnsChart.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/TierTurnsChart.tsx index 5b9b8563baa..e55ebc07656 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/TierTurnsChart.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/TierTurnsChart.tsx @@ -133,7 +133,7 @@ const TierTurnsChart: React.FC = ({ view, autoRouters }) => {slice.tier} {Math.round((100 * slice.turns) / total).toLocaleString()}%

{slice.models.length > 0 && ( -

{slice.models.join(", ")}

+

{slice.models.join(", ")}

)}
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 fa4e8daf9d0..3fce928ed67 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 @@ -223,7 +223,7 @@ const CostTrackingSettings: React.FC = ({ userID, use ) : (
= ({ userID, use ) : (
+
Token Pricing:{" "} {result.input_cost_per_token && ( Input ${formatNumberWithCommas(result.input_cost_per_token * 1_000_000, 2)}/1M @@ -294,17 +294,17 @@ const MultiCostResults: React.FC = ({ multiResult, timePe {record.error ? ( - - + - ) : ( {formatCost(record.cost_per_request)} )} {record.error ? ( - - + - ) : ( 0 ? "text-warning" : "text-muted-foreground/70"}`} + className={`font-mono text-sm ${(record.margin_cost_per_request ?? 0) > 0 ? "text-warning" : "text-muted-foreground"}`} > {formatCost(record.margin_cost_per_request)} @@ -312,7 +312,7 @@ const MultiCostResults: React.FC = ({ multiResult, timePe {record.error ? ( - - + - ) : ( {formatCost(periodCost)} )} @@ -325,7 +325,7 @@ const MultiCostResults: React.FC = ({ multiResult, timePe aria-expanded={isExpanded} aria-label={`${isExpanded ? "Hide" : "Show"} cost breakdown for ${record.model}`} onClick={() => toggleExpanded(record.id)} - className="text-muted-foreground/70 hover:text-muted-foreground" + className="text-muted-foreground hover:text-foreground" > {isExpanded ? : } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_margin_table.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_margin_table.tsx index fa322d1b4d1..04823ac4aa0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_margin_table.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_margin_table.tsx @@ -138,7 +138,7 @@ const ProviderMarginTable: React.FC = ({ autoFocus /> % - + + + $ 150 ? "text-destructive" : row.original.avgLatency > 50 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 cc2d2051a24..c9de758f50d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx @@ -358,7 +358,7 @@ function GuardrailCard({ @@ -727,7 +727,7 @@ function DetailPanel({
@@ -218,7 +220,7 @@ function SubmissionRulesPanel({ requiredFields, onChange, onSave, isSaving }: Su
{field.label}
-
{field.description}
+
{field.description}
); @@ -235,7 +237,7 @@ function SubmissionRulesPanel({ requiredFields, onChange, onSave, isSaving }: Su await onSave(); setExpanded(false); }} - className="px-4 py-1.5 text-sm font-medium text-white bg-info hover:bg-info/80 disabled:opacity-50 rounded-md transition-colors" + className="px-4 py-1.5 text-sm font-medium text-info-foreground bg-info hover:bg-info/80 disabled:opacity-50 rounded-md transition-colors" > {isSaving ? "Saving…" : "Save Rules"} @@ -296,11 +298,11 @@ function MCPServerCard({ server, onApprove, onReject, requiredFields }: MCPServe )} {server.url && (
- + {server.url}
)} -
+
Transport: {server.transport ?? "sse"} @@ -322,7 +324,7 @@ function MCPServerCard({ server, onApprove, onReject, requiredFields }: MCPServe @@ -341,7 +343,7 @@ function MCPServerCard({ server, onApprove, onReject, requiredFields }: MCPServe @@ -367,7 +369,11 @@ function MCPServerCard({ server, onApprove, onReject, requiredFields }: MCPServe allPassed ? "bg-success" : "bg-destructive" }`} > - {allPassed ? : } + {allPassed ? ( + + ) : ( + + )}
@@ -383,7 +389,7 @@ function MCPServerCard({ server, onApprove, onReject, requiredFields }: MCPServe @@ -392,7 +398,7 @@ function MCPServerCard({ server, onApprove, onReject, requiredFields }: MCPServe @@ -567,13 +573,13 @@ export function MCPSubmissionsTab({ accessToken }: MCPSubmissionsTabProps) {
- + setSearch(e.target.value)} - className="w-full pl-9 pr-4 py-2 border border-border rounded-md text-sm text-foreground placeholder-muted-foreground/70 focus:outline-hidden focus:ring-1 focus:ring-ring focus:border-info" + className="w-full pl-9 pr-4 py-2 border border-border rounded-md text-sm text-foreground placeholder:text-muted-foreground focus:outline-hidden focus:ring-1 focus:ring-ring focus:border-info" />
setSearchPrompt(e.target.value)} placeholder="Search prompts..." - className="w-full border border-border rounded-lg pl-8 pr-3 py-1.5 text-xs placeholder:text-muted-foreground/70 focus:outline-hidden focus:ring-2 focus:ring-blue-500/20 focus:border-info" + className="w-full border border-border rounded-lg pl-8 pr-3 py-1.5 text-xs placeholder:text-muted-foreground focus:outline-hidden focus:ring-2 focus:ring-blue-500/20 focus:border-info" />
@@ -883,7 +883,7 @@ export default function ComplianceUI({ > Select All - · + · @@ -985,7 +985,7 @@ export default function ComplianceUI({ Required columns:{" "} prompt,{" "} expected_result{" "} - (fail or pass) + (fail or pass)

Optional columns:{" "} @@ -1049,14 +1049,14 @@ export default function ComplianceUI({ className="w-full flex items-center gap-2 px-3 py-2.5 text-left bg-muted hover:bg-accent transition-colors rounded-lg border border-border" > {isExpanded ? ( - + ) : ( - + )}

{fw.name} - {fwPromptCount} prompts + {fwPromptCount} prompts
{fwSelectedCount > 0 && ( @@ -1092,9 +1092,9 @@ export default function ComplianceUI({ className="w-full flex items-center gap-1.5 px-2.5 py-2 text-left hover:bg-accent transition-colors" > {isCatExpanded ? ( - + ) : ( - + )} {category.name} - + {category.prompts.length} {selectedInCat > 0 && ( @@ -1118,7 +1118,7 @@ export default function ComplianceUI({ {isCatExpanded && (
-

+

{category.description}

) : ( -

+

No policies or guardrails selected — select above to test against specific rules.

)} @@ -1245,7 +1245,7 @@ export default function ComplianceUI({
- +

Type a prompt below to quickly test it.

@@ -1254,10 +1254,10 @@ export default function ComplianceUI({ {quickTestMessages.map((msg) => (

{msg.type === "system" && ( @@ -1284,7 +1284,7 @@ export default function ComplianceUI({ {isQuickTesting && (

- +
)} @@ -1300,23 +1300,23 @@ export default function ComplianceUI({ onKeyDown={handleQuickTestKeyDown} placeholder="Enter text to test..." rows={3} - className="w-full px-3 pt-3 pb-1 text-sm text-foreground placeholder:text-muted-foreground/70 focus:outline-hidden resize-none" + className="w-full px-3 pt-3 pb-1 text-sm text-foreground placeholder:text-muted-foreground focus:outline-hidden resize-none" />
- + Press Enter to submit ·{" "} Shift+Enter for new line - {quickTestInput.length} + {quickTestInput.length}