From 2c4fed10351660bce488fafa8a4ed017ef0fd0cd Mon Sep 17 00:00:00 2001 From: epistoteles Date: Thu, 21 May 2026 21:19:05 +0200 Subject: [PATCH 001/538] Add missing Databricks model pricing data Add cost data for 14 Databricks models that were missing from the model_prices_and_context_window.json file: - databricks-gpt-5-4, gpt-5-2, gpt-5-4-nano, gpt-5-4-mini - databricks-gpt-5-2-codex, gpt-5-3-codex - databricks-gpt-5-1-codex-mini, gpt-5-1-codex-max - databricks-gemini-3-pro, gemini-3-flash, gemini-3-1-pro, gemini-3-1-flash-lite - databricks-claude-sonnet-4-6, claude-opus-4-6 Pricing sourced from: https://www.databricks.com/product/pricing/proprietary-foundation-model-serving --- model_prices_and_context_window.json | 227 +++++++++++++++++++++++++++ 1 file changed, 227 insertions(+) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 31a5993a240..6ff71018f18 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -11243,6 +11243,26 @@ "supports_minimal_reasoning_effort": true, "supports_tool_choice": true }, + "databricks/databricks-claude-opus-4-6": { + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 2.5000010000000002e-05, + "output_dbu_cost_per_token": 0.000357143, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_minimal_reasoning_effort": true, + "supports_tool_choice": true + }, "databricks/databricks-claude-sonnet-4": { "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, @@ -11300,6 +11320,25 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "databricks/databricks-claude-sonnet-4-6": { + "input_cost_per_token": 2.9999900000000002e-06, + "input_dbu_cost_per_token": 4.2857e-05, + "litellm_provider": "databricks", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 1.5000020000000002e-05, + "output_dbu_cost_per_token": 0.000214286, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "databricks/databricks-gemini-2-5-flash": { "input_cost_per_token": 3.0001999999999996e-07, "input_dbu_cost_per_token": 4.285999999999999e-06, @@ -11334,6 +11373,74 @@ "supports_function_calling": true, "supports_tool_choice": true }, + "databricks/databricks-gemini-3-1-flash-lite": { + "input_cost_per_token": 3.1248e-07, + "input_dbu_cost_per_token": 4.464e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 1.87502e-06, + "output_dbu_cost_per_token": 2.6786e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "databricks/databricks-gemini-3-1-pro": { + "input_cost_per_token": 2.49998e-06, + "input_dbu_cost_per_token": 3.5714e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 1.5000020000000002e-05, + "output_dbu_cost_per_token": 0.000214286, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "databricks/databricks-gemini-3-flash": { + "input_cost_per_token": 6.2503e-07, + "input_dbu_cost_per_token": 8.929e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 3.74997e-06, + "output_dbu_cost_per_token": 5.3571e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "databricks/databricks-gemini-3-pro": { + "input_cost_per_token": 2.49998e-06, + "input_dbu_cost_per_token": 3.5714e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 1.5000020000000002e-05, + "output_dbu_cost_per_token": 0.000214286, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_function_calling": true, + "supports_tool_choice": true + }, "databricks/databricks-gemma-3-12b": { "input_cost_per_token": 1.5000999999999998e-07, "input_dbu_cost_per_token": 2.1429999999999996e-06, @@ -11379,6 +11486,126 @@ "output_dbu_cost_per_token": 0.000142857, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" }, + "databricks/databricks-gpt-5-1-codex-max": { + "input_cost_per_token": 1.24999e-06, + "input_dbu_cost_per_token": 1.7857e-05, + "litellm_provider": "databricks", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 9.999990000000002e-06, + "output_dbu_cost_per_token": 0.000142857, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + }, + "databricks/databricks-gpt-5-1-codex-mini": { + "input_cost_per_token": 2.4997e-07, + "input_dbu_cost_per_token": 3.571e-06, + "litellm_provider": "databricks", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 1.99997e-06, + "output_dbu_cost_per_token": 2.8571e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + }, + "databricks/databricks-gpt-5-2": { + "input_cost_per_token": 1.75e-06, + "input_dbu_cost_per_token": 2.5e-05, + "litellm_provider": "databricks", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 1.4e-05, + "output_dbu_cost_per_token": 0.0002, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + }, + "databricks/databricks-gpt-5-2-codex": { + "input_cost_per_token": 1.75e-06, + "input_dbu_cost_per_token": 2.5e-05, + "litellm_provider": "databricks", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 1.4e-05, + "output_dbu_cost_per_token": 0.0002, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + }, + "databricks/databricks-gpt-5-3-codex": { + "input_cost_per_token": 1.75e-06, + "input_dbu_cost_per_token": 2.5e-05, + "litellm_provider": "databricks", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 1.4e-05, + "output_dbu_cost_per_token": 0.0002, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + }, + "databricks/databricks-gpt-5-4": { + "input_cost_per_token": 2.49998e-06, + "input_dbu_cost_per_token": 3.5714e-05, + "litellm_provider": "databricks", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 1.5000020000000002e-05, + "output_dbu_cost_per_token": 0.000214286, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + }, + "databricks/databricks-gpt-5-4-mini": { + "input_cost_per_token": 7.4998e-07, + "input_dbu_cost_per_token": 1.0714e-05, + "litellm_provider": "databricks", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 4.50002e-06, + "output_dbu_cost_per_token": 6.4286e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + }, + "databricks/databricks-gpt-5-4-nano": { + "input_cost_per_token": 1.9999e-07, + "input_dbu_cost_per_token": 2.857e-06, + "litellm_provider": "databricks", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 1.24999e-06, + "output_dbu_cost_per_token": 1.7857e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + }, "databricks/databricks-gpt-5-mini": { "input_cost_per_token": 2.4997000000000006e-07, "input_dbu_cost_per_token": 3.571e-06, 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 002/538] 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 c62e1238d648607f9544a98bcba4b4ec8ad5e369 Mon Sep 17 00:00:00 2001 From: Chenlu Ji Date: Mon, 6 Jul 2026 22:31:48 -0700 Subject: [PATCH 003/538] feat(tinyfish): surface response headers + top-level response extras Follow-up to #31411 (superseded and merged as #31997). Two related fixes so LiteLLM callers see what TinyFish actually returns, plus small correctness cleanups. ## Response headers surfaced on _hidden_params TinyFish sets useful response headers (x-request-id on every response, retry-after and x-ratelimit-limit on 429s). Previously these were only accessible via BaseLLMException.headers on error paths; on the success path they were dropped entirely. Fix: stash headers on both LiteLLM-conventional channels, matching the pattern used by Gemini / Volcengine / Manus / ChatGPT / OpenAI-responses providers. - `_hidden_params["headers"]` -- raw dict from httpx, all keys lowercased. - `_hidden_params["additional_headers"]` -- passed through process_response_headers, which prefixes any x-litellm-* provider header with `llm_provider-` so downstream LiteLLM code that trusts bare x-litellm-* markers can't be spoofed (values still survive under the prefixed key for observability). ## Top-level response extras (query, total_results, page, future fields) transform_search_response was building a fresh SearchResponse from just `results`, silently dropping every top-level field TinyFish's response carries beyond `results` / `object`. Fix: mutate parsed.results to its truncated slice and return the same SearchResponse instance rather than reconstructing. Every field pydantic populated during model_validate -- declared attributes AND extras (query, total_results, page, parameter_warnings, and any future TinyFish additions) -- survives regardless of which storage bucket holds it. Robust against upstream schema evolution: if LiteLLM later promotes a field from extras to declared, this code needs no change. ## Code cleanup - List-valued custom params JSON-encoded on the wire (matching the existing dict handling), so callers can pass a natural Python list for JSON-array wire params. - URL-encodable-params adapter accepts float in addition to str / int / bool; server-side rejection of a wrong-typed float now surfaces cleanly with `TinyFish Search:` attribution + docs link. - Assorted comment / docstring / test-fixture hygiene (no logic changes). ## Tests 70 unit + integration tests pass locally. Live-tested against production TinyFish with 6 diverse queries (basic / max_results / country=US / language=ja / domain filter / fetch={"format":"html"}) -- all 6 pass every expected-behavior check. --- .../llms/tinyfish/search/transformation.py | 74 +++++--- tests/search_tests/test_tinyfish_search.py | 61 ++++++- .../llms/tinyfish/test_tinyfish_search.py | 165 +++++++++++++++--- 3 files changed, 250 insertions(+), 50 deletions(-) diff --git a/litellm/llms/tinyfish/search/transformation.py b/litellm/llms/tinyfish/search/transformation.py index cef5f9cd02e..aea0dfe8b8e 100644 --- a/litellm/llms/tinyfish/search/transformation.py +++ b/litellm/llms/tinyfish/search/transformation.py @@ -14,6 +14,7 @@ import httpx from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.search.transformation import ( @@ -22,13 +23,13 @@ from litellm.llms.base_llm.search.transformation import ( ) from litellm.secret_managers.main import get_secret_str -_UrlEncodableParams = TypeAdapter(dict[str, str | int | bool]) +_UrlEncodableParams = TypeAdapter(dict[str, str | int | float | bool]) _StrList = TypeAdapter(list[str]) _StrFrozenSet = TypeAdapter(frozenset[str]) _TINYFISH_PARAMS_KEY = "_tinyfish_params" _TINYFISH_DOCS_URL = "https://docs.tinyfish.ai/search-api" -_TINYFISH_RESULT_CAP = 10 # TinyFish's natural per-page SERP ceiling +_TINYFISH_RESULT_CAP = 10 # Client-side truncation cap for max_results class TinyfishSearchConfig(BaseSearchConfig): @@ -94,16 +95,16 @@ class TinyfishSearchConfig(BaseSearchConfig): TinyFish equivalents: - ``query`` (str or list[str]) → ``query`` (list joined by spaces) - ``country`` → ``location`` - - ``search_domain_filter`` (list[str]) → folded into the query as - ``() (site:a OR site:b ...)`` (TinyFish has no first-class - field today; see ML-2084 for the planned ``include_domains``) + - ``search_domain_filter`` (list[str]) → folded into the query using + search operators - ``max_results`` → not sent on the wire; stashed on ``self._caller_max_results`` for client-side response truncation (TinyFish doesn't honor it server-side) - ``max_tokens_per_page`` → silently dropped (no TinyFish equivalent) Any other ``optional_params`` keys are forwarded to TinyFish as-is. - dict/list values are JSON-encoded so they survive ``urlencode``. + dict and list values are JSON-encoded so structured payloads survive + ``urlencode``. Returns: ``{_TINYFISH_PARAMS_KEY: }``. @@ -144,14 +145,12 @@ class TinyfishSearchConfig(BaseSearchConfig): supported_perplexity = _StrFrozenSet.validate_python(raw_supported) for param, value in optional_params.items(): if param not in supported_perplexity and param not in request_data: - # `fetch` expects a JSON-encoded object on the wire; accept the - # natural Python dict form and serialize here so callers don't - # have to pre-stringify. - if isinstance(value, dict): + # Serialize dicts/lists as JSON so structured params survive urlencode. + if isinstance(value, (dict, list)): value = json.dumps(value, separators=(",", ":")) # `urlencode` would render Python bool as "True"/"False" - # (capitalized). ux-labs validators require lowercase - # "true"/"false" (e.g. `include_thumbnail`); normalize here. + # (capitalized). TinyFish Search's bool params require lowercase + # "true"/"false" strings on the wire; normalize here. elif isinstance(value, bool): value = "true" if value else "false" request_data[param] = value @@ -167,17 +166,35 @@ class TinyfishSearchConfig(BaseSearchConfig): """ Transform a TinyFish response to LiteLLM's unified ``SearchResponse``. - Mappings (per-result): - - ``title`` → ``SearchResult.title`` (defaults to ``""`` if missing/null) - - ``url`` → ``SearchResult.url`` (defaults to ``""``) - - ``snippet`` → ``SearchResult.snippet`` (defaults to ``""``) - - all other per-result fields (``position``, ``site_name``, - ``thumbnail_url``, ``fetch``, ``fetch_error``, ...) ride through as - extras on ``SearchResult`` via its ``extra="allow"`` config. + Per-result field handling: + - ``title``, ``url``, ``snippet`` are declared on ``SearchResult`` and + populated by ``SearchResponse.model_validate`` when present. Missing + or ``None`` values are defaulted to ``""`` beforehand by + ``_default_missing_result_fields`` so a degraded result flows through + instead of failing the whole call. + - All undeclared per-result fields (``position``, ``site_name``, and + any others TinyFish returns) ride through as extras via + ``SearchResult``'s ``extra="allow"`` config — accessible as + attributes on the result object or enumerable via + ``result.model_extra``. - Top-level ``parameter_warnings`` (see ML-2085) is read when present and - each entry is re-fired via ``verbose_logger.warning``. Absent or - malformed entries are silently skipped — never throws. + Top-level ``parameter_warnings`` is read when present and each entry + is re-fired via ``verbose_logger.warning``. Absent or malformed + entries are silently skipped — never throws. + + Top-level extras (``query``, ``total_results``, ``page``, and any + future TinyFish additions) ride through via + ``SearchResponse.extra="allow"``. The validated response is returned + in place after truncating ``results`` to the caller's ``max_results``, + so every field pydantic populated survives regardless of which + storage bucket (declared attribute or ``__pydantic_extra__``) holds it. + + TinyFish response headers (e.g. ``x-request-id``, ``retry-after``, + ``x-ratelimit-limit`` — httpx normalizes header names to lowercase) + are stashed on ``response._hidden_params["headers"]`` (raw) and + ``response._hidden_params["additional_headers"]`` (sanitized via + ``process_response_headers``) so callers can correlate a search with + server-side logs. Error paths routed through ``self._wrap_error`` for uniform ``"TinyFish Search: . See for details."`` wrapping: @@ -223,7 +240,12 @@ class TinyfishSearchConfig(BaseSearchConfig): _emit_parameter_warnings(parsed) max_results = self._caller_max_results or _TINYFISH_RESULT_CAP - return SearchResponse(results=list(parsed.results[:max_results])) + # Truncate in place so all pydantic-populated fields survive — declared and extras. + parsed.results = list(parsed.results[:max_results]) + raw_headers = dict(raw_response.headers) + parsed._hidden_params["headers"] = raw_headers + parsed._hidden_params["additional_headers"] = process_response_headers(raw_headers) + return parsed def _wrap_error( self, @@ -243,9 +265,9 @@ class TinyfishSearchConfig(BaseSearchConfig): carry the ``TinyFish Search:`` prefix — the bare error already names the host in the URL, so attribution is implicit there. """ - # ux-labs frontend wraps every error body as {"error": {"code", "message", "details"?}}. + # TinyFish Search wraps every error body as {"error": {"code", "message", "details"?}}. # Best-effort unwrap to surface the inner message; fall back to the raw body - # for non-ux-labs responses (CDN HTML pages, other JSON envelopes, plain text). + # for other envelope shapes (CDN HTML pages, other JSON envelopes, plain text). inner_message = error_message try: body: object = json.loads(error_message) # any-ok: json.loads -> Any @@ -290,7 +312,7 @@ def _default_missing_result_fields(raw_json: object) -> None: def _emit_parameter_warnings(parsed: SearchResponse) -> None: - """Re-fire TinyFish-side ``parameter_warnings`` (see ML-2085) as warnings. + """Re-fire TinyFish-side ``parameter_warnings`` as warnings. Defensive: skip silently on any shape we don't recognize so a malformed entry (or an early/partial rollout of the field) never throws. diff --git a/tests/search_tests/test_tinyfish_search.py b/tests/search_tests/test_tinyfish_search.py index aca28544513..becb8287a29 100644 --- a/tests/search_tests/test_tinyfish_search.py +++ b/tests/search_tests/test_tinyfish_search.py @@ -35,11 +35,16 @@ MOCK_TINYFISH_RESPONSE = { def _make_mock_response( - json_data: dict, status_code: int = 200, request_url: str | None = None + json_data: dict, + status_code: int = 200, + request_url: str | None = None, + headers: dict | None = None, ) -> MagicMock: mock = MagicMock() mock.status_code = status_code mock.json.return_value = json_data + # httpx.Headers normalizes keys to lowercase — mirror production behavior. + mock.headers = httpx.Headers(headers or {}) if request_url: mock.request = MagicMock() mock.request.url = httpx.URL(request_url) @@ -163,7 +168,7 @@ class TestTinyfishSearch: @pytest.mark.asyncio async def test_fetch_param_round_trip(self): - # End-to-end check: caller passes `fetch=...` (JSON-encoded tf-fetch + # End-to-end check: caller passes `fetch=...` (JSON-encoded fetch # config); param reaches TinyFish on the request side and the nested # `fetch` object on each result surfaces back to the SearchResult on the # response side. No LiteLLM-side support code is required. @@ -235,6 +240,58 @@ class TestTinyfishSearch: assert result.results[0].title == "Result 0" assert result.results[2].title == "Result 2" + @pytest.mark.asyncio + async def test_top_level_extras_surface_end_to_end(self): + # Envelope extras (`query`, `total_results`, `page`) must survive the + # full asearch dispatch — proves LiteLLM's entry-point plumbing outside + # our transformer doesn't accidentally strip them. + os.environ["TINYFISH_API_KEY"] = "sk-tinyfish-test" + + mock_response = _make_mock_response(MOCK_TINYFISH_RESPONSE) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get", + new_callable=AsyncMock, + ) as mock_get: + mock_get.return_value = mock_response + + response = await litellm.asearch( + query="web automation tools", + search_provider="tinyfish", + ) + + assert getattr(response, "query", None) == "web automation tools" + assert getattr(response, "total_results", None) == 2 + assert getattr(response, "page", None) == 0 + + @pytest.mark.asyncio + async def test_response_headers_surface_end_to_end(self): + # Response headers must land on `_hidden_params` after the full + # asearch dispatch (both raw and sanitized channels). + os.environ["TINYFISH_API_KEY"] = "sk-tinyfish-test" + + mock_response = _make_mock_response( + MOCK_TINYFISH_RESPONSE, + headers={"X-Request-ID": "req-e2e-1"}, + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get", + new_callable=AsyncMock, + ) as mock_get: + mock_get.return_value = mock_response + + response = await litellm.asearch( + query="test", + search_provider="tinyfish", + ) + + raw = response._hidden_params["headers"] + add = response._hidden_params["additional_headers"] + # httpx lowercases; both channels agree on the value. + assert raw["x-request-id"] == "req-e2e-1" + assert add["llm_provider-x-request-id"] == "req-e2e-1" + @pytest.mark.asyncio async def test_empty_results(self): os.environ["TINYFISH_API_KEY"] = "sk-tinyfish-test" diff --git a/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py b/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py index 58363e3baea..2dcccb8ea7e 100644 --- a/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py +++ b/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py @@ -47,7 +47,9 @@ def _make_mock_response( mock = MagicMock() mock.status_code = status_code - mock.headers = headers or {} + # httpx.Headers normalizes keys to lowercase — mirror production so tests + # assert what callers actually see. + mock.headers = httpx.Headers(headers or {}) if json_data is not None: mock.json.return_value = json_data mock.text = text if text is not None else _json.dumps(json_data) @@ -222,7 +224,7 @@ class TestTransformSearchRequest: assert param not in result["_tinyfish_params"] def test_arbitrary_param_passed_through(self): - # `fetch` is a TinyFish-specific param (JSON-encoded tf-fetch config). + # `fetch` is a TinyFish-specific param (JSON-encoded fetch config). # The passthrough loop should forward it verbatim without LiteLLM needing # to know about it. config = TinyfishSearchConfig() @@ -237,26 +239,49 @@ class TestTransformSearchRequest: config = TinyfishSearchConfig() result = config.transform_search_request( query="test", - optional_params={"fetch": {"format": "html", "fetch_path": "fast"}}, - ) - assert ( - result["_tinyfish_params"]["fetch"] - == '{"format":"html","fetch_path":"fast"}' + optional_params={"fetch": {"format": "html"}}, ) + assert result["_tinyfish_params"]["fetch"] == '{"format":"html"}' def test_bool_param_serialized_as_lowercase(self): - # urlencode renders Python bool as capitalized "True"/"False"; ux-labs - # rejects those (e.g. include_thumbnail must be literal "true"/"false"). - # Normalize before passing through. + # urlencode renders Python bool as capitalized "True"/"False"; TinyFish + # Search's bool params require lowercase "true"/"false" strings on the + # wire. Normalize before passing through. config = TinyfishSearchConfig() true_result = config.transform_search_request( - query="test", optional_params={"include_thumbnail": True} + query="test", optional_params={"some_bool_param": True} ) false_result = config.transform_search_request( - query="test", optional_params={"include_thumbnail": False} + query="test", optional_params={"some_bool_param": False} ) - assert true_result["_tinyfish_params"]["include_thumbnail"] == "true" - assert false_result["_tinyfish_params"]["include_thumbnail"] == "false" + assert true_result["_tinyfish_params"]["some_bool_param"] == "true" + assert false_result["_tinyfish_params"]["some_bool_param"] == "false" + + def test_float_param_passes_through(self): + # Float values pass the urlencode adapter and land on the wire as + # their decimal string form. If TinyFish's server rejects a float + # for a param it expects as int, the server's 400 response is + # attributed via _wrap_error (`TinyFish Search: ...`) — better than + # a client-side pydantic ValidationError with no context. + config = TinyfishSearchConfig() + result = config.transform_search_request( + query="test", + optional_params={"some_float_param": 0.5}, + ) + assert result["_tinyfish_params"]["some_float_param"] == 0.5 + + def test_list_param_auto_json_encoded(self): + # TinyFish Search's JSON-array params arrive on the wire as JSON- + # encoded strings. Accept the natural Python list form and serialize + # so the caller doesn't have to pre-stringify. Params whose wire + # format is a plain comma-separated string are the caller's + # responsibility to pass as a Python str. + config = TinyfishSearchConfig() + result = config.transform_search_request( + query="test", + optional_params={"some_list_param": ["a.example", "b.example"]}, + ) + assert result["_tinyfish_params"]["some_list_param"] == '["a.example","b.example"]' def test_pre_stringified_param_passed_unchanged(self): # If the caller already JSON-encoded, don't re-encode. @@ -422,10 +447,104 @@ class TestTransformSearchResponse: assert getattr(first, "position", None) == 1 assert getattr(first, "site_name", None) == "tinyfish.ai" + def test_top_level_extras_flow_through(self): + # TinyFish returns `query`, `total_results`, `page` at the envelope + # level. These must ride through to the caller via SearchResponse's + # extra="allow" so pagination logic, echo checks, etc. work. + config = TinyfishSearchConfig() + mock_response = _make_mock_response(MOCK_TINYFISH_RESPONSE) + result = config.transform_search_response( + raw_response=mock_response, logging_obj=None + ) + assert getattr(result, "query", None) == "web automation tools" + assert getattr(result, "total_results", None) == 2 + assert getattr(result, "page", None) == 0 + + def test_top_level_future_extras_flow_through(self): + # Any future TinyFish top-level field must ride through unchanged + # (design contract: no LiteLLM code change needed for new fields). + config = TinyfishSearchConfig() + body = { + "results": [ + {"title": "x", "url": "https://x", "snippet": "x"}, + ], + "query": "test", + "example_int_extra": 123, # hypothetical future field + "example_str_extra": "value", # hypothetical future field + "example_id_extra": "abc-def", # hypothetical future field + } + result = config.transform_search_response( + raw_response=_make_mock_response(body), logging_obj=None + ) + assert getattr(result, "example_int_extra", None) == 123 + assert getattr(result, "example_str_extra", None) == "value" + assert getattr(result, "example_id_extra", None) == "abc-def" + + def test_response_headers_stashed_on_hidden_params(self): + # TinyFish Search sets X-Request-ID on every success response. Confirm it + # lands on both `_hidden_params["headers"]` (raw) and + # `_hidden_params["additional_headers"]` (sanitized/prefixed). + # httpx.Headers lowercases every key, so assertions use lowercase. + config = TinyfishSearchConfig() + mock_response = _make_mock_response( + MOCK_TINYFISH_RESPONSE, + headers={"X-Request-ID": "req-abc-123", "Content-Type": "application/json"}, + ) + result = config.transform_search_response( + raw_response=mock_response, logging_obj=None + ) + # Raw copy — httpx has normalized keys to lowercase. + assert result._hidden_params["headers"]["x-request-id"] == "req-abc-123" + # process_response_headers prefixes non-OpenAI-standard keys with "llm_provider-". + assert result._hidden_params["additional_headers"]["llm_provider-x-request-id"] == "req-abc-123" + + def test_response_headers_future_headers_flow_through(self): + # "Accept extra": any header TinyFish Search adds later must ride + # through without a LiteLLM code change. + config = TinyfishSearchConfig() + mock_response = _make_mock_response( + MOCK_TINYFISH_RESPONSE, + headers={ + "X-Request-ID": "req-1", + "X-Example-Header-A": "value-a", # hypothetical future header + "X-Example-Header-B": "value-b", # hypothetical future header + }, + ) + result = config.transform_search_response( + raw_response=mock_response, logging_obj=None + ) + raw = result._hidden_params["headers"] + # httpx lowercases header names on read. + assert raw["x-example-header-a"] == "value-a" + assert raw["x-example-header-b"] == "value-b" + + def test_response_headers_strips_x_litellm_spoof(self): + # A provider setting `x-litellm-*` in its response must not be able to + # spoof LiteLLM-internal markers via _hidden_params["additional_headers"]. + # The raw copy preserves the header (opt-in debug view); the sanitized + # copy prefixes it with `llm_provider-` so bare `x-litellm-*` markers + # can't be spoofed (values still survive under the prefixed key for + # observability). + config = TinyfishSearchConfig() + mock_response = _make_mock_response( + MOCK_TINYFISH_RESPONSE, + headers={"x-litellm-attempted-fallbacks": "spoofed", "X-Request-ID": "r1"}, + ) + result = config.transform_search_response( + raw_response=mock_response, logging_obj=None + ) + # Raw view still has the spoof. + assert result._hidden_params["headers"]["x-litellm-attempted-fallbacks"] == "spoofed" + # Sanitized view: the spoof survives only under the llm_provider- prefix + # (never under the bare x-litellm-* key that LiteLLM downstream trusts). + additional = result._hidden_params["additional_headers"] + assert "x-litellm-attempted-fallbacks" not in additional + assert additional.get("llm_provider-x-litellm-attempted-fallbacks") == "spoofed" + def test_fetch_field_rides_through_to_search_result(self): - # Mirrors browser-search's per-result `fetch` nested object (see - # api/src/parser.rs SearchResult.fetch). Confirms `fetch=...` requests - # surface their content to LiteLLM callers without provider changes. + # Mirrors TinyFish Search's per-result `fetch` nested object. + # Confirms `fetch=...` requests surface their content to LiteLLM + # callers without provider changes. config = TinyfishSearchConfig() fetched = { "results": [ @@ -568,7 +687,7 @@ class TestTransformSearchResponse: class TestErrorHandling: def test_4xx_response_raises_with_attribution_and_unwrapped_message(self): - # Reproduces ux-labs' error envelope shape for an INVALID_INPUT response. + # Reproduces TinyFish Search's error envelope shape for an INVALID_INPUT response. config = TinyfishSearchConfig() body = { "error": { @@ -590,7 +709,7 @@ class TestErrorHandling: def test_429_preserves_status_code_and_headers(self): config = TinyfishSearchConfig() - body = {"error": {"code": "RATE_LIMIT_EXCEEDED", "message": "60 rpm"}} + body = {"error": {"code": "RATE_LIMIT_EXCEEDED", "message": "rate limit exceeded"}} mock_response = _make_mock_response( body, status_code=429, headers={"Retry-After": "60"} ) @@ -600,10 +719,12 @@ class TestErrorHandling: ) assert getattr(exc_info.value, "status_code", None) == 429 headers = getattr(exc_info.value, "headers", {}) or {} - assert headers.get("Retry-After") == "60" + # httpx lowercases; the exception carries the same dict shape. + assert headers.get("retry-after") == "60" - def test_5xx_with_non_ux_labs_body_falls_back_to_raw_text(self): - # Cloudflare-style JSON or any other envelope: unwrap fails, fall back to raw. + def test_5xx_with_non_tinyfish_envelope_shape_falls_back_to_raw_text(self): + # A JSON body that doesn't match TinyFish Search's error envelope shape: + # unwrap fails, fall back to the raw body text. config = TinyfishSearchConfig() body = {"errors": [{"code": "10000", "message": "Internal"}]} mock_response = _make_mock_response(body, status_code=502) From 39bd5fb8b7923f94712bac1bd9f1f6a93adb50ad Mon Sep 17 00:00:00 2001 From: Sujith Date: Tue, 14 Jul 2026 15:18:16 +0530 Subject: [PATCH 004/538] fix(main): forward store and prompt_cache_key params on chat completions (#33184) --- litellm/main.py | 8 ++++ litellm/utils.py | 2 + tests/test_litellm/test_main.py | 84 +++++++++++++++++++++++++++++++++ 3 files changed, 94 insertions(+) diff --git a/litellm/main.py b/litellm/main.py index 7d457d9cdd1..81b02082abd 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -435,6 +435,8 @@ async def acompletion( verbosity: Optional[Literal["low", "medium", "high"]] = None, safety_identifier: Optional[str] = None, service_tier: Optional[str] = None, + store: Optional[bool] = None, + prompt_cache_key: Optional[str] = None, # set api_base, api_version, api_key base_url: Optional[str] = None, api_version: Optional[str] = None, @@ -585,6 +587,8 @@ async def acompletion( "verbosity": verbosity, "safety_identifier": safety_identifier, "service_tier": service_tier, + "store": store, + "prompt_cache_key": prompt_cache_key, "extra_headers": extra_headers, "acompletion": True, # assuming this is a required parameter "thinking": thinking, @@ -4828,6 +4832,8 @@ def completion( # type: ignore extra_headers: Optional[dict] = None, safety_identifier: Optional[str] = None, service_tier: Optional[str] = None, + store: Optional[bool] = None, + prompt_cache_key: Optional[str] = None, # soon to be deprecated params by OpenAI functions: Optional[List] = None, function_call: Optional[str] = None, @@ -5249,6 +5255,8 @@ def completion( # type: ignore ), "safety_identifier": safety_identifier, "service_tier": service_tier, + "store": store, + "prompt_cache_key": prompt_cache_key, "allowed_openai_params": kwargs.get("allowed_openai_params"), "base_model": base_model, } diff --git a/litellm/utils.py b/litellm/utils.py index 18b89ee0d13..b16ecdf88be 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3791,6 +3791,8 @@ def get_optional_params( thinking: Optional[AnthropicThinkingParam] = None, web_search_options: Optional[OpenAIWebSearchOptions] = None, safety_identifier: Optional[str] = None, + store: Optional[bool] = None, + prompt_cache_key: Optional[str] = None, base_model: Optional[str] = None, **kwargs, ): diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 28cf4fa0744..0624b590df7 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -2081,3 +2081,87 @@ def test_stream_chunk_builder_text_completion_combines_text_and_usage(): assert response.usage.prompt_tokens > 0 assert response.usage.completion_tokens > 0 assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens + + +def test_completion_forwards_store_and_prompt_cache_key_to_openai(): + """ + Regression test for https://github.com/BerriAI/litellm/issues/33184 + + store and prompt_cache_key are documented OpenAI chat completion params that + were accepted as supported but silently dropped before the provider request + was built, because they were not named parameters of completion() and + get_optional_params() the way safety_identifier is. + """ + from openai import OpenAI + + client = OpenAI(api_key="fake-api-key") + + with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + try: + litellm.completion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + store=False, + prompt_cache_key="test-cache-key", + client=client, + ) + except Exception as e: + print(e) + + mock_client.assert_called_once() + request_body = mock_client.call_args.kwargs + assert request_body["store"] is False + assert request_body["prompt_cache_key"] == "test-cache-key" + + +@pytest.mark.asyncio +async def test_acompletion_forwards_store_and_prompt_cache_key_to_openai(): + """ + Async variant of the store/prompt_cache_key forwarding regression test for + https://github.com/BerriAI/litellm/issues/33184 + """ + from openai import AsyncOpenAI + + client = AsyncOpenAI(api_key="fake-api-key") + + with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + try: + await litellm.acompletion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + store=False, + prompt_cache_key="test-cache-key", + client=client, + ) + except Exception as e: + print(e) + + mock_client.assert_called_once() + request_body = mock_client.call_args.kwargs + assert request_body["store"] is False + assert request_body["prompt_cache_key"] == "test-cache-key" + + +def test_completion_omits_store_and_prompt_cache_key_when_not_passed(): + """ + When store and prompt_cache_key are not passed, they must not appear in the + outbound request body (guards against always forwarding None defaults). + """ + from openai import OpenAI + + client = OpenAI(api_key="fake-api-key") + + with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + try: + litellm.completion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + client=client, + ) + except Exception as e: + print(e) + + mock_client.assert_called_once() + request_body = mock_client.call_args.kwargs + assert "store" not in request_body + assert "prompt_cache_key" not in request_body From 4eaa70440a247cee2767bd16e7dc830da559105a Mon Sep 17 00:00:00 2001 From: Sujith Date: Tue, 14 Jul 2026 15:49:46 +0530 Subject: [PATCH 005/538] fix(main): use PEP 604 unions for new store and prompt_cache_key params --- litellm/main.py | 8 ++++---- litellm/utils.py | 4 ++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 81b02082abd..3b13e04bbe3 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -435,8 +435,8 @@ async def acompletion( verbosity: Optional[Literal["low", "medium", "high"]] = None, safety_identifier: Optional[str] = None, service_tier: Optional[str] = None, - store: Optional[bool] = None, - prompt_cache_key: Optional[str] = None, + store: bool | None = None, + prompt_cache_key: str | None = None, # set api_base, api_version, api_key base_url: Optional[str] = None, api_version: Optional[str] = None, @@ -4832,8 +4832,8 @@ def completion( # type: ignore extra_headers: Optional[dict] = None, safety_identifier: Optional[str] = None, service_tier: Optional[str] = None, - store: Optional[bool] = None, - prompt_cache_key: Optional[str] = None, + store: bool | None = None, + prompt_cache_key: str | None = None, # soon to be deprecated params by OpenAI functions: Optional[List] = None, function_call: Optional[str] = None, diff --git a/litellm/utils.py b/litellm/utils.py index b16ecdf88be..bc5b2447761 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3791,8 +3791,8 @@ def get_optional_params( thinking: Optional[AnthropicThinkingParam] = None, web_search_options: Optional[OpenAIWebSearchOptions] = None, safety_identifier: Optional[str] = None, - store: Optional[bool] = None, - prompt_cache_key: Optional[str] = None, + store: bool | None = None, + prompt_cache_key: str | None = None, base_model: Optional[str] = None, **kwargs, ): From 371fa670d6f79dfd579945e2d357f5b978d9af21 Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 17 Jul 2026 19:33:28 +0000 Subject: [PATCH 006/538] fix(proxy): forward Bedrock event-stream content-type on unbuffered passthrough Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 14 +++-- .../proxy/test_common_request_processing.py | 55 +++++++++++++++++++ 2 files changed, 65 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index c7c9397d850..24831a64410 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1766,6 +1766,7 @@ class ProxyBaseLLMRequestProcessing: return StreamingResponse( content=generator, # type: ignore[arg-type] status_code=status.HTTP_200_OK, + media_type=self._passthrough_event_stream_media_type(), headers=custom_headers, ) else: @@ -2216,10 +2217,15 @@ class ProxyBaseLLMRequestProcessing: def _passthrough_event_stream_media_type(self) -> Optional[str]: """ - Content-type for a buffered passthrough event-stream response, resolved - from the provider handler so the proxy stays provider-agnostic. Mirrors - the upstream content-type the non-streaming path forwards, since the - buffered streaming generator carries no headers of its own. + Content-type for a passthrough event-stream response, resolved from the + provider handler so the proxy stays provider-agnostic. Mirrors the + upstream content-type the non-streaming path forwards, since the + streaming generator carries no headers of its own. Used for both the + buffered (guardrail-rewritten) and the unbuffered relay paths so + clients that enforce the event-stream content-type (e.g. Claude Code on + Bedrock invoke-with-response-stream) see the correct header instead of + Starlette's application/octet-stream default. Returns None for providers + with no event-stream media type, leaving the response default unchanged. """ from litellm.llms.pass_through.guardrail_translation.handler import ( LlmPassthroughRouteHandler, diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index ebfbb46053d..f1a3745e85b 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -4137,6 +4137,61 @@ class TestAllmPassthroughStreamingProviderGate: assert streamed == chunks mock_handler.assert_not_awaited() + @pytest.mark.asyncio + async def test_bedrock_invoke_stream_sets_event_stream_content_type(self, monkeypatch): + """ + Regression for LIT-4561. The unbuffered Bedrock event-stream relay + (invoke-with-response-stream, no post-call guardrail rewriting) must set + content-type: application/vnd.amazon.eventstream instead of leaving it to + Starlette's application/octet-stream default, which trips Claude Code's + content-type guard added in 2.1.208 + """ + processing_obj = self._build_processing_obj( + "bedrock", "model/us.anthropic.claude-sonnet-4-20250514-v1:0/invoke-with-response-stream" + ) + chunks = [b"raw-1", b"raw-2"] + + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=False, + ): + result = await self._run(processing_obj, monkeypatch, chunks) + + assert isinstance(result, StreamingResponse) + assert result.media_type == "application/vnd.amazon.eventstream" + assert result.headers["content-type"] == "application/vnd.amazon.eventstream" + streamed = [chunk async for chunk in result.body_iterator] + assert streamed == chunks + + @pytest.mark.asyncio + async def test_non_bedrock_stream_keeps_default_content_type(self, monkeypatch): + """ + A provider with no registered event-stream media type must not have one + forced onto its unbuffered stream, so the response default is unchanged + """ + processing_obj = self._build_processing_obj("anthropic") + chunks = [b"chunk-1", b"chunk-2"] + + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=False, + ): + result = await self._run(processing_obj, monkeypatch, chunks) + + assert isinstance(result, StreamingResponse) + assert result.media_type is None + assert result.headers.get("content-type") != "application/vnd.amazon.eventstream" + class TestResponseCostHeaderForTypedDictResponses: """ From 3b843708b097d6c4ff96d374900282ddf5142f0d Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 18 Jul 2026 20:17:57 +0000 Subject: [PATCH 007/538] fix(bedrock): degrade gracefully on malformed tool-call arguments split_concatenated_json_objects re-raised JSONDecodeError on genuinely malformed (non-concatenated) tool-call arguments, which propagated out of _convert_to_bedrock_tool_call_invoke and turned every replayed Bedrock conversation into a 500. Catch the decode error, keep whatever complete objects parsed, log a warning, and let the caller fall back to input={} so the conversation continues. Fixes #18667 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../prompt_templates/common_utils.py | 27 +++++++++--- ...ore_utils_prompt_templates_common_utils.py | 30 +++++++++++-- ...llm_core_utils_prompt_templates_factory.py | 43 +++++++++++++++++++ 3 files changed, 89 insertions(+), 11 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 538d5f650ef..db3856ce86b 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1679,16 +1679,19 @@ def split_concatenated_json_objects(raw: str) -> List[Dict[str, Any]]: This helper uses ``json.JSONDecoder.raw_decode()`` to walk the string and extract each JSON object individually. + The walk degrades gracefully: if the string is malformed or truncated + (e.g. a stream that ended mid-tool-call), whatever complete objects were + parsed before the bad tail are returned and the remainder is discarded + with a warning, rather than raising. The sole caller + (``_convert_to_bedrock_tool_call_invoke``) treats an empty result as + ``input={}`` so the conversation can continue instead of hard-failing. + Returns ------- list[dict] A list of parsed dicts – one per JSON object found. If *raw* is - empty or whitespace-only, an empty list is returned. - - Raises - ------ - json.JSONDecodeError - If the string contains text that cannot be parsed as JSON at all. + empty, whitespace-only, or wholly unparseable, an empty list is + returned. """ import json @@ -1708,7 +1711,17 @@ def split_concatenated_json_objects(raw: str) -> List[Dict[str, Any]]: if idx >= length: break - obj, end_idx = decoder.raw_decode(raw, idx) + try: + obj, end_idx = decoder.raw_decode(raw, idx) + except json.JSONDecodeError as e: + verbose_logger.warning( + "split_concatenated_json_objects: discarding unparseable tool-call " + "arguments tail after %d complete object(s); error=%s at char %d", + len(results), + e, + idx, + ) + break if isinstance(obj, dict): results.append(obj) else: diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 1b1db634ed2..6d14d283b84 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -251,10 +251,32 @@ def test_split_concatenated_json_non_dict_value(): assert result == [{}] -def test_split_concatenated_json_invalid_raises(): - """Completely invalid JSON raises JSONDecodeError.""" - with pytest.raises(json.JSONDecodeError): - split_concatenated_json_objects("not json at all") +def test_split_concatenated_json_wholly_invalid_returns_empty(): + """ + Wholly unparseable JSON degrades to an empty list instead of raising. + + Regression for https://github.com/BerriAI/litellm/issues/18667: a raise + here propagated out of `_convert_to_bedrock_tool_call_invoke` and turned + every replayed conversation into a 500. + """ + assert split_concatenated_json_objects("not json at all") == [] + + +def test_split_concatenated_json_malformed_object_returns_empty(): + """ + A single malformed object (missing comma between keys) degrades to an + empty list rather than raising `Expecting ',' delimiter`. + """ + assert split_concatenated_json_objects('{"location": "Boston" "unit": "celsius"}') == [] + + +def test_split_concatenated_json_salvages_prefix_before_truncated_tail(): + """ + Complete objects parsed before an unparseable/truncated tail are kept; + only the bad tail is discarded. + """ + result = split_concatenated_json_objects('{"a": 1}{"b": 2}{"c":') + assert result == [{"a": 1}, {"b": 2}] # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index bcda88ea609..17f680df4d9 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -2286,6 +2286,49 @@ def test_bedrock_tool_call_invoke_non_dict_arguments(): assert result[0]["toolUse"]["input"] == {} +def test_bedrock_tool_call_invoke_malformed_json_does_not_raise(): + """ + Regression for https://github.com/BerriAI/litellm/issues/18667. + + When the model emits malformed JSON in tool-call arguments (here a + missing comma between keys), replaying that history must NOT raise + `Unable to convert openai tool calls ... Expecting ',' delimiter`. + It degrades to an empty-object input so the conversation can continue. + """ + tool_calls = [ + { + "id": "toolu_abc123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location": "Boston" "unit": "celsius"}', + }, + } + ] + result = _convert_to_bedrock_tool_call_invoke(tool_calls) + assert len(result) == 1 + assert result[0]["toolUse"]["toolUseId"] == "toolu_abc123" + assert result[0]["toolUse"]["name"] == "get_weather" + assert result[0]["toolUse"]["input"] == {} + + +def test_bedrock_tool_call_invoke_salvages_valid_prefix_before_truncated_tail(): + """ + A valid leading object followed by a truncated tail keeps the valid + object rather than dropping everything or raising. + """ + tool_calls = [ + { + "id": "call_partial", + "type": "function", + "function": {"name": "shell", "arguments": '{"cmd": "ls"}{"cmd":'}, + } + ] + result = _convert_to_bedrock_tool_call_invoke(tool_calls) + assert len(result) == 1 + assert result[0]["toolUse"]["input"] == {"cmd": "ls"} + + def test_make_valid_bedrock_tool_name_preserves_hyphens(): assert make_valid_bedrock_tool_name("my-tool") == "my-tool" assert ( From c2d8a4e4263e45bee96e94f7091071140bf79d83 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 18 Jul 2026 20:28:35 +0000 Subject: [PATCH 008/538] chore(bedrock): clarify tool-call decode warning to avoid double char reference Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/prompt_templates/common_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index db3856ce86b..9dcbca954b3 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1716,10 +1716,10 @@ def split_concatenated_json_objects(raw: str) -> List[Dict[str, Any]]: except json.JSONDecodeError as e: verbose_logger.warning( "split_concatenated_json_objects: discarding unparseable tool-call " - "arguments tail after %d complete object(s); error=%s at char %d", + "arguments tail after %d complete object(s); decode_start=%d error=%s", len(results), - e, idx, + e, ) break if isinstance(obj, dict): From b61484e6c919b2d8718249ef10889029192fa5b5 Mon Sep 17 00:00:00 2001 From: CrypticDriver <107245892+CrypticDriver@users.noreply.github.com> Date: Tue, 21 Jul 2026 11:31:50 +0000 Subject: [PATCH 009/538] feat: add Amazon Bedrock AgentCore Web Search as a native search provider Adds 'agentcore' to SearchProviders, backed by an AgentCore Gateway web-search connector target (MCP tools/call over Streamable HTTP). Web Search on Amazon Bedrock AgentCore is an AWS-managed web index (GA June 2026). Exposing it as a native search provider lets Bedrock users enable Claude Code / Anthropic-native WebSearch through websearch_interception with a pure-YAML config and AWS-native auth, keeping the whole search path inside AWS. Implementation: - New AgentCoreSearchConfig (litellm/llms/bedrock/search/) reusing BaseAWSLLM credential resolution. Auth follows the gateway's inbound authorizer type: AWS_IAM gateways get a SigV4-signed request (explicit aws_access_key_id/aws_secret_access_key params or the default credential chain); CUSTOM_JWT gateways get an OAuth2 bearer token via api_key / AGENTCORE_GATEWAY_TOKEN - SigV4 signing region is derived from the gateway URL so callers don't need aws_region_name to match their default region - Adds an optional sign_request() hook to BaseSearchConfig (no-op by default) and teaches the search HTTP handler to send a signed body verbatim, mirroring the existing anthropic_messages/chat pattern - Handles both plain-JSON and SSE-framed MCP responses, propagates MCP errors, truncates queries to the 200-char gateway limit Tested: - 13 unit tests: payload/signing, explicit AKSK passthrough, bearer token via api_key and env, query truncation, SSE frames, MCP error propagation, region derivation - Verified end-to-end against real AWS_IAM and CUSTOM_JWT gateways, including full Claude Code CLI WebSearch round-trips through the proxy with websearch_interception --- .../llms/base_llm/search/transformation.py | 23 ++ litellm/llms/bedrock/search/__init__.py | 0 litellm/llms/bedrock/search/transformation.py | 256 ++++++++++++++++++ litellm/llms/custom_httpx/llm_http_handler.py | 34 +++ .../agentcore_websearch_config.yaml | 39 +++ litellm/types/utils.py | 1 + litellm/utils.py | 2 + tests/search_tests/test_agentcore_search.py | 231 ++++++++++++++++ 8 files changed, 586 insertions(+) create mode 100644 litellm/llms/bedrock/search/__init__.py create mode 100644 litellm/llms/bedrock/search/transformation.py create mode 100644 litellm/proxy/example_config_yaml/agentcore_websearch_config.yaml create mode 100644 tests/search_tests/test_agentcore_search.py diff --git a/litellm/llms/base_llm/search/transformation.py b/litellm/llms/base_llm/search/transformation.py index fdfac6f5f9f..7a93cf43ca7 100644 --- a/litellm/llms/base_llm/search/transformation.py +++ b/litellm/llms/base_llm/search/transformation.py @@ -178,6 +178,29 @@ class BaseSearchConfig: """ return headers + def sign_request( + self, + headers: dict, + optional_params: dict, + request_data: Union[dict, list[dict]], + api_base: str, + api_key: str | None = None, + ) -> tuple[dict, bytes | None]: + """ + OPTIONAL + + Sign the request. Providers like Bedrock AgentCore need to SigV4-sign + the request before sending it to the API. + + For all other providers, this is a no-op and we just return the headers. + + Returns: + Tuple of (headers, signed_json_body). When signed_json_body is not + None, the handler MUST send it verbatim as the request body — + re-serializing the payload would invalidate the signature. + """ + return headers, None + def get_complete_url( self, api_base: Optional[str], diff --git a/litellm/llms/bedrock/search/__init__.py b/litellm/llms/bedrock/search/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py new file mode 100644 index 00000000000..16d671f26b0 --- /dev/null +++ b/litellm/llms/bedrock/search/transformation.py @@ -0,0 +1,256 @@ +""" +Calls an Amazon Bedrock AgentCore Gateway web-search target (MCP protocol) to search the web. + +Web Search on Amazon Bedrock AgentCore exposes Amazon's managed web index through +an AgentCore Gateway MCP endpoint. + +AWS docs: https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/gateway-target-connector-web-search-tool.html + +Authentication (matches the gateway's inbound authorizer type): +- AWS_IAM gateway: the request is SigV4-signed. Credentials come from explicit + params (aws_access_key_id / aws_secret_access_key / aws_session_token / + aws_region_name — also settable in a proxy search_tools entry) or the + standard AWS credential chain (env / profile / IRSA / assumed role) +- CUSTOM_JWT gateway: pass the OAuth2 bearer token (e.g. Cognito + client_credentials) as api_key, or set AGENTCORE_GATEWAY_TOKEN + +Setup: + 1. Create an AgentCore Gateway with a web-search connector target + 2. Set AGENTCORE_GATEWAY_URL (or pass api_base) to the gateway MCP endpoint, e.g. + https://.gateway.bedrock-agentcore..amazonaws.com/mcp + 3. AWS_IAM: ensure the credentials allow bedrock-agentcore:InvokeGateway + CUSTOM_JWT: set AGENTCORE_GATEWAY_TOKEN (or pass api_key) + +Usage: + response = litellm.search( + query="latest AI developments", + search_provider="agentcore", + max_results=5, + aws_access_key_id="...", # optional — omit to use the default chain + aws_secret_access_key="...", + ) +""" + +import json +import re +from typing import Union + +import httpx + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.search.transformation import ( + BaseSearchConfig, + SearchResponse, + SearchResult, +) +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.common_utils import BedrockError +from litellm.secret_managers.main import get_secret_str + +# AgentCore web-search rejects queries longer than 200 characters +AGENTCORE_MAX_QUERY_LENGTH = 200 + +# Default MCP tool name for a gateway web-search connector target: +# "___". Override with AGENTCORE_SEARCH_TOOL_NAME +# or optional_params["tool_name"] when the target uses a custom name. +AGENTCORE_DEFAULT_TOOL_NAME = "web-search-tool___WebSearch" + + +class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): + def __init__(self) -> None: + BaseSearchConfig.__init__(self) + BaseAWSLLM.__init__(self) + + @staticmethod + def ui_friendly_name() -> str: + return "Web Search on Amazon Bedrock" + + def validate_environment( + self, + headers: dict, + api_key: str | None = None, + api_base: str | None = None, + **kwargs, + ) -> dict: + """ + Set MCP transport headers. Per the MCP Streamable HTTP transport spec, + the client MUST accept both application/json and text/event-stream. + + Authentication itself happens in sign_request(): bearer token for + CUSTOM_JWT gateways, AWS SigV4 for AWS_IAM gateways. + """ + headers["Content-Type"] = "application/json" + headers["Accept"] = "application/json, text/event-stream" + return headers + + def get_complete_url( + self, + api_base: str | None, + optional_params: dict, + data: Union[dict, list[dict]] | None = None, + **kwargs, + ) -> str: + api_base = api_base or get_secret_str("AGENTCORE_GATEWAY_URL") + if not api_base: + raise ValueError( + "AGENTCORE_GATEWAY_URL is not set. Set it to your AgentCore Gateway MCP " + "endpoint (https://.gateway.bedrock-agentcore." + ".amazonaws.com/mcp) or pass api_base." + ) + return api_base + + def transform_search_request( + self, + query: Union[str, list[str]], + optional_params: dict, + **kwargs, + ) -> dict: + """ + Transform Search request to an MCP tools/call request. + + Args: + query: Search query (string or list of strings). AgentCore only + supports single string queries; lists are joined with spaces. + optional_params: Optional parameters for the request + - max_results: Maximum number of results (1-25), default 10 + - tool_name: Override the MCP tool name of the gateway target + + Returns: + Dict with the JSON-RPC 2.0 request body + """ + if isinstance(query, list): + query = " ".join(query) + query = query[:AGENTCORE_MAX_QUERY_LENGTH] + + tool_name = ( + optional_params.get("tool_name") + or get_secret_str("AGENTCORE_SEARCH_TOOL_NAME") + or AGENTCORE_DEFAULT_TOOL_NAME + ) + + arguments: dict[str, Union[str, int]] = {"query": query} + if "max_results" in optional_params: + arguments["maxResults"] = optional_params["max_results"] + + return { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": tool_name, "arguments": arguments}, + } + + def sign_request( + self, + headers: dict, + optional_params: dict, + request_data: Union[dict, list[dict]], + api_base: str, + api_key: str | None = None, + ) -> tuple[dict, bytes | None]: + """ + Authenticate the MCP request. + + CUSTOM_JWT gateways: attach the caller's OAuth2 bearer token (api_key + or AGENTCORE_GATEWAY_TOKEN) — no AWS credentials involved. + + AWS_IAM gateways: SigV4-sign with the bedrock-agentcore service name. + """ + if not isinstance(request_data, dict): + raise ValueError("AgentCore search expects a single dict request body") + + bearer_token = api_key or get_secret_str("AGENTCORE_GATEWAY_TOKEN") + if bearer_token: + headers["Authorization"] = f"Bearer {bearer_token}" + return headers, json.dumps(request_data).encode() + + # The signing region must match the gateway's region — derive it from + # the gateway URL so callers don't have to set aws_region_name to a + # region different from their default. + signing_params = dict(optional_params) + if signing_params.get("aws_region_name") is None: + match = re.search( + r"\.gateway\.bedrock-agentcore\.([a-z0-9-]+)\.amazonaws\.com", + api_base, + ) + if match: + signing_params["aws_region_name"] = match.group(1) + + return self._sign_request( + service_name="bedrock-agentcore", + headers=headers, + optional_params=signing_params, + request_data=request_data, + api_base=api_base, + ) + + def transform_search_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + **kwargs, + ) -> SearchResponse: + """ + Transform an MCP tools/call response to LiteLLM unified SearchResponse. + + The gateway returns JSON-RPC (as plain JSON or a single-message SSE + stream) whose result.content[] text blocks contain a JSON list of + {title, url, date/publishedDate, text} entries. + """ + response_json = self._parse_mcp_body(raw_response) + + if "error" in response_json: + raise BedrockError( + status_code=raw_response.status_code if raw_response.status_code >= 400 else 502, + message=f"AgentCore gateway MCP error: {response_json['error']}", + ) + + results: list[SearchResult] = [] + for block in response_json.get("result", {}).get("content", []): + if block.get("type") != "text": + continue + try: + parsed = json.loads(block["text"]) + except (json.JSONDecodeError, TypeError): + continue + items = parsed.get("results", []) if isinstance(parsed, dict) else parsed + for item in items: + if not isinstance(item, dict): + continue + results.append( + SearchResult( + title=item.get("title") or "", + url=item.get("url") or "", + snippet=item.get("text") or item.get("snippet") or "", + date=item.get("publishedDate") or item.get("date"), + last_updated=None, + ) + ) + + return SearchResponse(results=results, object="search") + + @staticmethod + def _parse_mcp_body(raw_response: httpx.Response) -> dict: + """Parse a JSON or SSE-framed (Streamable HTTP transport) MCP response.""" + text = raw_response.text + if text.lstrip().startswith(("event:", "data:")): + for line in text.splitlines(): + if line.startswith("data:"): + return json.loads(line[len("data:") :].strip()) + raise BedrockError( + status_code=502, + message=f"AgentCore gateway returned SSE without a data frame: {text[:200]}", + ) + return raw_response.json() + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict, + ) -> Exception: + return BaseLLMException( + status_code=status_code, + message=error_message, + headers=headers, + ) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 3e6f9ee08ee..aa33865df43 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1737,6 +1737,15 @@ class BaseLLMHTTPHandler: api_key=api_key, ) + # Sign the request if the provider requires it (e.g. AWS SigV4) + headers, signed_json_body = provider_config.sign_request( + headers=headers, + optional_params=optional_params, + request_data=data, + api_base=complete_url, + api_key=api_key, + ) + ## LOGGING logging_obj.pre_call( input=query if isinstance(query, str) else str(query), @@ -1762,6 +1771,14 @@ class BaseLLMHTTPHandler: url=complete_url, headers=headers, ) + elif signed_json_body is not None: + # Send the signed body verbatim — re-serializing would break the signature + response = client.post( + url=complete_url, + headers=headers, + data=signed_json_body, + timeout=timeout, + ) else: # Make POST request with JSON data response = client.post( @@ -1821,6 +1838,15 @@ class BaseLLMHTTPHandler: api_key=api_key, ) + # Sign the request if the provider requires it (e.g. AWS SigV4) + headers, signed_json_body = provider_config.sign_request( + headers=headers, + optional_params=optional_params, + request_data=data, + api_base=complete_url, + api_key=api_key, + ) + ## LOGGING logging_obj.pre_call( input=query if isinstance(query, str) else str(query), @@ -1851,6 +1877,14 @@ class BaseLLMHTTPHandler: url=complete_url, headers=headers, ) + elif signed_json_body is not None: + # Send the signed body verbatim — re-serializing would break the signature + response = await async_httpx_client.post( + url=complete_url, + headers=headers, + data=signed_json_body, + timeout=timeout, + ) else: # Make async POST request with JSON data response = await async_httpx_client.post( diff --git a/litellm/proxy/example_config_yaml/agentcore_websearch_config.yaml b/litellm/proxy/example_config_yaml/agentcore_websearch_config.yaml new file mode 100644 index 00000000000..f2c5a460bf0 --- /dev/null +++ b/litellm/proxy/example_config_yaml/agentcore_websearch_config.yaml @@ -0,0 +1,39 @@ +# Claude Code / Anthropic-native web search on Bedrock, backed by +# Amazon Bedrock AgentCore Web Search (AWS-managed web index, no third-party +# search API). See litellm/llms/bedrock/search/transformation.py for details. + +model_list: + - model_name: claude-sonnet + litellm_params: + model: bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0 + aws_region_name: us-east-1 + +search_tools: + - search_tool_name: agentcore-search + litellm_params: + search_provider: agentcore + # Your AgentCore Gateway MCP endpoint (gateway must have a `web-search` + # connector target). Alternatively set the AGENTCORE_GATEWAY_URL env var. + api_base: https://.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp + + # The gateway exposes the connector as "___WebSearch". + # Default is "web-search-tool___WebSearch", matching the target name used + # in the AWS docs' boto3/CLI setup examples. Set this ONLY if your target + # was created with a different name (misconfiguration surfaces as an MCP + # "tool not found" error): + # tool_name: MyWebSearchTarget___WebSearch + + # AWS_IAM gateway (default): SigV4-signed. Omit keys to use the standard + # AWS credential chain (env / profile / IRSA / instance role), or set them + # explicitly: + # aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID + # aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY + + # CUSTOM_JWT gateway alternative — OAuth2 bearer token instead of SigV4: + # api_key: os.environ/AGENTCORE_GATEWAY_TOKEN + +litellm_settings: + callbacks: ["websearch_interception"] + websearch_interception_params: + enabled_providers: ["bedrock"] + search_tool_name: agentcore-search diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ec8a9336ca7..088d2193055 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3489,6 +3489,7 @@ class SearchProviders(str, Enum): YOU_COM = "you_com" APISERPENT = "apiserpent" TINYFISH = "tinyfish" + AGENTCORE = "agentcore" # Create a set of all search provider values for quick lookup diff --git a/litellm/utils.py b/litellm/utils.py index 174bed09396..80a1f2b991f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8860,6 +8860,7 @@ class ProviderConfigManager: from litellm.llms.apiserpent.search.transformation import ( APISerpentSearchConfig, ) + from litellm.llms.bedrock.search.transformation import AgentCoreSearchConfig from litellm.llms.brave.search.transformation import BraveSearchConfig from litellm.llms.dataforseo.search.transformation import DataForSEOSearchConfig from litellm.llms.duckduckgo.search.transformation import DuckDuckGoSearchConfig @@ -8897,6 +8898,7 @@ class ProviderConfigManager: SearchProviders.YOU_COM: YouComSearchConfig, SearchProviders.APISERPENT: APISerpentSearchConfig, SearchProviders.TINYFISH: TinyfishSearchConfig, + SearchProviders.AGENTCORE: AgentCoreSearchConfig, } config_class = PROVIDER_TO_CONFIG_MAP.get(provider, None) if config_class is None: diff --git a/tests/search_tests/test_agentcore_search.py b/tests/search_tests/test_agentcore_search.py new file mode 100644 index 00000000000..5d2e4d1c3cd --- /dev/null +++ b/tests/search_tests/test_agentcore_search.py @@ -0,0 +1,231 @@ +""" +Tests for Amazon Bedrock AgentCore Web Search integration. +""" + +import json +import os +import sys +import pytest +from unittest.mock import AsyncMock, patch, MagicMock + +sys.path.insert(0, os.path.abspath("../..")) + +import litellm +from litellm.llms.bedrock.search.transformation import AgentCoreSearchConfig + +GATEWAY_URL = "https://testgateway-abc123.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp" + +MCP_RESULTS = [ + { + "title": "Test Result 1", + "url": "https://example.com/1", + "text": "Snippet for result 1", + "publishedDate": "2026-06-16", + }, + { + "title": "Test Result 2", + "url": "https://example.com/2", + "text": "Snippet for result 2", + }, +] + + +def _mcp_response_body() -> dict: + return { + "jsonrpc": "2.0", + "id": 1, + "result": {"content": [{"type": "text", "text": json.dumps(MCP_RESULTS)}]}, + } + + +def _make_mock_response(json_body: dict = None, text: str = None) -> MagicMock: + mock_response = MagicMock() + mock_response.status_code = 200 + if text is not None: + mock_response.text = text + else: + mock_response.text = json.dumps(json_body) + mock_response.json.return_value = json_body + return mock_response + + +class TestAgentCoreSearch: + """ + Tests for AgentCore Web Search functionality with mocked network/signing. + """ + + @pytest.mark.asyncio + async def test_agentcore_search_request_payload(self): + """Validates the MCP tools/call payload and SigV4 signing without real AWS calls.""" + os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL + + mock_response = _make_mock_response(_mcp_response_body()) + + with ( + patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post, + patch.object( + AgentCoreSearchConfig, + "_sign_request", + return_value=( + {"Authorization": "AWS4-HMAC-SHA256 test", "Content-Type": "application/json"}, + json.dumps({"signed": True}).encode(), + ), + ) as mock_sign, + ): + mock_post.return_value = mock_response + + response = await litellm.asearch( + query="latest developments in AI", + search_provider="agentcore", + max_results=5, + ) + + mock_post.assert_called_once() + call_kwargs = mock_post.call_args.kwargs + assert call_kwargs["url"] == GATEWAY_URL + # Signed body must be sent verbatim + assert call_kwargs["data"] == json.dumps({"signed": True}).encode() + assert "json" not in call_kwargs + + # Signing was invoked with the MCP request + mock_sign.assert_called_once() + sign_kwargs = mock_sign.call_args.kwargs + request_data = sign_kwargs["request_data"] + assert request_data["method"] == "tools/call" + assert request_data["params"]["name"] == "web-search-tool___WebSearch" + assert request_data["params"]["arguments"]["query"] == "latest developments in AI" + assert request_data["params"]["arguments"]["maxResults"] == 5 + assert sign_kwargs["service_name"] == "bedrock-agentcore" + + assert len(response.results) == 2 + assert response.results[0].title == "Test Result 1" + assert response.results[0].url == "https://example.com/1" + assert response.results[0].snippet == "Snippet for result 1" + assert response.results[0].date == "2026-06-16" + + def test_transform_search_request_query_truncation(self): + """AgentCore rejects queries > 200 chars; the request must truncate.""" + config = AgentCoreSearchConfig() + long_query = "a" * 300 + data = config.transform_search_request(query=long_query, optional_params={}) + assert len(data["params"]["arguments"]["query"]) == 200 + + def test_transform_search_request_joins_list_queries(self): + config = AgentCoreSearchConfig() + data = config.transform_search_request(query=["foo", "bar"], optional_params={}) + assert data["params"]["arguments"]["query"] == "foo bar" + + def test_transform_search_request_custom_tool_name(self): + config = AgentCoreSearchConfig() + data = config.transform_search_request(query="q", optional_params={"tool_name": "my-target___WebSearch"}) + assert data["params"]["name"] == "my-target___WebSearch" + + def test_get_complete_url_requires_gateway_url(self): + config = AgentCoreSearchConfig() + os.environ.pop("AGENTCORE_GATEWAY_URL", None) + with pytest.raises(ValueError, match="AGENTCORE_GATEWAY_URL"): + config.get_complete_url(api_base=None, optional_params={}) + + def test_get_complete_url_prefers_api_base(self): + config = AgentCoreSearchConfig() + assert config.get_complete_url(api_base=GATEWAY_URL, optional_params={}) == GATEWAY_URL + + def test_validate_environment_sets_mcp_headers(self): + """MCP Streamable HTTP requires accepting both JSON and SSE.""" + config = AgentCoreSearchConfig() + headers = config.validate_environment(headers={}) + assert headers["Accept"] == "application/json, text/event-stream" + assert headers["Content-Type"] == "application/json" + + def test_transform_search_response_parses_sse_frame(self): + """Gateway may answer with an SSE-framed JSON-RPC message.""" + config = AgentCoreSearchConfig() + body = _mcp_response_body() + sse_text = f"event: message\ndata: {json.dumps(body)}\n\n" + mock_response = _make_mock_response(text=sse_text) + + response = config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + assert len(response.results) == 2 + assert response.results[1].url == "https://example.com/2" + + def test_transform_search_response_raises_on_mcp_error(self): + config = AgentCoreSearchConfig() + mock_response = _make_mock_response( + {"jsonrpc": "2.0", "id": 1, "error": {"code": -32601, "message": "tool not found"}} + ) + with pytest.raises(Exception, match="tool not found"): + config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + + def test_sign_request_uses_bearer_token_when_api_key_set(self): + """CUSTOM_JWT gateways: api_key is sent as a bearer token, no SigV4.""" + config = AgentCoreSearchConfig() + request_data = {"jsonrpc": "2.0", "id": 1} + + headers, signed_body = config.sign_request( + headers={"Content-Type": "application/json"}, + optional_params={}, + request_data=request_data, + api_base=GATEWAY_URL, + api_key="test-jwt-token", + ) + assert headers["Authorization"] == "Bearer test-jwt-token" + assert signed_body == json.dumps(request_data).encode() + + def test_sign_request_uses_bearer_token_from_env(self): + config = AgentCoreSearchConfig() + os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" + try: + headers, _ = config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base=GATEWAY_URL, + ) + assert headers["Authorization"] == "Bearer env-jwt-token" + finally: + os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) + + def test_sign_request_passes_explicit_aws_credentials(self): + """Explicit aws_* params (e.g. from a proxy search_tools entry) reach the signer.""" + config = AgentCoreSearchConfig() + + with patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign: + config.sign_request( + headers={}, + optional_params={ + "aws_access_key_id": "AKIATEST", + "aws_secret_access_key": "secret", + "aws_session_token": "token", + }, + request_data={"jsonrpc": "2.0"}, + api_base=GATEWAY_URL, + ) + passed = mock_base_sign.call_args.kwargs["optional_params"] + assert passed["aws_access_key_id"] == "AKIATEST" + assert passed["aws_secret_access_key"] == "secret" + assert passed["aws_session_token"] == "token" + + def test_sign_request_derives_region_from_gateway_url(self): + """Signing region must come from the gateway URL, not the caller's default region.""" + config = AgentCoreSearchConfig() + eu_url = "https://gw-x.gateway.bedrock-agentcore.eu-central-1.amazonaws.com/mcp" + + with patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign: + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base=eu_url, + ) + assert mock_base_sign.call_args.kwargs["optional_params"]["aws_region_name"] == "eu-central-1" From ebdad6e3ef5fa35daee25cd8bad9ae5b6e54f353 Mon Sep 17 00:00:00 2001 From: CrypticDriver <107245892+CrypticDriver@users.noreply.github.com> Date: Wed, 22 Jul 2026 02:26:25 +0000 Subject: [PATCH 010/538] fix: address bot review findings (auth hardening, SSE parsing, defaults) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Refuse to send the server-managed AGENTCORE_GATEWAY_TOKEN to a caller-supplied api_base (reuses resolve_server_api_key's trusted-host guard) — closes the token-exfiltration path via /search_tools/test_connection - Disable BaseAWSLLM's AWS_BEARER_TOKEN_BEDROCK fallback when signing: that token is a Bedrock Runtime credential and must not reach an AgentCore gateway - Parse SSE responses per spec: join multi-line data fields, iterate events, and return the JSON-RPC response (result/error) instead of the first data line — progress notifications no longer shadow the result - Validate tool_name ends with ___WebSearch so a caller-supplied name cannot invoke unrelated tools on the same gateway with the proxy's credentials - Send the documented maxResults default (10) explicitly instead of leaving it to the gateway - Custom gateway hostnames: raise a clear error when no signing region can be derived and none is configured, instead of signing for a guessed region - 7 new unit tests covering each fix (20 total) --- litellm/llms/bedrock/search/transformation.py | 93 ++++++++++++++++--- tests/search_tests/test_agentcore_search.py | 92 ++++++++++++++++++ 2 files changed, 170 insertions(+), 15 deletions(-) diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index 16d671f26b0..3bae84f0026 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -51,11 +51,20 @@ from litellm.secret_managers.main import get_secret_str # AgentCore web-search rejects queries longer than 200 characters AGENTCORE_MAX_QUERY_LENGTH = 200 +# The provider contract documents a default of 10 results — send it explicitly +# so the gateway can't silently apply a different default. +AGENTCORE_DEFAULT_MAX_RESULTS = 10 + # Default MCP tool name for a gateway web-search connector target: # "___". Override with AGENTCORE_SEARCH_TOOL_NAME # or optional_params["tool_name"] when the target uses a custom name. AGENTCORE_DEFAULT_TOOL_NAME = "web-search-tool___WebSearch" +# All web-search connector tools share this suffix; rejecting other names keeps +# a caller-supplied tool_name from invoking unrelated tools on the same gateway +# with the proxy's credentials. +AGENTCORE_TOOL_NAME_SUFFIX = "___WebSearch" + class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): def __init__(self) -> None: @@ -128,10 +137,15 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): or get_secret_str("AGENTCORE_SEARCH_TOOL_NAME") or AGENTCORE_DEFAULT_TOOL_NAME ) + if not tool_name.endswith(AGENTCORE_TOOL_NAME_SUFFIX): + raise ValueError( + f"Invalid AgentCore search tool_name '{tool_name}': must end with " + f"'{AGENTCORE_TOOL_NAME_SUFFIX}' (a web-search connector tool). " + "Other gateway tools cannot be invoked through this provider." + ) arguments: dict[str, Union[str, int]] = {"query": query} - if "max_results" in optional_params: - arguments["maxResults"] = optional_params["max_results"] + arguments["maxResults"] = optional_params.get("max_results", AGENTCORE_DEFAULT_MAX_RESULTS) return { "jsonrpc": "2.0", @@ -159,14 +173,28 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): if not isinstance(request_data, dict): raise ValueError("AgentCore search expects a single dict request body") - bearer_token = api_key or get_secret_str("AGENTCORE_GATEWAY_TOKEN") + # Server-managed token fallback is gated on the request targeting the + # operator-configured gateway host — otherwise an authenticated caller + # could point api_base at their own server (e.g. via + # /search_tools/test_connection) and receive AGENTCORE_GATEWAY_TOKEN. + bearer_token = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("AGENTCORE_GATEWAY_TOKEN",), + base_env_var="AGENTCORE_GATEWAY_URL", + default_api_base=None, + ) if bearer_token: headers["Authorization"] = f"Bearer {bearer_token}" return headers, json.dumps(request_data).encode() # The signing region must match the gateway's region — derive it from - # the gateway URL so callers don't have to set aws_region_name to a - # region different from their default. + # standard gateway hostnames so callers don't have to set + # aws_region_name to a region different from their default. Custom or + # private hostnames can't be parsed: fall back to an explicitly + # configured region (param or AWS env vars), and error out rather than + # silently signing for a guessed region the gateway would reject with + # a confusing auth error. signing_params = dict(optional_params) if signing_params.get("aws_region_name") is None: match = re.search( @@ -175,13 +203,23 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): ) if match: signing_params["aws_region_name"] = match.group(1) + elif not any(get_secret_str(var) for var in ("AWS_REGION", "AWS_REGION_NAME", "AWS_DEFAULT_REGION")): + raise ValueError( + f"Cannot derive the SigV4 signing region from api_base '{api_base}'. " + "Set aws_region_name (or the AWS_REGION env var) to the gateway's " + "region when using a custom hostname." + ) + # api_key="" (not None, but falsy) disables BaseAWSLLM's fallback to the + # AWS_BEARER_TOKEN_BEDROCK env var: that token is a Bedrock Runtime + # credential and must not be sent to an AgentCore gateway. return self._sign_request( service_name="bedrock-agentcore", headers=headers, optional_params=signing_params, request_data=request_data, api_base=api_base, + api_key="", ) def transform_search_response( @@ -231,17 +269,42 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): @staticmethod def _parse_mcp_body(raw_response: httpx.Response) -> dict: - """Parse a JSON or SSE-framed (Streamable HTTP transport) MCP response.""" + """ + Parse a JSON or SSE-framed (Streamable HTTP transport) MCP response. + + Per the SSE spec, an event's data is the concatenation of all its + ``data:`` lines (joined with newlines), and a stream may carry several + events (e.g. progress notifications before the JSON-RPC response). + Return the event whose payload carries the ``id``-matched JSON-RPC + response — i.e. one containing ``result`` or ``error``. + """ text = raw_response.text - if text.lstrip().startswith(("event:", "data:")): - for line in text.splitlines(): - if line.startswith("data:"): - return json.loads(line[len("data:") :].strip()) - raise BedrockError( - status_code=502, - message=f"AgentCore gateway returned SSE without a data frame: {text[:200]}", - ) - return raw_response.json() + if not text.lstrip().startswith(("event:", "data:", ":", "id:", "retry:")): + return raw_response.json() + + last_parsed: dict | None = None + data_lines: list[str] = [] + # Trailing sentinel flushes the final event even without a blank line + for line in text.splitlines() + [""]: + if line.startswith("data:"): + data_lines.append(line[len("data:") :].lstrip()) + continue + if line == "" and data_lines: + try: + parsed = json.loads("\n".join(data_lines)) + except json.JSONDecodeError: + parsed = None + data_lines = [] + if isinstance(parsed, dict): + last_parsed = parsed + if "result" in parsed or "error" in parsed: + return parsed + if last_parsed is not None: + return last_parsed + raise BedrockError( + status_code=502, + message=f"AgentCore gateway returned SSE without a JSON data frame: {text[:200]}", + ) def get_error_class( self, diff --git a/tests/search_tests/test_agentcore_search.py b/tests/search_tests/test_agentcore_search.py index 5d2e4d1c3cd..d5f6ab3e82f 100644 --- a/tests/search_tests/test_agentcore_search.py +++ b/tests/search_tests/test_agentcore_search.py @@ -123,6 +123,18 @@ class TestAgentCoreSearch: data = config.transform_search_request(query="q", optional_params={"tool_name": "my-target___WebSearch"}) assert data["params"]["name"] == "my-target___WebSearch" + def test_transform_search_request_rejects_non_websearch_tool_name(self): + """A caller-supplied tool_name must not reach other tools on the gateway.""" + config = AgentCoreSearchConfig() + with pytest.raises(ValueError, match="must end with"): + config.transform_search_request(query="q", optional_params={"tool_name": "admin-target___DeleteUser"}) + + def test_transform_search_request_sends_documented_default_max_results(self): + """The documented default of 10 is sent explicitly, not left to the gateway.""" + config = AgentCoreSearchConfig() + data = config.transform_search_request(query="q", optional_params={}) + assert data["params"]["arguments"]["maxResults"] == 10 + def test_get_complete_url_requires_gateway_url(self): config = AgentCoreSearchConfig() os.environ.pop("AGENTCORE_GATEWAY_URL", None) @@ -151,6 +163,29 @@ class TestAgentCoreSearch: assert len(response.results) == 2 assert response.results[1].url == "https://example.com/2" + def test_transform_search_response_parses_multiline_sse_data(self): + """SSE data may be split across several data: lines (joined per spec).""" + config = AgentCoreSearchConfig() + pretty = json.dumps(_mcp_response_body(), indent=2) + sse_text = "event: message\n" + "\n".join(f"data: {line}" for line in pretty.splitlines()) + "\n\n" + mock_response = _make_mock_response(text=sse_text) + + response = config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + assert len(response.results) == 2 + + def test_transform_search_response_skips_progress_events(self): + """A progress notification before the JSON-RPC result must not shadow it.""" + config = AgentCoreSearchConfig() + progress = {"jsonrpc": "2.0", "method": "notifications/progress", "params": {"progress": 1}} + sse_text = ( + f"event: message\ndata: {json.dumps(progress)}\n\n" + f"event: message\ndata: {json.dumps(_mcp_response_body())}\n\n" + ) + mock_response = _make_mock_response(text=sse_text) + + response = config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + assert len(response.results) == 2 + def test_transform_search_response_raises_on_mcp_error(self): config = AgentCoreSearchConfig() mock_response = _make_mock_response( @@ -175,8 +210,10 @@ class TestAgentCoreSearch: assert signed_body == json.dumps(request_data).encode() def test_sign_request_uses_bearer_token_from_env(self): + """Server token is attached when the request targets the configured gateway host.""" config = AgentCoreSearchConfig() os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" + os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL try: headers, _ = config.sign_request( headers={}, @@ -187,6 +224,61 @@ class TestAgentCoreSearch: assert headers["Authorization"] == "Bearer env-jwt-token" finally: os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) + os.environ.pop("AGENTCORE_GATEWAY_URL", None) + + def test_sign_request_refuses_server_token_to_untrusted_host(self): + """Server-managed token must not be sent to a caller-chosen api_base.""" + config = AgentCoreSearchConfig() + os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" + os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL + try: + with pytest.raises(ValueError, match="Refusing to send"): + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base="https://attacker.example.com/mcp", + ) + finally: + os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) + os.environ.pop("AGENTCORE_GATEWAY_URL", None) + + def test_sign_request_does_not_leak_bedrock_bearer_token(self): + """AWS_BEARER_TOKEN_BEDROCK is a Bedrock Runtime credential — it must not + replace SigV4 on requests to an AgentCore gateway.""" + config = AgentCoreSearchConfig() + + with patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign: + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base=GATEWAY_URL, + ) + # api_key="" (falsy, not None) disables the base class's + # AWS_BEARER_TOKEN_BEDROCK env fallback. + assert mock_base_sign.call_args.kwargs["api_key"] == "" + + def test_sign_request_custom_hostname_requires_region(self): + """Non-standard hostnames can't yield a signing region — require it explicitly.""" + config = AgentCoreSearchConfig() + saved = {var: os.environ.pop(var, None) for var in ("AWS_REGION", "AWS_REGION_NAME", "AWS_DEFAULT_REGION")} + try: + with pytest.raises(ValueError, match="signing region"): + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base="https://gateway.internal.example.com/mcp", + ) + finally: + for var, val in saved.items(): + if val is not None: + os.environ[var] = val def test_sign_request_passes_explicit_aws_credentials(self): """Explicit aws_* params (e.g. from a proxy search_tools entry) reach the signer.""" From 2f342dc12d709666ca78f9d989313715d66ec216 Mon Sep 17 00:00:00 2001 From: CrypticDriver <107245892+CrypticDriver@users.noreply.github.com> Date: Wed, 22 Jul 2026 07:27:10 +0000 Subject: [PATCH 011/538] fix: honor AWS shared-config region for custom gateway hostnames MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous check only consulted AWS_REGION* env vars before rejecting custom hostnames, breaking deployments that configure their region via the AWS shared config (profile). Resolve through boto3's session (env vars + shared config) and only error when that chain yields nothing — never sign with a silently guessed region. --- litellm/llms/bedrock/search/transformation.py | 31 +++++++++++------ tests/search_tests/test_agentcore_search.py | 34 +++++++++++++++---- 2 files changed, 47 insertions(+), 18 deletions(-) diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index 3bae84f0026..9ab651f952c 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -190,11 +190,11 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): # The signing region must match the gateway's region — derive it from # standard gateway hostnames so callers don't have to set - # aws_region_name to a region different from their default. Custom or - # private hostnames can't be parsed: fall back to an explicitly - # configured region (param or AWS env vars), and error out rather than - # silently signing for a guessed region the gateway would reject with - # a confusing auth error. + # aws_region_name to a region different from their default. For custom + # or private hostnames, defer to BaseAWSLLM's normal region resolution + # (params, env vars, AWS shared config / profile); only error out when + # that chain yields nothing, rather than silently signing for a guessed + # region the gateway would reject with a confusing auth error. signing_params = dict(optional_params) if signing_params.get("aws_region_name") is None: match = re.search( @@ -203,12 +203,21 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): ) if match: signing_params["aws_region_name"] = match.group(1) - elif not any(get_secret_str(var) for var in ("AWS_REGION", "AWS_REGION_NAME", "AWS_DEFAULT_REGION")): - raise ValueError( - f"Cannot derive the SigV4 signing region from api_base '{api_base}'. " - "Set aws_region_name (or the AWS_REGION env var) to the gateway's " - "region when using a custom hostname." - ) + else: + # boto3's session resolution covers env vars AND the AWS shared + # config (profile region) — unlike BaseAWSLLM's helper, which + # silently defaults to us-west-2 when nothing is configured. + import boto3 + + configured_region = boto3.Session().region_name + if configured_region: + signing_params["aws_region_name"] = configured_region + else: + raise ValueError( + f"Cannot derive the SigV4 signing region from api_base '{api_base}' " + "or the AWS configuration chain. Set aws_region_name (or AWS_REGION / " + "a profile region) to the gateway's region when using a custom hostname." + ) # api_key="" (not None, but falsy) disables BaseAWSLLM's fallback to the # AWS_BEARER_TOKEN_BEDROCK env var: that token is a Bedrock Runtime diff --git a/tests/search_tests/test_agentcore_search.py b/tests/search_tests/test_agentcore_search.py index d5f6ab3e82f..c041bfc8d50 100644 --- a/tests/search_tests/test_agentcore_search.py +++ b/tests/search_tests/test_agentcore_search.py @@ -264,10 +264,12 @@ class TestAgentCoreSearch: assert mock_base_sign.call_args.kwargs["api_key"] == "" def test_sign_request_custom_hostname_requires_region(self): - """Non-standard hostnames can't yield a signing region — require it explicitly.""" + """Custom hostname + empty AWS config chain → clear error, no guessed region.""" config = AgentCoreSearchConfig() - saved = {var: os.environ.pop(var, None) for var in ("AWS_REGION", "AWS_REGION_NAME", "AWS_DEFAULT_REGION")} - try: + + mock_session = MagicMock() + mock_session.region_name = None # nothing configured anywhere + with patch("boto3.Session", return_value=mock_session): with pytest.raises(ValueError, match="signing region"): config.sign_request( headers={}, @@ -275,10 +277,28 @@ class TestAgentCoreSearch: request_data={"jsonrpc": "2.0"}, api_base="https://gateway.internal.example.com/mcp", ) - finally: - for var, val in saved.items(): - if val is not None: - os.environ[var] = val + + def test_sign_request_custom_hostname_uses_shared_config_region(self): + """Custom hostname + region from AWS shared config (profile) must be honored.""" + config = AgentCoreSearchConfig() + + mock_session = MagicMock() + mock_session.region_name = "eu-west-1" # e.g. from ~/.aws/config profile + with ( + patch("boto3.Session", return_value=mock_session), + patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign, + ): + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base="https://gateway.internal.example.com/mcp", + ) + assert mock_base_sign.call_args.kwargs["optional_params"]["aws_region_name"] == "eu-west-1" def test_sign_request_passes_explicit_aws_credentials(self): """Explicit aws_* params (e.g. from a proxy search_tools entry) reach the signer.""" From b43441814b8aaf51b3fab143424f6b89b80bf259 Mon Sep 17 00:00:00 2001 From: CrypticDriver <107245892+CrypticDriver@users.noreply.github.com> Date: Thu, 23 Jul 2026 15:31:10 +0000 Subject: [PATCH 012/538] test: mirror AgentCore search tests into tests/test_litellm for coverage Coverage collection runs against the sharded tests/test_litellm tree, so the provider tests living only in tests/search_tests were invisible to codecov (patch coverage reported ~31% despite the suite). Mirror them as tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py and add edge-case tests (malformed MCP content blocks, SSE without a JSON frame, notification-only streams, list request body, error-class mapping). transformation.py line coverage: 99% (26 tests x2 trees). --- tests/search_tests/test_agentcore_search.py | 55 +++ .../test_agentcore_search_transformation.py | 400 ++++++++++++++++++ 2 files changed, 455 insertions(+) create mode 100644 tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py diff --git a/tests/search_tests/test_agentcore_search.py b/tests/search_tests/test_agentcore_search.py index c041bfc8d50..578d2b0f63e 100644 --- a/tests/search_tests/test_agentcore_search.py +++ b/tests/search_tests/test_agentcore_search.py @@ -341,3 +341,58 @@ class TestAgentCoreSearch: api_base=eu_url, ) assert mock_base_sign.call_args.kwargs["optional_params"]["aws_region_name"] == "eu-central-1" + + +class TestAgentCoreSearchEdgeCases: + """Branch coverage for response parsing and error mapping.""" + + def test_transform_search_response_skips_non_text_and_bad_json_blocks(self): + """Non-text blocks and unparseable text blocks are skipped, not fatal.""" + config = AgentCoreSearchConfig() + body = { + "jsonrpc": "2.0", + "id": 1, + "result": { + "content": [ + {"type": "image", "data": "..."}, + {"type": "text", "text": "not-json"}, + {"type": "text", "text": json.dumps(["scalar", {"title": "T", "url": "u", "text": "s"}])}, + ] + }, + } + mock_response = _make_mock_response(body) + + response = config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + # only the one dict item survives; non-dict list entries are skipped + assert len(response.results) == 1 + assert response.results[0].title == "T" + + def test_parse_mcp_body_sse_without_json_frame_raises(self): + """An SSE stream carrying no parseable JSON object is a 502.""" + config = AgentCoreSearchConfig() + mock_response = _make_mock_response(text="event: ping\ndata: not-json\n\n") + with pytest.raises(Exception, match="SSE without a JSON data frame"): + config._parse_mcp_body(mock_response) + + def test_parse_mcp_body_returns_last_event_when_no_result_frame(self): + """A stream of only notifications returns the last parsed event.""" + config = AgentCoreSearchConfig() + note = {"jsonrpc": "2.0", "method": "notifications/progress"} + mock_response = _make_mock_response(text=f"data: {json.dumps(note)}\n\n") + assert config._parse_mcp_body(mock_response) == note + + def test_sign_request_rejects_list_request_body(self): + config = AgentCoreSearchConfig() + with pytest.raises(ValueError, match="single dict"): + config.sign_request( + headers={}, + optional_params={}, + request_data=[{"jsonrpc": "2.0"}], + api_base=GATEWAY_URL, + ) + + def test_get_error_class_maps_status_and_message(self): + config = AgentCoreSearchConfig() + err = config.get_error_class(error_message="boom", status_code=503, headers={}) + assert getattr(err, "status_code", None) == 503 + assert "boom" in str(err) diff --git a/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py b/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py new file mode 100644 index 00000000000..6bbaf66d3b3 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py @@ -0,0 +1,400 @@ +""" +Tests for Amazon Bedrock AgentCore Web Search integration. + +Mirror of tests/search_tests/test_agentcore_search.py placed in the +test_litellm tree so the AgentCoreSearchConfig transformation is exercised by +the sharded CI (coverage collection runs against this tree). +""" + +import json +import os + +import pytest +from unittest.mock import AsyncMock, patch, MagicMock + +import litellm +from litellm.llms.bedrock.search.transformation import AgentCoreSearchConfig + +GATEWAY_URL = "https://testgateway-abc123.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp" + +MCP_RESULTS = [ + { + "title": "Test Result 1", + "url": "https://example.com/1", + "text": "Snippet for result 1", + "publishedDate": "2026-06-16", + }, + { + "title": "Test Result 2", + "url": "https://example.com/2", + "text": "Snippet for result 2", + }, +] + + +def _mcp_response_body() -> dict: + return { + "jsonrpc": "2.0", + "id": 1, + "result": {"content": [{"type": "text", "text": json.dumps(MCP_RESULTS)}]}, + } + + +def _make_mock_response(json_body: dict = None, text: str = None) -> MagicMock: + mock_response = MagicMock() + mock_response.status_code = 200 + if text is not None: + mock_response.text = text + else: + mock_response.text = json.dumps(json_body) + mock_response.json.return_value = json_body + return mock_response + + +class TestAgentCoreSearch: + """ + Tests for AgentCore Web Search functionality with mocked network/signing. + """ + + @pytest.mark.asyncio + async def test_agentcore_search_request_payload(self): + """Validates the MCP tools/call payload and SigV4 signing without real AWS calls.""" + os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL + + mock_response = _make_mock_response(_mcp_response_body()) + + with ( + patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post, + patch.object( + AgentCoreSearchConfig, + "_sign_request", + return_value=( + {"Authorization": "AWS4-HMAC-SHA256 test", "Content-Type": "application/json"}, + json.dumps({"signed": True}).encode(), + ), + ) as mock_sign, + ): + mock_post.return_value = mock_response + + response = await litellm.asearch( + query="latest developments in AI", + search_provider="agentcore", + max_results=5, + ) + + mock_post.assert_called_once() + call_kwargs = mock_post.call_args.kwargs + assert call_kwargs["url"] == GATEWAY_URL + # Signed body must be sent verbatim + assert call_kwargs["data"] == json.dumps({"signed": True}).encode() + assert "json" not in call_kwargs + + # Signing was invoked with the MCP request + mock_sign.assert_called_once() + sign_kwargs = mock_sign.call_args.kwargs + request_data = sign_kwargs["request_data"] + assert request_data["method"] == "tools/call" + assert request_data["params"]["name"] == "web-search-tool___WebSearch" + assert request_data["params"]["arguments"]["query"] == "latest developments in AI" + assert request_data["params"]["arguments"]["maxResults"] == 5 + assert sign_kwargs["service_name"] == "bedrock-agentcore" + + assert len(response.results) == 2 + assert response.results[0].title == "Test Result 1" + assert response.results[0].url == "https://example.com/1" + assert response.results[0].snippet == "Snippet for result 1" + assert response.results[0].date == "2026-06-16" + + def test_transform_search_request_query_truncation(self): + """AgentCore rejects queries > 200 chars; the request must truncate.""" + config = AgentCoreSearchConfig() + long_query = "a" * 300 + data = config.transform_search_request(query=long_query, optional_params={}) + assert len(data["params"]["arguments"]["query"]) == 200 + + def test_transform_search_request_joins_list_queries(self): + config = AgentCoreSearchConfig() + data = config.transform_search_request(query=["foo", "bar"], optional_params={}) + assert data["params"]["arguments"]["query"] == "foo bar" + + def test_transform_search_request_custom_tool_name(self): + config = AgentCoreSearchConfig() + data = config.transform_search_request(query="q", optional_params={"tool_name": "my-target___WebSearch"}) + assert data["params"]["name"] == "my-target___WebSearch" + + def test_transform_search_request_rejects_non_websearch_tool_name(self): + """A caller-supplied tool_name must not reach other tools on the gateway.""" + config = AgentCoreSearchConfig() + with pytest.raises(ValueError, match="must end with"): + config.transform_search_request(query="q", optional_params={"tool_name": "admin-target___DeleteUser"}) + + def test_transform_search_request_sends_documented_default_max_results(self): + """The documented default of 10 is sent explicitly, not left to the gateway.""" + config = AgentCoreSearchConfig() + data = config.transform_search_request(query="q", optional_params={}) + assert data["params"]["arguments"]["maxResults"] == 10 + + def test_get_complete_url_requires_gateway_url(self): + config = AgentCoreSearchConfig() + os.environ.pop("AGENTCORE_GATEWAY_URL", None) + with pytest.raises(ValueError, match="AGENTCORE_GATEWAY_URL"): + config.get_complete_url(api_base=None, optional_params={}) + + def test_get_complete_url_prefers_api_base(self): + config = AgentCoreSearchConfig() + assert config.get_complete_url(api_base=GATEWAY_URL, optional_params={}) == GATEWAY_URL + + def test_validate_environment_sets_mcp_headers(self): + """MCP Streamable HTTP requires accepting both JSON and SSE.""" + config = AgentCoreSearchConfig() + headers = config.validate_environment(headers={}) + assert headers["Accept"] == "application/json, text/event-stream" + assert headers["Content-Type"] == "application/json" + + def test_transform_search_response_parses_sse_frame(self): + """Gateway may answer with an SSE-framed JSON-RPC message.""" + config = AgentCoreSearchConfig() + body = _mcp_response_body() + sse_text = f"event: message\ndata: {json.dumps(body)}\n\n" + mock_response = _make_mock_response(text=sse_text) + + response = config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + assert len(response.results) == 2 + assert response.results[1].url == "https://example.com/2" + + def test_transform_search_response_parses_multiline_sse_data(self): + """SSE data may be split across several data: lines (joined per spec).""" + config = AgentCoreSearchConfig() + pretty = json.dumps(_mcp_response_body(), indent=2) + sse_text = "event: message\n" + "\n".join(f"data: {line}" for line in pretty.splitlines()) + "\n\n" + mock_response = _make_mock_response(text=sse_text) + + response = config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + assert len(response.results) == 2 + + def test_transform_search_response_skips_progress_events(self): + """A progress notification before the JSON-RPC result must not shadow it.""" + config = AgentCoreSearchConfig() + progress = {"jsonrpc": "2.0", "method": "notifications/progress", "params": {"progress": 1}} + sse_text = ( + f"event: message\ndata: {json.dumps(progress)}\n\n" + f"event: message\ndata: {json.dumps(_mcp_response_body())}\n\n" + ) + mock_response = _make_mock_response(text=sse_text) + + response = config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + assert len(response.results) == 2 + + def test_transform_search_response_raises_on_mcp_error(self): + config = AgentCoreSearchConfig() + mock_response = _make_mock_response( + {"jsonrpc": "2.0", "id": 1, "error": {"code": -32601, "message": "tool not found"}} + ) + with pytest.raises(Exception, match="tool not found"): + config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + + def test_sign_request_uses_bearer_token_when_api_key_set(self): + """CUSTOM_JWT gateways: api_key is sent as a bearer token, no SigV4.""" + config = AgentCoreSearchConfig() + request_data = {"jsonrpc": "2.0", "id": 1} + + headers, signed_body = config.sign_request( + headers={"Content-Type": "application/json"}, + optional_params={}, + request_data=request_data, + api_base=GATEWAY_URL, + api_key="test-jwt-token", + ) + assert headers["Authorization"] == "Bearer test-jwt-token" + assert signed_body == json.dumps(request_data).encode() + + def test_sign_request_uses_bearer_token_from_env(self): + """Server token is attached when the request targets the configured gateway host.""" + config = AgentCoreSearchConfig() + os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" + os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL + try: + headers, _ = config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base=GATEWAY_URL, + ) + assert headers["Authorization"] == "Bearer env-jwt-token" + finally: + os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) + os.environ.pop("AGENTCORE_GATEWAY_URL", None) + + def test_sign_request_refuses_server_token_to_untrusted_host(self): + """Server-managed token must not be sent to a caller-chosen api_base.""" + config = AgentCoreSearchConfig() + os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" + os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL + try: + with pytest.raises(ValueError, match="Refusing to send"): + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base="https://attacker.example.com/mcp", + ) + finally: + os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) + os.environ.pop("AGENTCORE_GATEWAY_URL", None) + + def test_sign_request_does_not_leak_bedrock_bearer_token(self): + """AWS_BEARER_TOKEN_BEDROCK is a Bedrock Runtime credential — it must not + replace SigV4 on requests to an AgentCore gateway.""" + config = AgentCoreSearchConfig() + + with patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign: + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base=GATEWAY_URL, + ) + # api_key="" (falsy, not None) disables the base class's + # AWS_BEARER_TOKEN_BEDROCK env fallback. + assert mock_base_sign.call_args.kwargs["api_key"] == "" + + def test_sign_request_custom_hostname_requires_region(self): + """Custom hostname + empty AWS config chain → clear error, no guessed region.""" + config = AgentCoreSearchConfig() + + mock_session = MagicMock() + mock_session.region_name = None # nothing configured anywhere + with patch("boto3.Session", return_value=mock_session): + with pytest.raises(ValueError, match="signing region"): + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base="https://gateway.internal.example.com/mcp", + ) + + def test_sign_request_custom_hostname_uses_shared_config_region(self): + """Custom hostname + region from AWS shared config (profile) must be honored.""" + config = AgentCoreSearchConfig() + + mock_session = MagicMock() + mock_session.region_name = "eu-west-1" # e.g. from ~/.aws/config profile + with ( + patch("boto3.Session", return_value=mock_session), + patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign, + ): + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base="https://gateway.internal.example.com/mcp", + ) + assert mock_base_sign.call_args.kwargs["optional_params"]["aws_region_name"] == "eu-west-1" + + def test_sign_request_passes_explicit_aws_credentials(self): + """Explicit aws_* params (e.g. from a proxy search_tools entry) reach the signer.""" + config = AgentCoreSearchConfig() + + with patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign: + config.sign_request( + headers={}, + optional_params={ + "aws_access_key_id": "AKIATEST", + "aws_secret_access_key": "secret", + "aws_session_token": "token", + }, + request_data={"jsonrpc": "2.0"}, + api_base=GATEWAY_URL, + ) + passed = mock_base_sign.call_args.kwargs["optional_params"] + assert passed["aws_access_key_id"] == "AKIATEST" + assert passed["aws_secret_access_key"] == "secret" + assert passed["aws_session_token"] == "token" + + def test_sign_request_derives_region_from_gateway_url(self): + """Signing region must come from the gateway URL, not the caller's default region.""" + config = AgentCoreSearchConfig() + eu_url = "https://gw-x.gateway.bedrock-agentcore.eu-central-1.amazonaws.com/mcp" + + with patch.object( + AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM + "_sign_request", + return_value=({}, b"{}"), + ) as mock_base_sign: + config.sign_request( + headers={}, + optional_params={}, + request_data={"jsonrpc": "2.0"}, + api_base=eu_url, + ) + assert mock_base_sign.call_args.kwargs["optional_params"]["aws_region_name"] == "eu-central-1" + + +class TestAgentCoreSearchEdgeCases: + """Branch coverage for response parsing and error mapping.""" + + def test_transform_search_response_skips_non_text_and_bad_json_blocks(self): + """Non-text blocks and unparseable text blocks are skipped, not fatal.""" + config = AgentCoreSearchConfig() + body = { + "jsonrpc": "2.0", + "id": 1, + "result": { + "content": [ + {"type": "image", "data": "..."}, + {"type": "text", "text": "not-json"}, + {"type": "text", "text": json.dumps(["scalar", {"title": "T", "url": "u", "text": "s"}])}, + ] + }, + } + mock_response = _make_mock_response(body) + + response = config.transform_search_response(raw_response=mock_response, logging_obj=MagicMock()) + # only the one dict item survives; non-dict list entries are skipped + assert len(response.results) == 1 + assert response.results[0].title == "T" + + def test_parse_mcp_body_sse_without_json_frame_raises(self): + """An SSE stream carrying no parseable JSON object is a 502.""" + config = AgentCoreSearchConfig() + mock_response = _make_mock_response(text="event: ping\ndata: not-json\n\n") + with pytest.raises(Exception, match="SSE without a JSON data frame"): + config._parse_mcp_body(mock_response) + + def test_parse_mcp_body_returns_last_event_when_no_result_frame(self): + """A stream of only notifications returns the last parsed event.""" + config = AgentCoreSearchConfig() + note = {"jsonrpc": "2.0", "method": "notifications/progress"} + mock_response = _make_mock_response(text=f"data: {json.dumps(note)}\n\n") + assert config._parse_mcp_body(mock_response) == note + + def test_sign_request_rejects_list_request_body(self): + config = AgentCoreSearchConfig() + with pytest.raises(ValueError, match="single dict"): + config.sign_request( + headers={}, + optional_params={}, + request_data=[{"jsonrpc": "2.0"}], + api_base=GATEWAY_URL, + ) + + def test_get_error_class_maps_status_and_message(self): + config = AgentCoreSearchConfig() + err = config.get_error_class(error_message="boom", status_code=503, headers={}) + assert getattr(err, "status_code", None) == 503 + assert "boom" in str(err) From e1629b77dbce0e9eaddbdb726e13ed7edd614ba9 Mon Sep 17 00:00:00 2001 From: CrypticDriver <107245892+CrypticDriver@users.noreply.github.com> Date: Sun, 26 Jul 2026 07:26:46 +0000 Subject: [PATCH 013/538] fix(interactions): sync queued status enum from #34318 to unblock CI on stale daily branch --- litellm/types/interactions/generated.py | 2 ++ tests/test_litellm/interactions/test_openapi_compliance.py | 1 + 2 files changed, 3 insertions(+) diff --git a/litellm/types/interactions/generated.py b/litellm/types/interactions/generated.py index 793cc02ff17..4a1ef5ed696 100644 --- a/litellm/types/interactions/generated.py +++ b/litellm/types/interactions/generated.py @@ -173,6 +173,7 @@ class Status1(Enum): cancelled = "cancelled" incomplete = "incomplete" budget_exceeded = "budget_exceeded" + queued = "queued" class InteractionStatusUpdate(BaseModel): @@ -341,6 +342,7 @@ class Status3(Enum): CANCELLED = "cancelled" INCOMPLETE = "incomplete" BUDGET_EXCEEDED = "budget_exceeded" + QUEUED = "queued" class ModelOption(RootModel[str]): diff --git a/tests/test_litellm/interactions/test_openapi_compliance.py b/tests/test_litellm/interactions/test_openapi_compliance.py index 209e99895db..11b08fa45a8 100644 --- a/tests/test_litellm/interactions/test_openapi_compliance.py +++ b/tests/test_litellm/interactions/test_openapi_compliance.py @@ -194,6 +194,7 @@ class TestResponseCompliance: "cancelled", "incomplete", "budget_exceeded", + "queued", ] assert status_prop["enum"] == expected_statuses print(f"✓ Status enum values: {expected_statuses}") From 7e1f44f0cc15e639d8fe95fc8831c5a0b25cc76b Mon Sep 17 00:00:00 2001 From: mateo Date: Thu, 30 Jul 2026 02:38:28 +0000 Subject: [PATCH 014/538] 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/538] 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/538] 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 9dbe61aa6d8915f40a030ee11f522345afb9cc6a Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Fri, 31 Jul 2026 16:49:07 +0530 Subject: [PATCH 017/538] feat(proxy): add project-level ITPM and OTPM quotas Add model_itpm_limit and model_otpm_limit to project create and update requests, storing both quota maps in project metadata without a database migration Reserve input and output tokens independently before provider dispatch, expose separate project rate-limit headers, and reconcile counters across successful calls, failures, retries, fallbacks, streaming, caching, and cancellation Harden token estimation for pre-tokenized embeddings, multimodal inputs, Responses API requests, native Gemini requests, multiple candidates, and conflicting output-cap aliases Reject negative output caps, preserve conservative reservations when usage is missing or zero, bind reconciliation and refunds to the reservation window, prevent double refunds or negative counters, update generated API types, and add regression coverage --- litellm/proxy/_types.py | 8 +- litellm/proxy/auth/auth_utils.py | 2 +- .../proxy/hooks/dynamic_rate_limiter_v3.py | 4 +- .../hooks/parallel_request_limiter_v3.py | 1561 ++++++++++- .../hooks/test_parallel_request_limiter_v3.py | 132 +- .../proxy/hooks/test_tpm_concurrent.py | 2444 ++++++++++++++++- tests/test_litellm/proxy/test_proxy_types.py | 18 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 16 + 8 files changed, 4095 insertions(+), 90 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7d6829aca70..b4aab95e2d4 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1,7 +1,7 @@ import enum import json import os -from collections.abc import Callable +from collections.abc import Callable, Mapping from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal, Union @@ -2903,6 +2903,8 @@ class NewProjectRequest(LiteLLM_BudgetTable): models: list[str] = [] model_rpm_limit: dict | None = None model_tpm_limit: dict | None = None + model_itpm_limit: Mapping[str, int] | None = None + model_otpm_limit: Mapping[str, int] | None = None blocked: bool = False object_permission: LiteLLM_ObjectPermissionBase | None = None @@ -2935,6 +2937,8 @@ class UpdateProjectRequest(LiteLLM_BudgetTable): models: list[str] | None = None model_rpm_limit: dict | None = None model_tpm_limit: dict | None = None + model_itpm_limit: Mapping[str, int] | None = None + model_otpm_limit: Mapping[str, int] | None = None blocked: bool | None = None budget_id: str | None = None object_permission: LiteLLM_ObjectPermissionBase | None = None @@ -4072,6 +4076,8 @@ class PassThroughEndpointLoggingTypedDict(TypedDict): LiteLLM_ManagementEndpoint_MetadataFields: Final = [ "model_rpm_limit", "model_tpm_limit", + "model_itpm_limit", + "model_otpm_limit", "mcp_rpm_limit", "tag_rpm_limit", "rpm_limit_type", diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 87db8ed3b5e..21c41e08ecc 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -994,7 +994,7 @@ def get_key_model_tpm_limit( def get_model_rate_limit_from_metadata( user_api_key_dict: UserAPIKeyAuth, metadata_accessor_key: Literal["team_metadata", "organization_metadata", "project_metadata"], - rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"], + rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit", "model_itpm_limit", "model_otpm_limit"], ) -> dict[str, int] | None: if getattr(user_api_key_dict, metadata_accessor_key): return getattr(user_api_key_dict, metadata_accessor_key).get(rate_limit_key) diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 4492f42782c..de8834449de 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -454,7 +454,9 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): parent_otel_span=user_api_key_dict.parent_otel_span, ) - verbose_proxy_logger.debug("Atomic check+increment response: %s", json.dumps(atomic_response, indent=2)) + verbose_proxy_logger.debug( + "Atomic check+increment response: %s", json.dumps(atomic_response, indent=2, default=list) + ) if atomic_response["overall_code"] == "OVER_LIMIT": resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(model) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 2725da1ee12..e6cb54393ab 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -8,11 +8,22 @@ import asyncio import binascii import os import uuid -from collections.abc import Callable +from collections.abc import Callable, Mapping, Sequence, Set from contextvars import ContextVar from dataclasses import dataclass, field from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Final, + Literal, + Protocol, + TypedDict, + Union, + cast, +) + +from typing_extensions import NotRequired from litellm import DualCache from litellm._logging import verbose_proxy_logger @@ -34,11 +45,12 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ( ) from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit from litellm.types.caching import RedisPipelineIncrementOperation -from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject +from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage from litellm.types.utils import ( CallTypes, EmbeddingResponse, ModelResponse, + RerankResponse, TextCompletionResponse, Usage, ) @@ -109,7 +121,8 @@ CHECK_AND_INCREMENT_BY_N_SCRIPT: Final = """ -- ARGV[(i-1)*4 + 3] = ttl_seconds (counter TTL when window resets) -- ARGV[(i-1)*4 + 4] = window_size_seconds (sliding-window length) -- --- Return on success: { 0, new_counter_1, new_counter_2, ... } +-- Return on success: +-- { 0, new_counter_1, window_start_1, new_counter_2, window_start_2, ... } -- Return on over-limit: { 1, descriptor_index, current_counter, limit } local time_reply = redis.call('TIME') local now = tonumber(time_reply[1]) @@ -146,7 +159,7 @@ for i = 1, descriptor_count do return { 1, i, current_counter, limit } end - descriptor_state[i] = { window_expired, current_counter } + descriptor_state[i] = { window_expired, current_counter, window_start } end -- Pass 2: all checks passed. Apply increments. @@ -160,8 +173,10 @@ for i = 1, descriptor_count do local window_size = tonumber(ARGV[arg_base + 3]) local window_expired = descriptor_state[i][1] + local active_window_start if window_expired then + active_window_start = now redis.call('SET', window_key, tostring(now)) redis.call('SET', counter_key, increment) redis.call('EXPIRE', window_key, window_size) @@ -170,6 +185,7 @@ for i = 1, descriptor_count do end table.insert(results, increment) else + active_window_start = tonumber(descriptor_state[i][3]) local new_counter = redis.call('INCRBY', counter_key, increment) local current_ttl = redis.call('TTL', counter_key) if current_ttl == -1 and ttl > 0 then @@ -177,11 +193,39 @@ for i = 1, descriptor_count do end table.insert(results, new_counter) end + table.insert(results, active_window_start) end return results """ +WINDOW_GUARDED_TOKEN_INCREMENT_SCRIPT: Final = """ +local results = {} +for i = 1, #KEYS, 2 do + local window_key = KEYS[i] + local counter_key = KEYS[i + 1] + local arg_base = ((i - 1) / 2) * 3 + 1 + local expected_window_start = ARGV[arg_base] + local increment = tonumber(ARGV[arg_base + 1]) + local ttl = tonumber(ARGV[arg_base + 2]) + local active_window_start = redis.call('GET', window_key) + + if active_window_start and active_window_start == expected_window_start then + local new_counter = redis.call('INCRBY', counter_key, increment) + local current_ttl = redis.call('TTL', counter_key) + if current_ttl == -1 and ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, 1) + table.insert(results, new_counter) + else + table.insert(results, 0) + table.insert(results, tonumber(redis.call('GET', counter_key) or 0)) + end +end +return results +""" + PARALLEL_ACQUIRE_SCRIPT: Final = """ -- Atomic check-and-acquire for the max_parallel_requests concurrency gauge. -- Each gauge key is a sorted set of per-request slot ids scored by acquire @@ -286,6 +330,38 @@ DEFAULT_CHARS_PER_TOKEN: Final = 4 # (baseline floor) and to the smallest configured TPM limit (capped floor for # small per-tenant TPM caps). _TPM_FLOOR_FRACTION: Final = 4 +# Both embeddings and the Responses API put their prompt in data["input"], +# but only embeddings have no output tokens. Every "is this an embedding" +# check on data["input"] must exclude these call types, or a Responses call +# gets misclassified as an embedding and skips output-token reservation/caps. +RESPONSES_API_CALL_TYPES: Final = ("aresponses", "responses") +EMBEDDING_API_CALL_TYPES: Final = ("aembedding", "embedding") +TEXT_COMPLETION_API_CALL_TYPES: Final = ("atext_completion", "text_completion") +RERANK_API_CALL_TYPES: Final = (CallTypes.rerank.value, CallTypes.arerank.value) +GOOGLE_GENAI_NATIVE_CALL_TYPES: Final = ( + CallTypes.generate_content.value, + CallTypes.agenerate_content.value, + CallTypes.generate_content_stream.value, + CallTypes.agenerate_content_stream.value, +) +RESPONSES_API_MIN_OUTPUT_TOKENS: Final = 16 +# litellm.token_counter has no per-type handling for "input_audio" content +# blocks (unlike images, which use use_default_image_token_count) -- it +# silently contributes 0 tokens for them. When the block carries a base64 +# payload, the estimate is derived from the decoded byte count; when the +# block is a reference without a payload (or the payload is missing), this +# flat per-block floor is used instead. +DEFAULT_AUDIO_TOKEN_ESTIMATE: Final = 300 +# Conservative bytes-per-token assumption for size-based audio estimation: +# equivalent to 8 kHz mono PCM-16 (16 000 bytes/s) at 10 tokens/s. Choosing +# the lowest reasonable bitrate means we never under-reserve for higher- +# quality audio recorded at the same wall-clock duration. +_AUDIO_BYTES_PER_TOKEN: Final = 1600 +# Descriptor "key" values for project-scoped ITPM/OTPM. Distinct from +# "model_per_project" (the combined-TPM descriptor) so both can be enforced +# on the same project+model simultaneously without colliding on cache keys. +PROJECT_ITPM_DESCRIPTOR_KEY: Final = "model_per_project_itpm" +PROJECT_OTPM_DESCRIPTOR_KEY: Final = "model_per_project_otpm" # How long an acquired slot counts toward the in-flight total before it is # considered leaked (worker crashed without any release callback firing) and # pruned. Also the longest request duration the gauge can track: a request @@ -328,6 +404,13 @@ class RateLimitStatus(TypedDict): class RateLimitResponse(TypedDict): overall_code: str statuses: list[RateLimitStatus] + reservation_windows: NotRequired[frozenset[tuple[str, str, Literal["redis", "local"]]]] + + +class ReservationAwareIncrementOperation(RedisPipelineIncrementOperation): + window_key: NotRequired[str] + expected_window_start: NotRequired[str] + reservation_backend: NotRequired[Literal["redis", "local"]] class RateLimitResponseWithDescriptors(TypedDict): @@ -335,6 +418,10 @@ class RateLimitResponseWithDescriptors(TypedDict): response: RateLimitResponse +class _RateLimitDescriptorSink(Protocol): + def append(self, descriptor: RateLimitDescriptor, /) -> None: ... + + @dataclass(slots=True) class RequestRateLimiterStash: """ @@ -364,6 +451,16 @@ class RequestRateLimiterStash: reserved_tokens: int = 0 reserved_model: str | None = None reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset) + itpm_reserved_tokens: int = 0 + itpm_reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset) + itpm_reserved_window_identities: frozenset[tuple[str, str, Literal["redis", "local"]]] = field( + default_factory=frozenset + ) + otpm_reserved_tokens: int = 0 + otpm_reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset) + otpm_reserved_window_identities: frozenset[tuple[str, str, Literal["redis", "local"]]] = field( + default_factory=frozenset + ) reservation_released: bool = False @@ -426,6 +523,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self.check_and_increment_by_n_script = ( self.internal_usage_cache.dual_cache.redis_cache.async_register_script(CHECK_AND_INCREMENT_BY_N_SCRIPT) ) + self.window_guarded_token_increment_script = ( + self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + WINDOW_GUARDED_TOKEN_INCREMENT_SCRIPT + ) + ) self.parallel_acquire_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( PARALLEL_ACQUIRE_SCRIPT ) @@ -439,6 +541,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self.batch_rate_limiter_script = None self.token_increment_script = None self.check_and_increment_by_n_script = None + self.window_guarded_token_increment_script = None self.parallel_acquire_script = None self.parallel_release_script = None self.parallel_count_script = None @@ -505,18 +608,163 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return baseline return min(baseline, max(1, min_configured_tpm_limit // _TPM_FLOOR_FRACTION)) + @staticmethod + def _is_embedding_request(data: object, call_type: str | None) -> bool: + if call_type in EMBEDDING_API_CALL_TYPES: + return True + if call_type in RESPONSES_API_CALL_TYPES: + return False + if call_type: + return False + if not isinstance(data, dict): + return False + return data.get("input") is not None + + @staticmethod + def _translate_google_genai_native_request( + data: object, + call_type: str | None, + ) -> Mapping[str, object] | None: + contents = data.get("contents") if isinstance(data, dict) else None + if ( + not isinstance(data, dict) + or call_type not in GOOGLE_GENAI_NATIVE_CALL_TYPES + or not isinstance(contents, (dict, list)) + ): + return None + from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter + + config = data.get("config") if "config" in data else data.get("generationConfig") + translated_request = GoogleGenAIAdapter().translate_generate_content_to_completion( + model=data.get("model") if isinstance(data.get("model"), str) else "", + contents=contents, + config=config if isinstance(config, dict) else None, + systemInstruction=data.get("systemInstruction"), + system_instruction=data.get("system_instruction"), + tools=data.get("tools"), + toolConfig=data.get("toolConfig"), + tool_config=data.get("tool_config"), + ) + return translated_request + + @staticmethod + def _get_explicit_output_cap(data: object, call_type: str | None) -> int | None: + if not isinstance(data, dict): + return None + if call_type in GOOGLE_GENAI_NATIVE_CALL_TYPES: + config = data.get("config") if "config" in data else data.get("generationConfig") + values = tuple( + int(config[field]) + for field in ("maxOutputTokens", "max_output_tokens") + if isinstance(config, dict) and isinstance(config.get(field), (int, float, str)) + ) + return max(values, default=None) + if call_type in RESPONSES_API_CALL_TYPES: + value = data.get("max_output_tokens") + if value is None: + return None + if not isinstance(value, (int, float, str)): + return None + return max(RESPONSES_API_MIN_OUTPUT_TOKENS, int(value)) + if call_type in EMBEDDING_API_CALL_TYPES: + return None + fields = ( + ("max_tokens", "max_completion_tokens") + if call_type + else ("max_tokens", "max_completion_tokens", "max_output_tokens") + ) + values = tuple(int(data[field]) for field in fields if isinstance(data.get(field), (int, float, str))) + return max(values, default=None) + + @classmethod + def _has_explicit_output_cap(cls, data: object, call_type: str | None) -> bool: + """Whether the caller explicitly set an output-token cap. + + Checked via ``is not None`` (not truthiness) so an explicit 0 -- + a legitimate zero-output request -- counts as explicit. + """ + return cls._get_explicit_output_cap(data, call_type) is not None + + @staticmethod + def _get_output_candidate_count(data: object, call_type: str | None = None) -> int: + if not isinstance(data, dict): + return 1 + config = ( + (data.get("config") if "config" in data else data.get("generationConfig")) + if call_type in GOOGLE_GENAI_NATIVE_CALL_TYPES + else None + ) + candidate_values = ( + data.get("n"), + data.get("best_of"), + config.get("candidateCount") if isinstance(config, dict) else None, + config.get("candidate_count") if isinstance(config, dict) else None, + ) + candidate_count = 1 + for value in candidate_values: + try: + candidate_count = max(candidate_count, int(value or 1)) + except (TypeError, ValueError): + continue + return candidate_count + + @staticmethod + def _apply_implicit_output_cap( + data: object, + min_configured_limit: int | None, + call_type: str | None, + ) -> None: + """Hard-cap generation length when the request has no explicit cap. + + Guards against an unbounded response overshooting a small TPM/OTPM + budget before post-call reconciliation runs. Skips requests that + already set an explicit cap and embeddings, which have no generation + budget. The Responses API only honors ``max_output_tokens`` (its + underlying chat-completion transformation ignores ``max_tokens``), so + the cap must be written to that field for Responses call types. + """ + if not isinstance(data, dict): + return + capped_floor = _PROXY_MaxParallelRequestsHandler_v3._no_max_tokens_output_floor(min_configured_limit) + if call_type in RESPONSES_API_CALL_TYPES: + capped_floor = max(capped_floor, RESPONSES_API_MIN_OUTPUT_TOKENS) + baseline_floor = DEFAULT_MAX_TOKENS_ESTIMATE // _TPM_FLOOR_FRACTION + is_embedding = _PROXY_MaxParallelRequestsHandler_v3._is_embedding_request(data, call_type) + if ( + capped_floor >= baseline_floor + or _PROXY_MaxParallelRequestsHandler_v3._has_explicit_output_cap(data, call_type) + or is_embedding + ): + return + if call_type in GOOGLE_GENAI_NATIVE_CALL_TYPES: + config_field = "config" if "config" in data or "generationConfig" not in data else "generationConfig" + config = data.get(config_field) + if config is None or isinstance(config, dict): + data[config_field] = { # mutable-ok: downstream native routing requires a mutable request config + **(config or {}), # mutable-ok: downstream native routing requires a mutable request config + "maxOutputTokens": capped_floor, + } + return + cap_field = "max_output_tokens" if call_type in RESPONSES_API_CALL_TYPES else "max_tokens" + existing_cap = data.get(cap_field) + if existing_cap is None or capped_floor < existing_cap: + data[cap_field] = capped_floor + def _estimate_tokens_for_request( self, data: dict, model: str | None = None, min_configured_tpm_limit: int | None = None, + call_type: str | None = None, ) -> int: """ Estimate total tokens this request will consume so we can reserve them upfront (input + output budget): estimated = input_tokens + max_tokens. - Supports chat (messages), completions (prompt), and embeddings (input). + Supports chat (messages), completions (prompt), embeddings (input), + and the Responses API (also `input`, disambiguated from embeddings + via ``call_type``). ``min_configured_tpm_limit`` is the smallest ``tokens_per_unit`` among the TPM-bearing descriptors this request will be charged against. When @@ -524,34 +772,87 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): fraction of that limit so small TPM caps remain usable. Omit to preserve the unconstrained floor. """ - messages = data.get("messages") - prompt: Final = data.get("prompt") - input_text: Final = data.get("input") # embeddings + estimated_input_tokens, max_tokens_estimate = self._estimate_input_and_output_tokens( + data=data, + min_configured_tpm_limit=min_configured_tpm_limit, + call_type=call_type, + ) + total_estimated: Final = estimated_input_tokens + max_tokens_estimate + + verbose_proxy_logger.debug( + "TPM reservation estimate: input=%s, max_tokens=%s, total=%s", + estimated_input_tokens, + max_tokens_estimate, + total_estimated, + ) + + return total_estimated + + def _estimate_input_and_output_tokens( + self, + data: object, + min_configured_tpm_limit: int | None = None, + call_type: str | None = None, + ) -> tuple[int, int]: + """ + Estimate input tokens and output (max_tokens) budget separately, so + callers needing independent ITPM/OTPM reservations (rather than one + combined TPM reservation) can use each half on its own. + + ``min_configured_tpm_limit`` is the smallest ``tokens_per_unit`` among + the TPM-bearing descriptors this request will be charged against. When + provided, the no-``max_tokens`` output-budget floor is capped at a + fraction of that limit so small TPM caps remain usable. Omit to + preserve the unconstrained floor. + + ``call_type`` disambiguates embeddings from the Responses API: both + put their prompt in ``data["input"]``, but only embeddings have no + output tokens. Unset (the default) preserves the historical + "any `input` means zero output" behavior for callers that don't have + a call type to pass. + """ + if not isinstance(data, dict): + return 0, 0 + translated_data: Final = self._translate_google_genai_native_request(data, call_type) + estimable_data: Final = translated_data if translated_data is not None else data + messages = estimable_data.get("messages") + prompt = estimable_data.get("prompt") + input_text = estimable_data.get("input") + + if call_type in RESPONSES_API_CALL_TYPES or call_type in EMBEDDING_API_CALL_TYPES: + messages = None + prompt = None + elif call_type in TEXT_COMPLETION_API_CALL_TYPES: + messages = None + input_text = None + elif call_type: + prompt = None + input_text = None match (messages, prompt, input_text): - case (messages, _, _) if messages: - total_chars = len(get_str_from_messages(messages)) - case (_, str() as p, _): - total_chars = len(p) - case (_, list() as p, _): - total_chars = sum(len(str(item)) for item in p) - case (_, _, str() as t): - total_chars = len(t) - case (_, _, list() as t): - total_chars = sum(len(str(item)) for item in t) + case (selected_messages, _, _) if selected_messages: + total_chars = len(get_str_from_messages(selected_messages)) + case (_, str() as selected_prompt, _): + total_chars = len(selected_prompt) + case (_, list() as selected_prompt, _): + total_chars = sum(len(str(item)) for item in selected_prompt) + case (_, _, str() as selected_input): + total_chars = len(selected_input) + case (_, _, list() as selected_input): + total_chars = sum(len(str(item)) for item in selected_input) case _: total_chars = 0 estimated_input_tokens: Final = max(1, total_chars // DEFAULT_CHARS_PER_TOKEN) if total_chars > 0 else 0 - explicit_max_tokens: Final = data.get("max_tokens") or data.get("max_completion_tokens") + explicit_max_tokens: Final = self._get_explicit_output_cap(data, call_type) + is_embedding: Final = self._is_embedding_request(data, call_type) - match (explicit_max_tokens, input_text): - case (mt, _) if mt is not None: - max_tokens_estimate = int(mt) - case (_, embeddings_input) if embeddings_input: - # Embeddings have no output tokens + match (explicit_max_tokens, is_embedding): + case (_, True): max_tokens_estimate = 0 + case (mt, _) if mt is not None: + max_tokens_estimate = mt case _ if total_chars == 0: # Fully contentless request (no messages, prompt, or input). # Don't apply the conservative output-budget floor here — it @@ -566,20 +867,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # the smallest TPM limit this request will be charged against, # so a small per-tenant TPM cap can't be tripped by the floor # alone. - output_floor: Final = self._no_max_tokens_output_floor(min_configured_tpm_limit) + output_floor = self._no_max_tokens_output_floor(min_configured_tpm_limit) + if call_type in RESPONSES_API_CALL_TYPES: + output_floor = max(output_floor, RESPONSES_API_MIN_OUTPUT_TOKENS) max_tokens_estimate = max(estimated_input_tokens, output_floor) - total_estimated: Final = estimated_input_tokens + max_tokens_estimate - - verbose_proxy_logger.debug( - "TPM reservation estimate: input=%s, max_tokens=%s (explicit=%s), total=%s", - estimated_input_tokens, - max_tokens_estimate, - explicit_max_tokens is not None, - total_estimated, - ) - - return total_estimated + max_tokens_estimate *= self._get_output_candidate_count(data, call_type) + return estimated_input_tokens, max_tokens_estimate def _is_redis_cluster(self) -> bool: """ @@ -1367,6 +1661,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): descriptor i, refund descriptors 0..i-1's increments. On Lua failure mid-loop, refund applied increments and fall back to in-memory. """ + if not descriptor_groups: + return RateLimitResponse(overall_code="OK", statuses=[]) applied: Final[list[list[dict[str, Any]]]] = [] statuses: Final[list[RateLimitStatus]] = [] @@ -1400,10 +1696,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if response["overall_code"] == "OVER_LIMIT": await self._refund_applied_descriptor_groups(applied) return response + if len(descriptor_groups) == 1: + return response applied.append(meta) statuses.extend(response["statuses"]) - return RateLimitResponse(overall_code="OK", statuses=statuses) + return RateLimitResponse( + overall_code="OK", + statuses=statuses, + reservation_windows=frozenset(), + ) async def _refund_applied_descriptor_groups( self, @@ -1471,7 +1773,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) statuses: Final[list[RateLimitStatus]] = [] - for meta, new_counter in zip(per_counter_meta, raw[1:]): + for index, meta in enumerate(per_counter_meta): + new_counter = raw[1 + index * 2] statuses.append( RateLimitStatus( code="OK", @@ -1481,7 +1784,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): descriptor_key=meta["descriptor_key"], ) ) - return RateLimitResponse(overall_code="OK", statuses=statuses) + return RateLimitResponse( + overall_code="OK", + statuses=statuses, + reservation_windows=frozenset( + ( + meta["counter_key"], + str(int(raw[2 + index * 2])), + "redis", + ) + for index, meta in enumerate(per_counter_meta) + ), + ) async def _atomic_check_and_increment_in_memory( self, @@ -1539,7 +1853,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ], ) - descriptor_state.append({"window_expired": window_expired, "current": current_counter}) + descriptor_state.append( + { + "window_expired": window_expired, + "current": current_counter, + "window_start": str(now_int if window_expired else int(window_start)), + } + ) # Pass 2: apply increments. statuses: Final[list[RateLimitStatus]] = [] @@ -1569,7 +1889,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): descriptor_key=meta["descriptor_key"], ) ) - return RateLimitResponse(overall_code="OK", statuses=statuses) + return RateLimitResponse( + overall_code="OK", + statuses=statuses, + reservation_windows=frozenset( + (meta["counter_key"], state["window_start"], "local") + for meta, state in zip(per_counter_meta, descriptor_state) + ), + ) async def reserve_tpm_tokens( self, @@ -1586,6 +1913,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): TPM-only descriptor/increment list and delegates the all-or-nothing atomicity (Lua on Redis, asyncio-locked DualCache otherwise) to the shared primitive. + + Excludes project ITPM/OTPM descriptors -- those are reserved + separately (different estimate per bucket) via ``reserve_io_tokens``. """ tpm_descriptors: Final[list[RateLimitDescriptor]] = [ RateLimitDescriptor( @@ -1597,7 +1927,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ), ) for d in descriptors - if (d.get("rate_limit") or {}).get("tokens_per_unit") is not None + if d["key"] not in (PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY) + and (d.get("rate_limit") or {}).get("tokens_per_unit") is not None # mutable-ok: optional descriptor ] if not tpm_descriptors: return RateLimitResponse(overall_code="OK", statuses=[]) @@ -1611,6 +1942,142 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=parent_otel_span, ) + async def _refund_reserved_tokens( + self, + scopes: Sequence[tuple[str, str]], + amount: int, + reservation_windows: frozenset[tuple[str, str, Literal["redis", "local"]]] = frozenset(), + parent_otel_span: Span | None = None, + ) -> None: + """ + Directly decrement previously-reserved token counters for ``scopes`` + by ``amount``. Used to roll back a reservation that already + succeeded once a *different* bucket in the same request turns out to + be over its limit (e.g. ITPM reserved fine, OTPM then hits its + limit -- the ITPM reservation must not be left inflated). + """ + if amount <= 0 or not scopes: + return + if not reservation_windows: + await self.async_increment_tokens_with_ttl_preservation( + pipeline_operations=self._build_reservation_aware_tpm_ops( + targets=scopes, + reserved_scopes=frozenset(scopes), + actual_tokens=0, + reserved_tokens=amount, + ), + parent_otel_span=parent_otel_span, + ) + return + pipeline_operations: Final = self._build_project_reservation_ops( + targets=scopes, + reserved_scopes=frozenset(scopes), + actual_tokens=0, + reserved_tokens=amount, + reservation_window_identities=reservation_windows, + ) + await self.async_increment_reservation_aware_tokens( + pipeline_operations=pipeline_operations, + parent_otel_span=parent_otel_span, + ) + + async def reserve_io_tokens( + self, + descriptors: Sequence[RateLimitDescriptor], + estimated_input_tokens: int, + estimated_output_tokens: int, + parent_otel_span: Span | None = None, + ) -> tuple[RateLimitResponse, int, int]: + """ + Reserve ``estimated_input_tokens`` against project ITPM descriptors + and ``estimated_output_tokens`` against project OTPM descriptors. + + ITPM and OTPM are reserved from different-sized estimates, so unlike + same-size TPM descriptors they can't share a single + ``atomic_check_and_increment_by_n`` call -- each bucket gets its own + all-or-nothing atomic call. If the OTPM reservation is over limit + after ITPM already succeeded, the ITPM reservation this call made is + rolled back before returning, so a partial reservation never leaks. + + Returns ``(response, itpm_reserved, otpm_reserved)`` -- the latter two + are the amounts actually reserved (0 if that bucket wasn't + configured, or if the reservation failed), for the caller to stash + for post-call reconciliation. + """ + itpm_descriptors = [ # mutable-ok: atomic limiter API requires lists + d for d in descriptors if d["key"] == PROJECT_ITPM_DESCRIPTOR_KEY + ] + otpm_descriptors = [ # mutable-ok: atomic limiter API requires lists + d for d in descriptors if d["key"] == PROJECT_OTPM_DESCRIPTOR_KEY + ] + + if not itpm_descriptors and not otpm_descriptors: + return RateLimitResponse(overall_code="OK", statuses=[]), 0, 0 # mutable-ok: response contract uses a list + + itpm_response: RateLimitResponse | None = None + itpm_reserved = 0 + + if itpm_descriptors: + itpm_response = await self.atomic_check_and_increment_by_n( + descriptors=itpm_descriptors, + increments=[ # mutable-ok: atomic limiter API requires mutable increment records + {"tokens": estimated_input_tokens} # mutable-ok: atomic limiter increment record + for _ in itpm_descriptors + ], + parent_otel_span=parent_otel_span, + ) + if itpm_response["overall_code"] == "OVER_LIMIT": + return itpm_response, 0, 0 + itpm_reserved = estimated_input_tokens + + if otpm_descriptors: + otpm_response = await self.atomic_check_and_increment_by_n( + descriptors=otpm_descriptors, + increments=[ # mutable-ok: atomic limiter API requires mutable increment records + {"tokens": estimated_output_tokens} # mutable-ok: atomic limiter increment record + for _ in otpm_descriptors + ], + parent_otel_span=parent_otel_span, + ) + if otpm_response["overall_code"] == "OVER_LIMIT": + if itpm_reserved > 0: + await self._refund_reserved_tokens( + scopes=[ # mutable-ok: reservation rollback accepts collected scopes + (d["key"], d["value"]) for d in itpm_descriptors + ], + amount=itpm_reserved, + reservation_windows=itpm_response.get("reservation_windows", frozenset()), + parent_otel_span=parent_otel_span, + ) + return otpm_response, 0, 0 + statuses = ( + [ # mutable-ok: response contract uses a list + *itpm_response["statuses"], + *otpm_response["statuses"], + ] + if itpm_response is not None + else otpm_response["statuses"] + ) + return ( + RateLimitResponse( + overall_code="OK", + statuses=statuses, + reservation_windows=( + ( + itpm_response.get("reservation_windows", frozenset()) + if itpm_response is not None + else frozenset() + ) + | otpm_response.get("reservation_windows", frozenset()) + ), + ), + itpm_reserved, + estimated_output_tokens, + ) + + assert itpm_response is not None + return itpm_response, itpm_reserved, 0 + def create_organization_rate_limit_descriptor( self, user_api_key_dict: UserAPIKeyAuth, requested_model: str | None = None ) -> list[RateLimitDescriptor]: @@ -2317,6 +2784,62 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) + def _add_project_io_token_rate_limit_descriptors_from_metadata( + self, + user_api_key_dict: UserAPIKeyAuth, + requested_model: str | None, + descriptors: _RateLimitDescriptorSink, + ) -> None: + """Add project-scoped ITPM/OTPM descriptors from project_metadata. + + Enforced independently of, and alongside, the combined ``model_per_project`` + TPM descriptor above -- these give Bedrock Mantle-style separate input/output + token quotas at the project level. + """ + if requested_model is None or user_api_key_dict.project_id is None: + return + + itpm_limit_for_project_model = ( + get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_itpm_limit") + or {} # mutable-ok: metadata helper returns an optional mapping + ) + otpm_limit_for_project_model = ( + get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_otpm_limit") + or {} # mutable-ok: metadata helper returns an optional mapping + ) + + model_itpm_limit = itpm_limit_for_project_model.get(requested_model) + model_otpm_limit = otpm_limit_for_project_model.get(requested_model) + + if model_itpm_limit is None and model_otpm_limit is None: + return + + descriptor_value = f"{user_api_key_dict.project_id}:{requested_model}" + if model_itpm_limit is not None: + descriptors.append( + RateLimitDescriptor( + key=PROJECT_ITPM_DESCRIPTOR_KEY, + value=descriptor_value, + rate_limit={ # mutable-ok: descriptor TypedDict requires a runtime dict + "requests_per_unit": None, + "tokens_per_unit": model_itpm_limit, + "window_size": self.window_size, + }, + ) + ) + if model_otpm_limit is not None: + descriptors.append( + RateLimitDescriptor( + key=PROJECT_OTPM_DESCRIPTOR_KEY, + value=descriptor_value, + rate_limit={ # mutable-ok: descriptor TypedDict requires a runtime dict + "requests_per_unit": None, + "tokens_per_unit": model_otpm_limit, + "window_size": self.window_size, + }, + ) + ) + def _handle_rate_limit_error( self, response: RateLimitResponse, @@ -2361,6 +2884,396 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): llm_provider=llm_provider, ) + @staticmethod + def _estimate_audio_block_tokens(block: object) -> int: + """ + Token estimate for one ``input_audio`` content block. + + When the block carries a base64 ``data`` payload, the estimate comes + from the decoded byte count (``len(b64) * 3 // 4 // _AUDIO_BYTES_PER_TOKEN``), + assuming the lowest reasonable audio bitrate so we never under-reserve + for higher-quality recordings of the same duration. + + When no payload is present (reference-only block or missing ``data``), + falls back to ``DEFAULT_AUDIO_TOKEN_ESTIMATE``. + """ + if not isinstance(block, dict): + return DEFAULT_AUDIO_TOKEN_ESTIMATE + input_audio = block.get("input_audio") + b64_data = input_audio.get("data") if isinstance(input_audio, dict) else None + if b64_data and isinstance(b64_data, str): + decoded_bytes = len(b64_data) * 3 // 4 + return max(decoded_bytes // _AUDIO_BYTES_PER_TOKEN, DEFAULT_AUDIO_TOKEN_ESTIMATE) + return DEFAULT_AUDIO_TOKEN_ESTIMATE + + @classmethod + def _estimate_audio_content_tokens(cls, messages: object) -> int: + """ + Sum of per-block audio token estimates across all ``messages``. + Returns 0 when there are no ``input_audio`` blocks, which the caller + uses to skip the (relatively expensive) strip pass. + """ + if not isinstance(messages, list): + return 0 + total = 0 + for message in messages: + content = message.get("content") if isinstance(message, dict) else None + if not isinstance(content, list): + continue + total += sum( + cls._estimate_audio_block_tokens(block) + for block in content + if isinstance(block, dict) and block.get("type") == "input_audio" + ) + return total + + @staticmethod + def _strip_audio_content_blocks(messages: object) -> object: + """ + Drop ``input_audio`` content blocks before passing ``messages`` to + ``token_counter``, which raises ``ValueError`` on them (no per-type + handling, unlike images). The audio contribution is added back + separately via ``DEFAULT_AUDIO_TOKEN_ESTIMATE`` so the rest of the + message (text/images/tools) still gets counted accurately instead of + the whole call falling back to the cheap char-count estimate. + """ + if not isinstance(messages, list): + return messages + sanitized = [] # mutable-ok: token_counter requires a list of message dicts + for message in messages: + if not isinstance(message, dict): + sanitized.append(message) + continue + content = message.get("content") + if not isinstance(content, list): + sanitized.append(message) + continue + filtered_content = [ # mutable-ok: token_counter requires list content blocks + block for block in content if not (isinstance(block, dict) and block.get("type") == "input_audio") + ] + sanitized.append( # mutable-ok: token_counter requires mutable message dicts + {**message, "content": filtered_content} # mutable-ok: token_counter requires message dicts + ) + return sanitized + + @staticmethod + def _responses_input_to_chat_messages(data: object) -> Sequence[object]: + """ + Convert a Responses API ``input`` (string or list of input items) into + chat-completion-style messages via the standard LiteLLM transformation + (the same one guardrails use, e.g. ``purview_dlp.py``), so multimodal + ``input_image``/``input_text`` content blocks get counted by + ``token_counter``'s ``messages`` path instead of silently contributing + zero tokens via its ``text`` path, which only joins plain strings. + """ + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + if not isinstance(data, dict): + return () + return LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=data.get("input") or "", + responses_api_request=data, + ) + + @classmethod + def _contains_responses_file_reference(cls, value: object) -> bool: + if isinstance(value, dict): + return value.get("type") == "input_file" or any( + cls._contains_responses_file_reference(child) for child in value.values() + ) + if isinstance(value, list): + return any(cls._contains_responses_file_reference(child) for child in value) + return False + + @classmethod + def _contains_unmeasurable_chat_media(cls, value: object) -> bool: + if isinstance(value, dict): + return value.get("type") in ("document", "file", "video_url") or any( + cls._contains_unmeasurable_chat_media(child) for child in value.values() + ) + if isinstance(value, list): + return any(cls._contains_unmeasurable_chat_media(child) for child in value) + return False + + @classmethod + def _contains_image_content(cls, value: object) -> bool: + if isinstance(value, dict): + media_type = value.get("media_type") or value.get("mime_type") + return ( + value.get("type") in ("image", "image_url", "input_image") + or (isinstance(media_type, str) and media_type.startswith("image/")) + or any(cls._contains_image_content(child) for child in value.values()) + ) + if isinstance(value, list): + return any(cls._contains_image_content(child) for child in value) + return False + + @classmethod + def _requires_conservative_responses_input_reservation(cls, data: object, call_type: str | None) -> bool: + if not isinstance(data, dict): + return False + return call_type in RESPONSES_API_CALL_TYPES and ( + data.get("previous_response_id") is not None or cls._contains_responses_file_reference(data.get("input")) + ) + + @staticmethod + def _count_pretokenized_embedding_input(value: object) -> int | None: + if not isinstance(value, list): + return None + if all(isinstance(token, int) for token in value): + return len(value) + if all( + isinstance(token_ids, list) and all(isinstance(token, int) for token in token_ids) for token_ids in value + ): + return sum(len(token_ids) for token_ids in value) + return None + + @staticmethod + def _rerank_input_to_text(data: Mapping[str, object]) -> str: + documents = data.get("documents") + document_items: Sequence[object] = documents if isinstance(documents, list) else () # pyright: ignore[reportUnknownVariableType] # rerank documents are validated runtime JSON + input_parts: tuple[object, ...] = ( # pyright: ignore[reportUnknownVariableType] # list narrowing preserves unknown JSON element types + data.get("query"), + *document_items, + ) + return "\n".join( + str(part) # pyright: ignore[reportUnknownArgumentType] # accepted document dicts have provider-defined fields + for part in input_parts # pyright: ignore[reportUnknownVariableType] # runtime JSON list elements remain unknown after list narrowing + if isinstance(part, (str, dict)) + ) + + def _estimate_precise_input_tokens(self, data: object, model: str | None, call_type: str | None = None) -> int: + """ + Model-aware input token estimate for the project ITPM reservation, + using ``litellm.token_counter`` -- the same approach the + deployment-level itpm/otpm check uses in + ``io_token_rate_limit_check.py``. Unlike the cheap char-count + estimate the combined-TPM path uses, this accounts for image/tool + content and derives per-``input_audio``-block estimates from the + base64 payload size (assuming the lowest reasonable bitrate so + longer recordings always reserve proportionally more), so a burst + of multimodal, tool-heavy, or audio-heavy requests can't each + reserve only the one-token floor and blow past ITPM before + post-call reconciliation catches up. + + For the Responses API, ``input`` is converted to chat messages first + (via ``_responses_input_to_chat_messages``) so its own multimodal + content blocks are counted the same way; ``token_counter``'s ``text`` + argument can only see plain strings in a list, not content blocks. + + Falls back to the cheap char-count estimate if ``token_counter`` + can't resolve a tokenizer for this model (e.g. an unrecognized + custom model name) or otherwise raises -- the audio add-on still + applies on top of the fallback. + """ + from litellm import token_counter + + if not isinstance(data, dict): + return 0 + selected_text = None + countable_tools = data.get("tools") + countable_tool_choice = data.get("tool_choice") + if call_type in RESPONSES_API_CALL_TYPES: + messages = self._responses_input_to_chat_messages(data) + elif (translated_request := self._translate_google_genai_native_request(data, call_type)) is not None: + messages = translated_request.get("messages") + countable_tools = translated_request.get("tools") + countable_tool_choice = translated_request.get("tool_choice") + elif self._is_embedding_request(data, call_type): + messages = None + selected_text = data.get("input") + pretokenized_input_tokens = self._count_pretokenized_embedding_input(selected_text) + if pretokenized_input_tokens is not None: + return pretokenized_input_tokens + elif call_type in RERANK_API_CALL_TYPES: + messages = None + selected_text = self._rerank_input_to_text(data) # pyright: ignore[reportUnknownArgumentType] # proxy request bodies are runtime-validated JSON + elif call_type in TEXT_COMPLETION_API_CALL_TYPES: + messages = None + selected_text = data.get("prompt") + else: + messages = data.get("messages") + if messages is None: + selected_text = data.get("prompt") + if messages is None and selected_text is None: + selected_text = data.get("input") + + audio_token_estimate = self._estimate_audio_content_tokens(messages) + countable_messages = self._strip_audio_content_blocks(messages) if audio_token_estimate > 0 else messages + + try: + estimate = max( + 0, + int( + token_counter( + model=model or "", + messages=countable_messages, + text=selected_text, + tools=countable_tools, + tool_choice=countable_tool_choice, + use_default_image_token_count=True, + ) + ), + ) + return estimate + audio_token_estimate + except Exception: # noqa: BLE001 - any tokenizer/model-resolution/transform failure degrades to the cheap estimate, never a 500 + if call_type in RERANK_API_CALL_TYPES and isinstance(selected_text, str): + return max(0, len(selected_text) // DEFAULT_CHARS_PER_TOKEN) + estimated_input_tokens, _ = self._estimate_input_and_output_tokens(data=data, call_type=call_type) + return estimated_input_tokens + audio_token_estimate + + async def _reserve_project_io_tokens_or_raise( + self, + descriptors: Sequence[RateLimitDescriptor], + data: object, + requested_model: str | None, + user_api_key_dict: UserAPIKeyAuth, + tpm_reservation_scopes: Sequence[tuple[str, str]], + tpm_reservation_amount: int, + call_type: str | None = None, + ) -> None: + """ + Reserve project-scoped ITPM/OTPM tokens (Bedrock Mantle-style + separate input/output token buckets), independently of -- and, when + both are configured, in addition to -- the combined-TPM reservation + the caller already made. Raises (via ``_handle_rate_limit_error``) on + an over-limit reservation, first rolling back the combined-TPM + reservation named by ``tpm_reservation_scopes``/``tpm_reservation_amount`` + if one was made, so a partial reservation never leaks. + """ + if not isinstance(data, dict): + return + stash = claim_request_stash_for_data(data) + io_token_descriptors = [ # mutable-ok: reservation API requires descriptor lists + d for d in descriptors if d["key"] in (PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY) + ] + if not io_token_descriptors: + return + + configured_otpm_limits = [ # mutable-ok: min calculation materializes validated limits + int(v) + for d in io_token_descriptors + if d["key"] == PROJECT_OTPM_DESCRIPTOR_KEY + for v in [ # mutable-ok: comprehension binds the optional descriptor value + (d.get("rate_limit") or {}).get( # mutable-ok: optional descriptor fallback + "tokens_per_unit" + ) + ] + if v is not None + ] + min_configured_otpm_limit = min(configured_otpm_limits) if configured_otpm_limits else None + configured_itpm_limits = [ # mutable-ok: min calculation materializes validated limits + int(v) + for d in io_token_descriptors + if d["key"] == PROJECT_ITPM_DESCRIPTOR_KEY + for v in [ # mutable-ok: comprehension binds the optional descriptor value + (d.get("rate_limit") or {}).get( # mutable-ok: optional descriptor fallback + "tokens_per_unit" + ) + ] + if v is not None + ] + min_configured_itpm_limit = min(configured_itpm_limits) if configured_itpm_limits else None + + _, estimated_output_tokens = self._estimate_input_and_output_tokens( + data=data, + min_configured_tpm_limit=min_configured_otpm_limit, + call_type=call_type, + ) + estimated_input_tokens = ( + min_configured_itpm_limit + if min_configured_itpm_limit is not None + and ( + self._requires_conservative_responses_input_reservation(data, call_type) + or self._contains_unmeasurable_chat_media(data.get("messages")) + or self._contains_image_content(data) + ) + else self._estimate_precise_input_tokens(data=data, model=requested_model, call_type=call_type) + ) + estimated_input_tokens = max(estimated_input_tokens, 1) + if not self._has_explicit_output_cap(data, call_type): + estimated_output_tokens = max(estimated_output_tokens, 1) + + # Hard-cap generation length so an unbounded response can't overshoot + # the OTPM budget before post-call reconciliation runs, mirroring the + # combined-TPM floor cap in the caller. + self._apply_implicit_output_cap( + data=data, + min_configured_limit=min_configured_otpm_limit, + call_type=call_type, + ) + + io_response, itpm_reserved, otpm_reserved = await self.reserve_io_tokens( + descriptors=io_token_descriptors, + estimated_input_tokens=estimated_input_tokens, + estimated_output_tokens=estimated_output_tokens, + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + + if io_response["overall_code"] == "OVER_LIMIT": + # A combined-TPM reservation may have already succeeded above for + # this same request; refund it too, or its counter stays inflated + # until the window's TTL expires. Mark it released so the + # ProxyRateLimitError we're about to raise doesn't get refunded + # a second time when async_post_call_failure_hook sees the same + # (still-stashed) reservation and refunds it again. + if tpm_reservation_amount > 0: + await self._refund_reserved_tokens( + scopes=tpm_reservation_scopes, + amount=tpm_reservation_amount, + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + stash.reservation_released = True + acquisition = stash.parallel_slot + if acquisition is not None: + await self._release_parallel_request_slots( + acquisition=acquisition, + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + stash.parallel_slot = None + self._handle_rate_limit_error( + response=io_response, + descriptors=descriptors, + requested_model=requested_model, + ) + + if itpm_reserved > 0: + itpm_scopes = [ # mutable-ok: request stash freezes the collected scopes + (d["key"], d["value"]) for d in io_token_descriptors if d["key"] == PROJECT_ITPM_DESCRIPTOR_KEY + ] + stash.itpm_reserved_tokens = itpm_reserved + stash.itpm_reserved_scopes = frozenset(itpm_scopes) + stash.itpm_reserved_window_identities = frozenset( + (counter_key, window_start, backend) + for counter_key, window_start, backend in io_response.get("reservation_windows", frozenset()) + if "model_per_project_itpm" in counter_key + ) + if otpm_reserved > 0: + otpm_scopes = [ # mutable-ok: request stash freezes the collected scopes + (d["key"], d["value"]) for d in io_token_descriptors if d["key"] == PROJECT_OTPM_DESCRIPTOR_KEY + ] + stash.otpm_reserved_tokens = otpm_reserved + stash.otpm_reserved_scopes = frozenset(otpm_scopes) + stash.otpm_reserved_window_identities = frozenset( + (counter_key, window_start, backend) + for counter_key, window_start, backend in io_response.get("reservation_windows", frozenset()) + if "model_per_project_otpm" in counter_key + ) + + if stash.rate_limit_response is not None: + stash.rate_limit_response["statuses"].extend(io_response["statuses"]) + elif io_response["statuses"]: + stash.rate_limit_response = io_response + + verbose_proxy_logger.debug( + "ITPM/OTPM tokens reserved: itpm=%s, otpm=%s for model %s", + itpm_reserved, + otpm_reserved, + requested_model, + ) + async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, @@ -2433,6 +3346,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): requested_model=requested_model, descriptors=descriptors, ) + self._add_project_io_token_rate_limit_descriptors_from_metadata( + user_api_key_dict=user_api_key_dict, + requested_model=requested_model, + descriptors=descriptors, + ) # Org Level Rate Limits descriptors.extend(self.create_organization_rate_limit_descriptor(user_api_key_dict, requested_model)) @@ -2489,28 +3407,33 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): configured_tpm_limits: Final = [ int(v) for d in descriptors + if d["key"] not in (PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY) for v in [(d.get("rate_limit") or {}).get("tokens_per_unit")] if v is not None ] has_tpm_limits: Final = bool(configured_tpm_limits) + # Populated on a successful combined-TPM reservation below, so the + # project ITPM/OTPM block further down can roll it back if a + # different bucket in the same request subsequently hits its + # limit. Stays empty/0 whenever no combined-TPM reservation was + # made (or it was over limit, in which case execution never + # reaches the ITPM/OTPM block -- `_handle_rate_limit_error` raises). + tpm_reservation_scopes: Sequence[tuple[str, str]] = () + tpm_reservation_amount = 0 + if has_tpm_limits and self.tpm_reservation_enabled: min_configured_tpm_limit: Final = min(configured_tpm_limits) # When the configured TPM cap is small enough to constrain the - # no-max_tokens floor, also hard-cap the model output via - # data["max_tokens"] so concurrent unbounded generations can't - # spend past the limit before post-call reconciliation runs. - # Skip when the request already sets max_tokens or has no - # generation budget at all (embeddings). - capped_floor: Final = self._no_max_tokens_output_floor(min_configured_tpm_limit) - baseline_floor: Final = DEFAULT_MAX_TOKENS_ESTIMATE // _TPM_FLOOR_FRACTION - has_explicit_max_tokens: Final = ( - data.get("max_tokens") is not None or data.get("max_completion_tokens") is not None + # no-max_tokens floor, also hard-cap the model output so + # concurrent unbounded generations can't spend past the limit + # before post-call reconciliation runs. + self._apply_implicit_output_cap( + data=data, + min_configured_limit=min_configured_tpm_limit, + call_type=call_type, ) - is_embedding: Final = data.get("input") is not None - if capped_floor < baseline_floor and not has_explicit_max_tokens and not is_embedding: - data["max_tokens"] = capped_floor # Floor at 1 token so contentless requests (/responses, # tool-call continuations, empty messages) still flow @@ -2524,6 +3447,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): data=data, model=requested_model, min_configured_tpm_limit=min_configured_tpm_limit, + call_type=call_type, ), 1, ) @@ -2557,8 +3481,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): stash.reserved_scopes = frozenset( (d["key"], d["value"]) for d in descriptors - if (d.get("rate_limit") or {}).get("tokens_per_unit") is not None + if d["key"] not in (PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY) + and (d.get("rate_limit") or {}).get( # mutable-ok: optional descriptor fallback + "tokens_per_unit" + ) + is not None ) + tpm_reservation_scopes = tuple(stash.reserved_scopes) + tpm_reservation_amount = estimated_tokens # Merge TPM statuses into the stored rate-limit response # so x-ratelimit-{key}-remaining-tokens / -limit-tokens @@ -2573,6 +3503,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): "TPM tokens reserved: %s for model %s", estimated_tokens, requested_model ) + await self._reserve_project_io_tokens_or_raise( + descriptors=descriptors, + data=data, + requested_model=requested_model, + user_api_key_dict=user_api_key_dict, + tpm_reservation_scopes=tpm_reservation_scopes, + tpm_reservation_amount=tpm_reservation_amount, + call_type=call_type, + ) + def _create_pipeline_operations( self, key: str, @@ -2751,6 +3691,113 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): litellm_parent_otel_span=parent_otel_span, ) + async def _apply_local_window_guarded_token_increments( + self, + operations: Sequence[ReservationAwareIncrementOperation], + parent_otel_span: Span | None = None, + ) -> None: + async with self._check_and_increment_lock: + for operation in operations: + window_key = operation.get("window_key") + expected_window_start = operation.get("expected_window_start") + if window_key is None or expected_window_start is None: + continue + active_window_start = await self.internal_usage_cache.async_get_cache( + key=window_key, + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + if active_window_start is None or str(active_window_start) != expected_window_start: + continue + current_counter = ( + await self.internal_usage_cache.async_get_cache( + key=operation["key"], + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + or 0 + ) + await self.internal_usage_cache.async_set_cache( + key=operation["key"], + value=float(current_counter) + operation["increment_value"], + ttl=operation["ttl"], + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + + async def _apply_redis_window_guarded_token_increments( + self, + operations: Sequence[ReservationAwareIncrementOperation], + parent_otel_span: Span | None = None, + ) -> None: + for operation in operations: + window_key = operation.get("window_key") + expected_window_start = operation.get("expected_window_start") + if window_key is None or expected_window_start is None: + continue + if self.window_guarded_token_increment_script is not None: + try: + await self.window_guarded_token_increment_script( + keys=[window_key, operation["key"]], + args=[ + expected_window_start, + operation["increment_value"], + operation["ttl"] or 0, + ], + ) + continue + except Exception as e: + verbose_proxy_logger.warning( + "Window-guarded token adjustment failed for %s: %s", + operation["key"], + e, + ) + if operation["increment_value"] > 0: + await self.internal_usage_cache.async_increment_cache( + key=operation["key"], + value=operation["increment_value"], + litellm_parent_otel_span=parent_otel_span, + ttl=operation["ttl"], + ) + + async def async_increment_reservation_aware_tokens( + self, + pipeline_operations: Sequence[ReservationAwareIncrementOperation], + parent_otel_span: Span | None = None, + ) -> None: + for operation in pipeline_operations: + if operation.get("window_key") is None or operation.get("expected_window_start") is None: + await self.internal_usage_cache.async_increment_cache( + key=operation["key"], + value=operation["increment_value"], + litellm_parent_otel_span=parent_otel_span, + ttl=operation["ttl"], + ) + local_guarded_operations: Final = tuple( + operation + for operation in pipeline_operations + if operation.get("window_key") is not None + and operation.get("expected_window_start") is not None + and operation.get("reservation_backend") == "local" + ) + redis_guarded_operations: Final = tuple( + operation + for operation in pipeline_operations + if operation.get("window_key") is not None + and operation.get("expected_window_start") is not None + and operation.get("reservation_backend") != "local" + ) + if local_guarded_operations: + await self._apply_local_window_guarded_token_increments( + operations=local_guarded_operations, + parent_otel_span=parent_otel_span, + ) + if redis_guarded_operations: + await self._apply_redis_window_guarded_token_increments( + operations=redis_guarded_operations, + parent_otel_span=parent_otel_span, + ) + def get_rate_limit_type(self) -> Literal["output", "input", "total"]: from litellm.proxy.proxy_server import general_settings @@ -2780,6 +3827,163 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): merged[f"{prefix}-limit-{status['rate_limit_type']}"] = status["current_limit"] return merged + @staticmethod + def _resolve_rerank_token_usage(response_obj: object) -> tuple[int, int, bool] | None: + if not isinstance(response_obj, RerankResponse) or response_obj.meta is None: + return None + + rerank_tokens = response_obj.meta.get("tokens") # pyright: ignore[reportUnknownMemberType] # TypedDict's optional generic metadata widens get overloads + if rerank_tokens is not None: + input_tokens = rerank_tokens.get("input_tokens") or 0 # pyright: ignore[reportUnknownMemberType] # token fields are typed integers despite the generic get overload + output_tokens = rerank_tokens.get("output_tokens") or 0 # pyright: ignore[reportUnknownMemberType] # token fields are typed integers despite the generic get overload + if input_tokens or output_tokens: + return max(0, input_tokens), max(0, output_tokens), True + + billed_units = response_obj.meta.get("billed_units") # pyright: ignore[reportUnknownMemberType] # TypedDict's optional generic metadata widens get overloads + if billed_units is not None: + total_tokens = billed_units.get("total_tokens") or 0 # pyright: ignore[reportUnknownMemberType] # billed total is a typed integer despite the generic get overload + if total_tokens: + return max(0, total_tokens), 0, True + return None + + def _resolve_io_token_reconcile_usage( + self, + response_obj: object, + ) -> tuple[int, int, bool]: + """ + Resolve ``(billable_input_tokens, completion_tokens, usage_resolved)`` + for ITPM/OTPM reconciliation. Cache-read tokens are excluded from + billable input -- Bedrock Mantle doesn't count them toward ITPM -- + but they're untouched everywhere else (cost/usage logging still sees + the full prompt token count). + """ + rerank_usage = self._resolve_rerank_token_usage(response_obj) + if rerank_usage is not None: + return rerank_usage + + usage: object | None = None + if isinstance(response_obj, (Usage, ResponseAPIUsage)): + usage = response_obj + elif isinstance( + response_obj, + (ModelResponse, EmbeddingResponse, TextCompletionResponse, BaseLiteLLMOpenAIResponseObject), + ): + usage = getattr(response_obj, "usage", None) + elif isinstance(response_obj, dict): + usage = response_obj.get("usage") + if usage is None and any( + key in response_obj + for key in ( + "prompt_tokens", + "completion_tokens", + "input_tokens", + "output_tokens", + ) + ): + usage = response_obj + + if isinstance(usage, Usage): + prompt_tokens = usage.prompt_tokens or 0 + completion_tokens = usage.completion_tokens or 0 + cached_tokens = 0 + if usage.prompt_tokens_details is not None: + cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0 + elif isinstance(usage, ResponseAPIUsage): + # Responses API usage uses input_tokens/output_tokens instead of + # prompt_tokens/completion_tokens. + prompt_tokens = usage.input_tokens or 0 + completion_tokens = usage.output_tokens or 0 + cached_tokens = 0 + if usage.input_tokens_details is not None: + cached_tokens = usage.input_tokens_details.cached_tokens or 0 + elif isinstance(usage, dict): + prompt_tokens = usage.get("prompt_tokens") or usage.get("input_tokens") or 0 + completion_tokens = usage.get("completion_tokens") or usage.get("output_tokens") or 0 + prompt_details = ( + usage.get("prompt_tokens_details") + or usage.get("input_tokens_details") + or {} # mutable-ok: usage details are optional mappings + ) + cached_tokens = ( + (prompt_details.get("cached_tokens", 0) or 0) if isinstance(prompt_details, dict) else 0 + ) or (usage.get("cache_read_input_tokens") or 0) + else: + return 0, 0, False + + if prompt_tokens == 0 and completion_tokens == 0: + return 0, 0, False + return max(0, prompt_tokens - cached_tokens), completion_tokens, True + + def _build_io_token_reservation_ops( + self, + kwargs: object, + response_obj: object, + ) -> list[RedisPipelineIncrementOperation] | tuple[ReservationAwareIncrementOperation, ...]: + """ + Reconcile project ITPM/OTPM reservations to actual usage on success: + ITPM to billable input tokens, OTPM to actual completion tokens. + Reuses ``_build_reservation_aware_tpm_ops``'s delta pattern -- ITPM/OTPM + are stored in the same ":tokens" cache bucket as combined TPM, just + under distinct scope keys, so the reservation-aware increment math is + identical; only the usage fields being reconciled against differ. + """ + if not isinstance(kwargs, dict): + return () + stash = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) + if stash is None: + return () + + itpm_reserved = stash.itpm_reserved_tokens + otpm_reserved = stash.otpm_reserved_tokens + if itpm_reserved <= 0 and otpm_reserved <= 0: + return () + + billable_input, completion_tokens, usage_resolved = self._resolve_io_token_reconcile_usage(response_obj) + if not usage_resolved: + billable_input, completion_tokens, usage_resolved = self._resolve_io_token_reconcile_usage( + kwargs.get("combined_usage_object") + ) + if not usage_resolved: + if not stash.reservation_released: + return () + billable_input = itpm_reserved + completion_tokens = otpm_reserved + + if stash.reservation_released or ( + not stash.itpm_reserved_window_identities and not stash.otpm_reserved_window_identities + ): + return self._build_reservation_aware_tpm_ops( + targets=tuple(stash.itpm_reserved_scopes), + reserved_scopes=frozenset() if stash.reservation_released else stash.itpm_reserved_scopes, + actual_tokens=billable_input, + reserved_tokens=0 if stash.reservation_released else itpm_reserved, + ) + self._build_reservation_aware_tpm_ops( + targets=tuple(stash.otpm_reserved_scopes), + reserved_scopes=frozenset() if stash.reservation_released else stash.otpm_reserved_scopes, + actual_tokens=completion_tokens, + reserved_tokens=0 if stash.reservation_released else otpm_reserved, + ) + + itpm_ops: Sequence[ReservationAwareIncrementOperation] = () + if itpm_reserved > 0: + itpm_ops = self._build_project_reservation_ops( + targets=tuple(stash.itpm_reserved_scopes), + reserved_scopes=frozenset() if stash.reservation_released else stash.itpm_reserved_scopes, + actual_tokens=billable_input, + reserved_tokens=itpm_reserved, + reservation_window_identities=stash.itpm_reserved_window_identities, + ) + otpm_ops: Sequence[ReservationAwareIncrementOperation] = () + if otpm_reserved > 0: + otpm_ops = self._build_project_reservation_ops( + targets=tuple(stash.otpm_reserved_scopes), + reserved_scopes=frozenset() if stash.reservation_released else stash.otpm_reserved_scopes, + actual_tokens=completion_tokens, + reserved_tokens=otpm_reserved, + reservation_window_identities=stash.otpm_reserved_window_identities, + ) + return tuple((*itpm_ops, *otpm_ops)) + def _collect_tpm_scope_targets( self, standard_logging_metadata: dict[str, Any], @@ -2844,8 +4048,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _build_reservation_aware_tpm_ops( self, - targets: list[tuple[str, str]], - reserved_scopes: frozenset[tuple[str, str]], + targets: Sequence[tuple[str, str]], + reserved_scopes: Set[tuple[str, str]], actual_tokens: int, reserved_tokens: int, ) -> list[RedisPipelineIncrementOperation]: @@ -2878,6 +4082,66 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) return ops + def _build_project_reservation_op( + self, + scope: tuple[str, str], + reserved_scopes: Set[tuple[str, str]], + actual_tokens: int, + reserved_tokens: int, + reservation_window_identities: frozenset[tuple[str, str, Literal["redis", "local"]]], + ) -> ReservationAwareIncrementOperation | None: + scope_key, scope_value = scope + is_reserved_scope: Final = scope in reserved_scopes + increment: Final = actual_tokens - reserved_tokens if is_reserved_scope else actual_tokens + if increment == 0: + return None + counter_key: Final = self.create_rate_limit_keys(scope_key, scope_value, "tokens") + window_identity: Final = next( + ( + (window_start, backend) + for identity_counter_key, window_start, backend in reservation_window_identities + if identity_counter_key == counter_key + ), + None, + ) + if not is_reserved_scope or window_identity is None: + return ReservationAwareIncrementOperation( + key=counter_key, + increment_value=increment, + ttl=self.window_size, + ) + return ReservationAwareIncrementOperation( + key=counter_key, + increment_value=increment, + ttl=self.window_size, + window_key=f"{{{scope_key}:{scope_value}}}:window", + expected_window_start=window_identity[0], + reservation_backend=window_identity[1], + ) + + def _build_project_reservation_ops( + self, + targets: Sequence[tuple[str, str]], + reserved_scopes: Set[tuple[str, str]], + actual_tokens: int, + reserved_tokens: int, + reservation_window_identities: frozenset[tuple[str, str, Literal["redis", "local"]]], + ) -> tuple[ReservationAwareIncrementOperation, ...]: + return tuple( + operation + for scope in targets + if ( + operation := self._build_project_reservation_op( + scope=scope, + reserved_scopes=reserved_scopes, + actual_tokens=actual_tokens, + reserved_tokens=reserved_tokens, + reservation_window_identities=reservation_window_identities, + ) + ) + is not None + ) + def _build_success_event_pipeline_operations( self, kwargs: Any, @@ -2994,12 +4258,26 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): response_obj=response_obj, rate_limit_type=rate_limit_type, ) - if pipeline_operations: await self.async_increment_tokens_with_ttl_preservation( pipeline_operations=pipeline_operations, parent_otel_span=litellm_parent_otel_span, ) + io_token_operations: Final = self._build_io_token_reservation_ops( + kwargs=kwargs, + response_obj=response_obj, + ) + if io_token_operations: + if isinstance(io_token_operations, list): + await self.async_increment_tokens_with_ttl_preservation( + pipeline_operations=io_token_operations, + parent_otel_span=litellm_parent_otel_span, + ) + else: + await self.async_increment_reservation_aware_tokens( + pipeline_operations=io_token_operations, + parent_otel_span=litellm_parent_otel_span, + ) except Exception as e: verbose_proxy_logger.exception("Error in rate limit success event: %s", e) @@ -3092,9 +4370,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # already released it (proxy-level rejection that also bubbles up # here as an LLM-error callback). max_parallel_requests is its # own counter and is always decremented per call. - reserved_tokens = 0 - if stash is not None and not stash.reservation_released: + if stash is None or stash.reservation_released: + reserved_tokens = 0 + itpm_reserved = 0 + otpm_reserved = 0 + else: reserved_tokens = stash.reserved_tokens + itpm_reserved = stash.itpm_reserved_tokens + otpm_reserved = stash.otpm_reserved_tokens + if stash is not None and reserved_tokens > 0: verbose_proxy_logger.debug("Releasing reserved TPM tokens on failure: %s", reserved_tokens) # Refund only against the scopes the reservation actually @@ -3111,12 +4395,64 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) + # Refund project ITPM/OTPM reservations the same way -- full + # refund, since a failed call has no billable usage to reconcile + # against. + itpm_operations: Final = ( + self._build_project_reservation_ops( + targets=tuple(stash.itpm_reserved_scopes), + reserved_scopes=stash.itpm_reserved_scopes, + actual_tokens=0, + reserved_tokens=itpm_reserved, + reservation_window_identities=stash.itpm_reserved_window_identities, + ) + if stash is not None and itpm_reserved > 0 and stash.itpm_reserved_window_identities + else self._build_reservation_aware_tpm_ops( + targets=tuple(stash.itpm_reserved_scopes), + reserved_scopes=stash.itpm_reserved_scopes, + actual_tokens=0, + reserved_tokens=itpm_reserved, + ) + if stash is not None and itpm_reserved > 0 + else () + ) + + otpm_operations: Final = ( + self._build_project_reservation_ops( + targets=tuple(stash.otpm_reserved_scopes), + reserved_scopes=stash.otpm_reserved_scopes, + actual_tokens=0, + reserved_tokens=otpm_reserved, + reservation_window_identities=stash.otpm_reserved_window_identities, + ) + if stash is not None and otpm_reserved > 0 and stash.otpm_reserved_window_identities + else self._build_reservation_aware_tpm_ops( + targets=tuple(stash.otpm_reserved_scopes), + reserved_scopes=stash.otpm_reserved_scopes, + actual_tokens=0, + reserved_tokens=otpm_reserved, + ) + if stash is not None and otpm_reserved > 0 + else () + ) + if pipeline_operations: await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( increment_list=pipeline_operations, litellm_parent_otel_span=litellm_parent_otel_span, ) - if stash is not None and reserved_tokens > 0: + for project_operations in (itpm_operations, otpm_operations): + if isinstance(project_operations, list): + await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( + increment_list=project_operations, + litellm_parent_otel_span=litellm_parent_otel_span, + ) + elif project_operations: + await self.async_increment_reservation_aware_tokens( + pipeline_operations=project_operations, + parent_otel_span=litellm_parent_otel_span, + ) + if stash is not None and (reserved_tokens > 0 or itpm_reserved > 0 or otpm_reserved > 0): stash.reservation_released = True except Exception as e: verbose_proxy_logger.exception("Error in rate limit failure event: %s", e) @@ -3194,19 +4530,19 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): traceback_str: str | None = None, ) -> None: """ - Release the parallel-request slot and any TPM reservation when the - request is rejected after the pre-call hook acquired them but before - the LLM call ran (e.g. a downstream guardrail/auth hook raised). - Without this, those resources are stranded — async_log_failure_event - is a litellm completion-level callback and never fires for proxy-side - rejections, so a leaked slot would occupy the gauge for the full - PARALLEL_REQUEST_SLOT_TTL_SECONDS. + Release the parallel-request slot and any TPM/ITPM/OTPM reservation + when the request is rejected after the pre-call hook acquired them + but before the LLM call ran (e.g. a downstream guardrail/auth hook + raised). Without this, those resources are stranded — + async_log_failure_event is a litellm completion-level callback and + never fires for proxy-side rejections, so a leaked slot would occupy + the gauge for the full PARALLEL_REQUEST_SLOT_TTL_SECONDS. Idempotent: the slot release clears the stashed acquisition (and slot - removal is a no-op ZREM on a second run), and the TPM refund is - guarded by the stash's ``reservation_released`` flag — if both this - hook and async_log_failure_event end up running in the same flow, only - the first release/refund applies. + removal is a no-op ZREM on a second run), and the TPM/ITPM/OTPM + refund is guarded by the stash's ``reservation_released`` flag — if + both this hook and async_log_failure_event end up running in the same + flow, only the first release/refund applies. """ try: stash: Final = get_request_stash() @@ -3222,23 +4558,80 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if stash.reservation_released: return reserved_tokens: Final = stash.reserved_tokens - if reserved_tokens <= 0: + itpm_reserved: Final = stash.itpm_reserved_tokens + otpm_reserved: Final = stash.otpm_reserved_tokens + if reserved_tokens <= 0 and itpm_reserved <= 0 and otpm_reserved <= 0: return - ops: Final = self._build_reservation_aware_tpm_ops( - targets=list(stash.reserved_scopes), - reserved_scopes=stash.reserved_scopes, - actual_tokens=0, - reserved_tokens=reserved_tokens, - ) - if ops: - verbose_proxy_logger.debug( - "Releasing reserved TPM tokens on proxy-level rejection: %s", reserved_tokens + combined_ops: Final = ( + self._build_reservation_aware_tpm_ops( + targets=tuple(stash.reserved_scopes), + reserved_scopes=stash.reserved_scopes, + actual_tokens=0, + reserved_tokens=reserved_tokens, ) + if reserved_tokens > 0 + else () + ) + itpm_ops: Final = ( + self._build_project_reservation_ops( + targets=tuple(stash.itpm_reserved_scopes), + reserved_scopes=stash.itpm_reserved_scopes, + actual_tokens=0, + reserved_tokens=itpm_reserved, + reservation_window_identities=stash.itpm_reserved_window_identities, + ) + if itpm_reserved > 0 and stash.itpm_reserved_window_identities + else self._build_reservation_aware_tpm_ops( + targets=tuple(stash.itpm_reserved_scopes), + reserved_scopes=stash.itpm_reserved_scopes, + actual_tokens=0, + reserved_tokens=itpm_reserved, + ) + if itpm_reserved > 0 + else () + ) + otpm_ops: Final = ( + self._build_project_reservation_ops( + targets=tuple(stash.otpm_reserved_scopes), + reserved_scopes=stash.otpm_reserved_scopes, + actual_tokens=0, + reserved_tokens=otpm_reserved, + reservation_window_identities=stash.otpm_reserved_window_identities, + ) + if otpm_reserved > 0 and stash.otpm_reserved_window_identities + else self._build_reservation_aware_tpm_ops( + targets=tuple(stash.otpm_reserved_scopes), + reserved_scopes=stash.otpm_reserved_scopes, + actual_tokens=0, + reserved_tokens=otpm_reserved, + ) + if otpm_reserved > 0 + else () + ) + if combined_ops or itpm_ops or otpm_ops: + verbose_proxy_logger.debug( + "Releasing reserved tokens on proxy-level rejection: tpm=%s, itpm=%s, otpm=%s", + reserved_tokens, + itpm_reserved, + otpm_reserved, + ) + if combined_ops: await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( - increment_list=ops, + increment_list=combined_ops, litellm_parent_otel_span=user_api_key_dict.parent_otel_span, ) + for project_ops in (itpm_ops, otpm_ops): + if isinstance(project_ops, list): + await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( + increment_list=project_ops, + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + ) + elif project_ops: + await self.async_increment_reservation_aware_tokens( + pipeline_operations=project_ops, + parent_otel_span=user_api_key_dict.parent_otel_span, + ) stash.reservation_released = True except Exception as e: verbose_proxy_logger.exception("Error releasing TPM reservation on post-call failure: %s", e) diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 1c29c287c3a..872b447ad2d 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -3114,6 +3114,136 @@ async def test_project_model_rate_limits_not_triggered_for_other_model_v3(): ), f"model_per_project should not be added for unrelated model, got: {descriptor_keys}" +@pytest.mark.asyncio +async def test_project_model_itpm_otpm_limits_enforced_v3(): + """ + Project-level model_itpm_limit/model_otpm_limit must produce distinct + Bedrock Mantle-style input and output token descriptors. + """ + _api_key = hash_token("sk-project-io-test") + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + captured_descriptors = [] + + async def mock_should_rate_limit(descriptors, **kwargs): + captured_descriptors.extend(descriptors) + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + project_id="proj-mantle", + project_metadata={ + "model_itpm_limit": {"bedrock_mantle/claude-opus": 20000000}, + "model_otpm_limit": {"bedrock_mantle/claude-opus": 4000000}, + }, + ) + + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "bedrock_mantle/claude-opus"}, + call_type="", + ) + + descriptor_keys = [d["key"] for d in captured_descriptors] + assert "model_per_project_itpm" in descriptor_keys + assert "model_per_project_otpm" in descriptor_keys + assert "model_per_project" not in descriptor_keys + + itpm_descriptor = next( + d for d in captured_descriptors if d["key"] == "model_per_project_itpm" + ) + otpm_descriptor = next( + d for d in captured_descriptors if d["key"] == "model_per_project_otpm" + ) + assert itpm_descriptor["value"] == "proj-mantle:bedrock_mantle/claude-opus" + assert itpm_descriptor["rate_limit"]["tokens_per_unit"] == 20000000 + assert otpm_descriptor["value"] == "proj-mantle:bedrock_mantle/claude-opus" + assert otpm_descriptor["rate_limit"]["tokens_per_unit"] == 4000000 + + +@pytest.mark.asyncio +async def test_project_model_itpm_otpm_limits_not_triggered_for_other_model_v3(): + """Split project limits must not apply to an unrelated model.""" + _api_key = hash_token("sk-project-io-test-2") + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + captured_descriptors = [] + + async def mock_should_rate_limit(descriptors, **kwargs): + captured_descriptors.extend(descriptors) + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + project_id="proj-mantle", + project_metadata={ + "model_itpm_limit": {"bedrock_mantle/claude-opus": 20000000}, + }, + ) + + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4"}, + call_type="", + ) + + descriptor_keys = [d["key"] for d in captured_descriptors] + assert "model_per_project_itpm" not in descriptor_keys + assert "model_per_project_otpm" not in descriptor_keys + + +@pytest.mark.asyncio +async def test_project_model_itpm_and_tpm_limits_coexist_v3(): + """Combined project TPM and split ITPM/OTPM limits are enforced together.""" + _api_key = hash_token("sk-project-io-test-3") + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + captured_descriptors = [] + + async def mock_should_rate_limit(descriptors, **kwargs): + captured_descriptors.extend(descriptors) + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + project_id="proj-mantle", + project_metadata={ + "model_tpm_limit": {"bedrock_mantle/claude-opus": 1000}, + "model_itpm_limit": {"bedrock_mantle/claude-opus": 20000000}, + "model_otpm_limit": {"bedrock_mantle/claude-opus": 4000000}, + }, + ) + + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "bedrock_mantle/claude-opus"}, + call_type="", + ) + + descriptor_keys = [d["key"] for d in captured_descriptors] + assert "model_per_project" in descriptor_keys + assert "model_per_project_itpm" in descriptor_keys + assert "model_per_project_otpm" in descriptor_keys + + @pytest.mark.asyncio async def test_pre_call_hook_keeps_internal_stash_out_of_request_body(): """Regression for #27001 / #35197: the limiter's per-request bookkeeping @@ -3190,7 +3320,7 @@ async def test_responses_route_body_untouched_by_pre_call_hook(caller_metadata): _api_key = hash_token("sk-responses-regression") user_api_key_dict = UserAPIKeyAuth( api_key=_api_key, - tpm_limit=1000, + tpm_limit=100000, rpm_limit=5, ) local_cache = DualCache() diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py index f7bd37b412a..66fc84ab1e0 100644 --- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py +++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py @@ -15,7 +15,7 @@ Redis. """ import asyncio -from datetime import datetime +from datetime import datetime, timedelta from typing import Any, Dict import pytest @@ -23,15 +23,25 @@ import pytest from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + PROJECT_ITPM_DESCRIPTOR_KEY, + PROJECT_OTPM_DESCRIPTOR_KEY, + _AUDIO_BYTES_PER_TOKEN, _PROXY_MaxParallelRequestsHandler_v3 as RateLimitHandler, ) from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _call_id_from_callback_kwargs, _request_stash, get_or_create_request_stash, get_request_stash, ) from litellm.proxy.utils import InternalUsageCache, hash_token -from litellm.types.utils import ModelResponse, Usage +from litellm.types.llms.openai import ( + InputTokensDetails, + ResponseAPIUsage, + ResponsesAPIResponse, +) +from litellm.types.rerank import RerankResponse +from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage @pytest.fixture @@ -582,6 +592,47 @@ async def test_estimate_tokens_uses_max_tokens_when_explicit(rate_limiter): assert estimate == 4 + 25 +@pytest.mark.asyncio +async def test_estimate_tokens_honors_explicit_zero_max_tokens(rate_limiter): + """ + Regression for a Greptile finding: explicit_max_tokens was resolved via + `data.get("max_tokens") or data.get("max_completion_tokens") or + data.get("max_output_tokens")`, so an explicit 0 in the first field was + falsy and fell through to the next field (or the no-max_tokens floor), + silently discarding a caller's explicit zero-output request. + """ + handler, _cache = rate_limiter + + estimate = handler._estimate_tokens_for_request( + data={ + "messages": [ + {"role": "user", "content": "abcd" * 4} + ], # 16 chars ~ 4 tokens + "max_tokens": 0, + } + ) + assert estimate == 4, ( + f"expected input-only reservation (4) for an explicit max_tokens=0, got {estimate}" + ) + + +@pytest.mark.asyncio +async def test_estimate_tokens_honors_explicit_zero_max_output_tokens_for_responses( + rate_limiter, +): + handler, _cache = rate_limiter + + estimate = handler._estimate_tokens_for_request( + data={ + "input": "describe this image in detail", # 29 chars ~ 7 tokens + "max_output_tokens": 0, + }, + min_configured_tpm_limit=40, + call_type="aresponses", + ) + assert estimate == 23 + + @pytest.mark.asyncio async def test_estimate_tokens_zero_for_empty_embeddings(rate_limiter): """Embeddings have no output budget — reservation should equal input only.""" @@ -1197,5 +1248,2394 @@ async def test_small_tpm_cap_preserves_explicit_max_tokens(rate_limiter): assert data["max_tokens"] == 500 +@pytest.mark.asyncio +async def test_project_otpm_reservation_prevents_concurrent_bypass(rate_limiter): + """ + Bedrock Mantle-style OTPM: with a 100 OTPM limit and 5 concurrent + requests each reserving 50+ output tokens, upfront reservation must + reject the late arrivals -- not let all 5 through. Exercises the + in-memory fallback in ``atomic_check_and_increment_by_n`` for the + project-scoped ITPM/OTPM descriptors specifically. + """ + handler, cache = rate_limiter + + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-otpm-bypass"), + project_id="proj-mantle-bypass", + project_metadata={ + "model_otpm_limit": {"bedrock_mantle/claude-opus": 100}, + }, + ) + + request_data = { + "model": "bedrock_mantle/claude-opus", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 50, + } + + async def make_request(request_id: int) -> Dict[str, Any]: + data = request_data.copy() + try: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="", + ) + return {"request_id": request_id, "success": True} + except Exception as e: + return { + "request_id": request_id, + "success": False, + "status_code": getattr(e, "status_code", None), + } + + results = await asyncio.gather(*[make_request(i) for i in range(5)]) + + successful = [r for r in results if r["success"]] + rate_limited = [ + r for r in results if not r["success"] and r.get("status_code") == 429 + ] + + assert len(rate_limited) > 0, ( + f"Expected some OTPM-rate-limited requests but all {len(successful)} succeeded." + ) + + +@pytest.mark.asyncio +async def test_project_otpm_rejects_multiple_completion_candidates(rate_limiter): + handler, cache = rate_limiter + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-otpm-multiple-candidates"), + project_id="proj-multiple-candidates", + project_metadata={ + "model_otpm_limit": {"bedrock_mantle/claude-opus": 500}, + }, + ) + data = { + "model": "bedrock_mantle/claude-opus", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 100, + "n": 10, + } + + with pytest.raises(Exception) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="acompletion", + ) + + assert getattr(exc_info.value, "status_code", None) == 429 + + +@pytest.mark.asyncio +async def test_project_otpm_reserves_largest_conflicting_output_cap(rate_limiter): + handler, cache = rate_limiter + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-otpm-conflicting-caps"), + project_id="proj-conflicting-caps", + project_metadata={ + "model_otpm_limit": {"bedrock_mantle/claude-opus": 50}, + }, + ) + data = { + "model": "bedrock_mantle/claude-opus", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 1, + "max_completion_tokens": 100, + } + + with pytest.raises(Exception) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="acompletion", + ) + + assert getattr(exc_info.value, "status_code", None) == 429 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call_type", + ["agenerate_content", "agenerate_content_stream"], +) +@pytest.mark.parametrize("config_field", ["config", "generationConfig"]) +async def test_project_otpm_rejects_google_genai_native_output_cap( + rate_limiter, + call_type, + config_field, +): + handler, cache = rate_limiter + model = "gemini/gemini-3-flash-preview" + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-google-genai-native-otpm"), + project_id="project-google-genai-native-otpm", + project_metadata={"model_otpm_limit": {model: 50}}, + ) + + with pytest.raises(Exception) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data={ + "model": model, + "contents": [{"role": "user", "parts": [{"text": "Hello"}]}], + config_field: {"maxOutputTokens": 100}, + }, + call_type=call_type, + ) + + assert getattr(exc_info.value, "status_code", None) == 429 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call_type", + ["agenerate_content", "agenerate_content_stream"], +) +@pytest.mark.parametrize("candidate_count_field", ["candidateCount", "candidate_count"]) +async def test_project_otpm_rejects_google_genai_native_candidate_count( + rate_limiter, + call_type, + candidate_count_field, +): + handler, cache = rate_limiter + model = "gemini/gemini-3-flash-preview" + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-google-genai-native-candidate-count"), + project_id="project-google-genai-native-candidate-count", + project_metadata={"model_otpm_limit": {model: 150}}, + ) + + with pytest.raises(Exception) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data={ + "model": model, + "contents": [{"role": "user", "parts": [{"text": "Hello"}]}], + "config": { + "maxOutputTokens": 50, + candidate_count_field: 4, + }, + }, + call_type=call_type, + ) + + assert getattr(exc_info.value, "status_code", None) == 429 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call_type", + ["agenerate_content", "agenerate_content_stream"], +) +@pytest.mark.parametrize("config_field", [None, "config", "generationConfig"]) +async def test_project_otpm_injects_google_genai_native_output_cap( + rate_limiter, + call_type, + config_field, +): + handler, cache = rate_limiter + model = "gemini/gemini-3-flash-preview" + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-google-genai-native-implicit-otpm"), + project_id="project-google-genai-native-implicit-otpm", + project_metadata={"model_otpm_limit": {model: 40}}, + ) + data = { + "model": model, + "contents": [{"role": "user", "parts": [{"text": "Hello"}]}], + } + if config_field is not None: + data[config_field] = {"temperature": 0} + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type=call_type, + ) + + stash = get_request_stash() + assert stash is not None + assert stash.otpm_reserved_tokens == 10 + expected_config_field = config_field or "config" + assert data[expected_config_field]["maxOutputTokens"] == 10 + assert "max_tokens" not in data + + +@pytest.mark.asyncio +async def test_project_otpm_over_limit_rolls_back_itpm_reservation(rate_limiter): + """ + When ITPM reserves fine but OTPM is then over limit, the ITPM + reservation this same pre-call already made must be rolled back -- + otherwise it leaks until the window's TTL, silently shrinking the ITPM + budget for every other request in that minute. + """ + handler, cache = rate_limiter + + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-otpm-rollback"), + project_id="proj-mantle-rollback", + project_metadata={ + "model_itpm_limit": {"bedrock_mantle/claude-opus": 1000000}, + "model_otpm_limit": {"bedrock_mantle/claude-opus": 10}, + }, + ) + + itpm_counter_key = handler.create_rate_limit_keys( + key="model_per_project_itpm", + value="proj-mantle-rollback:bedrock_mantle/claude-opus", + rate_limit_type="tokens", + ) + + data = { + "model": "bedrock_mantle/claude-opus", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 500, # blows past the 10-token OTPM limit + } + + with pytest.raises(Exception) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="", + ) + assert getattr(exc_info.value, "status_code", None) == 429 + + cached_value = await cache.async_get_cache(key=itpm_counter_key, local_only=True) + assert int(cached_value or 0) == 0, ( + f"ITPM reservation leaked after OTPM rejection: counter={cached_value}" + ) + + +@pytest.mark.asyncio +async def test_project_itpm_reconciled_on_success_excludes_cached_tokens(rate_limiter): + """ + On success, ITPM reconciles to billable input tokens (prompt_tokens + minus cached_tokens) -- not raw prompt_tokens. Cached prompt-read tokens + are free under Bedrock Mantle and must not count against the ITPM quota, + even though they still appear in usage/cost logging elsewhere. + """ + handler, _cache = rate_limiter + + itpm_scope = ("model_per_project_itpm", "proj-mantle:model") + otpm_scope = ("model_per_project_otpm", "proj-mantle:model") + + stash = get_or_create_request_stash() + stash.itpm_reserved_tokens = 100 + stash.itpm_reserved_scopes = frozenset({itpm_scope}) + stash.otpm_reserved_tokens = 60 + stash.otpm_reserved_scopes = frozenset({otpm_scope}) + mock_kwargs = {} + + mock_response = ModelResponse( + id="test", + object="chat.completion", + created=int(datetime.now().timestamp()), + model="bedrock_mantle/claude-opus", + usage=Usage( + prompt_tokens=80, + completion_tokens=40, + total_tokens=120, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=30), + ), + choices=[], + ) + + increments = [] + + async def mock_increment(increment_list, **kwargs): + for op in increment_list: + increments.append({"key": op["key"], "increment": op["increment_value"]}) + + handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( + mock_increment + ) + + await handler.async_log_success_event( + kwargs=mock_kwargs, + response_obj=mock_response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + itpm_adjustments = [i for i in increments if "model_per_project_itpm" in i["key"]] + otpm_adjustments = [i for i in increments if "model_per_project_otpm" in i["key"]] + + # billable_input = 80 - 30 cached = 50; delta = 50 - 100 reserved = -50 + assert any(i["increment"] == -50 for i in itpm_adjustments), ( + f"Expected a -50 ITPM adjustment (50 billable - 100 reserved), got: {itpm_adjustments}" + ) + # delta = 40 actual completion - 60 reserved = -20 + assert any(i["increment"] == -20 for i in otpm_adjustments), ( + f"Expected a -20 OTPM adjustment (40 actual - 60 reserved), got: {otpm_adjustments}" + ) + + +@pytest.mark.asyncio +async def test_project_reconciliation_does_not_decrement_later_window(): + current_time = datetime(2026, 8, 5, 12, 0, 0) + cache = DualCache() + handler = RateLimitHandler( + internal_usage_cache=InternalUsageCache(cache), + time_provider=lambda: current_time, + ) + handler.window_size = 60 + scope = (PROJECT_ITPM_DESCRIPTOR_KEY, "project:model") + descriptor = { + "key": scope[0], + "value": scope[1], + "rate_limit": {"tokens_per_unit": 1000, "window_size": 60}, + } + + reservation = await handler.atomic_check_and_increment_by_n( + descriptors=[descriptor], + increments=[{"tokens": 100}], + ) + counter_key = handler.create_rate_limit_keys(*scope, rate_limit_type="tokens") + window_identity = next( + identity + for identity in reservation["reservation_windows"] + if identity[0] == counter_key + ) + stash = get_or_create_request_stash() + stash.itpm_reserved_tokens = 100 + stash.itpm_reserved_scopes = frozenset({scope}) + stash.itpm_reserved_window_identities = frozenset( + {window_identity} + ) + + current_time += timedelta(seconds=61) + later_reservation = await handler.atomic_check_and_increment_by_n( + descriptors=[descriptor], + increments=[{"tokens": 20}], + ) + assert window_identity not in later_reservation["reservation_windows"] + + await handler.async_log_success_event( + kwargs={}, + response_obj=ModelResponse( + usage=Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10) + ), + start_time=current_time, + end_time=current_time, + ) + + assert float(await cache.async_get_cache(key=counter_key, local_only=True) or 0) == 20 + + +@pytest.mark.asyncio +async def test_project_reconciliation_decrements_its_active_window(rate_limiter): + handler, cache = rate_limiter + scope = (PROJECT_ITPM_DESCRIPTOR_KEY, "project:model") + descriptor = { + "key": scope[0], + "value": scope[1], + "rate_limit": {"tokens_per_unit": 1000, "window_size": 60}, + } + reservation = await handler.atomic_check_and_increment_by_n( + descriptors=[descriptor], + increments=[{"tokens": 100}], + ) + counter_key = handler.create_rate_limit_keys(*scope, rate_limit_type="tokens") + window_identity = next( + identity + for identity in reservation["reservation_windows"] + if identity[0] == counter_key + ) + stash = get_or_create_request_stash() + stash.itpm_reserved_tokens = 100 + stash.itpm_reserved_scopes = frozenset({scope}) + stash.itpm_reserved_window_identities = frozenset( + {window_identity} + ) + + await handler.async_log_success_event( + kwargs={}, + response_obj=ModelResponse( + usage=Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10) + ), + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert float(await cache.async_get_cache(key=counter_key, local_only=True) or 0) == 10 + + +@pytest.mark.asyncio +async def test_redis_window_guard_uses_reservation_identity_and_never_falls_back_negative( + rate_limiter, +): + handler, _cache = rate_limiter + calls = [] + + async def failing_guard(*, keys, args): + calls.append((keys, args)) + raise RuntimeError("redis unavailable") + + unguarded_calls = [] + + async def capture_unguarded(pipeline_operations, **_kwargs): + unguarded_calls.extend(pipeline_operations) + + handler.window_guarded_token_increment_script = failing_guard + handler.async_increment_tokens_with_ttl_preservation = capture_unguarded + await handler.async_increment_reservation_aware_tokens( + pipeline_operations=[ + { + "key": "{model_per_project_itpm:project:model}:tokens", + "increment_value": -90, + "ttl": 60, + "window_key": "{model_per_project_itpm:project:model}:window", + "expected_window_start": "1234", + "reservation_backend": "redis", + } + ] + ) + + assert calls == [ + ( + [ + "{model_per_project_itpm:project:model}:window", + "{model_per_project_itpm:project:model}:tokens", + ], + ["1234", -90, 60], + ) + ] + assert unguarded_calls == [] + + +@pytest.mark.asyncio +async def test_atomic_lua_response_carries_redis_window_identity(rate_limiter): + handler, _cache = rate_limiter + counter_key = "{model_per_project_itpm:project:model}:tokens" + meta = [ + { + "descriptor_key": PROJECT_ITPM_DESCRIPTOR_KEY, + "current_limit": 100, + "rate_limit_type": "tokens", + "counter_key": counter_key, + } + ] + + async def successful_reservation(*, keys, args): + return [0, 25, 1234] + + handler.check_and_increment_by_n_script = successful_reservation + assert await handler._atomic_lua_per_descriptor([]) == { + "overall_code": "OK", + "statuses": [], + } + + response = await handler._atomic_lua_per_descriptor( + descriptor_groups=[ + ( + [ + "{model_per_project_itpm:project:model}:window", + counter_key, + ], + [100, 25, 60, 60], + meta, + ) + ] + ) + + assert response["statuses"][0]["limit_remaining"] == 75 + assert response["reservation_windows"] == frozenset( + {(counter_key, "1234", "redis")} + ) + + +@pytest.mark.asyncio +async def test_project_itpm_otpm_released_on_failure(rate_limiter): + """On failure, the full ITPM and OTPM reservations must be refunded.""" + handler, _cache = rate_limiter + + itpm_scope = ("model_per_project_itpm", "proj-mantle:model") + otpm_scope = ("model_per_project_otpm", "proj-mantle:model") + + stash = get_or_create_request_stash() + stash.itpm_reserved_tokens = 100 + stash.itpm_reserved_scopes = frozenset({itpm_scope}) + stash.otpm_reserved_tokens = 60 + stash.otpm_reserved_scopes = frozenset({otpm_scope}) + mock_kwargs = {} + + increments = [] + + async def mock_increment(increment_list, **kwargs): + for op in increment_list: + increments.append({"key": op["key"], "increment": op["increment_value"]}) + + handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( + mock_increment + ) + + await handler.async_log_failure_event( + kwargs=mock_kwargs, + response_obj=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + itpm_releases = [i for i in increments if "model_per_project_itpm" in i["key"]] + otpm_releases = [i for i in increments if "model_per_project_otpm" in i["key"]] + + assert any(i["increment"] == -100 for i in itpm_releases), itpm_releases + assert any(i["increment"] == -60 for i in otpm_releases), otpm_releases + + +@pytest.mark.asyncio +async def test_proxy_rejection_refunds_itpm_otpm_by_their_own_amount_not_combined( + rate_limiter, +): + """ + Regression for a Greptile-flagged bug: when a project configures both a + combined model_tpm_limit and split model_itpm_limit/model_otpm_limit for + the same model, async_post_call_failure_hook's proxy-side refund path + used to decrement every token descriptor -- including the ITPM/OTPM + ones -- by the flat combined reservation amount, instead of each + bucket's own reserved amount. That drives the split counters negative + (or under-refunds them) instead of returning them to exactly zero. + """ + handler, cache = rate_limiter + + api_key = hash_token("sk-mixed-tpm-io") + user_api_key_dict = UserAPIKeyAuth( + api_key=api_key, + project_id="proj-mixed", + project_metadata={ + "model_tpm_limit": {"bedrock_mantle/claude-opus": 100000}, + "model_itpm_limit": {"bedrock_mantle/claude-opus": 100000}, + "model_otpm_limit": {"bedrock_mantle/claude-opus": 100000}, + }, + ) + + data = { + "model": "bedrock_mantle/claude-opus", + "messages": [ + {"role": "user", "content": "hello there, this is a test message"} + ], + "max_tokens": 60, + } + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="", + ) + + tpm_counter_key = handler.create_rate_limit_keys( + key="model_per_project", + value="proj-mixed:bedrock_mantle/claude-opus", + rate_limit_type="tokens", + ) + itpm_counter_key = handler.create_rate_limit_keys( + key="model_per_project_itpm", + value="proj-mixed:bedrock_mantle/claude-opus", + rate_limit_type="tokens", + ) + otpm_counter_key = handler.create_rate_limit_keys( + key="model_per_project_otpm", + value="proj-mixed:bedrock_mantle/claude-opus", + rate_limit_type="tokens", + ) + + tpm_reserved = int( + await cache.async_get_cache(key=tpm_counter_key, local_only=True) or 0 + ) + itpm_reserved = int( + await cache.async_get_cache(key=itpm_counter_key, local_only=True) or 0 + ) + otpm_reserved = int( + await cache.async_get_cache(key=otpm_counter_key, local_only=True) or 0 + ) + assert tpm_reserved > 0 and itpm_reserved > 0 and otpm_reserved > 0 + + await handler.async_post_call_failure_hook( + request_data=data, + original_exception=Exception("guardrail rejected"), + user_api_key_dict=user_api_key_dict, + ) + + tpm_after = int( + await cache.async_get_cache(key=tpm_counter_key, local_only=True) or 0 + ) + itpm_after = int( + await cache.async_get_cache(key=itpm_counter_key, local_only=True) or 0 + ) + otpm_after = int( + await cache.async_get_cache(key=otpm_counter_key, local_only=True) or 0 + ) + + assert tpm_after == 0, f"combined TPM counter leaked: {tpm_after}" + assert itpm_after == 0, ( + f"ITPM counter corrupted by combined-amount refund: {itpm_after}" + ) + assert otpm_after == 0, ( + f"OTPM counter corrupted by combined-amount refund: {otpm_after}" + ) + + +@pytest.mark.asyncio +async def test_proxy_rejection_refunds_itpm_otpm_only_reservation_with_no_combined_tpm( + rate_limiter, +): + """ + Regression for the second half of the same bug: with only + model_itpm_limit/model_otpm_limit configured (no model_tpm_limit), the + combined reserved_tokens is 0, and the proxy-side refund path used to + return immediately on that -- leaking the ITPM/OTPM reservations until + the rate-limit window's TTL expired. + """ + handler, cache = rate_limiter + + api_key = hash_token("sk-io-only") + user_api_key_dict = UserAPIKeyAuth( + api_key=api_key, + project_id="proj-io-only", + project_metadata={ + "model_itpm_limit": {"bedrock_mantle/claude-opus": 100000}, + "model_otpm_limit": {"bedrock_mantle/claude-opus": 100000}, + }, + ) + + data = { + "model": "bedrock_mantle/claude-opus", + "messages": [ + {"role": "user", "content": "hello there, this is a test message"} + ], + "max_tokens": 60, + } + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="", + ) + + itpm_counter_key = handler.create_rate_limit_keys( + key="model_per_project_itpm", + value="proj-io-only:bedrock_mantle/claude-opus", + rate_limit_type="tokens", + ) + otpm_counter_key = handler.create_rate_limit_keys( + key="model_per_project_otpm", + value="proj-io-only:bedrock_mantle/claude-opus", + rate_limit_type="tokens", + ) + assert ( + int(await cache.async_get_cache(key=itpm_counter_key, local_only=True) or 0) > 0 + ) + assert ( + int(await cache.async_get_cache(key=otpm_counter_key, local_only=True) or 0) > 0 + ) + + await handler.async_post_call_failure_hook( + request_data=data, + original_exception=Exception("guardrail rejected"), + user_api_key_dict=user_api_key_dict, + ) + + itpm_after = int( + await cache.async_get_cache(key=itpm_counter_key, local_only=True) or 0 + ) + otpm_after = int( + await cache.async_get_cache(key=otpm_counter_key, local_only=True) or 0 + ) + assert itpm_after == 0, ( + f"ITPM-only reservation leaked on proxy rejection: {itpm_after}" + ) + assert otpm_after == 0, ( + f"OTPM-only reservation leaked on proxy rejection: {otpm_after}" + ) + + +@pytest.mark.asyncio +async def test_otpm_rejection_does_not_double_refund_combined_tpm(rate_limiter): + """ + Regression for a High-severity review finding: when the project ITPM + reservation succeeds but OTPM is then over limit, + _reserve_project_io_tokens_or_raise rolls back the combined-TPM + reservation that already succeeded earlier in the same pre-call, then + raises. If it doesn't also mark that reservation released, + async_post_call_failure_hook -- which fires next in the real request + lifecycle, since raising from async_pre_call_hook triggers it -- sees + the same still-stashed reservation and refunds it a second time, + driving the combined TPM counter negative and letting a caller push + past the project's real TPM budget by repeatedly triggering OTPM + rejections. + """ + handler, cache = rate_limiter + + api_key = hash_token("sk-double-refund") + user_api_key_dict = UserAPIKeyAuth( + api_key=api_key, + project_id="proj-double-refund", + project_metadata={ + "model_tpm_limit": {"bedrock_mantle/claude-opus": 100000}, + "model_itpm_limit": {"bedrock_mantle/claude-opus": 100000}, + "model_otpm_limit": {"bedrock_mantle/claude-opus": 5}, + }, + ) + + data = { + "model": "bedrock_mantle/claude-opus", + "messages": [ + {"role": "user", "content": "hello there, this is a test message"} + ], + "max_tokens": 60, # blows past the 5-token OTPM limit + } + + tpm_counter_key = handler.create_rate_limit_keys( + key="model_per_project", + value="proj-double-refund:bedrock_mantle/claude-opus", + rate_limit_type="tokens", + ) + + with pytest.raises(Exception) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="", + ) + assert getattr(exc_info.value, "status_code", None) == 429 + + tpm_after_pre_call = int( + await cache.async_get_cache(key=tpm_counter_key, local_only=True) or 0 + ) + assert tpm_after_pre_call == 0, ( + f"combined TPM reservation not rolled back: {tpm_after_pre_call}" + ) + + # In the real request lifecycle, async_post_call_failure_hook fires next + # for a pre-call rejection. It must not refund the same reservation again. + await handler.async_post_call_failure_hook( + request_data=data, + original_exception=exc_info.value, + user_api_key_dict=user_api_key_dict, + ) + + tpm_after_failure_hook = int( + await cache.async_get_cache(key=tpm_counter_key, local_only=True) or 0 + ) + assert tpm_after_failure_hook == 0, ( + f"combined TPM counter went negative from a double refund: {tpm_after_failure_hook}" + ) + + +@pytest.mark.parametrize( + "embedding_input", + [ + list(range(51)), + [list(range(25)), list(range(26))], + ], +) +@pytest.mark.asyncio +async def test_project_itpm_rejects_pretokenized_embedding_input( + rate_limiter, + embedding_input, +): + handler, cache = rate_limiter + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-pretokenized-embedding-itpm"), + project_id="proj-pretokenized-embedding", + project_metadata={ + "model_itpm_limit": {"text-embedding-3-small": 50}, + }, + ) + data = { + "model": "text-embedding-3-small", + "input": embedding_input, + } + + with pytest.raises(Exception) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="aembedding", + ) + + assert getattr(exc_info.value, "status_code", None) == 429 + + +@pytest.mark.asyncio +async def test_responses_api_not_misclassified_as_embedding_for_output_estimate( + rate_limiter, +): + """ + Regression for a High-severity review finding: the Responses API also + puts its prompt in data["input"], the same field embeddings use, so the + output-token estimate treated every Responses call as an embedding and + reserved zero output tokens. call_type now disambiguates the two: the + same input-only payload gets zero output tokens for an embedding call + but a real floor for a Responses API call. + """ + handler, _cache = rate_limiter + + data = {"input": "describe this image in detail"} + + _, embedding_output_estimate = handler._estimate_input_and_output_tokens( + data=data, call_type="aembedding" + ) + assert embedding_output_estimate == 0 + + _, responses_output_estimate = handler._estimate_input_and_output_tokens( + data=data, call_type="aresponses" + ) + assert responses_output_estimate > 0, ( + "Responses API call was misclassified as an embedding and reserved zero output tokens" + ) + + +@pytest.mark.parametrize( + ("data", "call_type", "expected_output_tokens"), + [ + ( + { + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 100, + "n": 10, + }, + "acompletion", + 1000, + ), + ( + { + "prompt": "hello", + "max_tokens": 100, + "n": 2, + "best_of": 5, + }, + "text_completion", + 500, + ), + ( + { + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 100, + "n": 0, + "best_of": "invalid", + }, + "acompletion", + 100, + ), + ], +) +def test_output_estimate_accounts_for_completion_candidates( + rate_limiter, + data, + call_type, + expected_output_tokens, +): + handler, _cache = rate_limiter + + _, estimated_output_tokens = handler._estimate_input_and_output_tokens( + data=data, + call_type=call_type, + ) + + assert estimated_output_tokens == expected_output_tokens + + +@pytest.mark.asyncio +async def test_responses_api_usage_reconciles_using_input_output_tokens_fields( + rate_limiter, +): + """ + Regression for the other half of the same finding: ResponseAPIUsage + exposes input_tokens/output_tokens, not prompt_tokens/completion_tokens. + Before this fix, _resolve_io_token_reconcile_usage couldn't resolve + Responses API usage at all, so the reservation was silently kept as-is + instead of being trued up to the much larger actual usage. + """ + handler, _cache = rate_limiter + + itpm_scope = ("model_per_project_itpm", "proj-responses:model") + otpm_scope = ("model_per_project_otpm", "proj-responses:model") + + stash = get_or_create_request_stash() + stash.itpm_reserved_tokens = 10 + stash.itpm_reserved_scopes = frozenset({itpm_scope}) + stash.otpm_reserved_tokens = 10 + stash.otpm_reserved_scopes = frozenset({otpm_scope}) + mock_kwargs = {} + + mock_response = ResponsesAPIResponse( + id="resp_test", + created_at=int(datetime.now().timestamp()), + output=[], + usage=ResponseAPIUsage(input_tokens=80, output_tokens=400, total_tokens=480), + ) + + increments = [] + + async def mock_increment(increment_list, **kwargs): + for op in increment_list: + increments.append({"key": op["key"], "increment": op["increment_value"]}) + + handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( + mock_increment + ) + + await handler.async_log_success_event( + kwargs=mock_kwargs, + response_obj=mock_response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + itpm_adjustments = [i for i in increments if "model_per_project_itpm" in i["key"]] + otpm_adjustments = [i for i in increments if "model_per_project_otpm" in i["key"]] + + # delta = 80 actual input - 10 reserved = +70 + assert any(i["increment"] == 70 for i in itpm_adjustments), ( + f"ITPM reservation was never trued up to actual Responses API usage: {itpm_adjustments}" + ) + # delta = 400 actual output - 10 reserved = +390 + assert any(i["increment"] == 390 for i in otpm_adjustments), ( + f"OTPM reservation was never trued up to actual Responses API usage: {otpm_adjustments}" + ) + + +@pytest.mark.asyncio +async def test_itpm_reservation_accounts_for_audio_content_not_just_text(rate_limiter): + """ + Regression for the audio half of a Medium-severity review finding: + litellm.token_counter has no per-type handling for `input_audio` + content blocks (unlike images, which it does count via + use_default_image_token_count), so it silently contributes 0 tokens for + them. Without DEFAULT_AUDIO_TOKEN_ESTIMATE, a burst of audio-heavy + requests with minimal text would each reserve only the one-token floor + and blow past the project ITPM limit before post-call reconciliation. + """ + handler, cache = rate_limiter + + api_key = hash_token("sk-audio-itpm") + user_api_key_dict = UserAPIKeyAuth( + api_key=api_key, + project_id="proj-audio", + project_metadata={ + # Tighter than DEFAULT_AUDIO_TOKEN_ESTIMATE (300), but far bigger + # than the handful of tokens the bare text "hi" would cost. + "model_itpm_limit": {"bedrock_mantle/claude-opus": 50}, + }, + ) + + data = { + "model": "bedrock_mantle/claude-opus", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hi"}, + { + "type": "input_audio", + "input_audio": {"data": "base64-audio-bytes", "format": "wav"}, + }, + ], + } + ], + } + + with pytest.raises(Exception) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="", + ) + assert getattr(exc_info.value, "status_code", None) == 429, ( + "Expected the audio content to push the ITPM reservation over the " + "50-token limit; if this doesn't raise, audio content isn't being " + "counted again." + ) + + +def test_audio_token_estimate_scales_with_payload_size(): + """ + Regression for veria-ai Low finding: audio token reservation was flat + 300 per block regardless of duration. A short clip and a long clip both + reserved the same amount, letting a caller hide long audio in one block + to exhaust ITPM quota while reserving almost nothing. + + The estimate must now grow proportionally with the base64 payload size + (len(b64) * 3 // 4 // _AUDIO_BYTES_PER_TOKEN), floored at + DEFAULT_AUDIO_TOKEN_ESTIMATE so reference-only blocks and genuinely + short clips still get a non-trivial reservation. + + To exceed the floor the decoded payload must be > 300 * 1600 = 480 000 + bytes. We synthesise a fake b64-length string of 650 000 chars + (decoded ≈ 487 500 bytes → 304 tokens) to avoid actually allocating + and encoding ~480 kB of audio in every test run. + """ + large_b64 = "A" * 650_000 + very_large_b64 = "A" * 12_900_000 + small_b64 = "A" * 1_000 + + large_block = { + "type": "input_audio", + "input_audio": {"data": large_b64, "format": "wav"}, + } + small_block = { + "type": "input_audio", + "input_audio": {"data": small_b64, "format": "wav"}, + } + very_large_block = { + "type": "input_audio", + "input_audio": {"data": very_large_b64, "format": "wav"}, + } + no_data_block = {"type": "input_audio", "input_audio": {"format": "wav"}} + + large_estimate = RateLimitHandler._estimate_audio_block_tokens(large_block) + very_large_estimate = RateLimitHandler._estimate_audio_block_tokens( + very_large_block + ) + small_estimate = RateLimitHandler._estimate_audio_block_tokens(small_block) + no_data_estimate = RateLimitHandler._estimate_audio_block_tokens(no_data_block) + + assert large_estimate > small_estimate, ( + f"Large payload ({large_estimate}) must reserve more than small payload " + f"({small_estimate}); flat-rate bug is back" + ) + assert very_large_estimate == len(very_large_b64) * 3 // 4 // _AUDIO_BYTES_PER_TOKEN + assert very_large_estimate > 6_000 + assert no_data_estimate >= 300, ( + f"Reference-only block (no data) must use the DEFAULT_AUDIO_TOKEN_ESTIMATE floor; got {no_data_estimate}" + ) + assert small_estimate >= 300, ( + f"Small payload must be floored at DEFAULT_AUDIO_TOKEN_ESTIMATE=300; got {small_estimate}" + ) + + +@pytest.mark.asyncio +async def test_itpm_rejects_large_audio_payload_that_would_pass_flat_estimate( + rate_limiter, +): + """ + Regression: a caller placing a long audio clip in one block previously + reserved only 300 tokens (the flat estimate). With the size-proportional + estimate, the same clip now reserves proportionally more and must trip + the ITPM limit when the limit is tuned to exactly expose the difference. + + 1 100 000 b64 chars → decoded ≈ 825 000 bytes → 825 000 // 1600 ≈ 515 + tokens > the 400-token limit. The flat estimate (300) would have passed. + """ + handler, cache = rate_limiter + + large_b64 = "A" * 1_100_000 + + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-large-audio"), + project_id="proj-large-audio", + project_metadata={ + "model_itpm_limit": {"bedrock_mantle/claude-opus": 400}, + }, + ) + + data = { + "model": "bedrock_mantle/claude-opus", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "transcribe this"}, + { + "type": "input_audio", + "input_audio": {"data": large_b64, "format": "wav"}, + }, + ], + } + ], + } + + with pytest.raises(Exception) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="", + ) + assert getattr(exc_info.value, "status_code", None) == 429, ( + "Large audio payload must exceed the 400-token ITPM limit under the " + "size-proportional estimate; the old flat-rate estimate (300 tokens) " + "would have passed this limit silently" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "request_data"), + [ + ( + "acompletion", + { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "https://example.com/high-resolution.png", + "detail": "high", + }, + } + ], + } + ] + }, + ), + ( + "aresponses", + { + "input": [ + { + "role": "user", + "content": [ + { + "type": "input_image", + "image_url": "https://example.com/high-resolution.png", + "detail": "high", + } + ], + } + ] + }, + ), + ], +) +async def test_image_content_reserves_full_project_itpm( + rate_limiter, + call_type, + request_data, +): + handler, cache = rate_limiter + model = "bedrock_mantle/claude-opus" + project_itpm_limit = 1_000 + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-high-resolution-image"), + project_id="project-high-resolution-image", + project_metadata={"model_itpm_limit": {model: project_itpm_limit}}, + ) + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data={"model": model, **request_data}, + call_type=call_type, + ) + + stash = get_request_stash() + assert stash is not None + assert stash.itpm_reserved_tokens == project_itpm_limit + + +@pytest.mark.asyncio +async def test_itpm_otpm_reservation_is_kept_on_stream_disconnect(rate_limiter): + handler, cache = rate_limiter + + api_key = hash_token("sk-disconnect-test") + user_api_key_dict = UserAPIKeyAuth( + api_key=api_key, + project_id="proj-disconnect", + project_metadata={ + "model_itpm_limit": {"bedrock_mantle/claude-opus": 1000}, + "model_otpm_limit": {"bedrock_mantle/claude-opus": 500}, + }, + ) + + data: Dict[str, Any] = { + "model": "bedrock_mantle/claude-opus", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 50, + } + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="", + ) + + stash = get_request_stash() + assert stash is not None + assert stash.itpm_reserved_tokens > 0, ( + "pre-call hook must stash an ITPM reservation" + ) + assert stash.otpm_reserved_tokens > 0, ( + "pre-call hook must stash an OTPM reservation" + ) + + increment_calls: list[dict] = [] + + async def mock_increment(increment_list, litellm_parent_otel_span=None): + for op in increment_list: + increment_calls.append( + {"key": op["key"], "increment": op["increment_value"]} + ) + + handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( + mock_increment + ) + + await handler.async_release_max_parallel_requests_on_disconnect( + user_api_key_dict=user_api_key_dict + ) + + itpm_refunds = [ + c + for c in increment_calls + if "model_per_project_itpm" in c["key"] and c["increment"] < 0 + ] + otpm_refunds = [ + c + for c in increment_calls + if "model_per_project_otpm" in c["key"] and c["increment"] < 0 + ] + + assert not itpm_refunds + assert not otpm_refunds + assert stash.reservation_released is False + + +@pytest.mark.asyncio +async def test_responses_api_otpm_output_cap_applied_not_skipped_as_embedding( + rate_limiter, +): + """ + Regression for a Greptile P1 finding: _reserve_project_io_tokens_or_raise + classified any request with data["input"] set as an embedding (no output + tokens), which also misclassifies the Responses API -- it puts its prompt + in "input" too, but does generate output. That skipped the output cap + applied whenever the configured OTPM limit is small enough to need it, + letting an unbounded Responses generation blow past OTPM before + post-call reconciliation catches up. + + The cap must land on data["max_output_tokens"], not data["max_tokens"]: + the Responses-to-chat-completion transformation only reads + max_output_tokens, so a max_tokens cap is silently dropped before + provider dispatch (a second Greptile finding on the same code path). + """ + handler, cache = rate_limiter + + api_key = hash_token("sk-responses-otpm-cap") + user_api_key_dict = UserAPIKeyAuth( + api_key=api_key, + project_id="proj-responses-otpm", + project_metadata={ + "model_otpm_limit": {"bedrock_mantle/claude-opus": 40}, + }, + ) + + data: Dict[str, Any] = { + "model": "bedrock_mantle/claude-opus", + "input": "describe this image in detail", + } + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="aresponses", + ) + + assert data.get("max_output_tokens") is not None, ( + "Responses call was misclassified as an embedding and skipped the OTPM output cap" + ) + assert data["max_output_tokens"] == 16 + assert data.get("max_tokens") is None, ( + "OTPM output cap was written to max_tokens, which the Responses transformation ignores" + ) + + +@pytest.mark.asyncio +async def test_explicit_zero_output_responses_call_reserves_effective_provider_minimum( + rate_limiter, +): + handler, cache = rate_limiter + + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-responses-zero-output"), + project_id="proj-responses-zero-output", + project_metadata={ + "model_otpm_limit": {"bedrock_mantle/claude-opus": 5}, + }, + ) + + with pytest.raises(Exception) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data={ + "model": "bedrock_mantle/claude-opus", + "input": "describe this image in detail", + "max_output_tokens": 0, + }, + call_type="aresponses", + ) + + assert getattr(exc_info.value, "status_code", None) == 429 + + +@pytest.mark.asyncio +async def test_responses_api_combined_tpm_output_cap_applied_not_skipped_as_embedding( + rate_limiter, +): + """ + Regression for the same misclassification bug in the combined-TPM + output-cap block of async_pre_call_hook (a second, independent + `is_embedding = data.get("input") is not None` check). A project with + only a combined model_tpm_limit (no split itpm/otpm) configured small + enough to need the output cap must still apply it to a Responses call, + and must write it to max_output_tokens for the same reason as the OTPM + case above. + """ + handler, cache = rate_limiter + + api_key = hash_token("sk-responses-tpm-cap") + user_api_key_dict = UserAPIKeyAuth( + api_key=api_key, + project_id="proj-responses-tpm", + project_metadata={ + "model_tpm_limit": {"bedrock_mantle/claude-opus": 40}, + }, + ) + + data: Dict[str, Any] = { + "model": "bedrock_mantle/claude-opus", + "input": "describe this image in detail", + } + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="aresponses", + ) + + assert data.get("max_output_tokens") is not None, ( + "Responses call was misclassified as an embedding and skipped the combined-TPM output cap" + ) + assert data["max_output_tokens"] == 16 + assert data.get("max_tokens") is None, ( + "combined-TPM output cap was written to max_tokens, which the Responses transformation ignores" + ) + + +@pytest.mark.asyncio +async def test_responses_api_multimodal_input_counts_image_content(rate_limiter): + """ + Regression for a Low-severity veria-ai finding: the Responses API's + `input` is commonly a list of message/content-block dicts, but + litellm.token_counter's `text` argument only joins plain string entries + in a list and silently drops everything else -- so an `input_image` + block contributed ~0 tokens to the ITPM estimate instead of the real + image token count. _estimate_precise_input_tokens now converts Responses + `input` to chat messages first (via the standard + transform_responses_api_input_to_messages helper) so image content is + counted the same way a chat completion's image content already is. + """ + handler, _cache = rate_limiter + + text_only_estimate = handler._estimate_precise_input_tokens( + data={"input": "hi"}, + model="bedrock_mantle/claude-opus", + call_type="aresponses", + ) + + multimodal_estimate = handler._estimate_precise_input_tokens( + data={ + "input": [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "hi"}, + { + "type": "input_image", + "image_url": "https://example.com/some-image.png", + }, + ], + } + ], + }, + model="bedrock_mantle/claude-opus", + call_type="aresponses", + ) + + assert multimodal_estimate > text_only_estimate + 100, ( + "Responses API input_image content block was not counted; got " + f"text_only={text_only_estimate}, multimodal={multimodal_estimate}" + ) + + +@pytest.mark.asyncio +async def test_refund_reserved_tokens_noop_when_amount_zero(rate_limiter): + """_refund_reserved_tokens returns immediately without calling Redis when amount=0.""" + handler, _cache = rate_limiter + + calls = [] + + async def mock_increment(pipeline_operations, **kwargs): + calls.extend(pipeline_operations) + + handler.async_increment_tokens_with_ttl_preservation = mock_increment + + await handler._refund_reserved_tokens( + scopes=[("api_key", "sk-test")], + amount=0, + ) + + assert not calls, "No Redis ops expected when amount is zero" + + +@pytest.mark.asyncio +async def test_reserve_io_tokens_noop_when_no_itpm_otpm_descriptors(rate_limiter): + """reserve_io_tokens returns OK immediately when no ITPM/OTPM descriptors present.""" + handler, _cache = rate_limiter + + non_io_descriptor = { + "key": "api_key", + "value": "sk-test", + "rate_limit": {"tokens_per_unit": 1000, "window_size": 60}, + } + response, itpm_reserved, otpm_reserved = await handler.reserve_io_tokens( + descriptors=[non_io_descriptor], + estimated_input_tokens=50, + estimated_output_tokens=50, + ) + + assert response["overall_code"] == "OK" + assert itpm_reserved == 0 + assert otpm_reserved == 0 + + +@pytest.mark.asyncio +async def test_reserve_io_tokens_itpm_only_no_otpm(rate_limiter): + """When only ITPM descriptors are present (no OTPM), returns itpm_reserved with otpm=0.""" + handler, cache = rate_limiter + + itpm_descriptor = { + "key": PROJECT_ITPM_DESCRIPTOR_KEY, + "value": "proj-a:model", + "rate_limit": {"tokens_per_unit": 10000, "window_size": 60}, + } + response, itpm_reserved, otpm_reserved = await handler.reserve_io_tokens( + descriptors=[itpm_descriptor], + estimated_input_tokens=100, + estimated_output_tokens=50, + ) + + assert response["overall_code"] == "OK" + assert itpm_reserved == 100 + assert otpm_reserved == 0 + + +def test_strip_audio_content_blocks_passthrough_non_list_messages(): + """Non-list input is returned unchanged (early return on line 2605).""" + result = RateLimitHandler._strip_audio_content_blocks("not a list") + assert result == "not a list" + + +def test_strip_audio_content_blocks_passthrough_non_dict_message(): + """Non-dict entries in the message list are appended unchanged.""" + messages = ["plain string message"] + result = RateLimitHandler._strip_audio_content_blocks(messages) + assert result == ["plain string message"] + + +def test_strip_audio_content_blocks_passthrough_non_list_content(): + """Messages with non-list content (e.g. plain string) pass through unchanged.""" + messages = [{"role": "user", "content": "hello"}] + result = RateLimitHandler._strip_audio_content_blocks(messages) + assert result == [{"role": "user", "content": "hello"}] + + +@pytest.mark.asyncio +async def test_otpm_rejection_releases_stashed_parallel_slot(rate_limiter): + """ + When OTPM is over limit and a parallel slot was already acquired, the + disconnect cleanup path in _reserve_project_io_tokens_or_raise must + release that slot. Exercises lines 2773-2777. + """ + handler, cache = rate_limiter + + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-otpm-slot"), + project_id="proj-slot", + project_metadata={"model_otpm_limit": {"m": 5}}, + ) + + data: Dict[str, Any] = { + "model": "m", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 50, + } + + slot_released = [] + + async def mock_release(acquisition, parent_otel_span=None): + slot_released.append(acquisition) + + handler._release_parallel_request_slots = mock_release + + stash = get_or_create_request_stash() + stash.parallel_slot = { + "slot_id": "test-slot-id", + "counter_keys": ["some-key"], + } + + otpm_descriptor = { + "key": PROJECT_OTPM_DESCRIPTOR_KEY, + "value": "proj-slot:m", + "rate_limit": {"tokens_per_unit": 5, "window_size": 60}, + } + + with pytest.raises(Exception) as exc_info: + await handler._reserve_project_io_tokens_or_raise( + descriptors=[otpm_descriptor], + data=data, + requested_model="m", + user_api_key_dict=user_api_key_dict, + tpm_reservation_scopes=[], + tpm_reservation_amount=0, + ) + assert getattr(exc_info.value, "status_code", None) == 429 + assert slot_released, "Parallel slot must be released when OTPM rejects" + assert stash.parallel_slot is None + + +@pytest.mark.asyncio +async def test_itpm_only_status_stored_when_no_prior_rate_limit_response(rate_limiter): + """ + When only ITPM is configured (no combined TPM/RPM to pre-populate + the request stash), a successful ITPM reservation must store its status + there so post-call headers can read it. + """ + handler, cache = rate_limiter + + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-itpm-only-store"), + project_id="proj-store", + ) + + data: Dict[str, Any] = {"model": "m", "messages": []} + + itpm_descriptor = { + "key": PROJECT_ITPM_DESCRIPTOR_KEY, + "value": "proj-store:m", + "rate_limit": {"tokens_per_unit": 100000, "window_size": 60}, + } + + await handler._reserve_project_io_tokens_or_raise( + descriptors=[itpm_descriptor], + data=data, + requested_model="m", + user_api_key_dict=user_api_key_dict, + tpm_reservation_scopes=[], + tpm_reservation_amount=0, + ) + + stash = get_request_stash() + assert stash is not None + stored = stash.rate_limit_response + assert stored is not None, ( + "ITPM status must be stored in litellm_proxy_rate_limit_response" + ) + assert stored.get("statuses"), "Stored response must contain statuses" + + +def test_resolve_io_token_usage_responses_api_with_cached_tokens(rate_limiter): + """ + ResponsesAPIResponse whose usage.input_tokens_details.cached_tokens is set + subtracts the cached portion from billable input. Covers line 3501. + """ + handler, _cache = rate_limiter + + response_obj = ResponsesAPIResponse( + id="resp_cached", + created_at=int(datetime.now().timestamp()), + output=[], + usage=ResponseAPIUsage( + input_tokens=100, + output_tokens=50, + total_tokens=150, + input_tokens_details=InputTokensDetails(cached_tokens=25), + ), + ) + billable_input, completion_tokens, resolved = ( + handler._resolve_io_token_reconcile_usage(response_obj) + ) + + assert resolved is True + assert billable_input == 75, f"Expected 100 - 25 cached = 75, got {billable_input}" + assert completion_tokens == 50 + + +def test_resolve_io_token_usage_dict_format(rate_limiter): + """ + Dict-shaped usage on a ModelResponse (older SDK versions or raw dicts in + the usage field) is parsed correctly. Covers lines 3502-3506. + """ + handler, _cache = rate_limiter + + response_obj = ModelResponse.model_construct( + usage={ + "prompt_tokens": 80, + "completion_tokens": 40, + "prompt_tokens_details": {"cached_tokens": 20}, + } + ) + billable_input, completion_tokens, resolved = ( + handler._resolve_io_token_reconcile_usage(response_obj) + ) + + assert resolved is True + assert billable_input == 60, f"Expected 80 - 20 cached = 60, got {billable_input}" + assert completion_tokens == 40 + + +def test_resolve_io_token_usage_unknown_type_returns_unresolved(rate_limiter): + """ + A ModelResponse whose usage attribute is not a Usage, ResponseAPIUsage, + or dict (e.g. a plain int) returns (0, 0, False) so the reservation is + kept rather than guessed. Covers lines 3507-3508. + """ + handler, _cache = rate_limiter + + response_obj = ModelResponse.model_construct(usage=42) + billable_input, completion_tokens, resolved = ( + handler._resolve_io_token_reconcile_usage(response_obj) + ) + + assert resolved is False + assert billable_input == 0 + assert completion_tokens == 0 + + +@pytest.mark.parametrize( + ("combined_usage", "expected_increments"), + [ + (None, ()), + ( + Usage(prompt_tokens=40, completion_tokens=15, total_tokens=55), + (-60, -45), + ), + ], +) +def test_zero_usage_keeps_reservations_unless_measured_fallback_exists( + rate_limiter, + combined_usage, + expected_increments, +): + handler, _cache = rate_limiter + itpm_scope = (PROJECT_ITPM_DESCRIPTOR_KEY, "project:model") + otpm_scope = (PROJECT_OTPM_DESCRIPTOR_KEY, "project:model") + stash = get_or_create_request_stash() + stash.itpm_reserved_tokens = 100 + stash.itpm_reserved_scopes = frozenset({itpm_scope}) + stash.otpm_reserved_tokens = 60 + stash.otpm_reserved_scopes = frozenset({otpm_scope}) + kwargs = {} if combined_usage is None else {"combined_usage_object": combined_usage} + response_obj = ModelResponse( + usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) + ) + + operations = handler._build_io_token_reservation_ops(kwargs, response_obj) + + assert tuple(operation["increment_value"] for operation in operations) == expected_increments + + +@pytest.mark.parametrize( + ("usage", "expected_increments"), + [ + ( + Usage(prompt_tokens=40, completion_tokens=15, total_tokens=55), + (40, 15), + ), + ( + Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), + (100, 60), + ), + ], +) +def test_retry_success_charges_released_project_io_reservations( + rate_limiter, + usage, + expected_increments, +): + handler, _cache = rate_limiter + itpm_scope = (PROJECT_ITPM_DESCRIPTOR_KEY, "project:model") + otpm_scope = (PROJECT_OTPM_DESCRIPTOR_KEY, "project:model") + stash = get_or_create_request_stash() + stash.itpm_reserved_tokens = 100 + stash.itpm_reserved_scopes = frozenset({itpm_scope}) + stash.otpm_reserved_tokens = 60 + stash.otpm_reserved_scopes = frozenset({otpm_scope}) + stash.reservation_released = True + + operations = handler._build_io_token_reservation_ops( + {}, + ModelResponse(usage=usage), + ) + + assert tuple(operation["increment_value"] for operation in operations) == expected_increments + + +@pytest.mark.asyncio +async def test_build_io_token_reservation_ops_skips_unresolvable_usage(rate_limiter): + """ + When response_obj has no parseable usage, _build_io_token_reservation_ops + returns [] to keep the reservation as-is rather than zeroing it out on a + bad guess. Covers line 3538. + """ + handler, _cache = rate_limiter + + itpm_scope = ("model_per_project_itpm", "proj-b:model") + stash = get_or_create_request_stash() + stash.itpm_reserved_tokens = 50 + stash.itpm_reserved_scopes = frozenset({itpm_scope}) + mock_kwargs = {} + + ops = handler._build_io_token_reservation_ops( + kwargs=mock_kwargs, + response_obj=object(), + ) + + assert not ops, f"Expected empty ops for unresolvable usage, got {ops}" + + +@pytest.mark.asyncio +async def test_post_call_failure_skips_rpm_only_descriptor_in_tpm_refund(rate_limiter): + """ + async_post_call_failure_hook skips descriptors without tokens_per_unit + (e.g. an RPM-only api_key scope) when building the combined-TPM refund ops, + so a key with rpm_limit but no tpm_limit doesn't receive a spurious refund + that would drive its counter negative. Covers the continue guard at line 4250. + """ + handler, cache = rate_limiter + + api_key = hash_token("sk-rpm-only-desc") + user_api_key_dict = UserAPIKeyAuth( + api_key=api_key, + rpm_limit=100, + project_id="proj-rpm-only-desc", + project_metadata={"model_tpm_limit": {"gpt-3.5-turbo": 100000}}, + ) + + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 20, + } + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="", + ) + + rpm_tokens_key = handler.create_rate_limit_keys( + key="api_key", value=api_key, rate_limit_type="tokens" + ) + + await handler.async_post_call_failure_hook( + request_data=data, + original_exception=Exception("rejected"), + user_api_key_dict=user_api_key_dict, + ) + + api_key_tokens_after = int( + await cache.async_get_cache(key=rpm_tokens_key, local_only=True) or 0 + ) + assert api_key_tokens_after >= 0, ( + f"RPM-only api_key scope must not receive a negative TPM refund; got {api_key_tokens_after}" + ) + + +@pytest.mark.asyncio +async def test_max_output_tokens_prevents_cap_injection(rate_limiter): + """ + Regression for veria-ai comment: when a Responses API request supplies + max_output_tokens (the canonical Responses output bound) but not max_tokens + or max_completion_tokens, the has_explicit_max_tokens check was False, so + the code injected data["max_tokens"] = capped_floor and silently truncated + the response. + + With the fix, max_output_tokens is included in the explicit-cap check and + data["max_tokens"] must NOT be injected when max_output_tokens is already + set. + """ + handler, cache = rate_limiter + + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-max-output-tokens"), + project_id="proj-responses-max-output", + project_metadata={ + "model_otpm_limit": {"mock-model": 100}, + }, + ) + + data: dict = { + "model": "mock-model", + "input": "Summarise the document", + "max_output_tokens": 80, + "litellm_call_id": "test-max-output-tokens", + "metadata": {}, + } + + try: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="responses", + ) + except Exception: + pass + + assert "max_tokens" not in data, ( + "data['max_tokens'] must not be injected when max_output_tokens is already " + "set; the cap injection was overriding the caller's explicit output bound" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "request_data", "cap_field", "reserved_tokens"), + [ + ("aresponses", {"input": "hello", "max_tokens": 1}, "max_output_tokens", 16), + ( + "acompletion", + { + "messages": [{"role": "user", "content": "hello"}], + "max_output_tokens": 1, + }, + "max_tokens", + 10, + ), + ], +) +async def test_output_reservation_ignores_cap_fields_from_other_endpoints( + rate_limiter, + call_type, + request_data, + cap_field, + reserved_tokens, +): + handler, cache = rate_limiter + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token(f"sk-{call_type}"), + project_id=f"project-{call_type}", + project_metadata={"model_otpm_limit": {"model": 40}}, + ) + data = {"model": "model", **request_data} + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type=call_type, + ) + + assert data[cap_field] == reserved_tokens + stash = get_request_stash() + assert stash is not None + assert stash.otpm_reserved_tokens == reserved_tokens + + +def test_responses_input_is_counted_even_when_messages_is_present(rate_limiter): + handler, _cache = rate_limiter + small_estimate = handler._estimate_precise_input_tokens( + data={"input": "short", "messages": [{"role": "user", "content": "ignored"}]}, + model="", + call_type="aresponses", + ) + large_estimate = handler._estimate_precise_input_tokens( + data={"input": "large input " * 500, "messages": []}, + model="", + call_type="aresponses", + ) + + assert large_estimate > small_estimate + + +def test_anthropic_messages_usage_reconciles_split_project_quota(rate_limiter): + handler, _cache = rate_limiter + + billable_input, output_tokens, resolved = handler._resolve_io_token_reconcile_usage( + { + "usage": { + "input_tokens": 100, + "output_tokens": 25, + "cache_read_input_tokens": 30, + } + } + ) + + assert resolved is True + assert billable_input == 70 + assert output_tokens == 25 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_data", + [ + {"input": "continue", "previous_response_id": "resp-123"}, + { + "input": [ + { + "role": "user", + "content": [{"type": "input_file", "file_id": "file-123"}], + } + ] + }, + ], +) +async def test_unmeasurable_responses_input_reserves_full_project_itpm( + rate_limiter, request_data +): + handler, cache = rate_limiter + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-unmeasurable-input"), + project_id="project-unmeasurable-input", + project_metadata={"model_itpm_limit": {"model": 100}}, + ) + data = {"model": "model", **request_data} + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="aresponses", + ) + + stash = get_request_stash() + assert stash is not None + assert stash.itpm_reserved_tokens == 100 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "media_block", + [ + { + "type": "document", + "source": { + "type": "base64", + "media_type": "application/pdf", + "data": "dGVzdA==", + }, + }, + { + "type": "file", + "file": { + "filename": "document.pdf", + "file_data": "data:application/pdf;base64,dGVzdA==", + }, + }, + { + "type": "video_url", + "video_url": {"url": "https://example.com/video.mp4"}, + }, + ], +) +async def test_unmeasurable_chat_media_reserves_full_project_itpm( + rate_limiter, + media_block, +): + handler, cache = rate_limiter + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-unmeasurable-chat-media"), + project_id="project-unmeasurable-chat-media", + project_metadata={"model_itpm_limit": {"model": 100}}, + ) + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data={ + "model": "model", + "messages": [{"role": "user", "content": [media_block]}], + }, + call_type="acompletion", + ) + + stash = get_request_stash() + assert stash is not None + assert stash.itpm_reserved_tokens == 100 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call_type", + ["agenerate_content", "agenerate_content_stream"], +) +async def test_google_genai_native_contents_reserve_project_itpm( + rate_limiter, + call_type, +): + handler, cache = rate_limiter + model = "gemini/gemini-2.5-flash" + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-google-genai-native-itpm"), + project_id="project-google-genai-native-itpm", + project_metadata={"model_itpm_limit": {model: 10_000}}, + ) + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data={ + "model": model, + "contents": [ + { + "role": "user", + "parts": [{"text": "Gemini quota input " * 200}], + } + ], + }, + call_type=call_type, + ) + + stash = get_request_stash() + assert stash is not None + assert stash.itpm_reserved_tokens > 100 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["rerank", "arerank"]) +async def test_rerank_query_and_documents_enforce_project_itpm( + rate_limiter, + monkeypatch, + call_type, +): + handler, cache = rate_limiter + captured = {} + + def token_counter(**kwargs): + captured.update(kwargs) + return 101 + + monkeypatch.setattr("litellm.token_counter", token_counter) + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token(f"sk-{call_type}-itpm"), + project_id=f"project-{call_type}-itpm", + project_metadata={"model_itpm_limit": {"rerank-model": 100}}, + ) + + with pytest.raises(Exception) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data={ + "model": "rerank-model", + "query": "Which document is most relevant?", + "documents": ["first document", {"text": "second document"}], + }, + call_type=call_type, + ) + + assert getattr(exc_info.value, "status_code", None) == 429 + assert captured["text"] == ( + "Which document is most relevant?\n" + "first document\n" + "{'text': 'second document'}" + ) + + +def test_rerank_input_estimate_falls_back_to_character_count( + rate_limiter, + monkeypatch, +): + handler, _cache = rate_limiter + data = { + "query": "query text", + "documents": ["first document", "second document"], + } + + def token_counter(**_kwargs): + raise ValueError("tokenizer unavailable") + + monkeypatch.setattr("litellm.token_counter", token_counter) + rerank_text = handler._rerank_input_to_text(data) + + assert handler._estimate_precise_input_tokens( + data, + model="custom-rerank-model", + call_type="rerank", + ) == len(rerank_text) // 4 + + +@pytest.mark.parametrize( + ("response_obj", "expected"), + [ + ( + RerankResponse( + meta={"tokens": {"input_tokens": 42, "output_tokens": 3}} + ), + (42, 3, True), + ), + ( + RerankResponse( + meta={ + "tokens": {"input_tokens": 0, "output_tokens": 0}, + "billed_units": {"total_tokens": 57}, + } + ), + (57, 0, True), + ), + ( + RerankResponse( + meta={ + "tokens": {"input_tokens": 0, "output_tokens": 0}, + "billed_units": {"total_tokens": 0}, + } + ), + (0, 0, False), + ), + ], +) +def test_rerank_usage_reconciles_project_split_token_quota( + rate_limiter, + response_obj, + expected, +): + handler, _cache = rate_limiter + + assert handler._resolve_io_token_reconcile_usage(response_obj) == expected + + +def test_split_quota_helpers_handle_non_mapping_inputs(rate_limiter): + handler, _cache = rate_limiter + + assert _call_id_from_callback_kwargs(object()) is None + assert handler._is_embedding_request(object(), None) is False + assert handler._get_explicit_output_cap(object(), None) is None + assert handler._get_output_candidate_count(object()) == 1 + assert ( + handler._get_explicit_output_cap({"max_output_tokens": []}, "responses") is None + ) + assert handler._apply_implicit_output_cap(object(), 100, "responses") is None + assert handler._estimate_input_and_output_tokens(object()) == (0, 0) + assert handler._build_io_token_reservation_ops(object(), object()) == () + + +@pytest.mark.parametrize( + ("call_type", "data"), + [ + ( + "text_completion", + { + "messages": [{"role": "user", "content": "ignored"}], + "prompt": "abcd", + "input": "ignored", + "max_tokens": 1, + }, + ), + (None, {"prompt": "abcd", "max_tokens": 1}), + (None, {"prompt": ["abcd", "efgh"], "max_tokens": 1}), + ], +) +def test_split_token_estimate_selects_endpoint_input(rate_limiter, call_type, data): + handler, _cache = rate_limiter + + estimated_input, estimated_output = handler._estimate_input_and_output_tokens( + data=data, + call_type=call_type, + ) + + assert estimated_input > 0 + assert estimated_output == 1 + + +def test_split_quota_multimodal_guards_handle_non_mapping_inputs(rate_limiter): + handler, _cache = rate_limiter + + assert handler._estimate_audio_block_tokens( + object() + ) == handler._estimate_audio_block_tokens({}) + assert handler._contains_unmeasurable_chat_media(object()) is False + assert handler._contains_image_content(object()) is False + assert handler._contains_image_content( + {"inline_data": {"mime_type": "image/png", "data": "dGVzdA=="}} + ) + assert handler._responses_input_to_chat_messages(object()) == () + assert ( + handler._requires_conservative_responses_input_reservation( + object(), "responses" + ) + is False + ) + assert handler._estimate_precise_input_tokens(object(), model=None) == 0 + + +@pytest.mark.parametrize( + ("call_type", "data", "expected_text"), + [ + ("embedding", {"input": "embedding input"}, "embedding input"), + ( + "embedding", + {"input": ["first embedding", "second embedding"]}, + ["first embedding", "second embedding"], + ), + ("text_completion", {"prompt": "completion prompt"}, "completion prompt"), + ], +) +def test_precise_input_estimate_selects_endpoint_text( + rate_limiter, + monkeypatch, + call_type, + data, + expected_text, +): + handler, _cache = rate_limiter + captured = {} + + def token_counter(**kwargs): + captured.update(kwargs) + return 7 + + monkeypatch.setattr("litellm.token_counter", token_counter) + + assert ( + handler._estimate_precise_input_tokens(data, model="test", call_type=call_type) + == 7 + ) + assert captured["messages"] is None + assert captured["text"] == expected_text + + +@pytest.mark.asyncio +async def test_project_io_reservation_ignores_non_mapping_request_data(rate_limiter): + handler, _cache = rate_limiter + + await handler._reserve_project_io_tokens_or_raise( + descriptors=[], + data=object(), + requested_model=None, + user_api_key_dict=UserAPIKeyAuth(), + tpm_reservation_scopes=(), + tpm_reservation_amount=0, + ) + + +@pytest.mark.asyncio +async def test_streaming_combined_usage_reconciles_project_io_reservations( + rate_limiter, +): + handler, _cache = rate_limiter + itpm_scope = (PROJECT_ITPM_DESCRIPTOR_KEY, "project:model") + otpm_scope = (PROJECT_OTPM_DESCRIPTOR_KEY, "project:model") + stash = get_or_create_request_stash() + stash.itpm_reserved_tokens = 100 + stash.itpm_reserved_scopes = frozenset({itpm_scope}) + stash.otpm_reserved_tokens = 60 + stash.otpm_reserved_scopes = frozenset({otpm_scope}) + kwargs = { + "combined_usage_object": Usage( + prompt_tokens=40, + completion_tokens=15, + total_tokens=55, + ), + } + increments = [] + + async def capture_increments(increment_list, **_kwargs): + increments.extend(increment_list) + + handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( + capture_increments + ) + + await handler.async_log_success_event( + kwargs=kwargs, + response_obj={"response": "stream body"}, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + itpm_adjustments = [ + operation + for operation in increments + if PROJECT_ITPM_DESCRIPTOR_KEY in operation["key"] + ] + otpm_adjustments = [ + operation + for operation in increments + if PROJECT_OTPM_DESCRIPTOR_KEY in operation["key"] + ] + assert [operation["increment_value"] for operation in itpm_adjustments] == [-60] + assert [operation["increment_value"] for operation in otpm_adjustments] == [-45] + + +def test_aggregate_only_combined_usage_keeps_project_io_reservations(rate_limiter): + handler, _cache = rate_limiter + stash = get_or_create_request_stash() + stash.itpm_reserved_tokens = 100 + stash.itpm_reserved_scopes = frozenset( + {(PROJECT_ITPM_DESCRIPTOR_KEY, "project:model")} + ) + kwargs = { + "combined_usage_object": Usage(total_tokens=55), + } + + assert handler._build_io_token_reservation_ops(kwargs, object()) == () + + +def test_raw_split_usage_dict_reconciles_project_io_tokens(rate_limiter): + handler, _cache = rate_limiter + + assert handler._resolve_io_token_reconcile_usage( + { + "input_tokens": 30, + "output_tokens": 12, + "input_tokens_details": {"cached_tokens": 5}, + } + ) == (25, 12, True) + + +@pytest.mark.asyncio +async def test_post_call_success_hook_contains_header_merge_failures( + rate_limiter, monkeypatch +): + handler, _cache = rate_limiter + response = ModelResponse() + response._hidden_params = {} + + def raise_on_merge(**_kwargs): + raise RuntimeError("header merge failed") + + monkeypatch.setattr( + handler, + "_merge_ratelimit_statuses_into_additional_headers", + raise_on_merge, + ) + + await handler.async_post_call_success_hook( + data={ + "litellm_proxy_rate_limit_response": { + "overall_code": "OK", + "statuses": (), + } + }, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) diff --git a/tests/test_litellm/proxy/test_proxy_types.py b/tests/test_litellm/proxy/test_proxy_types.py index 5354de182a0..4a93e9ac7ba 100644 --- a/tests/test_litellm/proxy/test_proxy_types.py +++ b/tests/test_litellm/proxy/test_proxy_types.py @@ -159,3 +159,21 @@ def test_update_key_request_requires_key_or_key_alias(): by_alias = UpdateKeyRequest(key_alias="my-alias") assert by_alias.key is None assert by_alias.key_alias == "my-alias" + + +@pytest.mark.parametrize("request_type", ["new", "update"]) +def test_project_io_token_limits_are_stored_in_metadata(request_type): + from litellm.proxy._types import NewProjectRequest, UpdateProjectRequest + + limits = { + "model_itpm_limit": {"bedrock_mantle/openai.gpt-oss-120b": 20_000_000}, + "model_otpm_limit": {"bedrock_mantle/openai.gpt-oss-120b": 4_000_000}, + } + request = ( + NewProjectRequest(team_id="team-1", **limits) + if request_type == "new" + else UpdateProjectRequest(project_id="project-1", **limits) + ) + + assert request.metadata == limits + assert request.model_dump(exclude_none=True)["metadata"] == limits diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index df85decc676..fd3eb6b16df 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -28772,10 +28772,18 @@ export interface components { metadata?: { [key: string]: unknown; } | null; + /** Model Itpm Limit */ + model_itpm_limit?: { + [key: string]: number; + } | null; /** Model Max Budget */ model_max_budget?: { [key: string]: unknown; } | null; + /** Model Otpm Limit */ + model_otpm_limit?: { + [key: string]: number; + } | null; /** Model Rpm Limit */ model_rpm_limit?: { [key: string]: unknown; @@ -33623,10 +33631,18 @@ export interface components { metadata?: { [key: string]: unknown; } | null; + /** Model Itpm Limit */ + model_itpm_limit?: { + [key: string]: number; + } | null; /** Model Max Budget */ model_max_budget?: { [key: string]: unknown; } | null; + /** Model Otpm Limit */ + model_otpm_limit?: { + [key: string]: number; + } | null; /** Model Rpm Limit */ model_rpm_limit?: { [key: string]: unknown; From db061d6e3118806dd59340e900af482b15450326 Mon Sep 17 00:00:00 2001 From: ayaangazali Date: Thu, 23 Jul 2026 16:25:13 -0700 Subject: [PATCH 018/538] fix(azure_ai): strip non-OpenAI-spec message fields before request --- litellm/llms/azure_ai/chat/transformation.py | 21 +++++- .../chat/test_azure_ai_transformation.py | 72 +++++++++++++++++++ 2 files changed, 91 insertions(+), 2 deletions(-) diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 5540d79f667..683c05cfbaa 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -11,6 +11,7 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( _audio_or_image_in_message_content, convert_content_list_to_str, + filter_value_from_dict, ) from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj @@ -28,6 +29,13 @@ class AzureFoundryErrorStrings(str, enum.Enum): SET_EXTRA_PARAMETERS_TO_PASS_THROUGH = "Set extra-parameters to 'pass-through'" +NON_OPENAI_SPEC_MESSAGE_FIELDS = ( + "thinking_blocks", + "provider_specific_fields", + "cache_control", +) + + class AzureAIStudioConfig(OpenAIConfig): def get_supported_openai_params(self, model: str) -> list: model_supports_tool_choice = True # azure ai supports this by default @@ -167,10 +175,19 @@ class AzureAIStudioConfig(OpenAIConfig): ) -> list: """ - Azure AI Studio doesn't support content as a list. This handles: - 1. Transforms list content to a string. - 2. If message contains an image or audio, send as is (user-intended) + 1. Strips message fields that are not part of the OpenAI chat-completions + schema (thinking_blocks, provider_specific_fields, cache_control). + Azure AI Foundry backends set additionalProperties=false and reject + these with "Extra inputs are not permitted", which breaks multi-turn + Anthropic-format clients that echo thinking blocks back as history. + 2. Transforms list content to a string. + 3. If message contains an image or audio, send as is (user-intended) """ for message in messages: + message_dict = cast(dict, message) # cast-ok: TypedDict is a runtime dict stripped in place + for field in NON_OPENAI_SPEC_MESSAGE_FIELDS: + filter_value_from_dict(message_dict, field) + # Do nothing if the message contains an image or audio if _audio_or_image_in_message_content(message): continue diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py index 2e75039139c..80c4355b560 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -262,3 +262,75 @@ def test_drop_tool_level_extra_fields_strips_copilot_mcp_server_name(): assert "copilot_mcp_server_name" not in tool assert result["tools"][0]["type"] == "function" assert result["tools"][1]["function"]["name"] == "read_file" + + +def _find_key_anywhere(obj, key: str) -> bool: + if isinstance(obj, dict): + if key in obj: + return True + return any(_find_key_anywhere(v, key) for v in obj.values()) + if isinstance(obj, list): + return any(_find_key_anywhere(item, key) for item in obj) + return False + + +def test_azure_ai_strips_non_openai_spec_message_fields(): + """ + Regression for https://github.com/BerriAI/litellm/issues/33961. + + Azure AI Foundry backends set additionalProperties=false, so any message + field outside the OpenAI chat-completions schema causes a 400 "Extra inputs + are not permitted". Anthropic-format clients (e.g. Claude Code) echo prior + assistant turns back as history carrying thinking_blocks, a nested thought + signature at tool_calls[].function.provider_specific_fields, and Anthropic + cache_control annotations. transform_request must strip all of these before + the request reaches the upstream. + """ + config = AzureAIStudioConfig() + + messages = [ + {"role": "user", "content": "Read a file."}, + { + "role": "assistant", + "content": "I can help.", + "thinking_blocks": [ + { + "type": "thinking", + "thinking": "The user wants me to read a file.", + "signature": "", + "cache_control": {"type": "ephemeral"}, + } + ], + "provider_specific_fields": {"thought_signature": "sig-top"}, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "read_file", + "arguments": "{}", + "provider_specific_fields": {"thought_signature": "sig-nested"}, + }, + } + ], + }, + {"role": "user", "content": "go ahead"}, + ] + + request = config.transform_request( + model="fw-glm-5.2", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + transformed_messages = request["messages"] + + assert not _find_key_anywhere(transformed_messages, "thinking_blocks") + assert not _find_key_anywhere(transformed_messages, "provider_specific_fields") + assert not _find_key_anywhere(transformed_messages, "cache_control") + + assistant_message = transformed_messages[1] + assert assistant_message["content"] == "I can help." + assert assistant_message["tool_calls"][0]["function"]["name"] == "read_file" From 95bc890fcfc157247072752e4005686dd9413a54 Mon Sep 17 00:00:00 2001 From: ayaangazali Date: Thu, 30 Jul 2026 11:42:03 -0700 Subject: [PATCH 019/538] fix(azure_ai): strip non-spec message fields on a copy, not the caller's messages --- litellm/llms/azure_ai/chat/transformation.py | 7 ++- .../chat/test_azure_ai_transformation.py | 50 +++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 683c05cfbaa..067c89214ab 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -1,3 +1,4 @@ +import copy import enum import re from typing import Any, Final, cast @@ -182,9 +183,13 @@ class AzureAIStudioConfig(OpenAIConfig): Anthropic-format clients that echo thinking blocks back as history. 2. Transforms list content to a string. 3. If message contains an image or audio, send as is (user-intended) + + Operates on a deep copy so the caller's messages keep their thinking blocks + and provider metadata, which a fallback to another provider still needs. """ + messages = copy.deepcopy(messages) for message in messages: - message_dict = cast(dict, message) # cast-ok: TypedDict is a runtime dict stripped in place + message_dict = cast(dict, message) # cast-ok: TypedDict is a runtime dict stripped on our copy for field in NON_OPENAI_SPEC_MESSAGE_FIELDS: filter_value_from_dict(message_dict, field) diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py index 80c4355b560..beb7e9dfab0 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -334,3 +334,53 @@ def test_azure_ai_strips_non_openai_spec_message_fields(): assistant_message = transformed_messages[1] assert assistant_message["content"] == "I can help." assert assistant_message["tool_calls"][0]["function"]["name"] == "read_file" + + +def test_azure_ai_stripping_does_not_mutate_caller_messages(): + """ + The stripping must not touch the caller's messages. LiteLLM reuses the same + message objects when falling back to another provider, so stripping in place + would hand the fallback a conversation history with its thinking blocks and + provider metadata already destroyed. + """ + config = AzureAIStudioConfig() + + messages = [ + {"role": "user", "content": "Read a file."}, + { + "role": "assistant", + "content": "I can help.", + "thinking_blocks": [ + {"type": "thinking", "thinking": "Reading the file.", "signature": "sig"} + ], + "provider_specific_fields": {"thought_signature": "sig-top"}, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "read_file", + "arguments": "{}", + "provider_specific_fields": {"thought_signature": "sig-nested"}, + }, + } + ], + }, + ] + + request = config.transform_request( + model="fw-glm-5.2", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert not _find_key_anywhere(request["messages"], "thinking_blocks") + + original_assistant = messages[1] + assert original_assistant["thinking_blocks"][0]["thinking"] == "Reading the file." + assert original_assistant["provider_specific_fields"] == {"thought_signature": "sig-top"} + assert original_assistant["tool_calls"][0]["function"]["provider_specific_fields"] == { + "thought_signature": "sig-nested" + } From af2246c5b8d75bcaefb67d5615063183bf5e7502 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:47:05 +0000 Subject: [PATCH 020/538] fix(anthropic,bedrock): report provider thinking tokens instead of classifying them as text Resolves LIT-5244 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_chunk_builder_utils.py | 2 +- litellm/llms/anthropic/chat/transformation.py | 77 ++++++++++-- .../bedrock/chat/converse_transformation.py | 19 ++- litellm/llms/bedrock/chat/invoke_handler.py | 8 +- .../transformation.py | 34 ++++-- litellm/types/llms/anthropic.py | 6 + .../test_streaming_chunk_builder_utils.py | 46 +++++++ .../test_anthropic_chat_transformation.py | 112 ++++++++++++++++++ .../chat/test_converse_transformation.py | 81 +++++++++++++ .../test_reasoning_content_transformation.py | 101 ++++++++++++++++ 10 files changed, 462 insertions(+), 24 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index a2f9c80f577..fe51b5cc822 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -583,7 +583,7 @@ class ChunkProcessor: for choice in response.choices: if ( hasattr(cast(Choices, choice).message, "reasoning_content") - and cast(Choices, choice).message.reasoning_content is not None + and cast(Choices, choice).message.reasoning_content ): if reasoning_tokens is None: reasoning_tokens = 0 diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 1f9022bf28f..5c27535014a 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1,9 +1,11 @@ import json import re import time +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, NoReturn, cast import httpx +from pydantic import ValidationError import litellm from litellm.constants import ( @@ -38,6 +40,7 @@ from litellm.types.llms.anthropic import ( AnthropicMessagesTool, AnthropicMessagesToolChoice, AnthropicOutputSchema, + AnthropicOutputTokensDetails, AnthropicSystemMessageContent, AnthropicThinkingParam, AnthropicWebSearchTool, @@ -2104,6 +2107,66 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): compaction_blocks, ) + @staticmethod + def _thinking_tokens_from_usage(usage_object: Mapping[str, object]) -> int | None: + details: Final = usage_object.get("output_tokens_details") + if not isinstance(details, Mapping): + return None + try: + return AnthropicOutputTokensDetails.model_validate(details).thinking_tokens + except ValidationError: + return None + + @staticmethod + def _response_has_thinking_block(completion_response: Mapping[str, object] | None) -> bool: + if completion_response is None: + return False + content: Final = completion_response.get("content") + if not isinstance(content, list): + return False + return any( + isinstance(block, Mapping) and block.get("type") in ("thinking", "redacted_thinking") for block in content + ) + + def _build_completion_token_details( + self, + usage_object: Mapping[str, object], + iterations: Sequence[object] | None, + completion_tokens: int, + reasoning_content: str | None, + completion_response: Mapping[str, object] | None, + ) -> CompletionTokensDetailsWrapper: + reported_thinking_tokens: Final = ( + self._sum_iteration_thinking_tokens(iterations) + if iterations + else self._thinking_tokens_from_usage(usage_object) + ) + if reported_thinking_tokens is not None: + capped_reported: Final = min(max(0, reported_thinking_tokens), completion_tokens) + return CompletionTokensDetailsWrapper( + reasoning_tokens=capped_reported, + text_tokens=completion_tokens - capped_reported, + ) + if reasoning_content: + estimated: Final = min( + token_counter(text=reasoning_content, count_response_tokens=True), + completion_tokens, + ) + return CompletionTokensDetailsWrapper( + reasoning_tokens=max(0, estimated), + text_tokens=completion_tokens - max(0, estimated), + ) + if self._response_has_thinking_block(completion_response): + return CompletionTokensDetailsWrapper(reasoning_tokens=None, text_tokens=None) + return CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=completion_tokens) + + def _sum_iteration_thinking_tokens(self, iterations: Sequence[object]) -> int | None: + per_iteration: Final = tuple( + self._thinking_tokens_from_usage(iteration) for iteration in iterations if isinstance(iteration, Mapping) + ) + reported: Final = tuple(tokens for tokens in per_iteration if tokens is not None) + return sum(reported) if reported else None + def calculate_usage( self, usage_object: dict, @@ -2182,14 +2245,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): cache_creation_token_details=cache_creation_token_details, text_tokens=raw_input_tokens, ) - # Always populate completion_token_details, not just when there's reasoning_content - estimated_reasoning_tokens: Final = ( - token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0 - ) - reasoning_tokens: Final = min(estimated_reasoning_tokens, completion_tokens) - completion_token_details: Final = CompletionTokensDetailsWrapper( - reasoning_tokens=max(0, reasoning_tokens), - text_tokens=(completion_tokens - reasoning_tokens if reasoning_tokens > 0 else completion_tokens), + completion_token_details: Final = self._build_completion_token_details( + usage_object=_usage, + iterations=iterations, + completion_tokens=completion_tokens, + reasoning_content=reasoning_content, + completion_response=completion_response, ) total_tokens: Final = prompt_tokens + completion_tokens diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 91adff50a17..93feabe7222 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1764,6 +1764,7 @@ class AmazonConverseConfig(BaseConfig): self, usage: ConverseTokenUsageBlock, reasoning_content: str | None = None, + thinking_ran: bool = False, ) -> Usage: input_tokens = usage["inputTokens"] output_tokens: Final = usage["outputTokens"] @@ -1784,10 +1785,19 @@ class AmazonConverseConfig(BaseConfig): cache_creation_tokens=cache_creation_input_tokens, text_tokens=raw_input_tokens, ) - reasoning_tokens = token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0 - completion_tokens_details: Final = CompletionTokensDetailsWrapper( - reasoning_tokens=reasoning_tokens, - text_tokens=(output_tokens - reasoning_tokens if reasoning_tokens > 0 else output_tokens), + reasoning_tokens: Final = ( + token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0 + ) + completion_tokens_details: Final = ( + CompletionTokensDetailsWrapper( + reasoning_tokens=reasoning_tokens, + text_tokens=output_tokens - reasoning_tokens, + ) + if reasoning_tokens > 0 + else CompletionTokensDetailsWrapper( + reasoning_tokens=None if thinking_ran else 0, + text_tokens=None if thinking_ran else output_tokens, + ) ) openai_usage: Final = Usage( prompt_tokens=input_tokens, @@ -2184,6 +2194,7 @@ class AmazonConverseConfig(BaseConfig): usage: Final = self._transform_usage( completion_response["usage"], reasoning_content=chat_completion_message.get("reasoning_content"), + thinking_ran=reasoningContentBlocks is not None, ) ## HANDLE TOOL CALLS diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index a2bb179f72f..57510ff334d 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -330,6 +330,7 @@ class AWSEventStreamDecoder: self.response_id: str | None = None self.json_mode = json_mode self._current_tool_name: str | None = None + self._thinking_ran = False def check_empty_tool_call_args(self) -> bool: """ @@ -559,7 +560,12 @@ class AWSEventStreamDecoder: elif "stopReason" in chunk_data: finish_reason = map_finish_reason(chunk_data.get("stopReason", "stop")) elif "usage" in chunk_data: - usage = converse_config._transform_usage(chunk_data.get("usage", {})) + usage = converse_config._transform_usage( + chunk_data.get("usage", {}), + thinking_ran=self._thinking_ran, + ) + if thinking_blocks: + self._thinking_ran = True model_response_provider_specific_fields: Final = {} if "trace" in chunk_data: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 79e05545358..174a55aac85 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -4,7 +4,7 @@ Handles transforming from Responses API -> LiteLLM completion (Chat Completion import json import re -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Any, Final, Literal, cast from openai.types.chat.chat_completion_named_tool_choice_param import ( @@ -1745,6 +1745,12 @@ class LiteLLMCompletionResponsesConfig: output_items.append(item) return output_items + @staticmethod + def _encode_thinking_blocks(message: Message) -> str | None: + thinking_blocks: Final[Sequence[Mapping[str, object]]] = getattr(message, "thinking_blocks", None) or () + preserved: Final = tuple(block for block in thinking_blocks if block.get("signature") or block.get("data")) + return json.dumps(preserved, separators=(",", ":")) if preserved else None + @staticmethod def _extract_reasoning_output_items( chat_completion_response: ModelResponse, @@ -1753,23 +1759,31 @@ class LiteLLMCompletionResponsesConfig: for choice in choices: if hasattr(choice, "message") and choice.message: message = choice.message - if hasattr(message, "reasoning_content") and message.reasoning_content: + reasoning_content = getattr(message, "reasoning_content", None) or "" + encrypted_content = LiteLLMCompletionResponsesConfig._encode_thinking_blocks(message) + if reasoning_content or encrypted_content: # Only check the first choice for reasoning content return [ GenericResponseOutputItem( type="reasoning", - id=f"rs_{hash(str(message.reasoning_content))}", + id=f"rs_{hash(reasoning_content or encrypted_content)}", status=LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status( choice.finish_reason ), role="assistant", - content=[ - OutputText( - type="output_text", - text=message.reasoning_content, - annotations=[], - ) - ], + content=( + [ + OutputText( + type="output_text", + text=reasoning_content, + annotations=[], + ) + ] + if reasoning_content + # mutable-ok: GenericResponseOutputItem.content is typed as a list + else [] + ), + encrypted_content=encrypted_content, ) ] return [] diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 95f8db66eda..f111d3c6e56 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -626,6 +626,12 @@ class AnthropicResponseUsageBlock(BaseModel): output_tokens: int +class AnthropicOutputTokensDetails(BaseModel): + model_config = ConfigDict(extra="allow") + + thinking_tokens: Optional[int] = None + + AnthropicFinishReason = Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"] diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index 0114db381cf..15bbe476a06 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -1180,3 +1180,49 @@ def test_get_combined_tool_content_joins_many_custom_tool_input_fragments_in_ord assert isinstance(combined[1], ChatCompletionMessageCustomToolCall) assert combined[1].custom.name == "run_script" assert combined[1].custom.input == "".join(object_fragments) + + +def _reasoning_stream_chunk() -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-reasoning", + model="claude-opus-4-8", + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="10", role="assistant"))], + ) + + +def test_count_reasoning_tokens_returns_none_for_signature_only_thinking(): + from litellm.types.utils import Choices, Message, ModelResponse + + processor = ChunkProcessor(chunks=[_reasoning_stream_chunk()]) + response = ModelResponse( + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="10", role="assistant", reasoning_content=""), + ) + ] + ) + + assert processor.count_reasoning_tokens(response) is None + + +def test_count_reasoning_tokens_counts_visible_reasoning(): + from litellm.types.utils import Choices, Message, ModelResponse + + processor = ChunkProcessor(chunks=[_reasoning_stream_chunk()]) + response = ModelResponse( + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content="10", + role="assistant", + reasoning_content="let me count the primes under thirty", + ), + ) + ] + ) + + assert processor.count_reasoning_tokens(response) > 0 diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 94a4a3fc945..063b965dd47 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -119,6 +119,118 @@ def test_calculate_usage_clamps_text_tokens_when_reasoning_estimate_exceeds_outp assert usage.completion_tokens_details.text_tokens == 0 +def test_calculate_usage_prefers_provider_reported_thinking_tokens(): + config = AnthropicConfig() + + usage = config.calculate_usage( + usage_object={ + "input_tokens": 32, + "output_tokens": 421, + "output_tokens_details": {"thinking_tokens": 372}, + }, + reasoning_content="", + completion_response={ + "content": [ + {"type": "thinking", "thinking": "", "signature": "sig"}, + {"type": "text", "text": "10"}, + ] + }, + ) + + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 372 + assert usage.completion_tokens_details.text_tokens == 49 + + +def test_calculate_usage_provider_thinking_tokens_win_over_visible_reasoning_estimate(): + config = AnthropicConfig() + + usage = config.calculate_usage( + usage_object={ + "input_tokens": 50, + "output_tokens": 811, + "output_tokens_details": {"thinking_tokens": 747}, + }, + reasoning_content="short visible reasoning that tokenizes to far fewer than 747 tokens", + ) + + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 747 + assert usage.completion_tokens_details.text_tokens == 64 + + +def test_calculate_usage_sums_provider_thinking_tokens_across_iterations(): + config = AnthropicConfig() + + usage = config.calculate_usage( + usage_object={ + "input_tokens": 10, + "output_tokens": 300, + "iterations": [ + {"input_tokens": 5, "output_tokens": 100, "output_tokens_details": {"thinking_tokens": 60}}, + {"input_tokens": 5, "output_tokens": 200, "output_tokens_details": {"thinking_tokens": 90}}, + ], + }, + reasoning_content=None, + ) + + assert usage.completion_tokens == 300 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 150 + assert usage.completion_tokens_details.text_tokens == 150 + + +def test_calculate_usage_reports_unknown_split_when_thinking_ran_without_a_count(): + config = AnthropicConfig() + + usage = config.calculate_usage( + usage_object={"input_tokens": 32, "output_tokens": 580}, + reasoning_content="", + completion_response={ + "content": [ + {"type": "redacted_thinking", "data": "encrypted"}, + {"type": "text", "text": "10"}, + ] + }, + ) + + assert usage.completion_tokens == 580 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens is None + assert usage.completion_tokens_details.text_tokens is None + + +def test_calculate_usage_without_thinking_reports_all_output_as_text(): + config = AnthropicConfig() + + usage = config.calculate_usage( + usage_object={"input_tokens": 32, "output_tokens": 171}, + reasoning_content=None, + completion_response={"content": [{"type": "text", "text": "10"}]}, + ) + + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 0 + assert usage.completion_tokens_details.text_tokens == 171 + + +def test_calculate_usage_ignores_malformed_provider_thinking_tokens(): + config = AnthropicConfig() + + usage = config.calculate_usage( + usage_object={ + "input_tokens": 32, + "output_tokens": 100, + "output_tokens_details": {"thinking_tokens": "not-a-number"}, + }, + reasoning_content=None, + ) + + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 0 + assert usage.completion_tokens_details.text_tokens == 100 + + def test_calculate_usage_handles_mocked_output_tokens_with_reasoning_content(): config = AnthropicConfig() diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 6d318bb8729..1f759b58cf7 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -5934,3 +5934,84 @@ def test_adaptive_thinking_dropped_when_max_tokens_too_small_converse(): ) assert "thinking" not in optional_params + + +def test_converse_usage_reports_unknown_split_for_signature_only_thinking(): + config = AmazonConverseConfig() + + usage = config._transform_usage( + ConverseTokenUsageBlock(inputTokens=32, outputTokens=581, totalTokens=613), + reasoning_content="", + thinking_ran=True, + ) + + assert usage.completion_tokens == 581 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens is None + assert usage.completion_tokens_details.text_tokens is None + + +def test_converse_usage_estimates_split_for_visible_thinking(): + config = AmazonConverseConfig() + + usage = config._transform_usage( + ConverseTokenUsageBlock(inputTokens=32, outputTokens=581, totalTokens=613), + reasoning_content="Let me think about how many primes there are under thirty.", + thinking_ran=True, + ) + + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens > 0 + assert ( + usage.completion_tokens_details.reasoning_tokens + usage.completion_tokens_details.text_tokens + == usage.completion_tokens + ) + + +def test_converse_usage_without_thinking_reports_all_output_as_text(): + config = AmazonConverseConfig() + + usage = config._transform_usage(ConverseTokenUsageBlock(inputTokens=32, outputTokens=171, totalTokens=203)) + + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 0 + assert usage.completion_tokens_details.text_tokens == 171 + + +def test_converse_transform_response_signature_only_thinking_reports_unknown_split(): + config = AmazonConverseConfig() + raw_response = MagicMock(status_code=200) + raw_response.text = json.dumps( + { + "output": { + "message": { + "role": "assistant", + "content": [ + {"reasoningContent": {"reasoningText": {"text": "", "signature": "sig"}}}, + {"text": "10"}, + ], + } + }, + "stopReason": "end_turn", + "usage": {"inputTokens": 32, "outputTokens": 581, "totalTokens": 613}, + } + ) + raw_response.json.return_value = json.loads(raw_response.text) + + response = config._transform_response( + model="bedrock/global.anthropic.claude-opus-4-8", + response=raw_response, + model_response=ModelResponse(), + stream=False, + logging_obj=None, + optional_params={}, + api_key=None, + data={}, + messages=[], + encoding=None, + ) + + assert response.choices[0].message.reasoning_content == "" + + assert response.usage.completion_tokens_details.reasoning_tokens is None + assert response.usage.completion_tokens_details.text_tokens is None diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_content_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_content_transformation.py index 020b5de0a2a..3c1980152a7 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_content_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_content_transformation.py @@ -263,6 +263,107 @@ class TestReasoningContentFinalResponse: assert len(reasoning_items) == 1, "Should have exactly one reasoning item" assert reasoning_items[0].content[0].text == "Reasoning for first answer" + def test_signature_only_thinking_block_still_emits_reasoning_item(self): + response = ModelResponse( + id="test-id", + created=1234567890, + model="test-model", + object="chat.completion", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content="10", + role="assistant", + reasoning_content="", + thinking_blocks=[ + {"type": "thinking", "thinking": "", "signature": "signature-payload"} + ], + ), + ) + ], + ) + + responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Test input", + responses_api_request={}, + chat_completion_response=response, + ) + + reasoning_items = [ + item for item in responses_api_response.output if item.type == "reasoning" + ] + assert len(reasoning_items) == 1, "Signature-only thinking should still surface a reasoning item" + assert reasoning_items[0].content == [] + assert "signature-payload" in reasoning_items[0].encrypted_content + + def test_redacted_thinking_block_preserved_as_encrypted_content(self): + response = ModelResponse( + id="test-id", + created=1234567890, + model="test-model", + object="chat.completion", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content="10", + role="assistant", + thinking_blocks=[{"type": "redacted_thinking", "data": "redacted-payload"}], + ), + ) + ], + ) + + responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Test input", + responses_api_request={}, + chat_completion_response=response, + ) + + reasoning_items = [ + item for item in responses_api_response.output if item.type == "reasoning" + ] + assert len(reasoning_items) == 1 + assert "redacted-payload" in reasoning_items[0].encrypted_content + + def test_visible_thinking_keeps_text_and_signature(self): + response = ModelResponse( + id="test-id", + created=1234567890, + model="test-model", + object="chat.completion", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content="10", + role="assistant", + reasoning_content="counting the primes", + thinking_blocks=[ + {"type": "thinking", "thinking": "counting the primes", "signature": "sig"} + ], + ), + ) + ], + ) + + responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Test input", + responses_api_request={}, + chat_completion_response=response, + ) + + reasoning_items = [ + item for item in responses_api_response.output if item.type == "reasoning" + ] + assert len(reasoning_items) == 1 + assert reasoning_items[0].content[0].text == "counting the primes" + assert "sig" in reasoning_items[0].encrypted_content + def test_streaming_chunk_id_raw(): """Test that streaming chunk IDs are raw (not encoded) to match OpenAI format""" From 53ee9c8293d6d1aeb038eb1a674e5d8ad090dbef Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 21:34:36 +0000 Subject: [PATCH 021/538] fix(anthropic): fall back when only some compaction iterations report thinking tokens Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/anthropic/chat/transformation.py | 10 +++-- .../transformation.py | 21 ++++----- litellm/types/llms/anthropic.py | 2 +- .../test_anthropic_chat_transformation.py | 44 +++++++++++++++++++ 4 files changed, 60 insertions(+), 17 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index e0e11be356c..feb26b19981 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -2134,9 +2134,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): reasoning_content: str | None, completion_response: Mapping[str, object] | None, ) -> CompletionTokensDetailsWrapper: + iteration_thinking_tokens: Final = self._sum_iteration_thinking_tokens(iterations) if iterations else None reported_thinking_tokens: Final = ( - self._sum_iteration_thinking_tokens(iterations) - if iterations + iteration_thinking_tokens + if iteration_thinking_tokens is not None else self._thinking_tokens_from_usage(usage_object) ) if reported_thinking_tokens is not None: @@ -2160,10 +2161,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def _sum_iteration_thinking_tokens(self, iterations: Sequence[object]) -> int | None: per_iteration: Final = tuple( - self._thinking_tokens_from_usage(iteration) for iteration in iterations if isinstance(iteration, Mapping) + self._thinking_tokens_from_usage(iteration) if isinstance(iteration, Mapping) else None + for iteration in iterations ) reported: Final = tuple(tokens for tokens in per_iteration if tokens is not None) - return sum(reported) if reported else None + return sum(reported) if len(reported) == len(per_iteration) else None @staticmethod def is_anthropic_usage_object(usage_object: dict) -> bool: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index fa2ce0d1505..f0614f1cacf 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1763,18 +1763,15 @@ class LiteLLMCompletionResponsesConfig: choice.finish_reason ), role="assistant", - content=( - [ - OutputText( - type="output_text", - text=reasoning_content, - annotations=[], - ) - ] - if reasoning_content - # mutable-ok: GenericResponseOutputItem.content is typed as a list - else [] - ), + content=[ + OutputText( + type="output_text", + text=text, + annotations=[], + ) + for text in (reasoning_content,) + if text + ], encrypted_content=encrypted_content, ) ] diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 7de383f6f13..f6b256ad5df 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -612,7 +612,7 @@ class AnthropicResponseUsageBlock(BaseModel): class AnthropicOutputTokensDetails(BaseModel): model_config = ConfigDict(extra="allow") - thinking_tokens: Optional[int] = None + thinking_tokens: int | None = None AnthropicFinishReason = Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"] diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 828a9c30fb9..de62c990f11 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -180,6 +180,50 @@ def test_calculate_usage_sums_provider_thinking_tokens_across_iterations(): assert usage.completion_tokens_details.text_tokens == 150 +def test_calculate_usage_falls_back_when_only_some_iterations_report_thinking_tokens(): + config = AnthropicConfig() + + usage = config.calculate_usage( + usage_object={ + "input_tokens": 10, + "output_tokens": 300, + "output_tokens_details": {"thinking_tokens": 240}, + "iterations": [ + {"input_tokens": 5, "output_tokens": 100, "output_tokens_details": {"thinking_tokens": 60}}, + {"input_tokens": 5, "output_tokens": 200}, + ], + }, + reasoning_content=None, + ) + + assert usage.completion_tokens == 300 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 240 + assert usage.completion_tokens_details.text_tokens == 60 + + +def test_calculate_usage_reports_unknown_split_when_only_some_iterations_report_thinking_tokens(): + config = AnthropicConfig() + + usage = config.calculate_usage( + usage_object={ + "input_tokens": 10, + "output_tokens": 300, + "iterations": [ + {"input_tokens": 5, "output_tokens": 100, "output_tokens_details": {"thinking_tokens": 60}}, + {"input_tokens": 5, "output_tokens": 200}, + ], + }, + reasoning_content="", + completion_response={"content": [{"type": "thinking", "thinking": "", "signature": "sig"}]}, + ) + + assert usage.completion_tokens == 300 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens is None + assert usage.completion_tokens_details.text_tokens is None + + def test_calculate_usage_reports_unknown_split_when_thinking_ran_without_a_count(): config = AnthropicConfig() From d0c1d2be8a82723d458eaf133159ad8576bd9a8b Mon Sep 17 00:00:00 2001 From: Shivam Rawat Date: Fri, 7 Aug 2026 18:08:33 -0700 Subject: [PATCH 022/538] feat(key_management): let any authenticated user resolve a raw key via /key/info Possession of the raw sk- key already lets the holder call /key/info with the key itself as the bearer token, so resolving raw key -> key info for any authenticated caller discloses nothing new. Lookups by hashed token remain restricted to admins, the key's owner, and teammates --- .../key_management_endpoints.py | 7 + .../test_key_management_endpoints.py | 146 ++++++++++++++++++ 2 files changed, 153 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index e4def45892b..99631377541 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -6351,6 +6351,12 @@ async def _can_user_query_key_info( ) -> bool: """ Helper to check if the user has access to the key's info + + Any authenticated caller who presents the raw key value (i.e. a preimage of the + stored token hash) is allowed: possession of the raw key already grants the + ability to call /key/info with that key as the bearer token, so resolving + raw key -> key info discloses nothing new. Lookups by hashed token remain + restricted to admins, the key's owner, and the key's teammates. """ if ( ( @@ -6359,6 +6365,7 @@ async def _can_user_query_key_info( ) or user_api_key_dict.api_key == key or key_info.user_id == user_api_key_dict.user_id + or (key is not None and hash_token(token=key) == key_info.token) or await TeamMemberPermissionChecks.user_belongs_to_keys_team( user_api_key_dict=user_api_key_dict, existing_key_row=key_info, diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index cf9aa477112..fc349f8448d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -15323,3 +15323,149 @@ async def test_rotate_master_key_rotates_sso_identity_assertions( prisma_client=mock_prisma_client, new_master_key="sk-new-master-key", ) + + +@pytest.mark.asyncio +async def test_can_user_query_key_info_raw_key_possession_allows_any_user(): + """ + Any authenticated user who presents the raw sk- key value can query that + key's info: possessing the raw key already lets them call /key/info with + the key itself as the bearer token, so this discloses nothing new. + """ + from litellm.proxy._types import hash_token + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _can_user_query_key_info, + ) + + raw_key = "sk-raw-key-owned-by-someone-else" + key_info = LiteLLM_VerificationToken( + token=hash_token(raw_key), + user_id="key-owner", + team_id=None, + key_alias="prod-batch-alias", + ) + caller = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="unrelated-user", + api_key="hashed-caller-token", + ) + + assert ( + await _can_user_query_key_info( + user_api_key_dict=caller, + key=raw_key, + key_info=key_info, + ) + is True + ) + + +@pytest.mark.asyncio +async def test_can_user_query_key_info_hashed_token_still_forbidden(): + """ + Querying by hashed token (e.g. copied from spend logs) must stay + restricted to admins, the key's owner, and teammates. + """ + from litellm.proxy._types import hash_token + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _can_user_query_key_info, + ) + + raw_key = "sk-raw-key-owned-by-someone-else" + key_info = LiteLLM_VerificationToken( + token=hash_token(raw_key), + user_id="key-owner", + team_id=None, + key_alias="prod-batch-alias", + ) + caller = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="unrelated-user", + api_key="hashed-caller-token", + ) + + assert ( + await _can_user_query_key_info( + user_api_key_dict=caller, + key=hash_token(raw_key), + key_info=key_info, + ) + is False + ) + + +@pytest.mark.asyncio +async def test_info_key_fn_resolves_alias_from_raw_key_for_any_user(monkeypatch): + """ + End-to-end through /key/info: a non-admin user unrelated to the key can + resolve raw sk- key -> key_alias. + """ + from litellm.proxy._types import hash_token + from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn + + raw_key = "sk-raw-key-owned-by-someone-else" + key_row = LiteLLM_VerificationToken( + token=hash_token(raw_key), + user_id="key-owner", + team_id=None, + key_alias="prod-batch-alias", + ) + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", AsyncMock()) + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=key_row + ) + + caller = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="unrelated-user", + api_key="hashed-caller-token", + ) + + result = await info_key_fn(key=raw_key, user_api_key_dict=caller) + + assert result["info"]["key_alias"] == "prod-batch-alias" + assert "token" not in result["info"] + + find_unique_kwargs = ( + mock_prisma_client.db.litellm_verificationtoken.find_unique.call_args.kwargs + ) + assert find_unique_kwargs["where"] == {"token": hash_token(raw_key)} + + +@pytest.mark.asyncio +async def test_info_key_fn_hashed_lookup_still_403_for_unrelated_user(monkeypatch): + """ + End-to-end through /key/info: the same unrelated user querying by hashed + token still gets a 403. + """ + from litellm.proxy._types import hash_token + from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn + + raw_key = "sk-raw-key-owned-by-someone-else" + key_row = LiteLLM_VerificationToken( + token=hash_token(raw_key), + user_id="key-owner", + team_id=None, + key_alias="prod-batch-alias", + ) + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", AsyncMock()) + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=key_row + ) + + caller = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="unrelated-user", + api_key="hashed-caller-token", + ) + + with pytest.raises(ProxyException) as exc_info: + await info_key_fn(key=hash_token(raw_key), user_api_key_dict=caller) + + assert int(exc_info.value.code) == 403 From a14b2ab960d093fe1198caac649f8f9884b9a773 Mon Sep 17 00:00:00 2001 From: Shivam Rawat Date: Fri, 7 Aug 2026 18:47:27 -0700 Subject: [PATCH 023/538] fix(router): forward target_model_names on file uploads to litellm_proxy deployments Uploads for the Batch API through a deployment that points at a second LiteLLM proxy arrived downstream as bare multipart requests with no model or target_model_names, so the second proxy could not route them and fell back to files_settings or the wrong endpoint shape. The router now injects target_model_names into extra_body when the deployment provider is litellm_proxy, and litellm_proxy is registered as an OpenAI-compatible files/batches provider so the downstream call uses the deployment api_base and api_key over the OpenAI wire format. Resolves https://github.com/BerriAI/litellm/issues/36176 --- litellm/files/main.py | 9 +- litellm/router.py | 9 ++ litellm/types/utils.py | 1 + tests/test_litellm/test_router.py | 135 ++++++++++++++++++++++++++++++ 4 files changed, 151 insertions(+), 3 deletions(-) diff --git a/litellm/files/main.py b/litellm/files/main.py index 34421d13761..cf8826895aa 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -24,12 +24,15 @@ FileCreateProvider = Literal[ "vertex_ai", "bedrock", "hosted_vllm", + "litellm_proxy", "manus", "anthropic", ] -FileRetrieveProvider = Literal["openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "manus", "anthropic"] -FileDeleteProvider = Literal["openai", "azure", "gemini", "manus", "anthropic"] -FileListProvider = Literal["openai", "azure", "manus", "anthropic"] +FileRetrieveProvider = Literal[ + "openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic" +] +FileDeleteProvider = Literal["openai", "azure", "gemini", "litellm_proxy", "manus", "anthropic"] +FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic"] import litellm from litellm import get_secret_str from litellm.files.streaming import FileContentStreamingResponse diff --git a/litellm/router.py b/litellm/router.py index 39d5080a33d..b35506554d1 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -21,6 +21,7 @@ import traceback from collections import defaultdict from collections.abc import AsyncGenerator, Callable, Generator, Mapping, Sequence from functools import lru_cache +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeVar, Union, cast import anyio @@ -192,6 +193,7 @@ from litellm.types.utils import ( CustomPricingLiteLLMParams, GenericBudgetConfigType, LiteLLMBatch, + LlmProviders, ModelInfo, ModelResponseStream, StandardLoggingPayload, @@ -4919,6 +4921,13 @@ class Router: ) kwargs_copy["file"] = file + if custom_llm_provider == LlmProviders.LITELLM_PROXY.value: + kwargs_copy["extra_body"] = MappingProxyType( + { + **(kwargs_copy.get("extra_body") or MappingProxyType({})), + "target_model_names": stripped_model, + } + ) if ( "gcs_bucket_name" in data ): # TODO: Remove this once we have a better way to handle GCS bucket name: Problem is that we need to pass the gcs_bucket_name to the router for the create_file call but it doesn't show up there diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 4f824262867..20225d7dc49 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3706,6 +3706,7 @@ LlmProvidersSet: Final = {provider.value for provider in LlmProviders} OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: set[str] = { LlmProviders.OPENAI.value, LlmProviders.HOSTED_VLLM.value, + LlmProviders.LITELLM_PROXY.value, } ListBatchesSupportedProvider = Literal["openai", "azure", "hosted_vllm", "vertex_ai"] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 8f35597768f..92085a009af 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -522,6 +522,141 @@ async def test_async_router_acreate_file_uses_deployment_custom_llm_provider(): assert mock_acreate_file.call_args.kwargs["custom_llm_provider"] == "azure" +@pytest.mark.asyncio +async def test_async_router_acreate_file_forwards_target_model_names_to_litellm_proxy(): + """ + A deployment pointing at a second LiteLLM proxy (litellm_proxy provider) must forward + target_model_names downstream so the second proxy can route the upload to the right + deployment. Regression test for https://github.com/BerriAI/litellm/issues/36176 + """ + import json + from io import BytesIO + from unittest.mock import MagicMock, patch + + jsonl_file = BytesIO( + json.dumps({"body": {"model": "chained-batch", "messages": [{"role": "user", "content": "hi"}]}}).encode( + "utf-8" + ) + ) + jsonl_file.name = "test.jsonl" + + router = litellm.Router( + model_list=[ + { + "model_name": "chained-batch", + "litellm_params": { + "model": "litellm_proxy/gpt-4.1-batch", + "api_base": "http://localhost:4001/v1", + "api_key": "sk-proxy-b", + }, + }, + ], + ) + + with patch("litellm.acreate_file", return_value=MagicMock()) as mock_acreate_file: + await router.acreate_file( + model="chained-batch", + purpose="batch", + file=jsonl_file, + ) + + assert mock_acreate_file.call_count == 1 + call_kwargs = mock_acreate_file.call_args.kwargs + assert call_kwargs["custom_llm_provider"] == "litellm_proxy" + assert call_kwargs["extra_body"] == {"target_model_names": "gpt-4.1-batch"} + uploaded_line = json.loads(call_kwargs["file"].read().decode("utf-8").split("\n")[0]) + assert uploaded_line["body"]["model"] == "gpt-4.1-batch" + + +@pytest.mark.asyncio +async def test_async_router_acreate_file_does_not_inject_target_model_names_for_other_providers(): + """ + target_model_names is a LiteLLM proxy routing hint; it must not leak into uploads + sent to non-litellm_proxy providers. + """ + from unittest.mock import MagicMock, patch + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4.1-batch", + "litellm_params": {"model": "gpt-4.1"}, + }, + ], + ) + + with patch("litellm.acreate_file", return_value=MagicMock()) as mock_acreate_file: + await router.acreate_file( + model="gpt-4.1-batch", + purpose="batch", + file=MagicMock(), + ) + + assert mock_acreate_file.call_count == 1 + assert mock_acreate_file.call_args.kwargs.get("extra_body") is None + + +@pytest.mark.asyncio +async def test_async_router_acreate_file_litellm_proxy_sends_target_model_names_in_multipart_form(): + """ + End-to-end through litellm.acreate_file and the OpenAI SDK: the multipart form that + reaches the second proxy must carry target_model_names as a form field, since the + downstream /v1/files endpoint reads it via Form(). Would raise BadRequestError + (unsupported provider) before litellm_proxy was supported for files. + """ + import json + from io import BytesIO + + import httpx + import respx + + jsonl_file = BytesIO( + json.dumps({"body": {"model": "chained-batch", "messages": [{"role": "user", "content": "hi"}]}}).encode( + "utf-8" + ) + ) + jsonl_file.name = "test.jsonl" + + router = litellm.Router( + model_list=[ + { + "model_name": "chained-batch", + "litellm_params": { + "model": "litellm_proxy/gpt-4.1-batch", + "api_base": "http://localhost:4001/v1", + "api_key": "sk-proxy-b", + }, + }, + ], + ) + + file_object_json = { + "id": "file-abc123", + "object": "file", + "bytes": 100, + "created_at": 1700000000, + "filename": "test.jsonl", + "purpose": "batch", + "status": "processed", + } + + with respx.mock(assert_all_called=True) as respx_mock: + create_route = respx_mock.post("http://localhost:4001/v1/files").mock( + return_value=httpx.Response(200, json=file_object_json) + ) + response = await router.acreate_file( + model="chained-batch", + purpose="batch", + file=jsonl_file, + ) + + assert response.id == "file-abc123" + request_body = create_route.calls.last.request.content + assert b'name="target_model_names"' in request_body + assert b"gpt-4.1-batch" in request_body + assert b'name="purpose"' in request_body + + @pytest.mark.asyncio async def test_async_router_afile_content_uses_deployment_custom_llm_provider(): """ From 3c96030488914eeea8a67dfc4d50944b79cc4c3b Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 8 Aug 2026 02:20:35 +0000 Subject: [PATCH 024/538] fix(advisor): resolve the advisor sub-call through the proxy router Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../messages/interceptors/advisor.py | 79 +++++++- .../messages/test_advisor_orchestration.py | 190 ++++++++++++++++++ 2 files changed, 264 insertions(+), 5 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py b/litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py index dfae7b4f4cf..805f5625ad9 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py @@ -16,7 +16,7 @@ How it works: import uuid from collections.abc import AsyncIterator -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final, cast import litellm import litellm.constants as _c @@ -28,6 +28,9 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) +if TYPE_CHECKING: + from litellm.router import Router + ADVISOR_MAX_USES: Final[int] = _c.ADVISOR_MAX_USES ADVISOR_NATIVE_PROVIDERS: Final[frozenset] = _c.ADVISOR_NATIVE_PROVIDERS ADVISOR_TOOL_DESCRIPTION: Final[str] = _c.ADVISOR_TOOL_DESCRIPTION @@ -138,13 +141,10 @@ class AdvisorOrchestrationHandler(MessagesInterceptor): # --- Advisor sub-call (always non-streaming, no tools) --- try: - advisor_response: AnthropicMessagesResponse = await _call_messages_handler( + advisor_response: AnthropicMessagesResponse = await _call_advisor( model=advisor_model, messages=advisor_messages, - tools=None, - stream=False, max_tokens=max_tokens, - custom_llm_provider=None, # let litellm resolve from model name metadata={ **metadata_base, "advisor_sub_call": True, @@ -357,6 +357,75 @@ def _inject_max_uses_error( ] +def _resolve_advisor_router(advisor_model: str) -> "Router | None": + """Return the proxy router when it can resolve ``advisor_model``. + + The advisor sub-call must honor the proxy's ``model_list`` (and its + fallbacks / credentials) exactly like a direct call to that model group + would. Without this, provider resolution falls back to the bare model + name, which for a ``claude-*`` advisor model means the public Anthropic + API, bypassing the configured deployment entirely. + + Returns ``None`` for SDK callers (no proxy router) and for advisor models + the router doesn't know about, so those keep resolving through + ``litellm.anthropic_messages()`` provider inference. + """ + try: + from litellm.proxy.proxy_server import llm_router + except (ImportError, ModuleNotFoundError): + return None + if llm_router is None: + return None + if llm_router.get_model_list(model_name=advisor_model): + return llm_router + if llm_router.model_group_alias and advisor_model in llm_router.model_group_alias: + return llm_router + if llm_router.pattern_router.route(advisor_model) is not None: + return llm_router + return None + + +async def _call_advisor( + *, + model: str, + messages: list[dict], + max_tokens: int, + metadata: dict, + api_key: str | None, + api_base: str | None, +) -> AnthropicMessagesResponse: + """Run the advisor sub-call, through the proxy router when it applies. + + A caller-supplied ``api_key`` / ``api_base`` override is an explicit + request to bypass the configured deployment, so it keeps the direct + SDK-level path. + """ + router: Final = None if (api_key or api_base) else _resolve_advisor_router(model) + response: Final = ( + await router.aanthropic_messages( + model=model, + messages=messages, + tools=None, + stream=False, + max_tokens=max_tokens, + metadata=metadata, + ) + if router is not None + else await _call_messages_handler( + model=model, + messages=messages, + tools=None, + stream=False, + max_tokens=max_tokens, + custom_llm_provider=None, + metadata=metadata, + api_key=api_key, + api_base=api_base, + ) + ) + return cast(AnthropicMessagesResponse, response) # cast-ok: both /messages entry points are untyped + + async def _call_messages_handler( model: str, messages: list[dict], diff --git a/tests/test_litellm/llms/anthropic/messages/test_advisor_orchestration.py b/tests/test_litellm/llms/anthropic/messages/test_advisor_orchestration.py index 3d35e93167f..bcb1843f408 100644 --- a/tests/test_litellm/llms/anthropic/messages/test_advisor_orchestration.py +++ b/tests/test_litellm/llms/anthropic/messages/test_advisor_orchestration.py @@ -1041,3 +1041,193 @@ async def test_executor_failure_is_not_tagged(): ) assert is_advisor_orchestration_failure(exc_info.value) is False + + +# --------------------------------------------------------------------------- +# 15. The advisor sub-call resolves through the proxy router when the advisor +# model is configured in model_list, instead of dialing the public +# Anthropic API (regression for LIT-5307). +# --------------------------------------------------------------------------- + + +def _router_with_advisor_deployment(recorder, advisor_model="claude-opus-4-8"): + """Build a Router whose only deployment is the advisor model on Foundry. + + The recorder replaces ``litellm.anthropic_messages`` before construction + because Router binds it at init time, so the returned Router exercises the + real deployment-resolution path and records what it dispatched. + """ + import litellm + from litellm.router import Router + + with patch("litellm.anthropic_messages", new=recorder): + return Router( + model_list=[ + { + "model_name": advisor_model, + "litellm_params": { + "model": f"azure_ai/{advisor_model}", + "api_base": "http://127.0.0.1:1/foundry", + "api_key": "fake-foundry-key", + }, + } + ], + num_retries=0, + ) + + +@pytest.mark.asyncio +async def test_advisor_sub_call_routes_through_proxy_router(): + import litellm.proxy.proxy_server as proxy_server + from litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor import ( + AdvisorOrchestrationHandler, + ) + + router_calls = [] + + async def recorder(**kwargs): + router_calls.append(kwargs) + return _make_text_response("Use trial division.", model="claude-opus-4-8") + + router = _router_with_advisor_deployment(recorder) + + call_count = 0 + + async def mock_call(model, messages, tools, stream, max_tokens, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _make_advisor_tool_use_response() + return _make_text_response("Final answer.") + + with ( + patch( + "litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor._call_messages_handler", + side_effect=mock_call, + ), + patch.object(proxy_server, "llm_router", router), + ): + h = AdvisorOrchestrationHandler() + result = await h.handle( + model="executor-model", + messages=MESSAGES, + tools=[{**ADVISOR_TOOL, "model": "claude-opus-4-8"}], + stream=False, + max_tokens=512, + custom_llm_provider="azure_ai", + ) + + assert call_count == 2 + assert len(router_calls) == 1 + assert router_calls[0]["model"] == "azure_ai/claude-opus-4-8" + assert router_calls[0]["api_base"] == "http://127.0.0.1:1/foundry" + assert router_calls[0]["api_key"] == "fake-foundry-key" + assert "Final answer." in result["content"][0]["text"] + + +@pytest.mark.asyncio +async def test_advisor_sub_call_bypasses_router_for_unconfigured_model(): + """An advisor model the router doesn't know about keeps the SDK-level path.""" + import litellm.proxy.proxy_server as proxy_server + from litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor import ( + AdvisorOrchestrationHandler, + ) + + router_calls = [] + + async def recorder(**kwargs): + router_calls.append(kwargs) + return _make_text_response("should not be used") + + router = _router_with_advisor_deployment(recorder, advisor_model="some-other-model") + + call_count = 0 + + async def mock_call(model, messages, tools, stream, max_tokens, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _make_advisor_tool_use_response() + if tools is None: + return _make_text_response("Advice.", model="claude-opus-4-8") + return _make_text_response("Final answer.") + + with ( + patch( + "litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor._call_messages_handler", + side_effect=mock_call, + ), + patch.object(proxy_server, "llm_router", router), + ): + h = AdvisorOrchestrationHandler() + await h.handle( + model="executor-model", + messages=MESSAGES, + tools=[{**ADVISOR_TOOL, "model": "claude-opus-4-8"}], + stream=False, + max_tokens=512, + custom_llm_provider="azure_ai", + ) + + assert router_calls == [] + assert call_count == 3 + + +@pytest.mark.asyncio +async def test_advisor_sub_call_client_override_bypasses_router(): + """A caller-supplied api_key/api_base override must not be re-routed.""" + import litellm + import litellm.proxy.proxy_server as proxy_server + from litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor import ( + AdvisorOrchestrationHandler, + ) + + router_calls = [] + + async def recorder(**kwargs): + router_calls.append(kwargs) + return _make_text_response("should not be used") + + router = _router_with_advisor_deployment(recorder) + + sub_calls = [] + + async def mock_call(model, messages, tools, stream, max_tokens, **kwargs): + sub_calls.append({"model": model, "tools": tools, **kwargs}) + if len(sub_calls) == 1: + return _make_advisor_tool_use_response() + if tools is None: + return _make_text_response("Advice.", model="claude-opus-4-8") + return _make_text_response("Final answer.") + + advisor_tool = { + **ADVISOR_TOOL, + "model": "claude-opus-4-8", + "api_key": "client-key", + "api_base": "https://client.example.com", + } + + with ( + patch( + "litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor._call_messages_handler", + side_effect=mock_call, + ), + patch.object(proxy_server, "llm_router", router), + patch.dict(proxy_server.general_settings, {"allow_client_side_credentials": True}), + patch.object(litellm, "user_url_validation", False), + ): + h = AdvisorOrchestrationHandler() + await h.handle( + model="executor-model", + messages=MESSAGES, + tools=[advisor_tool], + stream=False, + max_tokens=512, + custom_llm_provider="azure_ai", + ) + + assert router_calls == [] + advisor_sub_calls = [c for c in sub_calls if c["tools"] is None] + assert len(advisor_sub_calls) == 1 + assert advisor_sub_calls[0]["api_key"] == "client-key" + assert advisor_sub_calls[0]["api_base"] == "https://client.example.com" From 45317d58641c0ebcb1a84868d79fbcc19fca7b02 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 8 Aug 2026 02:48:08 +0000 Subject: [PATCH 025/538] refactor(advisor): resolve advisor router once per request Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../messages/interceptors/advisor.py | 83 +++++++------------ 1 file changed, 30 insertions(+), 53 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py b/litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py index 805f5625ad9..adf96db61c2 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py @@ -16,7 +16,7 @@ How it works: import uuid from collections.abc import AsyncIterator -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final import litellm import litellm.constants as _c @@ -100,6 +100,14 @@ class AdvisorOrchestrationHandler(MessagesInterceptor): parent_request_id: Final[str] = str(kwargs.pop("litellm_call_id", None) or uuid.uuid4()) metadata_base: Final[dict] = dict(kwargs.pop("metadata", None) or {}) + advisor_metadata: Final = { + **metadata_base, + "advisor_sub_call": True, + "parent_request_id": parent_request_id, + } + advisor_router: Final = ( + None if (advisor_api_key or advisor_api_base) else _resolve_advisor_router(advisor_model) + ) iteration = 0 while True: @@ -141,17 +149,27 @@ class AdvisorOrchestrationHandler(MessagesInterceptor): # --- Advisor sub-call (always non-streaming, no tools) --- try: - advisor_response: AnthropicMessagesResponse = await _call_advisor( - model=advisor_model, - messages=advisor_messages, - max_tokens=max_tokens, - metadata={ - **metadata_base, - "advisor_sub_call": True, - "parent_request_id": parent_request_id, - }, - api_key=advisor_api_key, - api_base=advisor_api_base, + advisor_response: AnthropicMessagesResponse = ( + await advisor_router.aanthropic_messages( + model=advisor_model, + messages=advisor_messages, + tools=None, + stream=False, + max_tokens=max_tokens, + metadata=advisor_metadata, + ) + if advisor_router is not None + else await _call_messages_handler( + model=advisor_model, + messages=advisor_messages, + tools=None, + stream=False, + max_tokens=max_tokens, + custom_llm_provider=None, + metadata=advisor_metadata, + api_key=advisor_api_key, + api_base=advisor_api_base, + ) ) except Exception as advisor_sub_call_exception: mark_advisor_orchestration_failure(advisor_sub_call_exception) @@ -385,47 +403,6 @@ def _resolve_advisor_router(advisor_model: str) -> "Router | None": return None -async def _call_advisor( - *, - model: str, - messages: list[dict], - max_tokens: int, - metadata: dict, - api_key: str | None, - api_base: str | None, -) -> AnthropicMessagesResponse: - """Run the advisor sub-call, through the proxy router when it applies. - - A caller-supplied ``api_key`` / ``api_base`` override is an explicit - request to bypass the configured deployment, so it keeps the direct - SDK-level path. - """ - router: Final = None if (api_key or api_base) else _resolve_advisor_router(model) - response: Final = ( - await router.aanthropic_messages( - model=model, - messages=messages, - tools=None, - stream=False, - max_tokens=max_tokens, - metadata=metadata, - ) - if router is not None - else await _call_messages_handler( - model=model, - messages=messages, - tools=None, - stream=False, - max_tokens=max_tokens, - custom_llm_provider=None, - metadata=metadata, - api_key=api_key, - api_base=api_base, - ) - ) - return cast(AnthropicMessagesResponse, response) # cast-ok: both /messages entry points are untyped - - async def _call_messages_handler( model: str, messages: list[dict], From 292161f766ca7ac88cf3cdbef0c4b599a5576a88 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 7 Aug 2026 23:25:10 -0700 Subject: [PATCH 026/538] fix(proxy): read through to the DB on registry misses so just-created models, guardrails, and agents resolve on sibling replicas --- .../proxy/agent_endpoints/a2a_endpoints.py | 15 +- litellm/proxy/agent_endpoints/a2a_routing.py | 6 +- .../common_utils/registry_read_through.py | 143 +++++++++++ .../proxy/guardrails/guardrail_endpoints.py | 6 +- ...model_access_group_management_endpoints.py | 30 ++- litellm/proxy/route_llm_request.py | 232 ++++++++++-------- ruff.toml | 2 +- .../test_registry_read_through.py | 231 +++++++++++++++++ .../test_access_group_management.py | 96 ++++++++ .../proxy/test_route_a2a_models.py | 75 ++++++ .../proxy/test_route_llm_request.py | 121 +++++++++ 11 files changed, 843 insertions(+), 114 deletions(-) create mode 100644 litellm/proxy/common_utils/registry_read_through.py create mode 100644 tests/test_litellm/proxy/common_utils/test_registry_read_through.py diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 27780aeb994..a4e1ac126d9 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -152,14 +152,13 @@ def _jsonrpc_error( ) -def _get_agent(agent_id: str): +async def _get_agent(agent_id: str) -> "AgentResponse | None": """Look up an agent by ID or name. Returns None if not found.""" - from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + from litellm.proxy.common_utils.registry_read_through import ( + get_agent_with_read_through, + ) - agent = global_agent_registry.get_agent_by_id(agent_id=agent_id) - if agent is None: - agent = global_agent_registry.get_agent_by_name(agent_name=agent_id) - return agent + return await get_agent_with_read_through(agent_id) def _enforce_inbound_trace_id(agent: Any, request: Request) -> None: @@ -531,7 +530,7 @@ async def get_agent_card( ) try: - agent: Final = _get_agent(agent_id) + agent: Final = await _get_agent(agent_id) if agent is None: raise HTTPException(status_code=404, detail=f"Agent '{agent_id}' not found") @@ -645,7 +644,7 @@ async def invoke_agent_a2a( params.pop(key) # Find the agent - agent: Final = _get_agent(agent_id) + agent: Final = await _get_agent(agent_id) if agent is None: return _jsonrpc_error(request_id, -32000, f"Agent '{agent_id}' not found", 404) diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index 038b6b4a840..2228735d805 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -25,10 +25,12 @@ async def route_a2a_agent_request( Returns None if not an A2A request (allows normal routing to continue). """ # Import here to avoid circular imports - from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, ) + from litellm.proxy.common_utils.registry_read_through import ( + get_agent_with_read_through, + ) from litellm.proxy.route_llm_request import ( ROUTE_ENDPOINT_MAPPING, ProxyModelNotFoundError, @@ -44,7 +46,7 @@ async def route_a2a_agent_request( agent_name: Final = model_name[4:] # Look up agent in registry - agent: Final = global_agent_registry.get_agent_by_name(agent_name) + agent: Final = await get_agent_with_read_through(agent_name) if agent is None: verbose_proxy_logger.error("[A2A] Agent '%s' not found in registry", agent_name) route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type) diff --git a/litellm/proxy/common_utils/registry_read_through.py b/litellm/proxy/common_utils/registry_read_through.py new file mode 100644 index 00000000000..b78106205d4 --- /dev/null +++ b/litellm/proxy/common_utils/registry_read_through.py @@ -0,0 +1,143 @@ +"""Read-through recovery for in-memory registries in multi-replica deployments. + +A management write (POST /model/new, /guardrails, /v1/agents) lands on one +replica and reaches Postgres, but sibling replicas only refresh their in-memory +registries on the periodic config reload or the Redis config-sync resync, both +of which lag by seconds. A request that uses the new object immediately can +land on a sibling that has never heard of it and fail with a 400/404. + +On a registry miss, callers here fetch the missing object from the DB and load +it into the local registry before giving up. A short negative-result TTL keeps +repeated lookups of genuinely unknown names from hammering the DB. +""" + +import asyncio +from collections.abc import Awaitable, Callable +from typing import TYPE_CHECKING, Final + +from litellm._logging import verbose_proxy_logger +from litellm.caching.in_memory_cache import InMemoryCache + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.agents import AgentResponse + +READ_THROUGH_MISS_TTL_SECONDS: Final = 2.0 + + +class RegistryReadThrough: + __slots__ = ("_lock", "_miss_ttl_seconds", "_recent_misses", "_resync") + + def __init__( + self, + resync: Callable[[str], Awaitable[bool]], + miss_ttl_seconds: float = READ_THROUGH_MISS_TTL_SECONDS, + ) -> None: + self._resync = resync + self._miss_ttl_seconds = miss_ttl_seconds + self._lock = asyncio.Lock() + self._recent_misses = InMemoryCache(max_size_in_memory=1000) + + async def attempt(self, key: str) -> bool: + if self._recent_misses.get_cache(key) is not None: + return False + async with self._lock: + if self._recent_misses.get_cache(key) is not None: + return False + try: + found: Final = await self._resync(key) + except Exception as e: # noqa: BLE001 # a failed read-through must surface the original miss error, not a 500 + verbose_proxy_logger.warning("registry read-through for %r failed: %s", key, e) + return False + if not found: + self._recent_misses.set_cache(key, True, ttl=self._miss_ttl_seconds) + return found + + +def _db_backed_registries_enabled() -> bool: + from litellm.proxy import proxy_server + + return proxy_server.prisma_client is not None and proxy_server.store_model_in_db is True + + +async def _resync_model_deployments(model_name: str) -> bool: + from litellm.proxy import proxy_server + from litellm.repositories.model_repository import ModelRepository + + if not _db_backed_registries_enabled(): + return False + prisma_client: Final = proxy_server.prisma_client + assert prisma_client is not None + rows: Final = await ModelRepository(prisma_client).table.find_many( + where={"OR": [{"model_name": model_name}, {"model_id": model_name}]} + ) + if not rows: + return False + if proxy_server.llm_router is None: + await proxy_server.proxy_config.add_deployment( + prisma_client=prisma_client, proxy_logging_obj=proxy_server.proxy_logging_obj + ) + return proxy_server.llm_router is not None + proxy_server.proxy_config._add_deployment(db_models=rows) + proxy_server.llm_model_list = proxy_server.llm_router.get_model_list() + return True + + +async def _resync_guardrails(guardrail_name: str) -> bool: + from litellm.proxy import proxy_server + + if not _db_backed_registries_enabled(): + return False + prisma_client: Final = proxy_server.prisma_client + assert prisma_client is not None + await proxy_server.proxy_config._init_guardrails_in_db(prisma_client=prisma_client) + return _initialized_guardrail(guardrail_name) is not None + + +async def _resync_agents(agent_id_or_name: str) -> bool: + from litellm.proxy import proxy_server + + if not _db_backed_registries_enabled(): + return False + prisma_client: Final = proxy_server.prisma_client + assert prisma_client is not None + await proxy_server.proxy_config._init_agents_in_db(prisma_client=prisma_client) + return _agent_from_registry(agent_id_or_name) is not None + + +model_registry_read_through: Final = RegistryReadThrough(resync=_resync_model_deployments) +guardrail_registry_read_through: Final = RegistryReadThrough(resync=_resync_guardrails) +agent_registry_read_through: Final = RegistryReadThrough(resync=_resync_agents) + + +def _agent_from_registry(agent_id_or_name: str) -> "AgentResponse | None": + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + by_id: Final = global_agent_registry.get_agent_by_id(agent_id=agent_id_or_name) + if by_id is not None: + return by_id + return global_agent_registry.get_agent_by_name(agent_name=agent_id_or_name) + + +async def get_agent_with_read_through(agent_id_or_name: str) -> "AgentResponse | None": + agent: Final = _agent_from_registry(agent_id_or_name) + if agent is not None: + return agent + if not await agent_registry_read_through.attempt(agent_id_or_name): + return None + return _agent_from_registry(agent_id_or_name) + + +def _initialized_guardrail(guardrail_name: str) -> "CustomGuardrail | None": + from litellm.proxy.guardrails import guardrail_endpoints + + return guardrail_endpoints.GUARDRAIL_REGISTRY.get_initialized_guardrail_callback(guardrail_name=guardrail_name) + + +async def get_initialized_guardrail_with_read_through(guardrail_name: str) -> "CustomGuardrail | None": + active: Final = _initialized_guardrail(guardrail_name) + if active is not None: + return active + if not await guardrail_registry_read_through.attempt(guardrail_name): + return None + return _initialized_guardrail(guardrail_name) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 761d8aabc8a..dff70ccf68d 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -2244,8 +2244,12 @@ async def apply_guardrail( litellm_logging_obj = None start_time: Final = datetime.now(timezone.utc) + from litellm.proxy.common_utils.registry_read_through import ( + get_initialized_guardrail_with_read_through, + ) + try: - active_guardrail: Final[CustomGuardrail | None] = GUARDRAIL_REGISTRY.get_initialized_guardrail_callback( + active_guardrail: Final[CustomGuardrail | None] = await get_initialized_guardrail_with_read_through( guardrail_name=request.guardrail_name ) if active_guardrail is None: diff --git a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py index 7051f705a03..75d33c6c40a 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -7,7 +7,10 @@ Endpoints here: import json from collections.abc import Mapping, Sequence -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final + +if TYPE_CHECKING: + from litellm.router import Router from fastapi import APIRouter, Depends, HTTPException @@ -52,6 +55,23 @@ def validate_models_exist(model_names: list[str], llm_router) -> tuple[bool, lis return (len(missing) == 0, missing) +async def _missing_models_after_read_through( + model_names: Sequence[str], llm_router: "Router | None" +) -> tuple[str, ...]: + from litellm.proxy import proxy_server + from litellm.proxy.common_utils.registry_read_through import ( + model_registry_read_through, + ) + + _, missing = validate_models_exist(model_names=list(model_names), llm_router=llm_router) + if not missing: + return () + for name in missing: + await model_registry_read_through.attempt(name) + _, still_missing = validate_models_exist(model_names=list(model_names), llm_router=proxy_server.llm_router) + return tuple(still_missing) + + def add_access_group_to_deployment(model_info: dict[str, Any], access_group: str) -> tuple[dict[str, Any], bool]: """ Add an access group to a deployment's model_info. @@ -369,12 +389,12 @@ async def create_model_group( # Validate model_names exist in router (only if using model_names path) if not use_model_ids and has_model_names: assert data.model_names is not None - all_valid, missing_models = validate_models_exist( + missing_models: Final = await _missing_models_after_read_through( model_names=data.model_names, llm_router=llm_router, ) - if not all_valid: + if missing_models: raise HTTPException( status_code=400, detail={"error": f"Model(s) not found: {', '.join(missing_models)}"}, @@ -633,12 +653,12 @@ async def update_access_group( # Validation: Check if all new models exist (only if using model_names path) if not use_model_ids and has_model_names: assert data.model_names is not None - all_valid, missing_models = validate_models_exist( + missing_models: Final = await _missing_models_after_read_through( model_names=data.model_names, llm_router=llm_router, ) - if not all_valid: + if missing_models: raise HTTPException( status_code=400, detail={"error": f"Model(s) not found: {', '.join(missing_models)}"}, diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index dd8deed57f1..657e0cbcafa 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -313,112 +313,150 @@ async def add_shared_session_to_data(data: dict) -> None: pass +RouteType = Literal[ + "acompletion", + "atext_completion", + "aembedding", + "aimage_generation", + "aspeech", + "atranscription", + "amoderation", + "arerank", + "aresponses", + "aget_responses", + "adelete_responses", + "acancel_responses", + "acompact_responses", + "acreate_response_reply", + "alist_input_items", + "_arealtime", # private function for realtime API + "acreate_realtime_client_secret", + "arealtime_calls", + "acreate_realtime_transcription_session", + "_aresponses_websocket", # private function for responses WebSocket mode + "aimage_edit", + "agenerate_content", + "agenerate_content_stream", + "allm_passthrough_route", + "acreate_batch", + "aretrieve_batch", + "alist_batches", + "afile_content", + "afile_retrieve", + "acreate_fine_tuning_job", + "acancel_fine_tuning_job", + "alist_fine_tuning_jobs", + "aretrieve_fine_tuning_job", + "avector_store_search", + "avector_store_create", + "avector_store_retrieve", + "avector_store_list", + "avector_store_update", + "avector_store_delete", + "avector_store_file_create", + "avector_store_file_list", + "avector_store_file_retrieve", + "avector_store_file_content", + "avector_store_file_update", + "avector_store_file_delete", + "aocr", + "asearch", + "avideo_generation", + "avideo_list", + "avideo_status", + "avideo_content", + "avideo_remix", + "avideo_create_character", + "avideo_get_character", + "avideo_edit", + "avideo_extension", + "acreate_container", + "alist_containers", + "aretrieve_container", + "adelete_container", + "aupload_container_file", + "alist_container_files", + "aretrieve_container_file", + "adelete_container_file", + "aretrieve_container_file_content", + "acreate_skill", + "alist_skills", + "aget_skill", + "adelete_skill", + "aingest", + "anthropic_messages", + "acreate_interaction", + "aget_interaction", + "adelete_interaction", + "acancel_interaction", + "acreate_agent", + "alist_agents", + "aget_agent", + "adelete_agent", + "alist_agent_versions", + "asend_message", + "call_mcp_tool", + "acancel_batch", + "afile_delete", + "acreate_eval", + "alist_evals", + "aget_eval", + "aupdate_eval", + "adelete_eval", + "acancel_eval", + "acreate_run", + "alist_runs", + "aget_run", + "acancel_run", + "adelete_run", +] + + async def route_request( data: dict, llm_router: LitellmRouter | None, user_model: str | None, - route_type: Literal[ - "acompletion", - "atext_completion", - "aembedding", - "aimage_generation", - "aspeech", - "atranscription", - "amoderation", - "arerank", - "aresponses", - "aget_responses", - "adelete_responses", - "acancel_responses", - "acompact_responses", - "acreate_response_reply", - "alist_input_items", - "_arealtime", # private function for realtime API - "acreate_realtime_client_secret", - "arealtime_calls", - "acreate_realtime_transcription_session", - "_aresponses_websocket", # private function for responses WebSocket mode - "aimage_edit", - "agenerate_content", - "agenerate_content_stream", - "allm_passthrough_route", - "acreate_batch", - "aretrieve_batch", - "alist_batches", - "afile_content", - "afile_retrieve", - "acreate_fine_tuning_job", - "acancel_fine_tuning_job", - "alist_fine_tuning_jobs", - "aretrieve_fine_tuning_job", - "avector_store_search", - "avector_store_create", - "avector_store_retrieve", - "avector_store_list", - "avector_store_update", - "avector_store_delete", - "avector_store_file_create", - "avector_store_file_list", - "avector_store_file_retrieve", - "avector_store_file_content", - "avector_store_file_update", - "avector_store_file_delete", - "aocr", - "asearch", - "avideo_generation", - "avideo_list", - "avideo_status", - "avideo_content", - "avideo_remix", - "avideo_create_character", - "avideo_get_character", - "avideo_edit", - "avideo_extension", - "acreate_container", - "alist_containers", - "aretrieve_container", - "adelete_container", - "aupload_container_file", - "alist_container_files", - "aretrieve_container_file", - "adelete_container_file", - "aretrieve_container_file_content", - "acreate_skill", - "alist_skills", - "aget_skill", - "adelete_skill", - "aingest", - "anthropic_messages", - "acreate_interaction", - "aget_interaction", - "adelete_interaction", - "acancel_interaction", - "acreate_agent", - "alist_agents", - "aget_agent", - "adelete_agent", - "alist_agent_versions", - "asend_message", - "call_mcp_tool", - "acancel_batch", - "afile_delete", - "acreate_eval", - "alist_evals", - "aget_eval", - "aupdate_eval", - "adelete_eval", - "acancel_eval", - "acreate_run", - "alist_runs", - "aget_run", - "acancel_run", - "adelete_run", - ], + route_type: RouteType, user_api_key_dict: UserAPIKeyAuth | None = None, ): """ Common helper to route the request """ + try: + return await _route_request_single_attempt( + data=data, + llm_router=llm_router, + user_model=user_model, + route_type=route_type, + user_api_key_dict=user_api_key_dict, + ) + except ProxyModelNotFoundError: + requested_model: Final = data.get("model", "") + if not isinstance(requested_model, str) or not requested_model: + raise + from litellm.proxy import proxy_server + from litellm.proxy.common_utils.registry_read_through import ( + model_registry_read_through, + ) + + if not await model_registry_read_through.attempt(requested_model): + raise + return await _route_request_single_attempt( + data=data, + llm_router=proxy_server.llm_router, + user_model=user_model, + route_type=route_type, + user_api_key_dict=user_api_key_dict, + ) + + +async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited provider coroutines; the inferred union keeps route_request's callers typed + data: dict, # noqa: LIT001 # request body is the proxy-wide mutable dict contract shared with route_request + llm_router: LitellmRouter | None, + user_model: str | None, + route_type: RouteType, + user_api_key_dict: UserAPIKeyAuth | None = None, +): raise_if_required_body_param_missing(route_type=route_type, data=data) await add_shared_session_to_data(data) diff --git a/ruff.toml b/ruff.toml index 095e3e24c52..bd3f8334e94 100644 --- a/ruff.toml +++ b/ruff.toml @@ -6,7 +6,7 @@ lint.extend-select = ["T20", "PGH004", "RUF008", "RUF009", "RUF100"] # litellm's own ruff config both rely on suppressions this config can't see. lint.external = [ # Enforced by the strict-rule gate (scripts/ruff_strict_gate.py + ruff-strict.toml) - "C901", "TID251", + "ANN202", "C901", "TID251", # Enforced by upstream litellm's ruff config, but not run in this repo's CI "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405", ] diff --git a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py new file mode 100644 index 00000000000..e1c8f031579 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py @@ -0,0 +1,231 @@ +import asyncio +from typing import Final + +import pytest + +from litellm.proxy.common_utils.registry_read_through import RegistryReadThrough + + +class ResyncSpy: + def __init__(self, found: bool = True, error: Exception | None = None) -> None: + self.found = found + self.error = error + self.calls: list[str] = [] + + async def __call__(self, key: str) -> bool: + self.calls.append(key) + if self.error is not None: + raise self.error + return self.found + + +@pytest.mark.asyncio +async def test_attempt_returns_true_when_resync_finds_object(): + spy: Final = ResyncSpy(found=True) + read_through: Final = RegistryReadThrough(resync=spy) + + assert await read_through.attempt("new-model") is True + assert spy.calls == ["new-model"] + + +@pytest.mark.asyncio +async def test_attempt_found_key_is_not_negative_cached(): + spy: Final = ResyncSpy(found=True) + read_through: Final = RegistryReadThrough(resync=spy) + + assert await read_through.attempt("new-model") is True + assert await read_through.attempt("new-model") is True + assert spy.calls == ["new-model", "new-model"] + + +@pytest.mark.asyncio +async def test_missing_key_is_negative_cached_within_ttl(): + spy: Final = ResyncSpy(found=False) + read_through: Final = RegistryReadThrough(resync=spy, miss_ttl_seconds=60.0) + + assert await read_through.attempt("ghost-model") is False + assert await read_through.attempt("ghost-model") is False + assert spy.calls == ["ghost-model"] + + +@pytest.mark.asyncio +async def test_negative_cache_expires_and_resync_runs_again(): + spy: Final = ResyncSpy(found=False) + read_through: Final = RegistryReadThrough(resync=spy, miss_ttl_seconds=0.05) + + assert await read_through.attempt("ghost-model") is False + await asyncio.sleep(0.1) + assert await read_through.attempt("ghost-model") is False + assert spy.calls == ["ghost-model", "ghost-model"] + + +@pytest.mark.asyncio +async def test_resync_exception_returns_false_without_negative_caching(): + spy: Final = ResyncSpy(error=RuntimeError("db down")) + read_through: Final = RegistryReadThrough(resync=spy) + + assert await read_through.attempt("new-model") is False + assert await read_through.attempt("new-model") is False + assert spy.calls == ["new-model", "new-model"] + + +@pytest.mark.asyncio +async def test_concurrent_attempts_for_missing_key_resync_once(): + class SlowResyncSpy(ResyncSpy): + async def __call__(self, key: str) -> bool: + await asyncio.sleep(0.05) + return await super().__call__(key) + + spy: Final = SlowResyncSpy(found=False) + read_through: Final = RegistryReadThrough(resync=spy, miss_ttl_seconds=60.0) + + results: Final = await asyncio.gather(*(read_through.attempt("ghost-model") for _ in range(5))) + assert results == [False] * 5 + assert spy.calls == ["ghost-model"] + + +@pytest.mark.asyncio +async def test_distinct_keys_do_not_share_negative_cache(): + spy: Final = ResyncSpy(found=False) + read_through: Final = RegistryReadThrough(resync=spy, miss_ttl_seconds=60.0) + + assert await read_through.attempt("ghost-a") is False + assert await read_through.attempt("ghost-b") is False + assert spy.calls == ["ghost-a", "ghost-b"] + + +class FakeAgentRow: + def __init__(self, agent_id: str, agent_name: str) -> None: + self.agent_id = agent_id + self.agent_name = agent_name + self.object_permission = None + self.spend = 0.0 + + def __iter__(self): + return iter( + { + "agent_id": self.agent_id, + "agent_name": self.agent_name, + "agent_card_params": {"name": self.agent_name, "url": "http://db-agent"}, + "litellm_params": {}, + }.items() + ) + + +@pytest.fixture +def clean_agent_registry(): + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + original_agents: Final = list(global_agent_registry.agent_list) + original_config_agents: Final = getattr(global_agent_registry, "config_agents", ()) + global_agent_registry.agent_list = [] + global_agent_registry.config_agents = () + try: + yield global_agent_registry + finally: + global_agent_registry.agent_list = original_agents + global_agent_registry.config_agents = original_config_agents + + +@pytest.mark.asyncio +async def test_get_agent_with_read_through_recovers_agent_created_on_sibling_replica( + clean_agent_registry, monkeypatch +): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through + + agent_id: Final = "read-through-db-agent-id" + prisma_client: Final = MagicMock() + prisma_client.db.litellm_agentstable.find_many = AsyncMock( + return_value=[FakeAgentRow(agent_id, "read-through-db-agent")] + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + assert clean_agent_registry.get_agent_by_id(agent_id=agent_id) is None + agent: Final = await get_agent_with_read_through(agent_id) + + assert agent is not None + assert agent.agent_id == agent_id + + +@pytest.mark.asyncio +async def test_get_agent_with_read_through_returns_none_for_unknown_agent(clean_agent_registry, monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through + + prisma_client: Final = MagicMock() + prisma_client.db.litellm_agentstable.find_many = AsyncMock(return_value=[]) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + assert await get_agent_with_read_through("agent-nobody-created") is None + + +class FakeGuardrailRow: + def __init__(self, guardrail_id: str, guardrail_name: str) -> None: + self.guardrail_id = guardrail_id + self.guardrail_name = guardrail_name + + def __iter__(self): + return iter( + { + "guardrail_id": self.guardrail_id, + "guardrail_name": self.guardrail_name, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "default_on": True, + "blocked_words": [{"keyword": "secret", "action": "BLOCK"}], + }, + "guardrail_info": {}, + }.items() + ) + + +@pytest.mark.asyncio +async def test_get_guardrail_with_read_through_recovers_guardrail_created_on_sibling_replica(monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import ( + get_initialized_guardrail_with_read_through, + ) + from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + + guardrail_id: Final = "read-through-db-guardrail-id" + guardrail_name: Final = "read-through-db-guardrail" + prisma_client: Final = MagicMock() + prisma_client.db.litellm_guardrailstable.find_many = AsyncMock( + return_value=[FakeGuardrailRow(guardrail_id, guardrail_name)] + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + try: + guardrail: Final = await get_initialized_guardrail_with_read_through(guardrail_name=guardrail_name) + assert guardrail is not None + assert guardrail.guardrail_name == guardrail_name + finally: + IN_MEMORY_GUARDRAIL_HANDLER.delete_in_memory_guardrail(guardrail_id) + + +@pytest.mark.asyncio +async def test_get_guardrail_with_read_through_returns_none_for_unknown_guardrail(monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.registry_read_through import ( + get_initialized_guardrail_with_read_through, + ) + + prisma_client: Final = MagicMock() + prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[]) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + assert await get_initialized_guardrail_with_read_through(guardrail_name="guardrail-nobody-created") is None diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py index 3240ad20edb..1722d8c377b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py @@ -430,3 +430,99 @@ async def test_delete_access_group_ignores_models_that_were_already_dead(): assert response.models_updated == 1 mock_prisma.db.litellm_proxymodeltable.update.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_create_access_group_read_through_recovers_model_created_on_sibling_replica(): + """Regression: an access group referencing a model that another replica just wrote + to the DB must be created instead of 400ing until the periodic config reload.""" + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + create_model_group, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + NewModelGroupRequest, + ) + + from types import SimpleNamespace + + model_name = "e2e-ag-sibling-replica-model" + db_row = SimpleNamespace( + model_id=f"{model_name}-id", + model_name=model_name, + litellm_params={"model": "openai/gpt-4o", "api_key": "fake", "mock_response": "hi"}, + model_info={}, + blocked=False, + ) + + mock_router = Router( + model_list=[ + { + "model_name": "some-other-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}, + } + ] + ) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(side_effect=[[db_row], [], [db_row]]) + mock_prisma.db.litellm_proxymodeltable.update = AsyncMock() + + with ( + patch("litellm.proxy.proxy_server.llm_router", mock_router), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch( + "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache", + new=AsyncMock(return_value=None), + ), + ): + response = await create_model_group( + data=NewModelGroupRequest(access_group="replica-lag-group", model_names=[model_name]), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert response.models_updated == 1 + assert response.model_names == [model_name] + assert mock_prisma.db.litellm_proxymodeltable.find_many.await_args_list[0].kwargs["where"] == { + "OR": [{"model_name": model_name}, {"model_id": model_name}] + } + + +@pytest.mark.asyncio +async def test_create_access_group_model_missing_everywhere_still_400s(): + from fastapi import HTTPException + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + create_model_group, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + NewModelGroupRequest, + ) + + model_name = "e2e-ag-model-nobody-created" + mock_router = Router( + model_list=[ + { + "model_name": "some-other-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}, + } + ] + ) + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + + with ( + patch("litellm.proxy.proxy_server.llm_router", mock_router), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + ): + with pytest.raises(HTTPException) as exc_info: + await create_model_group( + data=NewModelGroupRequest(access_group="ghost-group", model_names=[model_name]), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert exc_info.value.status_code == 400 + assert model_name in str(exc_info.value.detail) diff --git a/tests/test_litellm/proxy/test_route_a2a_models.py b/tests/test_litellm/proxy/test_route_a2a_models.py index 616fa62cda5..22f99d03d21 100644 --- a/tests/test_litellm/proxy/test_route_a2a_models.py +++ b/tests/test_litellm/proxy/test_route_a2a_models.py @@ -49,6 +49,7 @@ async def test_route_a2a_model_bypasses_router(): ) mock_registry = Mock() + mock_registry.get_agent_by_id = Mock(return_value=None) mock_registry.get_agent_by_name = Mock(return_value=mock_agent) # Mock litellm.acompletion to verify it's called @@ -104,3 +105,77 @@ async def test_route_non_a2a_model_raises_error_if_not_in_router(): user_model=None, route_type="acompletion", ) + + +class _DbAgentRow: + def __init__(self, agent_id: str, agent_name: str) -> None: + self.agent_id = agent_id + self.agent_name = agent_name + self.object_permission = None + self.spend = 0.0 + + def __iter__(self): + return iter( + { + "agent_id": self.agent_id, + "agent_name": self.agent_name, + "agent_card_params": {"name": self.agent_name, "url": "http://sibling-db-agent.example.com"}, + "litellm_params": {}, + }.items() + ) + + +def _router_without_models(): + mock_router = Mock() + mock_router.model_names = [] + mock_router.deployment_names = [] + mock_router.has_model_id = Mock(return_value=False) + mock_router.model_group_alias = None + mock_router.router_general_settings = Mock(pass_through_all_models=False) + mock_router.default_deployment = None + mock_router.pattern_router = Mock(patterns=[]) + mock_router.map_team_model = Mock(return_value=None) + return mock_router + + +@pytest.mark.asyncio +async def test_route_a2a_model_read_through_recovers_agent_created_on_sibling_replica(monkeypatch): + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + agent_name = "a2a-sibling-replica-agent" + prisma_client = Mock() + prisma_client.db.litellm_agentstable.find_many = AsyncMock( + return_value=[_DbAgentRow("a2a-sibling-replica-agent-id", agent_name)] + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + + original_agents = list(global_agent_registry.agent_list) + original_config_agents = getattr(global_agent_registry, "config_agents", ()) + global_agent_registry.agent_list = [] + global_agent_registry.config_agents = () + + data = { + "model": f"a2a/{agent_name}", + "messages": [{"role": "user", "content": "Hello"}], + } + mock_acompletion = AsyncMock(return_value={"id": "read-through-response"}) + + try: + with patch("litellm.acompletion", mock_acompletion): + await route_request( + data=data, + llm_router=_router_without_models(), + user_model=None, + route_type="acompletion", + ) + finally: + global_agent_registry.agent_list = original_agents + global_agent_registry.config_agents = original_config_agents + + mock_acompletion.assert_called_once() + call_kwargs = mock_acompletion.call_args.kwargs + assert call_kwargs["model"] == f"a2a/{agent_name}" + assert call_kwargs["api_base"] == "http://sibling-db-agent.example.com" + prisma_client.db.litellm_agentstable.find_many.assert_awaited() diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index 3ae0e1e7d18..ebd52c448bc 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -1091,3 +1091,124 @@ async def test_route_request_rejects_chat_completion_without_messages(): assert exc_info.value.status_code == 400 assert exc_info.value.param == "messages" llm_router.acompletion.assert_not_called() + + +class FakeProxyModelTable: + def __init__(self, rows): + self.rows = rows + self.find_many_wheres = [] + + async def find_many(self, where=None, **kwargs): + self.find_many_wheres.append(where) + return list(self.rows) + + +def _fake_prisma_client_with_models(rows): + from types import SimpleNamespace + + table = FakeProxyModelTable(rows) + return SimpleNamespace(db=SimpleNamespace(litellm_proxymodeltable=table)), table + + +def _db_model_row(model_name: str, mock_response: str): + from types import SimpleNamespace + + return SimpleNamespace( + model_id=f"{model_name}-id", + model_name=model_name, + litellm_params={"model": "openai/gpt-4o", "api_key": "fake", "mock_response": mock_response}, + model_info={}, + blocked=False, + ) + + +@pytest.mark.asyncio +async def test_route_request_read_through_recovers_model_created_on_sibling_replica(monkeypatch): + """Regression: a model written to the DB by another replica must be served on + first request instead of 400ing until the periodic config reload.""" + import litellm + import litellm.proxy.proxy_server as proxy_server + + model_name = "e2e-sibling-replica-model" + router = litellm.Router( + model_list=[ + { + "model_name": "some-other-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}, + } + ] + ) + fake_prisma, table = _fake_prisma_client_with_models([_db_model_row(model_name, "hello-from-db")]) + monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "llm_router", router) + + llm_call = await route_request( + data={"model": model_name, "messages": [{"role": "user", "content": "hi"}]}, + llm_router=router, + user_model=None, + route_type="acompletion", + ) + response = await llm_call + + assert response.choices[0].message.content == "hello-from-db" + assert len(table.find_many_wheres) == 1 + assert table.find_many_wheres[0] == {"OR": [{"model_name": model_name}, {"model_id": model_name}]} + + +@pytest.mark.asyncio +async def test_route_request_unknown_model_raises_and_hits_db_once_within_ttl(monkeypatch): + import litellm + import litellm.proxy.proxy_server as proxy_server + + model_name = "e2e-model-nobody-created" + router = litellm.Router( + model_list=[ + { + "model_name": "some-other-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}, + } + ] + ) + fake_prisma, table = _fake_prisma_client_with_models([]) + monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "llm_router", router) + + data = {"model": model_name, "messages": [{"role": "user", "content": "hi"}]} + with pytest.raises(ProxyModelNotFoundError): + await route_request(data=data, llm_router=router, user_model=None, route_type="acompletion") + with pytest.raises(ProxyModelNotFoundError): + await route_request(data=data, llm_router=router, user_model=None, route_type="acompletion") + + assert len(table.find_many_wheres) == 1 + + +@pytest.mark.asyncio +async def test_route_request_read_through_disabled_without_store_model_in_db(monkeypatch): + import litellm + import litellm.proxy.proxy_server as proxy_server + + model_name = "e2e-config-only-proxy-model" + router = litellm.Router( + model_list=[ + { + "model_name": "some-other-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}, + } + ] + ) + fake_prisma, table = _fake_prisma_client_with_models([_db_model_row(model_name, "should-not-load")]) + monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server, "store_model_in_db", False) + monkeypatch.setattr(proxy_server, "llm_router", router) + + with pytest.raises(ProxyModelNotFoundError): + await route_request( + data={"model": model_name, "messages": [{"role": "user", "content": "hi"}]}, + llm_router=router, + user_model=None, + route_type="acompletion", + ) + + assert table.find_many_wheres == [] From b1d77bb5dbc4e2697783258385eb9ac8029fe5b4 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 9 Aug 2026 01:36:33 +0000 Subject: [PATCH 027/538] style: ruff format agentcore search transformation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/search/transformation.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index 4faaad6cc00..f6d852ffb0a 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -118,9 +118,7 @@ def _iter_sse_events(text: str) -> Iterator[Mapping[str, object]]: progress notifications before the JSON-RPC response. """ for chunk in _SSE_EVENT_SEPARATOR.split(text): - payload = "\n".join( - line[len("data:") :].lstrip() for line in chunk.splitlines() if line.startswith("data:") - ) + payload = "\n".join(line[len("data:") :].lstrip() for line in chunk.splitlines() if line.startswith("data:")) if not payload: continue try: From 15a6664171f2ba2b559db1ea333db44996a9a7df Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 9 Aug 2026 02:17:44 +0000 Subject: [PATCH 028/538] fix(search): keep signed auth headers out of logging callbacks Log the pre-signing headers in the search pre_call hook so SigV4 and bearer Authorization values are never handed to user-configured logger callbacks, and tighten sign_request's annotations. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/base_llm/search/transformation.py | 8 ++++---- litellm/llms/bedrock/search/transformation.py | 8 ++++---- litellm/llms/custom_httpx/llm_http_handler.py | 4 ++-- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/litellm/llms/base_llm/search/transformation.py b/litellm/llms/base_llm/search/transformation.py index fad4be538c7..59039d68ede 100644 --- a/litellm/llms/base_llm/search/transformation.py +++ b/litellm/llms/base_llm/search/transformation.py @@ -180,12 +180,12 @@ class BaseSearchConfig: def sign_request( self, - headers: dict, # mutable-ok: matches the request header dict every other hook on this base takes - optional_params: dict, # mutable-ok: matches the optional params dict every other hook on this base takes - request_data: dict | list[dict], # mutable-ok: matches transform_search_request's JSON body return type + headers: dict[str, str], # mutable-ok: matches the request header dict every other hook on this base takes + optional_params: dict[str, object], # mutable-ok: matches every other hook on this base + request_data: dict[str, object] | list[dict[str, object]], # mutable-ok: transform_search_request's body api_base: str, api_key: str | None = None, - ) -> tuple[dict, bytes | None]: # mutable-ok: the handler passes these headers straight to httpx + ) -> tuple[dict[str, str], bytes | None]: # mutable-ok: the handler passes these headers straight to httpx """ OPTIONAL diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index f6d852ffb0a..ca9759ed151 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -221,12 +221,12 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): def sign_request( self, - headers: dict, # mutable-ok: BaseSearchConfig hands providers the mutable request header dict - optional_params: dict, # mutable-ok: BaseSearchConfig passes optional params as a dict - request_data: dict | list[dict], # mutable-ok: BaseSearchConfig request bodies are JSON dicts + headers: dict[str, str], # mutable-ok: BaseSearchConfig hands providers the mutable request header dict + optional_params: dict[str, object], # mutable-ok: BaseSearchConfig passes optional params as a dict + request_data: dict[str, object] | list[dict[str, object]], # mutable-ok: request bodies are JSON dicts api_base: str, api_key: str | None = None, - ) -> tuple[dict, bytes | None]: # mutable-ok: BaseSearchConfig.sign_request returns httpx headers + ) -> tuple[dict[str, str], bytes | None]: # mutable-ok: BaseSearchConfig.sign_request returns httpx headers """ Authenticate the MCP request. diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 3ee79241d30..b2f8c21d83c 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1755,7 +1755,7 @@ class BaseLLMHTTPHandler: additional_args={ "complete_input_dict": data, "api_base": complete_url, - "headers": signed_headers, + "headers": headers, }, ) @@ -1848,7 +1848,7 @@ class BaseLLMHTTPHandler: additional_args={ "complete_input_dict": data, "api_base": complete_url, - "headers": signed_headers, + "headers": headers, }, ) From 25144fc03cebf686483acdae02bf3bd1ce72e3d3 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 9 Aug 2026 03:11:12 +0000 Subject: [PATCH 029/538] chore: retrigger ci after docs merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> From 90cd378a5943bfa739b5eb2d1f4c4234a76920ac Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 12 Aug 2026 01:04:52 +0000 Subject: [PATCH 030/538] fix(streaming): accept provider cost objects when propagating usage cost Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/streaming_handler.py | 19 ++++- .../test_streaming_handler.py | 74 +++++++++++++++++++ 2 files changed, 91 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 99b1c1a2ab7..48e58c578e6 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1789,6 +1789,20 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason = "tool_calls" return model_response + @staticmethod + def _resolve_provider_reported_cost(usage_cost: object) -> float | None: + """ + Providers report usage.cost either as a number or, for Perplexity, as a + breakdown object whose total lives under ``total_cost``. + """ + if isinstance(usage_cost, bool): + return None + if isinstance(usage_cost, (int, float)): + return float(usage_cost) + if isinstance(usage_cost, dict): + return CustomStreamWrapper._resolve_provider_reported_cost(usage_cost.get("total_cost")) + return None + @staticmethod def _propagate_usage_cost_to_hidden_params( response: "ModelResponse", @@ -1799,10 +1813,11 @@ class CustomStreamWrapper: calculator uses it instead of a token-based estimate. """ _usage: Final[Usage | None] = getattr(response, "usage", None) - if _usage is not None and hasattr(_usage, "cost") and _usage.cost is not None: + _cost: Final = CustomStreamWrapper._resolve_provider_reported_cost(getattr(_usage, "cost", None)) + if _cost is not None: if "additional_headers" not in response._hidden_params: response._hidden_params["additional_headers"] = {} - response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = float(_usage.cost) + response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = _cost def __next__(self) -> "ModelResponseStream": cache_hit = False diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 101935cac0a..dad4faa98b4 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -1676,6 +1676,80 @@ def test_openrouter_streaming_cost_propagates_to_hidden_params(): assert provider_cost == 0.00025 +def test_perplexity_streaming_dict_cost_propagates_to_hidden_params(): + """ + Regression: Perplexity reports usage.cost as a breakdown object, which used to + blow up the end of the stream with + `float() argument must be a string or a real number, not 'dict'`. + """ + import litellm + from litellm.cost_calculator import get_response_cost_from_hidden_params + + chunks = [ + ModelResponseStream( + id="chatcmpl-pplx", + created=1742056047, + model="perplexity/sonar", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="Hi", role="assistant"), + ) + ], + usage=None, + ), + ModelResponseStream( + id="chatcmpl-pplx", + created=1742056048, + model="perplexity/sonar", + choices=[ + StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="")) + ], + usage=None, + ), + ModelResponseStream( + id="chatcmpl-pplx", + created=1742056049, + model="perplexity/sonar", + choices=[ + StreamingChoices(finish_reason=None, index=0, delta=Delta(content="")) + ], + usage=Usage( + completion_tokens=18, + prompt_tokens=12, + total_tokens=30, + cost={ + "input_tokens_cost": 0.000012, + "output_tokens_cost": 0.000018, + "request_cost": 0.005, + "total_cost": 0.00503, + }, + ), + ), + ] + + complete_response = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "test"}] + ) + + assert complete_response is not None + + CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response) + + assert ( + get_response_cost_from_hidden_params(complete_response._hidden_params) + == 0.00503 + ) + + +def test_provider_reported_cost_ignores_unusable_shapes(): + assert CustomStreamWrapper._resolve_provider_reported_cost(None) is None + assert CustomStreamWrapper._resolve_provider_reported_cost({}) is None + assert CustomStreamWrapper._resolve_provider_reported_cost({"total_cost": None}) is None + assert CustomStreamWrapper._resolve_provider_reported_cost(0.5) == 0.5 + + def test_handle_special_delta_attributes( initialized_custom_stream_wrapper: CustomStreamWrapper, ): From 84c1df918d40d6f9ce97c32684cb27f50cb8830b Mon Sep 17 00:00:00 2001 From: Daniel Meismer Date: Tue, 11 Aug 2026 22:59:02 -0400 Subject: [PATCH 031/538] fix(mcp): decouple OAuth discovery from startup Register remote MCP servers without awaiting OAuth metadata, warm discovery in the background, and share bounded request-time retries with per-server cooldowns. Preserve the existing discovered-tool boundary for explicit server calls. Co-Authored-By: Codex --- .../mcp_server/discoverable_endpoints.py | 28 +- .../mcp_server/mcp_server_manager.py | 532 ++++++++++++--- .../proxy/_experimental/mcp_server/server.py | 2 + .../mcp_server/test_discoverable_endpoints.py | 134 +++- .../mcp_server/test_mcp_server.py | 243 ++++--- .../mcp_server/test_mcp_server_manager.py | 639 ++++++++++++++++-- 6 files changed, 1328 insertions(+), 250 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 693e3f8e47d..48e92c2643e 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1676,10 +1676,17 @@ async def authorize( lookup_name: Final[str | None] = mcp_server_name or client_id client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) mcp_server = ( - global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip) if lookup_name else None + await global_mcp_server_manager.get_resolved_mcp_server_by_name(lookup_name, client_ip=client_ip) + if lookup_name + else None ) if mcp_server is None and mcp_server_name is None: - mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) + unresolved_server: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) + mcp_server = ( + await global_mcp_server_manager.ensure_oauth_metadata_discovered(unresolved_server) + if unresolved_server is not None + else None + ) if mcp_server is None: raise HTTPException(status_code=404, detail="MCP server not found") _raise_if_not_oauth2(mcp_server) @@ -1757,9 +1764,14 @@ async def token_endpoint( lookup_name: Final = mcp_server_name or client_id client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) - mcp_server = global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip) + mcp_server = await global_mcp_server_manager.get_resolved_mcp_server_by_name(lookup_name, client_ip=client_ip) if mcp_server is None and mcp_server_name is None: - mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) + unresolved_server: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) + mcp_server = ( + await global_mcp_server_manager.ensure_oauth_metadata_discovered(unresolved_server) + if unresolved_server is not None + else None + ) if mcp_server is None: raise HTTPException(status_code=404, detail="MCP server not found") return await exchange_token_with_server( @@ -2558,9 +2570,10 @@ async def register_client(request: Request, mcp_server_name: str | None = None): return await register_aggregate_client(request=request, request_body=data) resolved: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if resolved: + resolved_server: Final = await global_mcp_server_manager.ensure_oauth_metadata_discovered(resolved) return await register_client_with_server( request=request, - mcp_server=resolved, + mcp_server=resolved_server, client_name=data.get("client_name", ""), grant_types=data.get("grant_types", []), response_types=data.get("response_types", []), @@ -2570,7 +2583,10 @@ async def register_client(request: Request, mcp_server_name: str | None = None): ) return dummy_return - mcp_server: Final = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name, client_ip=client_ip) + mcp_server: Final = await global_mcp_server_manager.get_resolved_mcp_server_by_name( + mcp_server_name, + client_ip=client_ip, + ) if mcp_server is None: return dummy_return return await register_client_with_server( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index c8ff6e262d2..cdaf3f2f206 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -13,8 +13,9 @@ import json import os import re import time -from collections.abc import AsyncIterator, Callable, Sequence +from collections.abc import AsyncIterator, Callable, Iterable, Sequence from contextlib import asynccontextmanager +from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypedDict, cast from urllib.parse import ParseResult, urlparse @@ -216,12 +217,43 @@ _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES: Final[tuple[MCPAuth, ...]] = ( ) -# OAuth discovery retry cooldown for servers whose endpoints stay unresolved. The base is one -# reload cadence so a transient upstream failure recovers immediately; the cap bounds the request -# amplification and log volume of a permanently broken configuration. +_MCP_OAUTH_DISCOVERY_ON_STARTUP_ENV: Final = "LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP" +_TRUE_ENV_VALUES: Final = frozenset(("1", "true", "yes", "on")) +_OAUTH_DISCOVERY_RETRY_DELAYS_SECONDS: Final = (0.05, 0.15) _OAUTH_DISCOVERY_RETRY_BASE_SECONDS: Final = 30.0 _OAUTH_DISCOVERY_RETRY_MAX_SECONDS: Final = 900.0 + +def _oauth_discovery_now() -> float: + return time.monotonic() + + +def _oauth_discovery_retry_delay(consecutive_failures: int) -> float: + backoff_multiplier: Final[int] = 1 << max(consecutive_failures - 1, 0) + return min( + _OAUTH_DISCOVERY_RETRY_BASE_SECONDS * backoff_multiplier, + _OAUTH_DISCOVERY_RETRY_MAX_SECONDS, + ) + + +def _mcp_oauth_discovery_on_startup_enabled() -> bool: + """Return whether remote MCP OAuth metadata is discovered during registration. + + Discovery is deferred until the first admitted request unless explicitly + enabled with ``1``, ``true``, ``yes``, or ``on``. + """ + value: Final = os.getenv(_MCP_OAUTH_DISCOVERY_ON_STARTUP_ENV) + return value is not None and value.strip().lower() in _TRUE_ENV_VALUES + + +def _requires_oauth_discovery( + server_url: str | None, + use_issuer_anchor: bool, + server: MCPServer, +) -> bool: + return _has_oauth_discovery_source(server_url, use_issuer_anchor) and _oauth_endpoints_unresolved(server) + + _StringList: TypeAlias = list[str] _StringMap: TypeAlias = dict[str, str] _ToolParamMap: TypeAlias = dict[str, list[str]] @@ -230,6 +262,34 @@ _InMemoryCacheDict: TypeAlias = dict[str, object] _ToolArguments: TypeAlias = dict[str, object] +@dataclass(frozen=True, slots=True) +class _OAuthDiscoveryResolved: + server: MCPServer + + +@dataclass(frozen=True, slots=True) +class _OAuthDiscoveryFailed: + server_id: str + timed_out: bool + + +@dataclass(frozen=True, slots=True) +class _OAuthDiscoveryStale: + server_id: str + + +_OAuthDiscoveryOutcome: TypeAlias = _OAuthDiscoveryResolved | _OAuthDiscoveryFailed | _OAuthDiscoveryStale + + +@dataclass(frozen=True, slots=True) +class _OAuthDiscoverySlot: + server_id: str + generation: int + task: asyncio.Task[_OAuthDiscoveryOutcome] | None = None + consecutive_failures: int = 0 + retry_not_before: float = 0.0 + + class MCPServerConfig(TypedDict, total=False): """Shape of a single ``mcp_servers`` entry in config.yaml, as consumed by :meth:`MCPServerManager.load_servers_from_config`. Every key is optional: YAML supplies @@ -620,6 +680,7 @@ def _warn_oauth_endpoints_unresolved( server_ref: str, server_url: str | None, discovery_attempted: bool, + discovery_deferred: bool = False, issuer_anchored: bool, metadata: MCPOAuthMetadata | None, needs_authorization_url: bool, @@ -638,7 +699,7 @@ def _warn_oauth_endpoints_unresolved( are needed (client_credentials never needs authorization_url; OBO needs only token_url); the issuer-anchored arm is excluded here because it has its own RFC 8414 §3.3 warning. """ - if issuer_anchored: + if discovery_deferred or issuer_anchored: return unresolved: Final = tuple( field @@ -1393,41 +1454,288 @@ class MCPServerManager: # empty result, or failure). Used to throttle re-probes for servers that do # not return instructions, and to apply a short cooldown after failures. self._upstream_initialize_instructions_probed_at: dict[str, float] = {} - # Per-server (consecutive failures, monotonic timestamp) for OAuth discovery retries, so a - # server whose endpoints never resolve backs off instead of re-running the full - # RFC 9728 -> 8414 chain, and re-logging its warning, on every reload forever. - self._oauth_discovery_retry_state: dict[ - str, tuple[int, float] - ] = {} # mutable-ok: retry cooldown cache, keyed per server and pruned on success + self._oauth_discovery_on_startup = _mcp_oauth_discovery_on_startup_enabled() + self._oauth_discovery_generation_counter = 0 + self._oauth_discovery_slots: tuple[_OAuthDiscoverySlot, ...] = () - def _oauth_discovery_retry_due(self, server_id: str) -> bool: - """Whether an unresolved server is due for another discovery attempt. + def _oauth_discovery_slot(self, server_id: str) -> _OAuthDiscoverySlot | None: + return next((slot for slot in self._oauth_discovery_slots if slot.server_id == server_id), None) - The reload fast-path exemption is what retries a failed discovery, so without a cooldown a - permanently unresolvable server re-runs the whole RFC 9728 -> RFC 8414 -> origin-fallback - chain and re-emits its unresolved-endpoints warning on every reload, per server, forever. - Delay doubles per consecutive failure from ``_OAUTH_DISCOVERY_RETRY_BASE_SECONDS`` up to - ``_OAUTH_DISCOVERY_RETRY_MAX_SECONDS``, so a transient outage still recovers on the next - reload while a broken configuration settles to one attempt per cap. - """ - state: Final = self._oauth_discovery_retry_state.get(server_id) - if state is None: - return True - failures, attempted_at = state - backoff_multiplier: Final[int] = 2 ** max(failures - 1, 0) - delay: Final = min( - _OAUTH_DISCOVERY_RETRY_BASE_SECONDS * backoff_multiplier, - _OAUTH_DISCOVERY_RETRY_MAX_SECONDS, + def _remove_oauth_discovery_slot(self, server_id: str) -> None: + self._oauth_discovery_slots = tuple(slot for slot in self._oauth_discovery_slots if slot.server_id != server_id) + + def _store_oauth_discovery_slot(self, slot: _OAuthDiscoverySlot) -> None: + self._oauth_discovery_slots = ( + *(existing for existing in self._oauth_discovery_slots if existing.server_id != slot.server_id), + slot, ) - return (time.monotonic() - attempted_at) >= delay - def _record_oauth_discovery_outcome(self, server: MCPServer) -> None: - """Advance or clear a server's retry cooldown after a rebuild resolved it or did not.""" - if not _oauth_endpoints_unresolved(server): - self._oauth_discovery_retry_state.pop(server.server_id, None) + def _set_oauth_discovery_deferred(self, server_id: str, discovery_deferred: bool) -> None: + previous: Final = self._oauth_discovery_slot(server_id) + self._remove_oauth_discovery_slot(server_id) + if previous is not None and previous.task is not None and not previous.task.done(): + previous.task.cancel() + if discovery_deferred: + self._oauth_discovery_generation_counter += 1 + self._store_oauth_discovery_slot( + _OAuthDiscoverySlot( + server_id=server_id, + generation=self._oauth_discovery_generation_counter, + ) + ) + + def _invalidate_oauth_discovery_state(self, server_id: str) -> None: + previous: Final = self._oauth_discovery_slot(server_id) + self._remove_oauth_discovery_slot(server_id) + if previous is not None and previous.task is not None and not previous.task.done(): + previous.task.cancel() + + def _registered_server(self, server: MCPServer) -> MCPServer: + return self.registry.get(server.server_id) or self.config_mcp_servers.get(server.server_id) or server + + async def _discover_oauth_metadata_for_server(self, server: MCPServer) -> MCPOAuthMetadata | None: + manual_issuer: Final = _blank_to_none(server.issuer) + manual_authorization_url: Final = _blank_to_none(server.authorization_url) + manual_token_url: Final = _blank_to_none(server.token_url) + is_discovery_auth_type: Final = server.auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES + use_issuer_anchor: Final = server.issuer_is_anchored + obo_needs_discovery: Final = self._obo_needs_endpoint_discovery( + server.auth_type, + server.token_exchange_endpoint, + manual_token_url, + ) + needs_authorization_url: Final = is_discovery_auth_type and server.oauth2_flow != "client_credentials" + needs_token_url: Final = is_discovery_auth_type or obo_needs_discovery + warn_on_empty_discovery: Final = _discovery_failure_leaves_needs_unresolved( + needs_authorization_url=needs_authorization_url, + needs_token_url=needs_token_url, + manual_authorization_url=manual_authorization_url, + manual_token_url=manual_token_url, + ) + metadata: Final = await ( + self._fetch_issuer_anchored_oauth_metadata(manual_issuer, server.url) + if use_issuer_anchor and manual_issuer is not None + else self._descovery_metadata( + server_url=server.url or "", + allow_origin_fallback=is_discovery_auth_type, + warn_when_no_metadata=warn_on_empty_discovery, + ) + ) + if use_issuer_anchor: + return metadata + gated_metadata: Final = ( + _restrict_discovery_to_corroborated_authorization_server( + metadata, + manual_authorization_url, + server.server_id, + server.is_dcr_bridge, + ) + if is_discovery_auth_type + else metadata + ) + _warn_oauth_endpoints_unresolved( + server_ref=server.alias or server.server_name or server.server_id, + server_url=server.url, + discovery_attempted=True, + issuer_anchored=False, + metadata=gated_metadata, + needs_authorization_url=needs_authorization_url, + needs_token_url=needs_token_url, + manual_authorization_url=manual_authorization_url, + manual_token_url=manual_token_url, + ) + return gated_metadata + + @staticmethod + def _merge_discovered_oauth_metadata(server: MCPServer, metadata: MCPOAuthMetadata | None) -> MCPServer: + if metadata is None: + return server + discovered_issuer: Final = metadata.discovered_issuer if not metadata.from_origin_fallback else None + resolved: Final = server.model_copy() + resolved.scopes = server.scopes or metadata.scopes + resolved.issuer = server.issuer or discovered_issuer + resolved.authorization_url = server.authorization_url or metadata.authorization_url + resolved.token_url = server.token_url or metadata.token_url + resolved.registration_url = server.registration_url or metadata.registration_url + return resolved + + def _oauth_discovery_slot_is_current(self, server_id: str, generation: int) -> bool: + slot: Final = self._oauth_discovery_slot(server_id) + return slot is not None and slot.generation == generation + + def _publish_resolved_oauth_server( + self, + server: MCPServer, + generation: int, + ) -> MCPServer | None: + if not self._oauth_discovery_slot_is_current(server.server_id, generation): + return None + if server.server_id in self.registry: + self.registry[server.server_id] = server + elif server.server_id in self.config_mcp_servers: + self.config_mcp_servers[server.server_id] = server + else: + return None + self._remove_oauth_discovery_slot(server.server_id) + return server + + async def _attempt_oauth_metadata_once( + self, + server: MCPServer, + generation: int, + ) -> _OAuthDiscoveryOutcome | None: + if not self._oauth_discovery_slot_is_current(server.server_id, generation): + return _OAuthDiscoveryStale(server_id=server.server_id) + current: Final = self._registered_server(server) + if not _oauth_endpoints_unresolved(current): + published: Final = self._publish_resolved_oauth_server(current, generation) + return ( + _OAuthDiscoveryResolved(server=published) + if published is not None + else _OAuthDiscoveryStale(server_id=server.server_id) + ) + metadata: Final = await self._discover_oauth_metadata_for_server(current) + if not self._oauth_discovery_slot_is_current(server.server_id, generation): + return _OAuthDiscoveryStale(server_id=server.server_id) + candidate: Final = self._merge_discovered_oauth_metadata(self._registered_server(server), metadata) + if _oauth_endpoints_unresolved(candidate): + return None + published_candidate: Final = self._publish_resolved_oauth_server(candidate, generation) + return ( + _OAuthDiscoveryResolved(server=published_candidate) + if published_candidate is not None + else _OAuthDiscoveryStale(server_id=server.server_id) + ) + + async def _attempt_oauth_metadata_resolution( + self, + server: MCPServer, + generation: int, + retry_delays: tuple[float, ...] = _OAUTH_DISCOVERY_RETRY_DELAYS_SECONDS, + ) -> _OAuthDiscoveryOutcome: + outcome: Final = await self._attempt_oauth_metadata_once(server, generation) + if outcome is not None: + return outcome + if not retry_delays: + return _OAuthDiscoveryFailed(server_id=server.server_id, timed_out=False) + await asyncio.sleep(retry_delays[0]) + return await self._attempt_oauth_metadata_resolution(server, generation, retry_delays[1:]) + + async def _run_oauth_metadata_resolution( + self, + server: MCPServer, + generation: int, + ) -> _OAuthDiscoveryOutcome: + try: + outcome: Final = await asyncio.wait_for( + self._attempt_oauth_metadata_resolution(server, generation), + timeout=MCP_METADATA_TIMEOUT, + ) + except asyncio.TimeoutError: + verbose_logger.warning( + "Deferred MCP OAuth discovery timed out after %ss for server %s", + MCP_METADATA_TIMEOUT, + server.server_id, + ) + failure: Final = _OAuthDiscoveryFailed(server_id=server.server_id, timed_out=True) + self._record_oauth_discovery_failure(server.server_id, generation) + return failure + if isinstance(outcome, _OAuthDiscoveryFailed): + self._record_oauth_discovery_failure(server.server_id, generation) + return outcome + + def _record_oauth_discovery_failure(self, server_id: str, generation: int) -> None: + slot: Final = self._oauth_discovery_slot(server_id) + if slot is None or slot.generation != generation: return - failures, _ = self._oauth_discovery_retry_state.get(server.server_id, (0, 0.0)) - self._oauth_discovery_retry_state[server.server_id] = (failures + 1, time.monotonic()) + consecutive_failures: Final = slot.consecutive_failures + 1 + self._store_oauth_discovery_slot( + replace( + slot, + consecutive_failures=consecutive_failures, + retry_not_before=_oauth_discovery_now() + _oauth_discovery_retry_delay(consecutive_failures), + ) + ) + + def _get_or_start_oauth_discovery_task( + self, + server: MCPServer, + ) -> tuple[asyncio.Task[_OAuthDiscoveryOutcome], int] | None: + slot: Final = self._oauth_discovery_slot(server.server_id) + if slot is None: + return None + if slot.task is not None: + if not slot.task.done() or _oauth_discovery_now() < slot.retry_not_before: + return slot.task, slot.generation + task: Final = asyncio.create_task( + self._run_oauth_metadata_resolution(self._registered_server(server), slot.generation) + ) + self._store_oauth_discovery_slot(replace(slot, task=task)) + return task, slot.generation + + def prime_oauth_metadata_discovery(self, server: MCPServer) -> None: + """Start best-effort OAuth metadata discovery for ``server``. + + The call returns immediately and never delays registration. It is a no-op + when the server has no deferred discovery slot. + + Args: + server: The registered MCP server to warm metadata for. + """ + self._get_or_start_oauth_discovery_task(server) + + def _prime_oauth_metadata_discovery_for_servers(self, servers: Iterable[MCPServer]) -> None: + for server in servers: + self.prime_oauth_metadata_discovery(server) + + def _reconcile_oauth_discovery_slots_for_servers(self, servers: Iterable[MCPServer]) -> None: + """Align retry slots after an atomic registry replacement.""" + for server in servers: + should_defer = bool(server.url) and _oauth_endpoints_unresolved(server) + has_slot = self._oauth_discovery_slot(server.server_id) is not None + if should_defer != has_slot: + self._set_oauth_discovery_deferred(server.server_id, should_defer) + + async def ensure_oauth_metadata_discovered(self, server: MCPServer) -> MCPServer: + """Join the bounded discovery task and return the resolved server. + + Concurrent callers share one task per server. A failed attempt remains + retryable after a per-server cooldown. + + Args: + server: The MCP server whose OAuth metadata must be resolved. + + Returns: + The resolved server, or the registered server when no discovery is + pending. + + Raises: + HTTPException: Status 503 when discovery times out or returns + incomplete metadata. + """ + acquisition: Final = self._get_or_start_oauth_discovery_task(server) + if acquisition is None: + return self._registered_server(server) + task, generation = acquisition + try: + outcome: Final = await asyncio.shield(task) + except asyncio.CancelledError: + if task.cancelled() and not self._oauth_discovery_slot_is_current(server.server_id, generation): + return await self.ensure_oauth_metadata_discovered(server) + raise + match outcome: + case _OAuthDiscoveryResolved(resolved_server): + return resolved_server + case _OAuthDiscoveryStale(): + return await self.ensure_oauth_metadata_discovered(server) + case _OAuthDiscoveryFailed(timed_out=timed_out): + current: Final = self._registered_server(server) + server_ref: Final = current.alias or current.server_name or current.name or current.server_id + reason: Final = "timed out" if timed_out else "returned incomplete metadata" + raise HTTPException( + status_code=503, + detail=f"OAuth metadata discovery {reason} for MCP server {server_ref!r}", + ) def _remember_upstream_initialize_instructions(self, server: MCPServer, client: MCPClient) -> None: raw: Final[str | None] = getattr(client, "_last_initialize_instructions", None) @@ -1618,7 +1926,8 @@ class MCPServerManager: manual_authorization_url=manual_authorization_url, manual_token_url=manual_token_url, ) - if not should_discover: + discovery_deferred = should_discover and not self._oauth_discovery_on_startup + if not should_discover or discovery_deferred: mcp_oauth_metadata = None elif use_issuer_anchor and manual_issuer is not None: mcp_oauth_metadata = await self._fetch_issuer_anchored_oauth_metadata(manual_issuer, server_url) @@ -1696,6 +2005,7 @@ class MCPServerManager: server_ref=server_name or server_id, server_url=server_url, discovery_attempted=should_discover, + discovery_deferred=discovery_deferred, issuer_anchored=use_issuer_anchor, metadata=gated_oauth_metadata, needs_authorization_url=needs_authorization_url, @@ -1774,6 +2084,10 @@ class MCPServerManager: self._assign_unique_short_prefix(new_server) _warn_internal_delegate_pkce_if_applicable(new_server, source="config") self.config_mcp_servers[server_id] = new_server + self._set_oauth_discovery_deferred( + server_id, + _requires_oauth_discovery(server_url, use_issuer_anchor, new_server), + ) # Check if this is an OpenAPI-based server spec_path = server_config.get("spec_path", None) @@ -1791,6 +2105,8 @@ class MCPServerManager: await self._hydrate_config_servers_dcr_clients() + self._prime_oauth_metadata_discovery_for_servers(self.config_mcp_servers.values()) + self.initialize_tool_name_to_mcp_server_name_mapping() async def _hydrate_config_servers_dcr_clients(self) -> None: @@ -1992,6 +2308,7 @@ class MCPServerManager: if evicted is not None: verbose_logger.debug("Removed MCP Server: %s", mcp_server.server_id or mcp_server.server_name) self._cleanup_server_tool_routing_artifacts(evicted) + self._invalidate_oauth_discovery_state(evicted.server_id) else: verbose_logger.warning("Server ID %s not found in registry", mcp_server.server_id) @@ -2023,7 +2340,7 @@ class MCPServerManager: use_issuer_anchor: bool, scopes: list[str] | None, token_exchange_endpoint: str | None, - ) -> MCPOAuthMetadata | None: + ) -> tuple[MCPOAuthMetadata | None, bool]: obo_needs_discovery = self._obo_needs_endpoint_discovery(auth_type, token_exchange_endpoint, manual_token_url) needs_authorization_url: Final = ( is_discovery_auth_type and getattr(mcp_server, "oauth2_flow", None) != "client_credentials" @@ -2039,7 +2356,8 @@ class MCPServerManager: needs_discovery: Final = _has_oauth_discovery_source(server_url, use_issuer_anchor) and ( (is_discovery_auth_type and not has_all_upstream_oauth_fields) or obo_needs_discovery ) - if not needs_discovery: + discovery_deferred: Final = needs_discovery and not self._oauth_discovery_on_startup + if not needs_discovery or discovery_deferred: mcp_oauth_metadata: MCPOAuthMetadata | None = None elif use_issuer_anchor and manual_issuer is not None: mcp_oauth_metadata = await self._fetch_issuer_anchored_oauth_metadata(manual_issuer, server_url) @@ -2050,7 +2368,7 @@ class MCPServerManager: warn_when_no_metadata=warn_on_empty_discovery, ) if use_issuer_anchor: - return mcp_oauth_metadata + return mcp_oauth_metadata, discovery_deferred gated_metadata: Final = ( _restrict_discovery_to_corroborated_authorization_server( mcp_oauth_metadata, @@ -2065,6 +2383,7 @@ class MCPServerManager: server_ref=mcp_server.alias or mcp_server.server_name or mcp_server.server_id, server_url=server_url, discovery_attempted=needs_discovery, + discovery_deferred=discovery_deferred, issuer_anchored=False, metadata=gated_metadata, needs_authorization_url=needs_authorization_url, @@ -2072,7 +2391,7 @@ class MCPServerManager: manual_authorization_url=manual_authorization_url, manual_token_url=manual_token_url, ) - return gated_metadata + return gated_metadata, discovery_deferred async def build_mcp_server_from_table( self, @@ -2177,7 +2496,7 @@ class MCPServerManager: manual_registration_url, mcp_server.alias or mcp_server.server_name or mcp_server.server_id, ) - gated_oauth_metadata: Final = await self._resolve_table_oauth_metadata( + gated_oauth_metadata, _ = await self._resolve_table_oauth_metadata( mcp_server=mcp_server, auth_type=auth_type, server_url=server_url, @@ -2283,6 +2602,10 @@ class MCPServerManager: max_concurrent_requests=getattr(mcp_server, "max_concurrent_requests", None), ) _warn_internal_delegate_pkce_if_applicable(new_server, source="database") + self._set_oauth_discovery_deferred( + new_server.server_id, + _requires_oauth_discovery(server_url, use_issuer_anchor, new_server), + ) return new_server async def _maybe_register_openapi_tools(self, server: MCPServer, *, initialize_mapping: bool = True): @@ -2316,6 +2639,7 @@ class MCPServerManager: self._assign_unique_short_prefix(new_server) self.registry[mcp_server.server_id] = new_server await self._maybe_register_openapi_tools(new_server) + self.prime_oauth_metadata_discovery(new_server) verbose_logger.debug("Added MCP Server: %s", new_server.name) except Exception as e: @@ -2332,6 +2656,7 @@ class MCPServerManager: evicted = self.registry.pop(mcp_server.server_name, None) if evicted is not None: self._cleanup_server_tool_routing_artifacts(evicted) + self._invalidate_oauth_discovery_state(evicted.server_id) return try: if mcp_server.server_id in self.registry: @@ -2350,6 +2675,7 @@ class MCPServerManager: self._assign_unique_short_prefix(new_server) self.registry[mcp_server.server_id] = new_server await self._maybe_register_openapi_tools(new_server) + self.prime_oauth_metadata_discovery(new_server) verbose_logger.debug("Updated MCP Server: %s", new_server.name) except Exception as e: @@ -3137,7 +3463,8 @@ class MCPServerManager: subject_token: Final = self._extract_bearer_token(oauth2_headers, None) if not subject_token: return - spec: Final = to_server_spec(server) + resolved_server: Final = await self.ensure_oauth_metadata_discovered(server) + spec: Final = to_server_spec(resolved_server) if spec is None or not isinstance(spec.config, TokenExchangeConfig): return match await self._cred_provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec): @@ -3146,7 +3473,7 @@ class MCPServerManager: case Error(err): if err.tag == "unauthorized": raise_token_exchange_challenge( - server, + resolved_server, root_path=get_server_root_path(), claims=err.unauthorized.claims, ) @@ -3182,8 +3509,9 @@ class MCPServerManager: Returns: Configured MCP client instance. """ - transport: Final = server.transport or MCPTransport.sse - spec = None if transport == MCPTransport.stdio else _to_server_spec_fail_closed(server) + resolved_server: Final = await self.ensure_oauth_metadata_discovered(server) + transport: Final = resolved_server.transport or MCPTransport.sse + spec = None if transport == MCPTransport.stdio else _to_server_spec_fail_closed(resolved_server) provider: Final = cred_provider or self._cred_provider # A caller-supplied per-request override (mcp_auth_header / x-mcp-*) defers to the v1 path # so it wins - except for the modes the v2 resolver owns per-caller (authorization_code's @@ -3202,16 +3530,20 @@ class MCPServerManager: ) ): spec = None - auth_value: Final = await resolve_mcp_auth(server, mcp_auth_header) if spec is None else None + auth_value: Final = await resolve_mcp_auth(resolved_server, mcp_auth_header) if spec is None else None # Create sampling and elicitation callbacks for this client - sampling_cb = _create_sampling_callback(user_api_key_auth=user_api_key_auth) if server.allow_sampling else None - elicitation_cb: Final = _create_elicitation_callback() if server.allow_elicitation else None + sampling_cb = ( + _create_sampling_callback(user_api_key_auth=user_api_key_auth) if resolved_server.allow_sampling else None + ) + elicitation_cb: Final = _create_elicitation_callback() if resolved_server.allow_elicitation else None # Handle stdio transport if transport == MCPTransport.stdio: resolved_env: Final = ( - stdio_env if stdio_env is not None else (dict(server.env) if server.env is not None else None) + stdio_env + if stdio_env is not None + else (dict(resolved_server.env) if resolved_server.env is not None else None) ) # Ensure npm-based STDIO MCP servers have a writable cache dir. @@ -3222,8 +3554,8 @@ class MCPServerManager: # Defense-in-depth: block commands not in the allowlist. # The Pydantic validator blocks new servers; this catches legacy # config/DB records predating the allowlist. - if server.command: - base_command: Final = os.path.basename(server.command) + if resolved_server.command: + base_command: Final = os.path.basename(resolved_server.command) # Strip .exe/.cmd/.bat/.com suffix for Windows compatibility base_command_no_ext = base_command.lower() for ext in [".exe", ".cmd", ".bat", ".com"]: @@ -3236,24 +3568,24 @@ class MCPServerManager: ): raise HTTPException( status_code=403, - detail=f"MCP stdio command '{server.command}' is not in the allowlist ({sorted(MCP_STDIO_ALLOWED_COMMANDS)}). " + detail=f"MCP stdio command '{resolved_server.command}' is not in the allowlist ({sorted(MCP_STDIO_ALLOWED_COMMANDS)}). " f"Add it to LITELLM_MCP_STDIO_EXTRA_COMMANDS to allow this command.", ) stdio_config: MCPStdioConfig | None = None - if server.command and server.args is not None: + if resolved_server.command and resolved_server.args is not None: stdio_config = MCPStdioConfig( - command=server.command, - args=server.args, + command=resolved_server.command, + args=resolved_server.args, env=resolved_env, ) return MCPClient( server_url="", # Not used for stdio transport_type=transport, - auth_type=server.auth_type, + auth_type=resolved_server.auth_type, auth_value=auth_value, - timeout=(server.timeout if server.timeout is not None else MCP_CLIENT_TIMEOUT), + timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), stdio_config=stdio_config, extra_headers=extra_headers, sampling_callback=sampling_cb, @@ -3261,7 +3593,7 @@ class MCPServerManager: ) else: # For HTTP/SSE transports - server_url: Final = server.url or "" + server_url: Final = resolved_server.url or "" if spec is not None: inbound_token = subject_token @@ -3271,7 +3603,7 @@ class MCPServerManager: if per_server_token is not None: inbound_token = per_server_token resolved_auth, extra_headers = await self._resolve_v2_auth( - server=server, + server=resolved_server, spec=spec, provider=provider, subject_token=inbound_token, @@ -3281,8 +3613,8 @@ class MCPServerManager: return MCPClient( server_url=server_url, transport_type=transport, - auth_type=server.auth_type, - timeout=(server.timeout if server.timeout is not None else MCP_CLIENT_TIMEOUT), + auth_type=resolved_server.auth_type, + timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), extra_headers=extra_headers, resolved_auth=resolved_auth, sampling_callback=sampling_cb, @@ -3291,23 +3623,23 @@ class MCPServerManager: # Create SigV4 auth if configured aws_auth = None - if server.auth_type == MCPAuth.aws_sigv4: + if resolved_server.auth_type == MCPAuth.aws_sigv4: aws_auth = MCPSigV4Auth( - aws_access_key_id=server.aws_access_key_id, - aws_secret_access_key=server.aws_secret_access_key, - aws_session_token=server.aws_session_token, - aws_region_name=server.aws_region_name, - aws_service_name=server.aws_service_name, - aws_role_name=server.aws_role_name, - aws_session_name=server.aws_session_name, + aws_access_key_id=resolved_server.aws_access_key_id, + aws_secret_access_key=resolved_server.aws_secret_access_key, + aws_session_token=resolved_server.aws_session_token, + aws_region_name=resolved_server.aws_region_name, + aws_service_name=resolved_server.aws_service_name, + aws_role_name=resolved_server.aws_role_name, + aws_session_name=resolved_server.aws_session_name, ) return MCPClient( server_url=server_url, transport_type=transport, - auth_type=server.auth_type, + auth_type=resolved_server.auth_type, auth_value=auth_value, - timeout=(server.timeout if server.timeout is not None else MCP_CLIENT_TIMEOUT), + timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), extra_headers=extra_headers, aws_auth=aws_auth, sampling_callback=sampling_cb, @@ -3763,7 +4095,10 @@ class MCPServerManager: ) -> tuple[MCPOAuthMetadata | None, tuple[str, ...]]: origin: Final = _redact_mcp_resource_url(server_url) or "" try: - client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) + client: Final = get_async_httpx_client( + llm_provider=httpxSpecialProvider.MCP, + params={"timeout": MCP_METADATA_TIMEOUT}, # mutable-ok: HTTP client factory requires a dict + ) response: Final = await client.get(server_url) response.raise_for_status() ( @@ -5347,6 +5682,8 @@ class MCPServerManager: Note: This now handles prefixed tool names """ for server in self.get_registry().values(): + if self._oauth_discovery_slot(server.server_id) is not None: + continue if server.needs_user_oauth_token: # Skip OAuth2 servers that rely on user-provided tokens continue @@ -5459,9 +5796,9 @@ class MCPServerManager: and existing_server.updated_at is not None and server.updated_at is not None and existing_server.updated_at == server.updated_at - and not ( - _oauth_endpoints_unresolved(existing_server) - and self._oauth_discovery_retry_due(server.server_id) + and ( + self._oauth_discovery_slot(server.server_id) is not None + or not _oauth_endpoints_unresolved(existing_server) ) ): # Re-use existing server instance to avoid re-running build_mcp_server_from_table() @@ -5480,7 +5817,6 @@ class MCPServerManager: # already-decrypted records add_server/update_server are handed. # Decrypt them while building the registry entry. new_server = await self.build_mcp_server_from_table(server, env_vars_are_encrypted=True) - self._record_oauth_discovery_outcome(new_server) # Carry the cached short_prefix from the previous registry entry # (if any) so the prefix is stable across reloads. if existing_server is not None and existing_server.short_prefix: @@ -5517,7 +5853,17 @@ class MCPServerManager: e, ) + dropped_registry_keys: Final = previous_registry.keys() - registered_registry.keys() + for registry_key in dropped_registry_keys: + self._invalidate_oauth_discovery_state(previous_registry[registry_key].server_id) + self.registry = registered_registry + # A discovery task may have published into ``previous_registry`` while + # this replacement was being staged. Reconcile every published entry + # synchronously after the swap so a lost publication cannot also leave + # the replacement unresolved with no retry slot. + self._reconcile_oauth_discovery_slots_for_servers(registered_registry.values()) + self._prime_oauth_metadata_discovery_for_servers(registered_registry.values()) if registered_openapi_tools: self.initialize_tool_name_to_mcp_server_name_mapping() @@ -5705,6 +6051,14 @@ class MCPServerManager: return server return None + async def get_resolved_mcp_server_by_name( + self, + server_name: str, + client_ip: str | None = None, + ) -> MCPServer | None: + server: Final = self.get_mcp_server_by_name(server_name, client_ip=client_ip) + return await self.ensure_oauth_metadata_discovered(server) if server is not None else None + def get_filtered_registry(self, client_ip: str | None = None) -> dict[str, MCPServer]: """ Get registry filtered by client IP access control. @@ -5799,21 +6153,19 @@ class MCPServerManager: should_skip_health_check = True if not should_skip_health_check: - resolved_static_headers: Final = await self._resolve_static_headers_with_env_vars( - server=server, - user_api_key_auth=None, - raise_on_missing=False, - ) - extra_headers: Final = dict(resolved_static_headers) if resolved_static_headers else {} - - client: Final = await self._create_mcp_client( - server=server, - mcp_auth_header=None, - extra_headers=extra_headers, - stdio_env=None, - ) - try: + resolved_static_headers: Final = await self._resolve_static_headers_with_env_vars( + server=server, + user_api_key_auth=None, + raise_on_missing=False, + ) + extra_headers: Final = dict(resolved_static_headers) if resolved_static_headers else {} + client: Final = await self._create_mcp_client( + server=server, + mcp_auth_header=None, + extra_headers=extra_headers, + stdio_env=None, + ) async def _noop(session): return "ok" diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 49a1f1314f0..e4ac40734dc 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3737,6 +3737,8 @@ if MCP_AVAILABLE: # preemptive challenge and let downstream authorization # return 403. continue + if server is not None: + server = await global_mcp_server_manager.ensure_oauth_metadata_discovered(server) if server and server.auth_type == MCPAuth.oauth2: # The challenge decision is per oauth2 sub-mode, not per header: # gateway-managed modes (M2M and interactive authorization_code) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 9bc84b43fc5..cde436e9794 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -41,6 +41,128 @@ def _mock_callback_request(base_url: str = "http://localhost:3000/"): return req +def _unresolved_oauth_server(): + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + return MCPServer( + server_id="cold-oauth-server", + name="cold_oauth_server", + server_name="cold_oauth_server", + alias="cold_oauth_server", + url="https://mcp.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + client_id="client-id", + ) + + +def _resolved_oauth_metadata(): + from litellm.types.mcp_server.mcp_server_manager import MCPOAuthMetadata + + return MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + registration_url="https://idp.example.com/register", + scopes=["mcp.read"], + ) + + +@pytest.mark.asyncio +async def test_authorize_resolves_cold_oauth_metadata(): + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + server = _unresolved_oauth_server() + global_mcp_server_manager.registry[server.server_id] = server + global_mcp_server_manager._set_oauth_discovery_deferred(server.server_id, True) + request = _mock_callback_request("https://litellm.example.com/") + expected = MagicMock() + + with ( + patch.object( + global_mcp_server_manager, + "_discover_oauth_metadata_for_server", + new=AsyncMock(return_value=_resolved_oauth_metadata()), + ) as discovery, + patch.object(discoverable_endpoints, "authorize_with_server", new=AsyncMock(return_value=expected)) as relay, + ): + response = await discoverable_endpoints.authorize( + request=request, + client_id="client-id", + mcp_server_name=server.server_name, + redirect_uri="http://127.0.0.1:60108/callback", + ) + + discovery.assert_awaited_once_with(server) + assert relay.await_args.kwargs["mcp_server"].authorization_url == "https://idp.example.com/authorize" + assert response is expected + + +@pytest.mark.asyncio +async def test_token_resolves_cold_oauth_metadata(): + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + server = _unresolved_oauth_server() + global_mcp_server_manager.registry[server.server_id] = server + global_mcp_server_manager._set_oauth_discovery_deferred(server.server_id, True) + request = _mock_callback_request("https://litellm.example.com/") + expected = MagicMock() + + with ( + patch.object( + global_mcp_server_manager, + "_discover_oauth_metadata_for_server", + new=AsyncMock(return_value=_resolved_oauth_metadata()), + ) as discovery, + patch.object( + discoverable_endpoints, "exchange_token_with_server", new=AsyncMock(return_value=expected) + ) as relay, + ): + response = await discoverable_endpoints.token_endpoint( + request=request, + grant_type="refresh_token", + client_id="client-id", + refresh_token="refresh-token", + mcp_server_name=server.server_name, + ) + + discovery.assert_awaited_once_with(server) + assert relay.await_args.kwargs["mcp_server"].token_url == "https://idp.example.com/token" + assert response is expected + + +@pytest.mark.asyncio +async def test_register_resolves_cold_oauth_metadata(): + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + server = _unresolved_oauth_server() + global_mcp_server_manager.registry[server.server_id] = server + global_mcp_server_manager._set_oauth_discovery_deferred(server.server_id, True) + request = _mock_callback_request("https://litellm.example.com/") + expected = MagicMock() + + with ( + patch.object( + global_mcp_server_manager, + "_discover_oauth_metadata_for_server", + new=AsyncMock(return_value=_resolved_oauth_metadata()), + ) as discovery, + patch.object(discoverable_endpoints, "_read_request_body", new=AsyncMock(return_value={})), + patch.object( + discoverable_endpoints, "register_client_with_server", new=AsyncMock(return_value=expected) + ) as relay, + ): + response = await discoverable_endpoints.register_client(request=request, mcp_server_name=server.server_name) + + discovery.assert_awaited_once_with(server) + assert relay.await_args.kwargs["mcp_server"].registration_url == "https://idp.example.com/register" + assert response is expected + + @pytest.fixture def trust_xff(): """Force ``IPAddressUtils.is_request_from_trusted_proxy`` to True. @@ -8457,6 +8579,8 @@ async def test_authorize_wall_names_the_issuer_for_anchored_servers(): assert "verify the Issuer" in detail_text assert "Servers with no url" not in detail_text assert "idp.example.com" not in detail_text + + def test_passthrough_authorization_code_round_trips_and_rejects_hostile_input(): """The passthrough gateway code seals and recovers the ephemeral DCR client and upstream code, and is total over hostile input: a raw upstream code opens to None, and a tampered or @@ -8865,7 +8989,9 @@ async def test_mint_ephemeral_dcr_client_unusable_registration_response_is_502(p ) from litellm.types.mcp import MCPAuth - server = _bridge_server(auth_type=MCPAuth.true_passthrough, dcr_bridge=None, server_id=server_id, server_name=server_id) + server = _bridge_server( + auth_type=MCPAuth.true_passthrough, dcr_bridge=None, server_id=server_id, server_name=server_id + ) mock_response = MagicMock() mock_response.text = json.dumps(payload) mock_response.raise_for_status = MagicMock() @@ -8943,8 +9069,6 @@ async def test_token_exchange_authenticates_with_the_sealed_clients_own_auth_met assert sent_body["client_secret"] == "mint-secret" - - # --------------------------------------------------------------------------- # LIT-4339: RFC 8707 resource indicators on the upstream OAuth legs # --------------------------------------------------------------------------- @@ -9197,7 +9321,9 @@ def test_upstream_resource_auto_keeps_the_path_because_it_identifies_the_server( sets ``upstream_resource`` explicitly instead of using ``auto``.""" from litellm.proxy._experimental.mcp_server.oauth_utils import resolve_upstream_resource - first = resolve_upstream_resource(_resource_server(url="https://gw.example.com/team-a/mcp", upstream_resource="auto")) + first = resolve_upstream_resource( + _resource_server(url="https://gw.example.com/team-a/mcp", upstream_resource="auto") + ) second = resolve_upstream_resource( _resource_server(url="https://gw.example.com/team-b/mcp", upstream_resource="auto") ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 850d01c6e34..0193f2c9152 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -21,7 +21,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.types.mcp import MCPAuth -from litellm.types.mcp_server.mcp_server_manager import MCPServer +from litellm.types.mcp_server.mcp_server_manager import MCPOAuthMetadata, MCPServer def _rendered_log_message(call): @@ -43,13 +43,19 @@ def cleanup_mcp_global_state(): global_mcp_server_manager, ) - # Clear before test + for slot in global_mcp_server_manager._oauth_discovery_slots: + if slot.task is not None and not slot.task.done(): + slot.task.cancel() global_mcp_server_manager.registry.clear() global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.clear() + global_mcp_server_manager._oauth_discovery_slots = () yield - # Clear after test + for slot in global_mcp_server_manager._oauth_discovery_slots: + if slot.task is not None and not slot.task.done(): + slot.task.cancel() global_mcp_server_manager.registry.clear() global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.clear() + global_mcp_server_manager._oauth_discovery_slots = () except ImportError: # MCP not available, skip cleanup yield @@ -1207,9 +1213,7 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing(): assert result.outcomes["failing2"].tag == "internal" # Verify failure logging for both servers - rendered_exceptions = [ - _rendered_log_message(c) for c in mock_logger.exception.call_args_list if c.args - ] + rendered_exceptions = [_rendered_log_message(c) for c in mock_logger.exception.call_args_list if c.args] assert ( "Error getting tools from server failing_server1: Server failing_server1 connection failed" in rendered_exceptions @@ -5632,13 +5636,17 @@ async def test_delegate_bad_token_gets_connect_time_401(): server = _delegate_auth_mcp_server() scope = _delegate_scope([(b"authorization", b"Bearer bogus-token")]) - with _patch_delegate_resolver(server, "delegate_test"), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", - new=AsyncMock(return_value=[server]), - ), patch( - "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", - new=AsyncMock(return_value=(401, 'Bearer realm="upstream", error="invalid_token"')), - ) as probe: + with ( + _patch_delegate_resolver(server, "delegate_test"), + patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, 'Bearer realm="upstream", error="invalid_token"')), + ) as probe, + ): with pytest.raises(HTTPException) as exc_info: await _check_passthrough_upstream_auth( scope=scope, @@ -5650,7 +5658,9 @@ async def test_delegate_bad_token_gets_connect_time_401(): assert exc_info.value.status_code == 401 challenge = exc_info.value.headers["www-authenticate"] assert 'error="invalid_token"' in challenge - assert 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/delegate_test"' in challenge + assert ( + 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/delegate_test"' in challenge + ) probe.assert_awaited_once() probe_url, probe_auth = probe.call_args.args assert probe_url == "http://upstream:9401/mcp" @@ -5668,13 +5678,17 @@ async def test_delegate_valid_token_passes_preflight(): server = _delegate_auth_mcp_server() scope = _delegate_scope([(b"authorization", b"Bearer good-token")]) - with _patch_delegate_resolver(server, "delegate_test"), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", - new=AsyncMock(return_value=[server]), - ), patch( - "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", - new=AsyncMock(return_value=(200, None)), - ) as probe: + with ( + _patch_delegate_resolver(server, "delegate_test"), + patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(200, None)), + ) as probe, + ): await _check_passthrough_upstream_auth( scope=scope, user_api_key_auth=UserAPIKeyAuth(), @@ -5698,12 +5712,16 @@ async def test_delegate_valid_token_forbidden_returns_403(): server = _delegate_auth_mcp_server() scope = _delegate_scope([(b"authorization", b"Bearer scoped-out-token")]) - with _patch_delegate_resolver(server, "delegate_test"), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", - new=AsyncMock(return_value=[server]), - ), patch( - "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", - new=AsyncMock(return_value=(403, None)), + with ( + _patch_delegate_resolver(server, "delegate_test"), + patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(403, None)), + ), ): with pytest.raises(HTTPException) as exc_info: await _check_passthrough_upstream_auth( @@ -5729,13 +5747,17 @@ async def test_delegate_tokenless_request_not_probed(): server = _delegate_auth_mcp_server() scope = _delegate_scope([(b"content-type", b"application/json")]) - with _patch_delegate_resolver(server, "delegate_test"), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", - new=AsyncMock(return_value=[server]), - ), patch( - "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", - new=AsyncMock(return_value=(401, None)), - ) as probe: + with ( + _patch_delegate_resolver(server, "delegate_test"), + patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, None)), + ) as probe, + ): await _check_passthrough_upstream_auth( scope=scope, user_api_key_auth=UserAPIKeyAuth(), @@ -5758,13 +5780,17 @@ async def test_delegate_preflight_skipped_on_multi_server_routes(): servers = [_delegate_auth_mcp_server("delegate-1"), _delegate_auth_mcp_server("delegate-2")] scope = _delegate_scope([(b"authorization", b"Bearer bogus-token")]) - with _patch_delegate_resolver(servers[0], "delegate_test", "other_server"), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", - new=AsyncMock(return_value=servers), - ), patch( - "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", - new=AsyncMock(return_value=(401, None)), - ) as probe: + with ( + _patch_delegate_resolver(servers[0], "delegate_test", "other_server"), + patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=servers), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, None)), + ) as probe, + ): await _check_passthrough_upstream_auth( scope=scope, user_api_key_auth=UserAPIKeyAuth(), @@ -5797,13 +5823,17 @@ async def test_bare_authorization_never_probes_passthrough_servers(): ) scope = _delegate_scope([(b"authorization", b"Bearer ambiguous-token")]) - with _patch_delegate_resolver(passthrough_server, "pt_server"), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", - new=AsyncMock(return_value=[passthrough_server]), - ), patch( - "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", - new=AsyncMock(return_value=(401, None)), - ) as probe: + with ( + _patch_delegate_resolver(passthrough_server, "pt_server"), + patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[passthrough_server]), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, None)), + ) as probe, + ): await _check_passthrough_upstream_auth( scope=scope, user_api_key_auth=UserAPIKeyAuth(), @@ -5839,13 +5869,17 @@ async def test_delegate_not_probed_when_named_only_via_server_id(): "headers": [(b"authorization", b"Bearer sk-litellm-proxy-key")], } - with _patch_delegate_resolver(server, "delegate_test"), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", - new=AsyncMock(return_value=[server]), - ), patch( - "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", - new=AsyncMock(return_value=(401, None)), - ) as probe: + with ( + _patch_delegate_resolver(server, "delegate_test"), + patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, None)), + ) as probe, + ): await _check_passthrough_upstream_auth( scope=scope, user_api_key_auth=UserAPIKeyAuth(user_id="u1", api_key="hashed-sk"), @@ -5893,12 +5927,16 @@ async def test_delegate_preflight_with_unpatched_probe(): server = _delegate_auth_mcp_server() - with _patch_delegate_resolver(server, "delegate_test"), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", - new=AsyncMock(return_value=[server]), - ), patch( - "litellm.proxy._experimental.mcp_server.server.get_async_httpx_client", - return_value=mock_client, + with ( + _patch_delegate_resolver(server, "delegate_test"), + patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.get_async_httpx_client", + return_value=mock_client, + ), ): with pytest.raises(HTTPException) as exc_info: await _check_passthrough_upstream_auth( @@ -5918,7 +5956,9 @@ async def test_delegate_preflight_with_unpatched_probe(): assert exc_info.value.status_code == 401 challenge = exc_info.value.headers["www-authenticate"] assert 'error="invalid_token"' in challenge - assert 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/delegate_test"' in challenge + assert ( + 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/delegate_test"' in challenge + ) probed_urls = [call.kwargs["url"] for call in mock_client.post.await_args_list] assert probed_urls == ["http://upstream:9401/mcp", "http://upstream:9401/mcp"] @@ -5943,12 +5983,16 @@ async def test_delegate_challenge_echoes_requested_alias(): "headers": [(b"authorization", b"Bearer bogus-token")], } - with _patch_delegate_resolver(server, "dt-alias"), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", - new=AsyncMock(return_value=[server]), - ), patch( - "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", - new=AsyncMock(return_value=(401, 'Bearer error="invalid_token"')), + with ( + _patch_delegate_resolver(server, "dt-alias"), + patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, 'Bearer error="invalid_token"')), + ), ): with pytest.raises(HTTPException) as exc_info: await _check_passthrough_upstream_auth( @@ -5975,13 +6019,17 @@ async def test_delegate_probe_not_fanned_out_to_access_group_members(): group_member = _delegate_auth_mcp_server() - with _patch_delegate_resolver(group_member, "delegate_test"), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", - new=AsyncMock(return_value=[group_member]), - ), patch( - "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", - new=AsyncMock(return_value=(401, None)), - ) as probe: + with ( + _patch_delegate_resolver(group_member, "delegate_test"), + patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[group_member]), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(401, None)), + ) as probe, + ): await _check_passthrough_upstream_auth( scope=_delegate_scope([(b"authorization", b"Bearer bogus-token")]), user_api_key_auth=UserAPIKeyAuth(), @@ -7889,6 +7937,38 @@ class TestPreemptive401ModeAware: client_ip=None, ) + @pytest.mark.asyncio + async def test_deferred_discovery_runs_before_delegate_challenge(self): + from litellm.proxy._experimental.mcp_server import server as server_module + + manager = server_module.global_mcp_server_manager + server = _make_oauth2_server( + "lazy_delegate", + oauth2_flow="authorization_code", + delegate_auth_to_upstream=True, + ) + manager.registry[server.server_id] = server + manager._set_oauth_discovery_deferred(server.server_id, True) + metadata = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + registration_url="https://idp.example.com/register", + ) + + with ( + patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)) as discovery, + pytest.raises(HTTPException) as exc, + ): + await self._run(server, None, has_stored_token=False) + + discovery.assert_awaited_once() + resolved = manager.registry[server.server_id] + assert resolved.authorization_url == "https://idp.example.com/authorize" + assert resolved.token_url == "https://idp.example.com/token" + assert resolved.registration_url == "https://idp.example.com/register" + assert manager._oauth_discovery_slot(server.server_id) is None + assert exc.value.status_code == 401 + @pytest.mark.asyncio async def test_gateway_managed_interactive_no_token_challenges_with_x_litellm_api_key(self): """No stored token, key in x-litellm-api-key (oauth2_headers empty): 401.""" @@ -8217,16 +8297,13 @@ class TestListFiltersHonorThePrefixBoundary: url="http://127.0.0.1:5115/mcp", transport=MCPTransport.http, ) - published = MCPTool( - name=f"{self.SERVER_ID}-read_wiki_contents", description="", inputSchema={"type": "object"} - ) + published = MCPTool(name=f"{self.SERVER_ID}-read_wiki_contents", description="", inputSchema={"type": "object"}) auth = UserAPIKeyAuth(api_key="sk-test") - with patch.object( - MCPRequestHandler, "get_allowed_tools_for_server", AsyncMock(return_value=grants) - ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager" - ) as mock_manager: + with ( + patch.object(MCPRequestHandler, "get_allowed_tools_for_server", AsyncMock(return_value=grants)), + patch("litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager") as mock_manager, + ): mock_manager.get_mcp_server_by_id.return_value = server listed = await filter_tools_by_key_team_permissions([published], self.SERVER_ID, auth) != [] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 54fd5242d5f..fe09303b298 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -2,11 +2,10 @@ import importlib import asyncio import json import logging -import time import os import sys from datetime import datetime -from typing import Any, Dict, Optional +from typing import Any, Dict, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -33,10 +32,12 @@ from mcp.types import ( ) from mcp.types import Tool as MCPTool +from litellm.constants import MCP_METADATA_TIMEOUT from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, _deserialize_json_dict, _flow_endpoints_missing, + _mcp_oauth_discovery_on_startup_enabled, _oauth_endpoints_unresolved, _deserialize_json_list, _normalize_mcp_server_cost_info, @@ -54,6 +55,7 @@ from litellm.proxy._types import ( MCPTransport, UserAPIKeyAuth, ) +from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.mcp import MCPAuth, MCPAuthType from litellm.types.mcp_server.mcp_server_manager import MCPOAuthMetadata, MCPServer @@ -72,6 +74,11 @@ def _reload_mcp_manager_module(): return reloaded +@pytest.fixture(autouse=True) +def enable_eager_mcp_oauth_discovery(monkeypatch): + monkeypatch.setenv("LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP", "1") + + class TestMCPServerManager: """Test MCP Server Manager stdio functionality""" @@ -428,6 +435,468 @@ class TestMCPServerManager: base.update(overrides) return {"m2mserver": base} + @pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "on"]) + def test_mcp_oauth_discovery_on_startup_true_values(self, value): + with patch.dict(os.environ, {"LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP": value}): + assert _mcp_oauth_discovery_on_startup_enabled() is True + + @pytest.mark.parametrize("value", ["0", "false", "FALSE", "no", "off", "", "invalid"]) + def test_mcp_oauth_discovery_on_startup_non_true_values(self, value): + with patch.dict(os.environ, {"LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP": value}): + assert _mcp_oauth_discovery_on_startup_enabled() is False + + def test_mcp_oauth_discovery_on_startup_defaults_to_disabled(self): + with patch.dict(os.environ, {}, clear=True): + assert _mcp_oauth_discovery_on_startup_enabled() is False + + @pytest.mark.asyncio + async def test_config_oauth_discovery_warmup_is_non_blocking_and_shared(self): + metadata = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + scopes=["mcp.read"], + ) + with patch.dict(os.environ, {}, clear=True): + manager = MCPServerManager() + + started = asyncio.Event() + release = asyncio.Event() + + async def discover(_server): + started.set() + await release.wait() + return metadata + + with ( + patch.object(manager, "_discover_oauth_metadata_for_server", side_effect=discover) as discovery, + patch.object(manager, "initialize_tool_name_to_mcp_server_name_mapping"), + ): + load_task = asyncio.create_task( + manager.load_servers_from_config( + self._oauth2_config( + oauth2_flow="authorization_code", + authorization_url=None, + token_url=None, + ) + ) + ) + await started.wait() + assert load_task.done() + await load_task + + server = next(iter(manager.config_mcp_servers.values())) + waiters = [asyncio.create_task(manager.ensure_oauth_metadata_discovered(server)) for _ in range(10)] + await asyncio.sleep(0) + release.set() + resolved = await asyncio.gather(*waiters) + + discovery.assert_awaited_once_with(server) + assert all(result is resolved[0] for result in resolved) + assert resolved[0].authorization_url == "https://idp.example.com/authorize" + assert resolved[0].token_url == "https://idp.example.com/token" + assert resolved[0].scopes == ["mcp.read"] + assert manager.config_mcp_servers[server.server_id] is resolved[0] + assert server.authorization_url is None + assert manager._oauth_discovery_slot(server.server_id) is None + + @pytest.mark.asyncio + async def test_table_oauth_discovery_can_be_deferred_until_first_use(self): + row = LiteLLM_MCPServerTable( + server_id="lazy-db-1", + alias="lazy_db", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + created_at=datetime.now(), + updated_at=datetime.now(), + ) + metadata = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + ) + with patch.dict(os.environ, {}, clear=True): + manager = MCPServerManager() + + discovery = AsyncMock(return_value=metadata) + with patch.object(manager, "_descovery_metadata", new=discovery): + server = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) + + discovery.assert_not_awaited() + assert manager._oauth_discovery_slot(server.server_id) is not None + manager.registry[server.server_id] = server + + with patch.object(manager, "_descovery_metadata", new=discovery): + resolved = await manager.ensure_oauth_metadata_discovered(server) + + discovery.assert_awaited_once() + assert resolved.authorization_url == "https://idp.example.com/authorize" + assert resolved.token_url == "https://idp.example.com/token" + + @pytest.mark.asyncio + async def test_lazy_oauth_discovery_failure_is_shared_and_retries_after_cooldown(self): + with patch.dict(os.environ, {"LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP": "false"}): + manager = MCPServerManager() + + metadata = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + ) + discovery = AsyncMock(side_effect=[None, None, None, metadata]) + discovery_clock: Final = MagicMock(return_value=100.0) + with ( + patch.object(manager, "_discover_oauth_metadata_for_server", new=discovery), + patch.object(manager, "initialize_tool_name_to_mcp_server_name_mapping"), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager._oauth_discovery_now", + new=discovery_clock, + ), + ): + await manager.load_servers_from_config( + self._oauth2_config( + oauth2_flow="authorization_code", + authorization_url=None, + token_url=None, + ) + ) + + server = next(iter(manager.config_mcp_servers.values())) + failures: Final = await asyncio.gather( + *(manager.ensure_oauth_metadata_discovered(server) for _ in range(10)), + return_exceptions=True, + ) + cooldown_failures: Final = await asyncio.gather( + *(manager.ensure_oauth_metadata_discovered(server) for _ in range(10)), + return_exceptions=True, + ) + discovery_clock.return_value = 130.0 + resolutions: Final = await asyncio.gather( + *(manager.ensure_oauth_metadata_discovered(server) for _ in range(10)) + ) + + assert discovery.await_count == 4 + assert all(isinstance(failure, HTTPException) and failure.status_code == 503 for failure in failures) + assert all(isinstance(failure, HTTPException) and failure.status_code == 503 for failure in cooldown_failures) + assert len({id(resolution) for resolution in resolutions}) == 1 + assert resolutions[0].authorization_url == "https://idp.example.com/authorize" + assert resolutions[0].token_url == "https://idp.example.com/token" + assert manager._oauth_discovery_slot(server.server_id) is None + + @pytest.mark.asyncio + async def test_lazy_oauth_discovery_timeout_is_bounded(self): + manager = MCPServerManager() + server = MCPServer( + server_id="lazy-timeout-1", + name="lazy_timeout", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + ) + manager.registry[server.server_id] = server + manager._set_oauth_discovery_deferred(server.server_id, True) + + async def never_returns(_server): + await asyncio.Future() + + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCP_METADATA_TIMEOUT", + 0.01, + ), + patch.object(manager, "_discover_oauth_metadata_for_server", side_effect=never_returns) as discovery, + ): + with pytest.raises(HTTPException) as exc: + await asyncio.wait_for(manager.ensure_oauth_metadata_discovered(server), timeout=0.2) + + assert exc.value.status_code == 503 + assert "timed out" in str(exc.value.detail) + discovery.assert_awaited_once_with(server) + assert manager._oauth_discovery_slot(server.server_id) is not None + + @pytest.mark.asyncio + async def test_cancelling_one_waiter_does_not_cancel_shared_discovery(self): + manager = MCPServerManager() + server = MCPServer( + server_id="lazy-cancel-1", + name="lazy_cancel", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + ) + manager.registry[server.server_id] = server + manager._set_oauth_discovery_deferred(server.server_id, True) + started = asyncio.Event() + release = asyncio.Event() + metadata = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + ) + + async def discover(_server): + started.set() + await release.wait() + return metadata + + with patch.object(manager, "_discover_oauth_metadata_for_server", side_effect=discover) as discovery: + cancelled_waiter = asyncio.create_task(manager.ensure_oauth_metadata_discovered(server)) + successful_waiter = asyncio.create_task(manager.ensure_oauth_metadata_discovered(server)) + await started.wait() + cancelled_waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled_waiter + release.set() + resolved = await successful_waiter + + discovery.assert_awaited_once_with(server) + assert resolved.authorization_url == "https://idp.example.com/authorize" + + @pytest.mark.asyncio + async def test_lazy_oauth_discovery_ignores_stale_registration_result(self): + manager = MCPServerManager() + old_server = MCPServer( + server_id="lazy-reload-1", + name="lazy_reload", + url="https://old.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + ) + replacement = old_server.model_copy(update={"url": "https://new.example.com/mcp"}) + manager.registry[old_server.server_id] = old_server + manager._set_oauth_discovery_deferred(old_server.server_id, True) + started = asyncio.Event() + metadata = MCPOAuthMetadata( + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + ) + + async def discover(candidate): + if candidate.url == old_server.url: + started.set() + await asyncio.Future() + return metadata + + with patch.object(manager, "_discover_oauth_metadata_for_server", side_effect=discover): + old_attempt = asyncio.create_task(manager.ensure_oauth_metadata_discovered(old_server)) + await started.wait() + manager.registry[replacement.server_id] = replacement + manager._set_oauth_discovery_deferred(replacement.server_id, True) + resolved = await old_attempt + + assert resolved is manager.registry[replacement.server_id] + assert resolved.url == replacement.url + assert old_server.authorization_url is None + assert old_server.token_url is None + assert replacement.authorization_url is None + assert replacement.token_url is None + assert resolved.authorization_url == "https://idp.example.com/authorize" + assert resolved.token_url == "https://idp.example.com/token" + assert manager._oauth_discovery_slot(replacement.server_id) is None + + def _assert_oauth_discovery_state_removed(self, manager, server_id): + assert manager._oauth_discovery_slot(server_id) is None + + @pytest.mark.asyncio + async def test_deactivated_server_clears_lazy_oauth_discovery_state(self): + manager = MCPServerManager() + server = MCPServer( + server_id="lazy-deactivated-1", + name="lazy_deactivated", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + ) + manager.registry[server.server_id] = server + manager._set_oauth_discovery_deferred(server.server_id, True) + record = LiteLLM_MCPServerTable( + server_id=server.server_id, + server_name=server.name, + url=server.url, + transport=MCPTransport.http, + approval_status="rejected", + ) + + await manager.update_server(record) + + assert manager.registry == {} + self._assert_oauth_discovery_state_removed(manager, server.server_id) + + @pytest.mark.asyncio + async def test_database_reload_drop_clears_lazy_oauth_discovery_state(self): + manager = MCPServerManager() + server = MCPServer( + server_id="lazy-dropped-1", + name="lazy_dropped", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + ) + manager.registry[server.server_id] = server + manager._set_oauth_discovery_deferred(server.server_id, True) + repository = MagicMock() + repository.table.find_many = AsyncMock(return_value=[]) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerRepository", + return_value=repository, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + ): + await manager.reload_servers_from_database() + + assert manager.registry == {} + self._assert_oauth_discovery_state_removed(manager, server.server_id) + + @pytest.mark.asyncio + async def test_database_reload_rearms_discovery_lost_to_registry_swap(self): + """A resolution published into the old registry while reload is staged + must leave the swapped-in unresolved entry with a fresh retry slot. + """ + manager = MCPServerManager() + stamp = datetime.now() + server = MCPServer( + server_id="lazy-swap-1", + name="lazy_swap", + server_name="lazy_swap", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + updated_at=stamp, + ) + manager.registry[server.server_id] = server + previous_registry = manager.registry + manager._set_oauth_discovery_deferred(server.server_id, True) + old_generation = manager._oauth_discovery_slot(server.server_id).generation + resolved = server.model_copy( + update={ + "authorization_url": "https://idp.example.com/authorize", + "token_url": "https://idp.example.com/token", + } + ) + row = LiteLLM_MCPServerTable( + server_id=server.server_id, + server_name=server.server_name, + url=server.url, + transport=server.transport, + auth_type=server.auth_type, + oauth2_flow=server.oauth2_flow, + updated_at=stamp, + ) + raw_row = MagicMock() + raw_row.model_dump.return_value = row.model_dump() + repository = MagicMock() + repository.table.find_many = AsyncMock(return_value=[raw_row]) + + async def publish_while_staged(*_args, **_kwargs): + assert manager.registry is previous_registry + assert manager._publish_resolved_oauth_server(resolved, old_generation) is resolved + + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerRepository", + return_value=repository, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch.object( + manager, + "_maybe_register_openapi_tools", + new=AsyncMock(side_effect=publish_while_staged), + ), + patch.object(manager, "_prime_oauth_metadata_discovery_for_servers"), + ): + await manager.reload_servers_from_database() + + assert previous_registry[server.server_id] is resolved + assert manager.registry[server.server_id] is server + retry_slot = manager._oauth_discovery_slot(server.server_id) + assert retry_slot is not None + assert retry_slot.generation > old_generation + + @pytest.mark.asyncio + async def test_lazy_oauth_discovery_preserves_manual_authorization_url_gate(self): + with patch.dict(os.environ, {"LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP": "false"}): + manager = MCPServerManager() + + metadata = MCPOAuthMetadata( + authorization_url="https://attacker.example.com/authorize", + token_url="https://attacker.example.com/token", + scopes=["mcp.read"], + ) + discovery = AsyncMock(return_value=metadata) + with ( + patch.object(manager, "_descovery_metadata", new=discovery), + patch.object(manager, "initialize_tool_name_to_mcp_server_name_mapping"), + ): + await manager.load_servers_from_config( + self._oauth2_config( + oauth2_flow="authorization_code", + authorization_url="https://idp.example.com/authorize", + token_url=None, + ) + ) + + server = next(iter(manager.config_mcp_servers.values())) + with ( + patch.object(manager, "_descovery_metadata", new=discovery), + pytest.raises(HTTPException) as exc, + ): + await manager.ensure_oauth_metadata_discovered(server) + + assert exc.value.status_code == 503 + assert manager.config_mcp_servers[server.server_id].authorization_url == "https://idp.example.com/authorize" + assert manager.config_mcp_servers[server.server_id].token_url is None + assert manager.config_mcp_servers[server.server_id].scopes is None + assert manager._oauth_discovery_slot(server.server_id) is not None + + @pytest.mark.asyncio + async def test_create_mcp_client_triggers_deferred_oauth_discovery(self): + manager = MCPServerManager() + server = MCPServer( + server_id="lazy-client-1", + name="lazy_client", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + ) + ensure_oauth_metadata_discovered: Final = AsyncMock(return_value=server) + + with ( + patch.object( + manager, + "ensure_oauth_metadata_discovered", + new=ensure_oauth_metadata_discovered, + ), + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient"), + ): + await manager._create_mcp_client(server) + + ensure_oauth_metadata_discovered.assert_awaited_once_with(server) + + @pytest.mark.asyncio + async def test_startup_tool_mapping_skips_servers_with_deferred_discovery(self): + manager = MCPServerManager() + server = MCPServer( + server_id="lazy-map-1", + name="lazy_map", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.true_passthrough, + ) + manager.registry[server.server_id] = server + manager._set_oauth_discovery_deferred(server.server_id, True) + + with patch.object(manager, "_get_tools_from_server", new=AsyncMock()) as get_tools: + await manager._initialize_tool_name_to_mcp_server_name_mapping() + + get_tools.assert_not_awaited() + @pytest.mark.asyncio async def test_load_servers_from_config_requires_oauth2_flow(self): """auth_type oauth2 without an explicit oauth2_flow is a config error: the @@ -1492,7 +1961,9 @@ class TestMCPServerManager: ) resource_rooted = AsyncMock(return_value=MCPOAuthMetadata(token_url="https://attacker.example.com/steal")) with ( - patch.object(manager, "_fetch_issuer_anchored_oauth_metadata", new=AsyncMock(return_value=resolved)) as anchored, + patch.object( + manager, "_fetch_issuer_anchored_oauth_metadata", new=AsyncMock(return_value=resolved) + ) as anchored, patch.object(manager, "_descovery_metadata", new=resource_rooted), ): built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) @@ -1526,7 +1997,9 @@ class TestMCPServerManager: patch.object( manager, "_fetch_single_authorization_server_metadata", new=AsyncMock(return_value=issuer_document) ) as issuer_fetch, - patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=resource_document)) as resource_fetch, + patch.object( + manager, "_descovery_metadata", new=AsyncMock(return_value=resource_document) + ) as resource_fetch, ): result = await manager._fetch_issuer_anchored_oauth_metadata( "https://idp.example.com", "https://up.example.com/mcp" @@ -1916,6 +2389,29 @@ class TestMCPServerManager: await manager.preflight_token_exchange(server=server, oauth2_headers=None, user_api_key_auth=None) assert resolved == ["good-subject"] + @pytest.mark.asyncio + async def test_preflight_token_exchange_skips_discovery_for_other_auth_modes(self): + """Preflight must not make unrelated auth modes depend on OAuth discovery.""" + manager = MCPServerManager() + server = MCPServer( + server_id="plain-preflight", + name="plain_preflight", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + manager.ensure_oauth_metadata_discovered = AsyncMock( + side_effect=AssertionError("non-token-exchange server was resolved") + ) + + await manager.preflight_token_exchange( + server=server, + oauth2_headers={"Authorization": "Bearer subject"}, + user_api_key_auth=None, + ) + + manager.ensure_oauth_metadata_discovered.assert_not_awaited() + @pytest.mark.asyncio async def test_call_regular_mcp_tool_passthrough_strips_authorization_when_admission_consumed_litellm_key( self, @@ -2816,7 +3312,7 @@ class TestMCPServerManager: patch( "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", return_value=mock_client, - ), + ) as get_client, patch.object( manager, "_attempt_well_known_discovery", @@ -2835,6 +3331,10 @@ class TestMCPServerManager: ): result = await manager._descovery_metadata("http://localhost:8001/mcp") + get_client.assert_called_once_with( + llm_provider=httpxSpecialProvider.MCP, + params={"timeout": MCP_METADATA_TIMEOUT}, + ) mock_well_known.assert_awaited_once_with("http://localhost:8001/mcp") mock_fetch_auth.assert_awaited_once_with( ["https://login.microsoftonline.com/test-tenant-id/v2.0"], @@ -3125,7 +3625,9 @@ class TestMCPServerManager: registration_url="https://discovered.example.com/register", ) - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): + async def fake_discovery( + server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False + ): assert server_url == "https://example.com/mcp" # oauth2 (browser flow) keeps the origin fallback; only OBO disables it. assert allow_origin_fallback is True @@ -3375,6 +3877,29 @@ class TestMCPServerManager: assert result.health_check_error == "Connection timeout" assert result.last_health_check is not None + @pytest.mark.asyncio + async def test_health_check_server_contains_client_creation_failure(self): + """Deferred discovery failures are reported unhealthy, not raised.""" + manager = MCPServerManager() + server = MCPServer( + server_id="discovery-failure", + name="discovery-failure", + transport=MCPTransport.http, + auth_type=None, + authentication_token="test-token", + url="https://up.example.com/mcp", + ) + manager.get_mcp_server_by_id = MagicMock(return_value=server) + manager._resolve_static_headers_with_env_vars = AsyncMock(return_value=None) + manager._create_mcp_client = AsyncMock( + side_effect=HTTPException(status_code=503, detail="OAuth discovery unavailable") + ) + + result = await manager.health_check_server(server.server_id) + + assert result.status == "unhealthy" + assert "OAuth discovery unavailable" in (result.health_check_error or "") + @pytest.mark.asyncio async def test_health_check_server_not_found(self): """Test health check for a server that doesn't exist""" @@ -4177,6 +4702,20 @@ class TestMCPServerManager: with pytest.raises(ValueError, match="Tool .* not found"): manager._resolve_mcp_server_for_tool_call("nonexistent", "ghost_tool") + def test_resolve_mcp_server_for_tool_call_unscoped_cached_tool_still_fails(self): + """Without an explicit server, an unmapped tool remains ambiguous.""" + manager = MCPServerManager() + manager.registry = { + "github": MCPServer( + server_id="github", + name="github", + transport=MCPTransport.http, + ) + } + + with pytest.raises(ValueError, match="Tool cached_tool not found"): + manager._resolve_mcp_server_for_tool_call("", "cached_tool") + def test_resolve_mcp_server_for_tool_call_unknown_tool_with_known_server(self): """Server-name match alone must not let unknown tools slip through. @@ -5520,7 +6059,9 @@ class TestMCPServerTimestamps: manager = MCPServerManager() calls: list[bool] = [] - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): + async def fake_discovery( + server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False + ): calls.append(allow_origin_fallback) return MCPOAuthMetadata( scopes=None, @@ -5555,7 +6096,9 @@ class TestMCPServerTimestamps: manager = MCPServerManager() calls: list[str] = [] - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): + async def fake_discovery( + server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False + ): calls.append(server_url) raise AssertionError("discovery must not run when token_exchange_endpoint is configured") @@ -5588,7 +6131,9 @@ class TestMCPServerTimestamps: lives on the in-memory registry entry only, for oauth2 and OBO alike.""" manager = MCPServerManager() - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): + async def fake_discovery( + server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False + ): return MCPOAuthMetadata( scopes=["mcp.read"], authorization_url="https://idp.example.com/authorize", @@ -5670,9 +6215,7 @@ class TestMCPServerTimestamps: assert _flow_endpoints_missing(MCPAuth.oauth2, "client_credentials", None, None) is True assert _flow_endpoints_missing(MCPAuth.oauth2_token_exchange, None, None, None) is True assert _flow_endpoints_missing(MCPAuth.oauth2_token_exchange, None, None, "https://idp/token") is False - assert ( - _flow_endpoints_missing(MCPAuth.oauth2_token_exchange, None, None, None, "https://idp/exchange") is False - ) + assert _flow_endpoints_missing(MCPAuth.oauth2_token_exchange, None, None, None, "https://idp/exchange") is False assert _flow_endpoints_missing(MCPAuth.api_key, None, None, None) is False def test_unresolved_check_uses_the_flow_judge_not_the_raw_column(self): @@ -5717,7 +6260,9 @@ class TestMCPServerTimestamps: registration_url=None, ) assert _oauth_endpoints_unresolved(relay_arm) is True - assert _oauth_endpoints_unresolved(relay_arm.model_copy(update={"registration_url": "https://idp/reg"})) is False + assert ( + _oauth_endpoints_unresolved(relay_arm.model_copy(update={"registration_url": "https://idp/reg"})) is False + ) assert _oauth_endpoints_unresolved(relay_arm.model_copy(update={"client_id": "admin-client"})) is False def test_entra_obo_without_scopes_is_unresolved(self): @@ -5738,50 +6283,6 @@ class TestMCPServerTimestamps: assert _oauth_endpoints_unresolved(entra.model_copy(update={"scopes": ["api://app/.default"]})) is False assert _oauth_endpoints_unresolved(entra.model_copy(update={"token_exchange_profile": "rfc8693"})) is False - def test_oauth_discovery_retry_backs_off_per_server(self): - """Without a cooldown the fast-path exemption re-runs the full discovery chain, and re-emits - the unresolved warning, on every reload forever for a server that can never resolve. Delay - doubles per consecutive failure up to the cap, a success clears the state so the next failure - starts from the base delay again, and the cooldown is per server.""" - manager = MCPServerManager() - - def unresolved(server_id): - return MCPServer( - server_id=server_id, - name=server_id, - url="https://up.example.com/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - oauth2_flow="authorization_code", - ) - - assert manager._oauth_discovery_retry_due("a") is True - - manager._record_oauth_discovery_outcome(unresolved("a")) - assert manager._oauth_discovery_retry_due("a") is False - assert manager._oauth_discovery_retry_due("b") is True, "cooldown must be per server" - - failures_before, _ = manager._oauth_discovery_retry_state["a"] - manager._record_oauth_discovery_outcome(unresolved("a")) - failures_after, _ = manager._oauth_discovery_retry_state["a"] - assert failures_after == failures_before + 1 - - # An elapsed cooldown lets the retry through, and the delay grows with the failure count - manager._oauth_discovery_retry_state["a"] = (1, time.monotonic() - 31.0) - assert manager._oauth_discovery_retry_due("a") is True - manager._oauth_discovery_retry_state["a"] = (5, time.monotonic() - 31.0) - assert manager._oauth_discovery_retry_due("a") is False - - resolved = unresolved("a").model_copy( - update={ - "authorization_url": "https://idp.example.com/authorize", - "token_url": "https://idp.example.com/token", - } - ) - manager._record_oauth_discovery_outcome(resolved) - assert "a" not in manager._oauth_discovery_retry_state - assert manager._oauth_discovery_retry_due("a") is True - @pytest.mark.asyncio async def test_reload_fast_path_retries_unresolved_oauth_servers(self): """A server whose discovery failed must not be pinned broken by the updated_at fast path: @@ -8574,7 +9075,9 @@ class TestOBOEndpointDiscovery: ) seen = [] - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): + async def fake_discovery( + server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False + ): seen.append((server_url, allow_origin_fallback)) return discovered @@ -8602,7 +9105,9 @@ class TestOBOEndpointDiscovery: async def test_config_obo_with_configured_endpoint_skips_discovery(self): manager = MCPServerManager() - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): + async def fake_discovery( + server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False + ): raise AssertionError("discovery must not run when the endpoint is configured") manager._descovery_metadata = fake_discovery # type: ignore[attr-defined] @@ -9007,7 +9512,9 @@ class TestUrllessIssuerDiscovery: ) resource_rooted = AsyncMock(return_value=None) with ( - patch.object(manager, "_fetch_issuer_anchored_oauth_metadata", new=AsyncMock(return_value=resolved)) as anchored, + patch.object( + manager, "_fetch_issuer_anchored_oauth_metadata", new=AsyncMock(return_value=resolved) + ) as anchored, patch.object(manager, "_descovery_metadata", new=resource_rooted), ): built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) @@ -9055,7 +9562,9 @@ class TestUrllessIssuerDiscovery: resolved = MCPOAuthMetadata(token_url="https://idp.example.com/token") resource_rooted = AsyncMock(return_value=None) with ( - patch.object(manager, "_fetch_issuer_anchored_oauth_metadata", new=AsyncMock(return_value=resolved)) as anchored, + patch.object( + manager, "_fetch_issuer_anchored_oauth_metadata", new=AsyncMock(return_value=resolved) + ) as anchored, patch.object(manager, "_descovery_metadata", new=resource_rooted), ): built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) @@ -9073,9 +9582,7 @@ class TestDiscoveryFailureLogging: def _connect_error_client(self, url: str) -> MagicMock: client = MagicMock() - client.get = AsyncMock( - side_effect=httpx.ConnectError(f"[Errno 8] nodename nor servname provided for {url}") - ) + client.get = AsyncMock(side_effect=httpx.ConnectError(f"[Errno 8] nodename nor servname provided for {url}")) return client @pytest.mark.asyncio @@ -9118,9 +9625,7 @@ class TestDiscoveryFailureLogging: manager = MCPServerManager() url = "https://real-host.example.com/mcp-typo" client = MagicMock() - client.get = AsyncMock( - return_value=httpx.Response(404, request=httpx.Request("GET", url)) - ) + client.get = AsyncMock(return_value=httpx.Response(404, request=httpx.Request("GET", url))) with ( patch( "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", From 79a6d2b8d264502040698ff1856671669d4b5795 Mon Sep 17 00:00:00 2001 From: RayJueWang <570828708@qq.com> Date: Tue, 28 Jul 2026 14:01:34 +0800 Subject: [PATCH 032/538] fix(proxy): retry spend updates on Postgres deadlock instead of dropping them Spend-update transactions increment non-idempotent counters (spend = spend + x) inside prisma interactive transactions. Every retry loop only caught DB_RETRY_SAFE_ERROR_TYPES (httpx.ConnectError); a Postgres deadlock (SQLSTATE 40P01, surfaced by prisma as transaction conflict code P2034) fell through to a bare except that re-raised immediately, so on multi-pod / high-concurrency deployments any pod that lost a deadlock silently dropped its increment. A deadlock is replay-safe even though the increment is non-idempotent: Postgres aborts and fully rolls back the victim transaction, so no partial spend is committed. Add PrismaDBExceptionHandler.is_deadlock_error and route every spend path (user, end-user/key, team, team_member, org, tag/agent via _update_entity_spend_in_db, and the daily-spend upsert) through a shared _handle_spend_update_failure that retries connection errors and deadlocks with randomized jitter backoff and re-raises everything else or on exhaustion. --- litellm/proxy/db/db_spend_update_writer.py | 143 ++++++-------- litellm/proxy/db/exception_handler.py | 12 ++ .../proxy/db/test_db_spend_update_writer.py | 184 ++++++++++++++++++ .../proxy/db/test_exception_handler.py | 33 ++++ 4 files changed, 293 insertions(+), 79 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index b2b72c1cac4..16389cda336 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1122,6 +1122,23 @@ class DBSpendUpdateWriter: except Exception as e: verbose_proxy_logger.debug("_flush_tool_discovery_queue error (non-blocking): %s", e) + @staticmethod + async def _handle_spend_update_failure( + e: Exception, + attempt: int, + n_retry_times: int, + start_time: float, + proxy_logging_obj: ProxyLogging, + ) -> None: + """Retry a failed spend-update transaction on connection errors or deadlocks, else re-raise.""" + from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler + from litellm.proxy.utils import _raise_failed_update_spend_exception + + is_retryable = isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) or PrismaDBExceptionHandler.is_deadlock_error(e) + if not is_retryable or attempt >= n_retry_times: + _raise_failed_update_spend_exception(e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj) + await asyncio.sleep(random.uniform(2**attempt, 2 ** (attempt + 1))) + async def _commit_spend_updates_to_db( self, prisma_client: PrismaClient, @@ -1133,10 +1150,7 @@ class DBSpendUpdateWriter: Commits all the spend `UPDATE` transactions to the Database """ - from litellm.proxy.utils import ( - ProxyUpdateSpend, - _raise_failed_update_spend_exception, - ) + from litellm.proxy.utils import ProxyUpdateSpend ### UPDATE USER TABLE ### user_list_transactions: Final = db_spend_update_transactions["user_list_transactions"] @@ -1156,18 +1170,13 @@ class DBSpendUpdateWriter: data={"spend": {"increment": response_cost}}, ) break - except DB_RETRY_SAFE_ERROR_TYPES as e: - if i >= n_retry_times: # If we've reached the maximum number of retries - _raise_failed_update_spend_exception( - e=e, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - # Optionally, sleep for a bit before retrying - await asyncio.sleep(2**i) # Exponential backoff except Exception as e: - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj + await self._handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, ) ### UPDATE END-USER TABLE ### @@ -1199,18 +1208,13 @@ class DBSpendUpdateWriter: }, ) break - except DB_RETRY_SAFE_ERROR_TYPES as e: - if i >= n_retry_times: # If we've reached the maximum number of retries - _raise_failed_update_spend_exception( - e=e, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - # Optionally, sleep for a bit before retrying - await asyncio.sleep(2**i) # Exponential backoff except Exception as e: - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj + await self._handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, ) ### UPDATE TEAM TABLE ### @@ -1232,18 +1236,13 @@ class DBSpendUpdateWriter: data={"spend": {"increment": response_cost}}, ) break - except DB_RETRY_SAFE_ERROR_TYPES as e: - if i >= n_retry_times: # If we've reached the maximum number of retries - _raise_failed_update_spend_exception( - e=e, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - # Optionally, sleep for a bit before retrying - await asyncio.sleep(2**i) # Exponential backoff except Exception as e: - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj + await self._handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, ) ### UPDATE TEAM Membership TABLE with spend ### @@ -1279,18 +1278,13 @@ class DBSpendUpdateWriter: ) # Transaction succeeded, break out of retry loop break - except DB_RETRY_SAFE_ERROR_TYPES as e: - if i >= n_retry_times: # If we've reached the maximum number of retries - _raise_failed_update_spend_exception( - e=e, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - # Optionally, sleep for a bit before retrying - await asyncio.sleep(2**i) # Exponential backoff except Exception as e: - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj + await self._handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, ) # Invalidate cache for updated team memberships @@ -1321,25 +1315,13 @@ class DBSpendUpdateWriter: data={"spend": {"increment": response_cost}}, ) break - except DB_RETRY_SAFE_ERROR_TYPES as e: - if i >= n_retry_times: # If we've reached the maximum number of retries - _raise_failed_update_spend_exception( - e=e, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - # Optionally, sleep for a bit before retrying - await asyncio.sleep( - # Sleep a random amount to avoid retrying and deadlocking again: when two transactions deadlock they are - # cancelled basically at the same time, so if they wait the same time they will also retry at the same time - # and thus they are more likely to deadlock again. - # Instead, we sleep a random amount so that they retry at slightly different times, lowering the chance of - # repeated deadlocks, and therefore of exceeding the retry limit. - random.uniform(2**i, 2 ** (i + 1)) - ) except Exception as e: - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj + await self._handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, ) ### UPDATE TAG TABLE ### @@ -1388,8 +1370,6 @@ class DBSpendUpdateWriter: prisma_client: Prisma client instance proxy_logging_obj: Proxy logging object """ - from litellm.proxy.utils import _raise_failed_update_spend_exception - verbose_proxy_logger.debug("%s Spend transactions: %s", entity_name, transactions) if transactions is not None and len(transactions.keys()) > 0: for i in range(n_retry_times + 1): @@ -1411,17 +1391,13 @@ class DBSpendUpdateWriter: data={"spend": {"increment": response_cost}}, ) break - except DB_RETRY_SAFE_ERROR_TYPES as e: - if i >= n_retry_times: - _raise_failed_update_spend_exception( - e=e, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - await asyncio.sleep(2**i) # Exponential backoff except Exception as e: - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj + await DBSpendUpdateWriter._handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, ) # fmt: off @@ -1590,7 +1566,16 @@ class DBSpendUpdateWriter: break - except DB_RETRY_SAFE_ERROR_TYPES as e: + except Exception as e: + from litellm.proxy.db.exception_handler import ( + PrismaDBExceptionHandler, + ) + + is_retryable = isinstance( + e, DB_RETRY_SAFE_ERROR_TYPES + ) or PrismaDBExceptionHandler.is_deadlock_error(e) + if not is_retryable: + raise if i >= n_retry_times: _raise_failed_update_spend_exception( e=e, diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index e0a21ceed26..91c7e576dff 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -166,6 +166,18 @@ class PrismaDBExceptionHandler: return True return False + @staticmethod + def is_deadlock_error(e: Exception) -> bool: + """True iff ``e`` is a Postgres deadlock (P2034 / 40P01) surfaced through prisma.""" + import prisma + + if not isinstance(e, prisma.errors.PrismaError): + return False + if getattr(e, "code", None) == "P2034": + return True + error_message = str(e).lower() + return "deadlock detected" in error_message or "40p01" in error_message + @staticmethod def is_prisma_engine_internal_error(e: Exception) -> bool: """True iff ``e`` is a non-``PrismaError`` exception raised from inside diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index ca7d5fcd273..49ef653d4e1 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2268,3 +2268,187 @@ async def test_daily_transaction_internal_call_keeps_spend_but_not_request_count assert internal["autorouter_savings_spend"] == 0.0 assert user_sent["api_requests"] == 1 assert user_sent["successful_requests"] == 1 + + +def _deadlock_error(): + from prisma.errors import RawQueryError + + return RawQueryError( + data={"user_facing_error": {"error_code": "P2034", "meta": {"table": "LiteLLM_VerificationToken"}}} + ) + + +def _empty_spend_transactions(**overrides): + base = { + "user_list_transactions": {}, + "end_user_list_transactions": {}, + "key_list_transactions": {}, + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "tag_list_transactions": {}, + "agent_list_transactions": {}, + } + return {**base, **overrides} + + +def _good_tx(mock_batcher): + tx = AsyncMock() + tx.__aenter__ = AsyncMock(return_value=tx) + tx.__aexit__ = AsyncMock(return_value=False) + tx.batch_ = MagicMock( + return_value=AsyncMock( + __aenter__=AsyncMock(return_value=mock_batcher), + __aexit__=AsyncMock(return_value=False), + ) + ) + return tx + + +def _failing_tx(error): + tx = MagicMock() + tx.__aenter__ = AsyncMock(side_effect=error) + tx.__aexit__ = AsyncMock(return_value=False) + return tx + + +@pytest.mark.asyncio +async def test_commit_spend_updates_retries_deadlock_then_commits(monkeypatch): + """Regression: a deadlock on the key-spend UPDATE is retried and commits the increment exactly once.""" + slept = [] + monkeypatch.setattr( + "litellm.proxy.db.db_spend_update_writer.asyncio.sleep", + AsyncMock(side_effect=lambda s: slept.append(s)), + ) + + mock_batcher = MagicMock() + mock_prisma_client = MagicMock() + mock_prisma_client.db.tx = MagicMock(side_effect=[_failing_tx(_deadlock_error()), _good_tx(mock_batcher)]) + + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + await DBSpendUpdateWriter()._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=3, + proxy_logging_obj=proxy_logging, + db_spend_update_transactions=_empty_spend_transactions(key_list_transactions={"sk-abc": 0.5}), + ) + + assert mock_prisma_client.db.tx.call_count == 2 + mock_batcher.litellm_verificationtoken.update_many.assert_called_once() + call_kwargs = mock_batcher.litellm_verificationtoken.update_many.call_args[1] + assert call_kwargs["where"] == {"token": "sk-abc"} + assert call_kwargs["data"]["spend"] == {"increment": 0.5} + assert len(slept) == 1 + proxy_logging.failure_handler.assert_not_called() + + +@pytest.mark.asyncio +async def test_commit_spend_updates_raises_after_exhausting_deadlock_retries(monkeypatch): + """A deadlock that never clears must surface after the retry budget is spent, not loop or swallow.""" + monkeypatch.setattr("litellm.proxy.db.db_spend_update_writer.asyncio.sleep", AsyncMock(return_value=None)) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.tx = MagicMock(side_effect=lambda *a, **k: _failing_tx(_deadlock_error())) + + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + from prisma.errors import RawQueryError + + with pytest.raises(RawQueryError): + await DBSpendUpdateWriter()._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=2, + proxy_logging_obj=proxy_logging, + db_spend_update_transactions=_empty_spend_transactions(key_list_transactions={"sk-abc": 0.5}), + ) + + assert mock_prisma_client.db.tx.call_count == 3 + + +@pytest.mark.asyncio +async def test_commit_spend_updates_does_not_retry_non_deadlock_data_error(monkeypatch): + """A non-retryable data-layer error raises on the first attempt, never retried against the increment.""" + monkeypatch.setattr("litellm.proxy.db.db_spend_update_writer.asyncio.sleep", AsyncMock(return_value=None)) + + from prisma.errors import UniqueViolationError + + non_deadlock = UniqueViolationError( + data={"user_facing_error": {"error_code": "P2002", "meta": {"table": "LiteLLM_VerificationToken"}}} + ) + mock_prisma_client = MagicMock() + mock_prisma_client.db.tx = MagicMock(side_effect=lambda *a, **k: _failing_tx(non_deadlock)) + + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + with pytest.raises(UniqueViolationError): + await DBSpendUpdateWriter()._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=3, + proxy_logging_obj=proxy_logging, + db_spend_update_transactions=_empty_spend_transactions(key_list_transactions={"sk-abc": 0.5}), + ) + + mock_prisma_client.db.tx.assert_called_once() + + +@pytest.mark.asyncio +async def test_update_daily_spend_retries_deadlock(monkeypatch): + """The daily-spend upsert path retries a deadlock on the bulk upsert and then drains successfully.""" + mock_prisma_client = MagicMock() + mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[_deadlock_error(), None]) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + monkeypatch.setattr("litellm.proxy.db.db_spend_update_writer.asyncio.sleep", AsyncMock(return_value=None)) + daily_spend_transactions = {"k1": _daily_txn()} + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=3, + prisma_client=mock_prisma_client, + proxy_logging_obj=proxy_logging, + daily_spend_transactions=daily_spend_transactions, + entity_type="user", + entity_id_field="user_id", + ) + + assert mock_prisma_client.db.execute_raw.call_count == 2 + assert daily_spend_transactions == {} + proxy_logging.failure_handler.assert_not_called() + + +@pytest.mark.parametrize( + "transactions_key, sample_key", + [ + ("user_list_transactions", "user-1"), + ("team_list_transactions", "team-1"), + ("team_member_list_transactions", "team_id::team-1::user_id::user-1"), + ("org_list_transactions", "org-1"), + ("tag_list_transactions", "tag-1"), + ("agent_list_transactions", "agent-1"), + ], +) +@pytest.mark.asyncio +async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkeypatch, transactions_key, sample_key): + """Every per-entity spend path, not just keys, retries a deadlock instead of dropping the increment.""" + monkeypatch.setattr("litellm.proxy.db.db_spend_update_writer.asyncio.sleep", AsyncMock(return_value=None)) + + mock_batcher = MagicMock() + mock_prisma_client = MagicMock() + mock_prisma_client.db.tx = MagicMock(side_effect=[_failing_tx(_deadlock_error()), _good_tx(mock_batcher)]) + + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + proxy_logging.call_details = {} + + await DBSpendUpdateWriter()._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=3, + proxy_logging_obj=proxy_logging, + db_spend_update_transactions=_empty_spend_transactions(**{transactions_key: {sample_key: 0.5}}), + ) + + assert mock_prisma_client.db.tx.call_count == 2 + proxy_logging.failure_handler.assert_not_called() diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index a188289bfce..474e571e592 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -549,3 +549,36 @@ def test_handle_db_exception_surfaces_a_permanent_fault_even_when_degraded_mode_ with pytest.raises(BinaryNotFoundError): PrismaDBExceptionHandler.handle_db_exception(BinaryNotFoundError("query engine binary not found")) + + +@pytest.mark.parametrize( + "error", + [ + RawQueryError(data={"user_facing_error": {"error_code": "P2034", "meta": {"table": "t"}}}), + PrismaError("Transaction failed due to a write conflict or a deadlock. Please retry your transaction"), + RawQueryError(data={"user_facing_error": {"message": "deadlock detected", "meta": {"table": "t"}}}), + RawQueryError( + data={"user_facing_error": {"message": "ERROR: 40P01: deadlock detected", "meta": {"table": "t"}}} + ), + ], +) +def test_is_deadlock_error_matches_postgres_deadlock(error): + """A Postgres deadlock surfaced through prisma (P2034 or 40P01 / "deadlock detected" text) is recognized.""" + assert PrismaDBExceptionHandler.is_deadlock_error(error) is True + + +@pytest.mark.parametrize( + "error", + [ + UniqueViolationError(data={"user_facing_error": {"error_code": "P2002", "meta": {"table": "t"}}}), + RecordNotFoundError(data={"user_facing_error": {"meta": {"table": "t"}}}), + PrismaError("validation failed on query"), + PrismaError("can't reach database server"), + httpx.ConnectError("connection refused"), + RuntimeError("deadlock detected"), + ValueError("40P01"), + ], +) +def test_is_deadlock_error_excludes_non_deadlocks(error): + """Non-deadlock prisma errors, connectivity failures, and non-prisma exceptions are not treated as deadlocks.""" + assert PrismaDBExceptionHandler.is_deadlock_error(error) is False From 16e6ad6fb73566b1cbf3aa04fe90f7ee4653a389 Mon Sep 17 00:00:00 2001 From: RayJueWang <570828708@qq.com> Date: Thu, 6 Aug 2026 16:27:33 +0800 Subject: [PATCH 033/538] fix(proxy): recognize P2034 write-conflict deadlock text in is_deadlock_error The prisma P2034 transaction conflict can surface only as the message "Transaction failed due to a write conflict or a deadlock" without the code being reachable on the raised object, so the message fallback in is_deadlock_error now matches that canonical wording in addition to 40P01 / deadlock detected. Fixes the proxy-infra unit test that asserts this exact prisma message is treated as a retryable deadlock. --- litellm/proxy/db/exception_handler.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 91c7e576dff..f7a39aaa50f 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -176,7 +176,11 @@ class PrismaDBExceptionHandler: if getattr(e, "code", None) == "P2034": return True error_message = str(e).lower() - return "deadlock detected" in error_message or "40p01" in error_message + return ( + "deadlock detected" in error_message + or "40p01" in error_message + or "write conflict or a deadlock" in error_message + ) @staticmethod def is_prisma_engine_internal_error(e: Exception) -> bool: From 28c1e431968d61ea7caf1d82e351aca765ec83a1 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 14 Aug 2026 00:01:16 -0700 Subject: [PATCH 034/538] feat(ui): standardize the Teams page header --- .../_components/AccessGroupsPage.tsx | 4 +- .../budgets/_components/budget_panel.tsx | 4 +- .../projects/_components/ProjectsPage.tsx | 4 +- .../src/components/Teams.test.tsx | 37 ++++++--- ui/litellm-dashboard/src/components/Teams.tsx | 47 +++++------ .../VirtualKeysPage/VirtualKeysTable.tsx | 4 +- .../shared/LegacyPageHeader.test.tsx | 33 ++++++++ .../components/shared/LegacyPageHeader.tsx | 25 ++++++ .../src/components/shared/PageHeader.test.tsx | 80 +++++++++++++++---- .../src/components/shared/PageHeader.tsx | 63 +++++++++++---- 10 files changed, 224 insertions(+), 77 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/shared/LegacyPageHeader.test.tsx create mode 100644 ui/litellm-dashboard/src/components/shared/LegacyPageHeader.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx index f37acb3d85a..aeff249fdd3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx @@ -3,7 +3,7 @@ import { useDeleteAccessGroup } from "@/app/(dashboard)/hooks/accessGroups/useDe import { Plus, SearchIcon, X } from "lucide-react"; import { useMemo, useState } from "react"; import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; -import { PageHeader } from "@/components/shared/PageHeader"; +import { LegacyPageHeader } from "@/components/shared/LegacyPageHeader"; import { Button } from "@/components/ui/button"; import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; import { AccessGroupDetail } from "./AccessGroupsDetailsPage"; @@ -61,7 +61,7 @@ export function AccessGroupsPage() { return (
- = ({ accessToken }) => { return (
- } title="Budgets" subtitle="Spend, TPM and RPM limits you can assign to customers." diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsPage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsPage.tsx index 91ad9f847f5..dc18a05edca 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsPage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsPage.tsx @@ -3,7 +3,7 @@ import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import { Plus, SearchIcon, X } from "lucide-react"; import { parseAsString, useQueryState } from "nuqs"; import { useMemo, useState } from "react"; -import { PageHeader } from "@/components/shared/PageHeader"; +import { LegacyPageHeader } from "@/components/shared/LegacyPageHeader"; import { Button } from "@/components/ui/button"; import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; import { CreateProjectModal } from "./ProjectModals/CreateProjectModal"; @@ -56,7 +56,7 @@ export function ProjectsPage() { return (
- { expect(onUrlUpdate.mock.calls.at(-1)![0].searchParams.has("team")).toBe(false); await waitFor(() => expect(screen.queryByTestId("team-info-view")).not.toBeInTheDocument()); }); + + it("should preserve the legacy inset for the team detail view", async () => { + renderWithQueryClient(, { + searchParams: "?team=team-from-url", + }); + + await waitFor(() => expect(mockTeamInfoView).toHaveBeenCalled()); + expect(screen.getByRole("main")).toHaveClass("px-12", "py-6"); + }); }); describe("Teams - Create Team CTA is grouped with the tabs on the left", () => { @@ -521,22 +530,28 @@ describe("Teams - Create Team CTA is grouped with the tabs on the left", () => { mockUseOrganizations.mockReturnValue({ data: [] }); }); - it("renders the Create Team button inside the tab bar, ahead of the tabs", () => { - const { container } = renderWithQueryClient(); + it("should render the Create Team button inside the tab bar, ahead of the tabs", () => { + renderWithQueryClient(); - const createButton = screen.getByTestId("create-team-button"); - const tabNav = container.querySelector(".ant-tabs-nav"); + const tabNav = screen.getByRole("tablist"); + const createButton = within(tabNav).getByTestId("create-team-button"); + const firstTab = within(tabNav).getByRole("tab", { name: "Your Teams" }); + const tabs = tabNav.closest(".ant-tabs"); - // The CTA lives in the tab bar's left slot, not the standalone page header. - expect(tabNav).not.toBeNull(); - expect(tabNav!.contains(createButton)).toBe(true); - - // It reads as the left end of the cluster: it precedes the first tab in DOM order. - const firstTab = screen.getByRole("tab", { name: "Your Teams" }); + expect(screen.getByRole("main")).toHaveClass("p-8"); + expect(within(tabNav).getByRole("separator")).toBeInTheDocument(); expect(createButton.compareDocumentPosition(firstTab) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); + expect(tabs).toHaveClass( + "[&>.ant-tabs-nav]:!mb-6", + "[&>.ant-tabs-nav]:before:!border-b-0", + "[&_.ant-tabs-ink-bar]:!h-0.5", + "[&_.ant-tabs-tab]:!py-[7px]", + "[&_.ant-tabs-tab+_.ant-tabs-tab]:!ml-[22px]", + "[&_.ant-tabs-tab-active]:font-semibold", + ); }); - it("omits the Create Team CTA for a role that cannot manage teams", () => { + it("should omit the Create Team CTA for a role that cannot manage teams", () => { renderWithQueryClient(); expect(screen.queryByTestId("create-team-button")).not.toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index becbe0e2b48..5b79b067412 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -6,7 +6,7 @@ import TeamSSOSettings from "@/components/TeamSSOSettings"; import { isProxyAdminRole } from "@/utils/roles"; import { InfoCircleOutlined } from "@ant-design/icons"; import { Accordion, AccordionBody, AccordionHeader, TextInput } from "@tremor/react"; -import { Button, Form, Input, Layout, Modal, Select, Switch, Tabs, theme, Tooltip, Typography } from "antd"; +import { Button, Form, Input, Layout, Modal, Select, Switch, Tabs, Tooltip, Typography } from "antd"; import { Plus, Users } from "lucide-react"; import React, { useEffect, useState } from "react"; import { useQuery, useQueryClient } from "@tanstack/react-query"; @@ -403,7 +403,6 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser return false; }; - const { token } = theme.useToken(); const { Text } = Typography; const { Content } = Layout; @@ -474,7 +473,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser ]; return ( - + {selectedTeamId ? ( = ({ accessToken, userID, userRole, premiumUser premiumUser={premiumUser} /> ) : ( - <> -
- } - title="Teams" - subtitle="Manage teams, members, and their access to models and budgets" + } + title="Teams" + subtitle="Manage teams, members, and their access to models and budgets" + primaryAction={ + canCreateOrManageTeams(userRole, userID, organizations) ? ( + setIsTeamModalVisible(true)} data-testid="create-team-button"> + + Create Team + + ) : undefined + } + tabs={({ leadingControls }) => ( + -
- - - setIsTeamModalVisible(true)} data-testid="create-team-button"> - - Create Team - -
-
- ) : undefined, - }} - /> - + )} + /> )} {canCreateOrManageTeams(userRole, userID, organizations) && ( diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index fa0360c0dda..b278b4f675d 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -12,7 +12,7 @@ import { DataTableToolbar, } from "@/components/shared/DataTable"; import { SearchSelect } from "@/components/shared/SearchSelect"; -import { PageHeader } from "@/components/shared/PageHeader"; +import { LegacyPageHeader } from "@/components/shared/LegacyPageHeader"; import { Input } from "@/components/ui/input"; import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; import { ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; @@ -172,7 +172,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { return (
- } title="Virtual Keys" subtitle="Every key that authenticates requests to the gateway." diff --git a/ui/litellm-dashboard/src/components/shared/LegacyPageHeader.test.tsx b/ui/litellm-dashboard/src/components/shared/LegacyPageHeader.test.tsx new file mode 100644 index 00000000000..a0081c1c7f1 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/LegacyPageHeader.test.tsx @@ -0,0 +1,33 @@ +import { renderWithProviders, screen } from "@/../tests/test-utils"; +import { describe, expect, it } from "vitest"; + +import { LegacyPageHeader } from "./LegacyPageHeader"; + +describe("LegacyPageHeader", () => { + it("should render the title as a heading", () => { + renderWithProviders(); + + expect(screen.getByRole("heading", { name: "Virtual Keys" })).toBeInTheDocument(); + }); + + it("should render the optional identity and actions", () => { + renderWithProviders( + Key icon} + actions={} + />, + ); + + expect(screen.getByText("Every key that authenticates requests")).toBeInTheDocument(); + expect(screen.getByText("Key icon")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Create New Key" })).toBeInTheDocument(); + }); + + it("should omit optional actions when none are provided", () => { + renderWithProviders(); + + expect(screen.queryByRole("button")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/LegacyPageHeader.tsx b/ui/litellm-dashboard/src/components/shared/LegacyPageHeader.tsx new file mode 100644 index 00000000000..43979ad00b4 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/LegacyPageHeader.tsx @@ -0,0 +1,25 @@ +"use client"; + +import * as React from "react"; + +interface LegacyPageHeaderProps { + title: React.ReactNode; + subtitle?: React.ReactNode; + icon?: React.ReactNode; + actions?: React.ReactNode; +} + +export function LegacyPageHeader({ title, subtitle, icon, actions }: LegacyPageHeaderProps) { + return ( +
+
+ {icon != null && {icon}} +
+

{title}

+ {subtitle != null &&

{subtitle}

} +
+
+ {actions != null &&
{actions}
} +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/shared/PageHeader.test.tsx b/ui/litellm-dashboard/src/components/shared/PageHeader.test.tsx index f7a313271da..3741542abad 100644 --- a/ui/litellm-dashboard/src/components/shared/PageHeader.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/PageHeader.test.tsx @@ -1,31 +1,77 @@ -import { render, screen } from "@testing-library/react"; +import { renderWithProviders, screen, within } from "@/../tests/test-utils"; import { describe, expect, it } from "vitest"; import { PageHeader } from "./PageHeader"; +const identity = { + icon: Teams icon, + title: "Teams", + subtitle: "Manage teams, members, and their access to models and budgets", +}; + describe("PageHeader", () => { - it("renders the title as a heading", () => { - render(); - expect(screen.getByRole("heading", { name: "Virtual Keys" })).toBeInTheDocument(); + it("should render the page identity", () => { + renderWithProviders(); + + expect(screen.getByRole("heading", { name: "Teams" })).toBeInTheDocument(); + expect(screen.getByText("Teams icon").parentElement).toHaveAttribute("aria-hidden", "true"); + expect(screen.getByText(identity.subtitle)).toBeInTheDocument(); }); - it("renders the subtitle, icon, and actions when provided", () => { - render( + it("should apply the standard title and subtext typography", () => { + renderWithProviders(); + + const icon = screen.getByText("Teams icon").parentElement; + expect(screen.getByRole("heading", { name: "Teams" })).toHaveClass("text-2xl", "font-semibold", "tracking-tight"); + expect(screen.getByText(identity.subtitle)).toHaveClass("mt-1.5", "text-sm", "text-muted-foreground"); + expect(icon).toHaveClass("size-5", "[&_svg]:size-5", "[&_svg]:stroke-[1.75]"); + expect(icon?.parentElement).toHaveClass("gap-2.5"); + }); + + it("should render the primary action, divider, tabs, and utilities in the standard control row", () => { + renderWithProviders( } - actions={} + {...identity} + primaryAction={} + tabs={ +
+ +
+ } + utilities={} />, ); - expect(screen.getByText("Every key that authenticates requests")).toBeInTheDocument(); - expect(screen.getByTestId("icon")).toBeInTheDocument(); - expect(screen.getByRole("button", { name: "Create New Key" })).toBeInTheDocument(); + + const controls = screen.getByRole("group", { name: "Page controls" }); + expect(controls).toHaveClass("mt-5", "h-9"); + expect(within(controls).getByRole("separator")).toHaveClass("mx-4", "h-6"); + expect(controls).toHaveTextContent("Create TeamYour TeamsRefresh"); }); - it("omits the optional slots when not provided", () => { - render(); - expect(screen.queryByRole("button")).not.toBeInTheDocument(); - expect(document.querySelector("p")).toBeNull(); + it("should omit the divider when tabs are absent", () => { + renderWithProviders(Create Team} />); + + expect(screen.queryByRole("separator")).not.toBeInTheDocument(); + }); + + it("should provide standard controls to an embedded tab shell", () => { + renderWithProviders( + Create Team} + tabs={({ leadingControls, utilities }) => ( +
+ {leadingControls} + + {utilities} +
+ )} + utilities={} + />, + ); + + const tabs = screen.getByRole("tablist"); + expect(within(tabs).getByRole("separator")).toBeInTheDocument(); + expect(tabs).toHaveTextContent("Create TeamYour TeamsRefresh"); }); }); diff --git a/ui/litellm-dashboard/src/components/shared/PageHeader.tsx b/ui/litellm-dashboard/src/components/shared/PageHeader.tsx index e314e8e8bc2..81092821efc 100644 --- a/ui/litellm-dashboard/src/components/shared/PageHeader.tsx +++ b/ui/litellm-dashboard/src/components/shared/PageHeader.tsx @@ -2,24 +2,57 @@ import * as React from "react"; -interface PageHeaderProps { - title: React.ReactNode; - subtitle?: React.ReactNode; - icon?: React.ReactNode; - actions?: React.ReactNode; +import { ToolbarSeparator } from "./ToolbarSeparator"; + +interface EmbeddedTabsSlots { + leadingControls: React.ReactNode; + utilities: React.ReactNode; } -export function PageHeader({ title, subtitle, icon, actions }: PageHeaderProps) { - return ( -
-
- {icon != null && {icon}} -
-

{title}

- {subtitle != null &&

{subtitle}

} -
+interface PageHeaderProps { + title: React.ReactNode; + subtitle: React.ReactNode; + icon: React.ReactNode; + primaryAction?: React.ReactNode; + tabs?: React.ReactNode | ((slots: EmbeddedTabsSlots) => React.ReactNode); + utilities?: React.ReactNode; +} + +export function PageHeader({ title, subtitle, icon, primaryAction, tabs, utilities }: PageHeaderProps) { + const leadingControls = + primaryAction == null ? null : ( +
+ {primaryAction} + {tabs != null && }
- {actions != null &&
{actions}
} + ); + const utilityControls = utilities == null ? null :
{utilities}
; + const hasControlRow = primaryAction != null || tabs != null || utilities != null; + + return ( +
+
+ +

{title}

+
+

{subtitle}

+ + {typeof tabs === "function" ? ( +
{tabs({ leadingControls, utilities: utilityControls })}
+ ) : ( + hasControlRow && ( +
+ {leadingControls} + {tabs} + {utilityControls != null &&
{utilityControls}
} +
+ ) + )}
); } From 312d12fe0c98409b07b65fa7a677db0981b60cb5 Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Fri, 14 Aug 2026 21:39:34 +0530 Subject: [PATCH 035/538] fix(proxy): enforce project ITPM/OTPM quota on every Responses WebSocket frame The connection-level pre-call hook only ran once per WebSocket connection, so a project caller could send unlimited high-token response.create frames after a single minimal reservation. Adds enforce_project_io_token_quota_for_frame to the v3 rate limiter and wires it into both the native and managed WebSocket handlers via a duck-typed litellm.callbacks lookup, so the SDK layer stays free of proxy imports. A rejected frame gets an error event; the connection stays open for the client to retry. Also fixes the RET504 and BLE001 strict-lint-budget violations the litellm_internal_staging merge introduced in parallel_request_limiter_v3.py, which were failing the lint check. --- litellm/llms/custom_httpx/llm_http_handler.py | 28 +++ .../hooks/parallel_request_limiter_v3.py | 40 ++++- litellm/responses/streaming_iterator.py | 112 +++++++++++- .../hooks/test_parallel_request_limiter_v3.py | 56 ++++++ .../test_responses_websocket_all_providers.py | 165 ++++++++++++++++++ 5 files changed, 396 insertions(+), 5 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 721b9545ac1..e4c3a29b956 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -252,6 +252,30 @@ def _has_pre_call_deployment_hook(logging_obj: LiteLLMLoggingObj) -> bool: return False +def _collect_ws_project_quota_callbacks() -> list: + """Duck-type discover proxy hooks exposing per-frame project ITPM/OTPM + enforcement, so the Responses WebSocket loop can charge every + ``response.create`` frame, not just the connection's first one. + + Uses duck-typing on ``litellm.callbacks`` (rather than importing the + proxy hook directly) to avoid a layering violation (SDK importing from + the proxy layer). + """ + try: + import litellm as _litellm + + return [ + cb for cb in _litellm.callbacks if callable(getattr(cb, "enforce_project_io_token_quota_for_frame", None)) + ] + except Exception as exc: # noqa: BLE001 - discovery must not block the connection + verbose_logger.warning( + "Responses WebSocket: failed to collect project quota callbacks — " + "per-frame ITPM/OTPM enforcement will be skipped. Error: %s", + exc, + ) + return [] + + class BaseLLMHTTPHandler: async def _make_common_async_call( self, @@ -6168,6 +6192,8 @@ class BaseLLMHTTPHandler: - Uses ManagedResponsesWebSocketHandler which makes HTTP streaming calls - Forwards events over the websocket connection """ + _ws_quota_callbacks: Final = _collect_ws_project_quota_callbacks() + if responses_api_provider_config is None or not responses_api_provider_config.supports_native_websocket(): from litellm.responses.streaming_iterator import ( ManagedResponsesWebSocketHandler, @@ -6184,6 +6210,7 @@ class BaseLLMHTTPHandler: timeout=timeout, custom_llm_provider=custom_llm_provider, first_message=first_message, + quota_callbacks=_ws_quota_callbacks, **kwargs, ) await handler.run() @@ -6304,6 +6331,7 @@ class BaseLLMHTTPHandler: first_message=first_message, guardrail_callbacks=_ws_guardrail_callbacks, output_guardrail_callbacks=_ws_output_guardrail_callbacks, + quota_callbacks=_ws_quota_callbacks, authorized_model=model, ) await streaming.bidirectional_forward() diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 1c7ae1ea9f8..7a97718247e 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -680,7 +680,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter config = data.get("config") if "config" in data else data.get("generationConfig") - translated_request = GoogleGenAIAdapter().translate_generate_content_to_completion( + return GoogleGenAIAdapter().translate_generate_content_to_completion( model=data.get("model") if isinstance(data.get("model"), str) else "", contents=contents, config=config if isinstance(config, dict) else None, @@ -690,7 +690,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): toolConfig=data.get("toolConfig"), tool_config=data.get("tool_config"), ) - return translated_request @staticmethod def _get_explicit_output_cap(data: object, call_type: str | None) -> int | None: @@ -2171,6 +2170,41 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): assert itpm_response is not None return itpm_response, itpm_reserved, 0 + async def enforce_project_io_token_quota_for_frame( + self, + user_api_key_dict: UserAPIKeyAuth, + requested_model: str | None, + estimated_input_tokens: int, + estimated_output_tokens: int, + ) -> None: + """Reserve one WebSocket ``response.create`` frame's tokens against + the caller's project ITPM/OTPM quota. + + The Responses WebSocket connection-level pre-call hook only runs once + per connection, but a connection accepts many ``response.create`` + frames over its lifetime. Without this, a project caller could send + unlimited high-token generations after a single minimal reservation. + There is no per-frame post-call hook to reconcile against, so -- + like the batch rate limiter -- this charges the estimate immediately + and never refunds it. + """ + descriptors: Final[list[RateLimitDescriptor]] = [] + self._add_project_io_token_rate_limit_descriptors_from_metadata( + user_api_key_dict=user_api_key_dict, + requested_model=requested_model, + descriptors=descriptors, + ) + if not descriptors: + return + response, _itpm_reserved, _otpm_reserved = await self.reserve_io_tokens( + descriptors=descriptors, + estimated_input_tokens=estimated_input_tokens, + estimated_output_tokens=estimated_output_tokens, + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + if response["overall_code"] == "OVER_LIMIT": + self._handle_rate_limit_error(response, descriptors, requested_model) + def create_organization_rate_limit_descriptor( self, user_api_key_dict: UserAPIKeyAuth, requested_model: str | None = None ) -> list[RateLimitDescriptor]: @@ -3858,7 +3892,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ], ) continue - except Exception as e: + except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the plain increment fallback, never a 500 verbose_proxy_logger.warning( "Window-guarded token adjustment failed for %s: %s", operation["key"], diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 25e5fcb6976..af561c9092b 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -20,7 +20,7 @@ from litellm.constants import ( LITELLM_MAX_STREAMING_DURATION_SECONDS, STREAM_SSE_DONE_STRING, ) -from litellm.exceptions import MidStreamFallbackError +from litellm.exceptions import MidStreamFallbackError, RateLimitError from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -1326,6 +1326,79 @@ def _build_synthetic_response_events( from litellm._logging import verbose_logger +# Conservative per-frame output-token floor used when a response.create +# frame omits max_output_tokens, so a project OTPM quota can't be bypassed +# by simply never declaring an output cap. +_FRAME_NO_MAX_OUTPUT_TOKENS_FLOOR: Final = 1024 + +# Rough chars-per-token ratio for estimating a frame's input tokens without +# resolving a real per-model tokenizer, matching the conservative estimate +# the proxy's own rate limiter uses for the same purpose. +_FRAME_CHARS_PER_TOKEN_ESTIMATE: Final = 4 + + +def _extract_frame_quota_estimate_inputs(msg_obj: Mapping[str, object]) -> tuple[int, int | None]: + """Extract a rough input-token count and any explicit max_output_tokens + from a ``response.create`` frame, handling both wire shapes: + flat: {"type": "response.create", "input": ..., "max_output_tokens": ...} + nested: {"type": "response.create", "response": {"input": ..., "max_output_tokens": ...}} + """ + nested: Final = msg_obj.get("response") + params: Final[Mapping[str, object]] = ( + nested if _is_json_object(nested) and nested else {k: v for k, v in msg_obj.items() if k != "type"} + ) + text_parts: list[str] = [] # mutable-ok: local accumulator built in one pass, not shared + + def _collect_text(value: object) -> None: + if isinstance(value, str): + text_parts.append(value) + elif _is_json_array(value): + for item in value: + if isinstance(item, str): + text_parts.append(item) + elif _is_json_object(item): + _collect_text(item.get("content")) + _collect_text(item.get("text")) + + _collect_text(params.get("input")) + _collect_text(params.get("instructions")) + total_chars: Final = sum(len(part) for part in text_parts) + estimated_input_tokens: Final = max(1, total_chars // _FRAME_CHARS_PER_TOKEN_ESTIMATE) if total_chars else 0 + + max_output_tokens = params.get("max_output_tokens") + return estimated_input_tokens, max_output_tokens if isinstance(max_output_tokens, int) else None + + +async def _enforce_frame_project_quota( + quota_callbacks: Sequence[Any], + user_api_key_dict: UserAPIKeyAuth | None, + model: str | None, + raw_message: str, +) -> None: + """Charge one response.create frame's estimated tokens against every + registered project ITPM/OTPM quota callback, in isolation from PII + masking / logging so a malformed frame still reaches those callbacks.""" + if not quota_callbacks: + return + try: + msg_obj = json.loads(raw_message) + except (json.JSONDecodeError, TypeError): + return + if not _is_json_object(msg_obj) or msg_obj.get("type") != "response.create": + return + estimated_input_tokens, explicit_max_output_tokens = _extract_frame_quota_estimate_inputs(msg_obj) + estimated_output_tokens = ( + explicit_max_output_tokens if explicit_max_output_tokens is not None else _FRAME_NO_MAX_OUTPUT_TOKENS_FLOOR + ) + for callback in quota_callbacks: + await callback.enforce_project_io_token_quota_for_frame( + user_api_key_dict=user_api_key_dict, + requested_model=model, + estimated_input_tokens=estimated_input_tokens, + estimated_output_tokens=estimated_output_tokens, + ) + + RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [ "response.created", "response.completed", @@ -1360,6 +1433,7 @@ class ResponsesWebSocketStreaming: first_message: str | None = None, guardrail_callbacks: list[Any] | None = None, output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None, + quota_callbacks: list[Any] | None = None, authorized_model: str | None = None, ): self.websocket = websocket @@ -1372,6 +1446,7 @@ class ResponsesWebSocketStreaming: self.first_message = first_message self.guardrail_callbacks: list[Any] = guardrail_callbacks or [] self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or [] + self.quota_callbacks: list[Any] = quota_callbacks or [] # Model name authorized at connection time; enforced on every # response.create frame to prevent deployment-substitution attacks. self.authorized_model: str | None = authorized_model @@ -1781,10 +1856,31 @@ class ResponsesWebSocketStreaming: return json.dumps(evt_obj) if modified else response_str + async def _enforce_or_reject_frame(self, message: str) -> bool: + """Run the per-frame project quota check. + + On rejection, sends an ``error`` event to the client and reports that + the frame must be dropped instead of forwarded, so the connection + stays open for the client to retry once the window resets. + """ + try: + await _enforce_frame_project_quota( + self.quota_callbacks, self.user_api_key_dict, self.authorized_model, message + ) + except RateLimitError as e: + try: + await self.websocket.send_text( + json.dumps({"type": "error", "error": {"type": "rate_limit_exceeded", "message": str(e)}}) + ) + except Exception: # noqa: BLE001, S110 - best-effort notification, client may already be gone + pass + return False + return True + async def client_to_backend(self) -> None: """Forward response.create events from client to backend.""" try: - if self.first_message is not None: + if self.first_message is not None and await self._enforce_or_reject_frame(self.first_message): masked_first: Final = await self._mask_response_create(self.first_message) self._store_input(masked_first) self._store_event(masked_first) @@ -1792,6 +1888,8 @@ class ResponsesWebSocketStreaming: while True: message = await self.websocket.receive_text() + if not await self._enforce_or_reject_frame(message): + continue masked = await self._mask_response_create(message) self._store_input(masked) self._store_event(masked) @@ -1871,6 +1969,7 @@ class ManagedResponsesWebSocketHandler: timeout: float | None = None, custom_llm_provider: str | None = None, first_message: str | None = None, + quota_callbacks: list[Any] | None = None, **kwargs: object, ) -> None: self.websocket = websocket @@ -1887,6 +1986,7 @@ class ManagedResponsesWebSocketHandler: self.custom_llm_provider = custom_llm_provider self._connection_provider = self._resolve_provider(model) or custom_llm_provider self.first_message = first_message + self.quota_callbacks: list[Any] = quota_callbacks or [] # Carry through safe pass-through kwargs (e.g. extra_headers) self.extra_kwargs: dict[str, object] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS} # In-memory session history: response_id → full accumulated message list. @@ -2292,6 +2392,14 @@ class ManagedResponsesWebSocketHandler: verbose_logger.debug("ManagedResponsesWS: error sending warmup ack: %s", exc) return + try: + await _enforce_frame_project_quota( + self.quota_callbacks, self.user_api_key_dict, self.model_group or self.model, raw_message + ) + except RateLimitError as e: + await self._send_error(str(e), error_type="rate_limit_exceeded") + return + call_kwargs: Final = self._build_base_call_kwargs(msg_obj) call_kwargs["stream"] = True diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index ac0396c11c1..4ec0183f33c 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -3246,6 +3246,62 @@ async def test_project_model_itpm_and_tpm_limits_coexist_v3(): assert "model_per_project_otpm" in descriptor_keys +@pytest.mark.asyncio +async def test_enforce_project_io_token_quota_for_frame_blocks_over_limit_otpm(): + """VERIA regression: the Responses WebSocket connection-level pre-call + hook only runs once, but a connection accepts many response.create + frames. enforce_project_io_token_quota_for_frame is the per-frame check + that closes that gap; it must reserve against the caller's project OTPM + limit and reject once a frame's estimated output tokens exceed it.""" + _api_key = hash_token("sk-ws-frame-otpm") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + project_id="proj-mantle-ws", + project_metadata={"model_otpm_limit": {"gpt-4o": 50}}, + ) + + await handler.enforce_project_io_token_quota_for_frame( + user_api_key_dict=user_api_key_dict, + requested_model="gpt-4o", + estimated_input_tokens=1, + estimated_output_tokens=30, + ) + + with pytest.raises(HTTPException) as exc: + await handler.enforce_project_io_token_quota_for_frame( + user_api_key_dict=user_api_key_dict, + requested_model="gpt-4o", + estimated_input_tokens=1, + estimated_output_tokens=30, + ) + + assert exc.value.status_code == 429 + assert "model_per_project_otpm" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_enforce_project_io_token_quota_for_frame_noop_without_project_limits(): + """A key with no project ITPM/OTPM configured must never be blocked by + the per-frame check (no descriptors to reserve against).""" + _api_key = hash_token("sk-ws-frame-no-limits") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key) + + await handler.enforce_project_io_token_quota_for_frame( + user_api_key_dict=user_api_key_dict, + requested_model="gpt-4o", + estimated_input_tokens=10_000_000, + estimated_output_tokens=10_000_000, + ) + + @pytest.mark.asyncio async def test_pre_call_hook_keeps_internal_stash_out_of_request_body(): """Regression for #27001 / #35197: the limiter's per-request bookkeeping diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index 4509abc7749..2d523bfdeb3 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -1030,6 +1030,171 @@ class TestWebSocketErrorHandling: assert "Invalid JSON" in error_event +class TestWebSocketProjectQuotaEnforcement: + """VERIA regression: the connection-level pre-call hook only runs once, + but a WebSocket connection accepts many response.create frames. Every + frame must be checked against any registered project ITPM/OTPM quota + callback, not just the first one.""" + + @pytest.mark.asyncio + async def test_managed_handler_blocks_frame_rejected_by_quota_callback(self, monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + import litellm + from litellm.exceptions import RateLimitError + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + aresponses_called = False + + async def fake_aresponses(*args, **kwargs): + nonlocal aresponses_called + aresponses_called = True + + monkeypatch.setattr(litellm, "aresponses", fake_aresponses) + + quota_callback = MagicMock() + quota_callback.enforce_project_io_token_quota_for_frame = AsyncMock( + side_effect=RateLimitError(message="project OTPM exceeded", llm_provider="", model="") + ) + + mock_websocket = MagicMock() + mock_websocket.send_text = AsyncMock() + mock_logging_obj = Logging( + model="test-model", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ) + handler = ManagedResponsesWebSocketHandler( + websocket=mock_websocket, + model="test-model", + logging_obj=mock_logging_obj, + quota_callbacks=[quota_callback], + ) + + await handler._process_response_create(json.dumps({"type": "response.create", "input": "hi"})) + + quota_callback.enforce_project_io_token_quota_for_frame.assert_awaited_once() + assert aresponses_called is False + mock_websocket.send_text.assert_called_once() + error_event = mock_websocket.send_text.call_args[0][0] + assert "rate_limit_exceeded" in error_event + + @pytest.mark.asyncio + async def test_managed_handler_forwards_frame_allowed_by_quota_callback(self, monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + aresponses_called = False + + async def fake_aresponses(*args, **kwargs): + nonlocal aresponses_called + aresponses_called = True + + async def _empty(): + return + yield + + return _empty() + + monkeypatch.setattr(litellm, "aresponses", fake_aresponses) + + quota_callback = MagicMock() + quota_callback.enforce_project_io_token_quota_for_frame = AsyncMock(return_value=None) + + mock_websocket = MagicMock() + mock_websocket.send_text = AsyncMock() + mock_logging_obj = Logging( + model="test-model", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ) + handler = ManagedResponsesWebSocketHandler( + websocket=mock_websocket, + model="test-model", + logging_obj=mock_logging_obj, + quota_callbacks=[quota_callback], + ) + + await handler._process_response_create(json.dumps({"type": "response.create", "input": "hi"})) + + quota_callback.enforce_project_io_token_quota_for_frame.assert_awaited_once() + assert aresponses_called is True + + @pytest.mark.asyncio + async def test_native_handler_blocks_frame_rejected_by_quota_callback(self): + from unittest.mock import AsyncMock, MagicMock + + from litellm.exceptions import RateLimitError + from litellm.responses.streaming_iterator import ResponsesWebSocketStreaming + + quota_callback = MagicMock() + quota_callback.enforce_project_io_token_quota_for_frame = AsyncMock( + side_effect=RateLimitError(message="project OTPM exceeded", llm_provider="", model="") + ) + + mock_backend_ws = MagicMock() + mock_backend_ws.send = AsyncMock() + mock_websocket = MagicMock() + mock_websocket.send_text = AsyncMock() + + handler = ResponsesWebSocketStreaming( + websocket=mock_websocket, + backend_ws=mock_backend_ws, + logging_obj=MagicMock(), + authorized_model="gpt-4o", + quota_callbacks=[quota_callback], + ) + + allowed = await handler._enforce_or_reject_frame( + json.dumps({"type": "response.create", "input": "hi"}) + ) + + assert allowed is False + mock_backend_ws.send.assert_not_called() + mock_websocket.send_text.assert_called_once() + assert "rate_limit_exceeded" in mock_websocket.send_text.call_args[0][0] + + @pytest.mark.asyncio + async def test_native_handler_forwards_frame_allowed_by_quota_callback(self): + from unittest.mock import AsyncMock, MagicMock + + from litellm.responses.streaming_iterator import ResponsesWebSocketStreaming + + quota_callback = MagicMock() + quota_callback.enforce_project_io_token_quota_for_frame = AsyncMock(return_value=None) + + handler = ResponsesWebSocketStreaming( + websocket=MagicMock(), + backend_ws=MagicMock(), + logging_obj=MagicMock(), + authorized_model="gpt-4o", + quota_callbacks=[quota_callback], + ) + + allowed = await handler._enforce_or_reject_frame( + json.dumps({"type": "response.create", "input": "hi"}) + ) + + assert allowed is True + quota_callback.enforce_project_io_token_quota_for_frame.assert_awaited_once() + + class TestNativeWebSocketGuardrails: @pytest.mark.asyncio async def test_response_create_injects_authorized_model(self): From 2b23295f82d29cac7eb97474ea0d53b992c3f511 Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Fri, 14 Aug 2026 23:31:18 +0530 Subject: [PATCH 036/538] fix(proxy): reconcile project quota reservations --- basedpyright-code-budget.json | 6 +- litellm/llms/custom_httpx/llm_http_handler.py | 26 +- litellm/proxy/hooks/batch_rate_limiter.py | 50 +- .../hooks/parallel_request_limiter_v3.py | 589 +++++++++--------- litellm/responses/streaming_iterator.py | 59 +- .../proxy/hooks/test_tpm_concurrent.py | 10 +- type-discipline-budget.json | 8 +- 7 files changed, 407 insertions(+), 341 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 521b4315e6e..67b0575a0cb 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5707 }, "reportMissingTypeArgument": { - "limit": 15642 + "limit": 15641 }, "reportMissingTypeStubs": { "limit": 40 @@ -108,7 +108,7 @@ "limit": 39237 }, "reportUnknownParameterType": { - "limit": 19969 + "limit": 19968 }, "reportUnknownVariableType": { "limit": 30881 @@ -132,7 +132,7 @@ "limit": 27 }, "reportUnusedClass": { - "limit": 23 + "limit": 22 }, "reportUnusedFunction": { "limit": 139 diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index e4c3a29b956..8627d1797a5 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2,7 +2,7 @@ import asyncio import json import os import ssl -from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping +from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence from contextlib import asynccontextmanager from functools import lru_cache from types import ModuleType @@ -69,6 +69,7 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.responses.streaming_iterator import ( BaseResponsesAPIStreamingIterator, MockResponsesAPIStreamingIterator, + ProjectQuotaCallback, ResponsesAPIStreamingIterator, ResponsesWebSocketStreaming, SyncResponsesAPIStreamingIterator, @@ -252,7 +253,7 @@ def _has_pre_call_deployment_hook(logging_obj: LiteLLMLoggingObj) -> bool: return False -def _collect_ws_project_quota_callbacks() -> list: +def _collect_ws_project_quota_callbacks() -> tuple[ProjectQuotaCallback, ...]: """Duck-type discover proxy hooks exposing per-frame project ITPM/OTPM enforcement, so the Responses WebSocket loop can charge every ``response.create`` frame, not just the connection's first one. @@ -261,19 +262,16 @@ def _collect_ws_project_quota_callbacks() -> list: proxy hook directly) to avoid a layering violation (SDK importing from the proxy layer). """ - try: - import litellm as _litellm + import litellm as _litellm - return [ - cb for cb in _litellm.callbacks if callable(getattr(cb, "enforce_project_io_token_quota_for_frame", None)) - ] - except Exception as exc: # noqa: BLE001 - discovery must not block the connection - verbose_logger.warning( - "Responses WebSocket: failed to collect project quota callbacks — " - "per-frame ITPM/OTPM enforcement will be skipped. Error: %s", - exc, - ) - return [] + callbacks: Final = cast( # cast-ok: callback registry is inspected before protocol use + Sequence[object], _litellm.callbacks + ) + return tuple( + cast(ProjectQuotaCallback, callback) # cast-ok: required callback method is callable + for callback in callbacks + if callable(getattr(callback, "enforce_project_io_token_quota_for_frame", None)) + ) class BaseLLMHTTPHandler: diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 3091b0f6973..efef246a7a6 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -18,11 +18,12 @@ Quick summary: """ import json -from collections.abc import Iterable +from collections.abc import Iterable, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn from fastapi import HTTPException -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -77,6 +78,9 @@ else: RateLimitDescriptor = dict[str, Any] +_BATCH_BODY_ADAPTER: Final = TypeAdapter(dict[str, object]) + + class BatchFileUsage(BaseModel): """ Internal model for batch file usage tracking, used for batch rate limiting @@ -214,7 +218,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): tpm_limit_type=None, model_has_failures=False, ) - self.parallel_request_limiter._add_project_io_token_rate_limit_descriptors_from_metadata( + self.parallel_request_limiter.add_project_io_token_rate_limit_descriptors_from_metadata( user_api_key_dict=user_api_key_dict, requested_model=self._get_batch_routing_model(data), descriptors=descriptors, @@ -314,7 +318,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): def _estimate_entry_output_tokens( self, - entry: dict, + entry: Mapping[str, object], min_configured_otpm_limit: int | None, ) -> int: """Conservative per-row output-token estimate for the project OTPM reservation. @@ -325,16 +329,21 @@ class _PROXY_BatchRateLimiter(CustomLogger): that omits ``max_tokens`` can't be used to bypass OTPM the way an unbounded streaming request could. """ - body: Final = entry.get("body", {}) or {} + raw_body: Final = entry.get("body") + body: Final[Mapping[str, object]] = ( + MappingProxyType(_BATCH_BODY_ADAPTER.validate_python(raw_body)) + if isinstance(raw_body, Mapping) + else MappingProxyType({}) # mutable-ok: immediately frozen empty fallback + ) if body.get("input") is not None and body.get("messages") is None and body.get("prompt") is None: return 0 # embeddings: no output tokens - explicit_cap = body.get("max_tokens", body.get("max_completion_tokens")) + explicit_cap: Final = body.get("max_tokens", body.get("max_completion_tokens")) if explicit_cap is not None: try: return max(0, int(explicit_cap)) except (TypeError, ValueError): pass - return self.parallel_request_limiter._no_max_tokens_output_floor(min_configured_otpm_limit) + return self.parallel_request_limiter.no_max_tokens_output_floor(min_configured_otpm_limit) @staticmethod def _has_applicable_batch_rate_limits( @@ -446,7 +455,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): f"Limit resets at: {reset_time_formatted}" ) else: # tokens - batch_token_count = ( + batch_token_count: Final = ( batch_usage.output_tokens if descriptor.get("key") == PROJECT_OTPM_DESCRIPTOR_KEY else batch_usage.total_tokens @@ -496,8 +505,8 @@ class _PROXY_BatchRateLimiter(CustomLogger): data=data, ) - increments: Final[list[dict[Literal["requests", "tokens"], int]]] = [ - { + increments: Final = [ # mutable-ok: atomic limiter API requires mutable increment records + { # mutable-ok: atomic limiter API requires mutable increment records "requests": batch_usage.request_count, "tokens": batch_usage.output_tokens if d.get("key") == PROJECT_OTPM_DESCRIPTOR_KEY @@ -530,7 +539,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", user_api_key_dict: UserAPIKeyAuth | None = None, data: dict | None = None, - descriptors: list["RateLimitDescriptor"] | None = None, + descriptors: Sequence["RateLimitDescriptor"] | None = None, ) -> BatchFileUsage: """ Count number of requests and tokens in a batch input file. @@ -545,13 +554,14 @@ class _PROXY_BatchRateLimiter(CustomLogger): Returns: BatchFileUsage with total_tokens, output_tokens, and request_count """ - otpm_limits: Final = [ + otpm_limits: Final = tuple( int(v) - for d in (descriptors or []) + for d in (descriptors or ()) if d.get("key") == PROJECT_OTPM_DESCRIPTOR_KEY - for v in [(d.get("rate_limit") or {}).get("tokens_per_unit")] + for rate_limit in (d.get("rate_limit"),) + for v in (rate_limit.get("tokens_per_unit") if rate_limit is not None else None,) if v is not None - ] + ) min_configured_otpm_limit: Final = min(otpm_limits) if otpm_limits else None try: # Check if this is a managed file (base64 encoded unified file ID) @@ -605,7 +615,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): # the counter can't measure is estimated, not hard-rejected. models: Final[set] = set() total_tokens = 0 - output_tokens = 0 + output_tokens = 0 # rebind-ok: accumulated per JSONL row in the loop below request_count = 0 for raw_line in _iter_batch_input_lines(file_content_bytes): request_count += 1 @@ -613,9 +623,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): entry = json.loads(raw_line) except Exception: total_tokens += _estimate_batch_entry_tokens(raw_line) - output_tokens += self.parallel_request_limiter._no_max_tokens_output_floor( - min_configured_otpm_limit - ) + output_tokens += self.parallel_request_limiter.no_max_tokens_output_floor(min_configured_otpm_limit) continue if isinstance(entry, dict): model = (entry.get("body") or {}).get("model") @@ -623,9 +631,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): models.add(model) output_tokens += self._estimate_entry_output_tokens(entry, min_configured_otpm_limit) else: - output_tokens += self.parallel_request_limiter._no_max_tokens_output_floor( - min_configured_otpm_limit - ) + output_tokens += self.parallel_request_limiter.no_max_tokens_output_floor(min_configured_otpm_limit) try: total_tokens += _count_entry_tokens(entry) except Exception: diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 7a97718247e..f858fd2af98 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -22,7 +22,7 @@ from typing import ( TypedDict, ) -from typing_extensions import NotRequired +from typing_extensions import NotRequired, ReadOnly from litellm import DualCache from litellm._logging import verbose_proxy_logger @@ -69,6 +69,7 @@ else: Span = Any InternalUsageCache = Any + BATCH_RATE_LIMITER_SCRIPT: Final = """ local results = {} local now = tonumber(ARGV[1]) @@ -413,7 +414,7 @@ class RateLimitStatus(TypedDict): class RateLimitResponse(TypedDict): overall_code: str statuses: list[RateLimitStatus] - reservation_windows: NotRequired[frozenset[tuple[str, str, Literal["redis", "local"]]]] + reservation_windows: NotRequired[ReadOnly[frozenset[tuple[str, str, Literal["redis", "local"]]]]] class ReservationAwareIncrementOperation(RedisPipelineIncrementOperation): @@ -452,6 +453,7 @@ class AtomicCounterMeta(TypedDict): class AtomicCounterState(TypedDict): window_expired: bool current: int + window_start: ReadOnly[str] DescriptorAtomicGroup: TypeAlias = tuple[list[str], list[int], list[AtomicCounterMeta]] @@ -639,7 +641,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return self._time_provider() @staticmethod - def _no_max_tokens_output_floor( + def no_max_tokens_output_floor( min_configured_tpm_limit: int | None, ) -> int: """Output-budget floor used when the request omits max_tokens. @@ -670,7 +672,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): data: object, call_type: str | None, ) -> Mapping[str, object] | None: - contents = data.get("contents") if isinstance(data, dict) else None + contents: Final = data.get("contents") if isinstance(data, dict) else None if ( not isinstance(data, dict) or call_type not in GOOGLE_GENAI_NATIVE_CALL_TYPES @@ -679,7 +681,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return None from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter - config = data.get("config") if "config" in data else data.get("generationConfig") + config: Final = data.get("config") if "config" in data else data.get("generationConfig") return GoogleGenAIAdapter().translate_generate_content_to_completion( model=data.get("model") if isinstance(data.get("model"), str) else "", contents=contents, @@ -696,15 +698,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if not isinstance(data, dict): return None if call_type in GOOGLE_GENAI_NATIVE_CALL_TYPES: - config = data.get("config") if "config" in data else data.get("generationConfig") - values = tuple( - int(config[field]) + config: Final = data.get("config") if "config" in data else data.get("generationConfig") + google_cap_values: Final = tuple( + int(raw_value) for field in ("maxOutputTokens", "max_output_tokens") - if isinstance(config, dict) and isinstance(config.get(field), (int, float, str)) + if isinstance(config, dict) + for raw_value in (config.get(field),) + if isinstance(raw_value, (int, float, str)) ) - return max(values, default=None) + return max(google_cap_values, default=None) if call_type in RESPONSES_API_CALL_TYPES: - value = data.get("max_output_tokens") + value: Final = data.get("max_output_tokens") if value is None: return None if not isinstance(value, (int, float, str)): @@ -712,13 +716,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return max(RESPONSES_API_MIN_OUTPUT_TOKENS, int(value)) if call_type in EMBEDDING_API_CALL_TYPES: return None - fields = ( + fields: Final = ( ("max_tokens", "max_completion_tokens") if call_type else ("max_tokens", "max_completion_tokens", "max_output_tokens") ) - values = tuple(int(data[field]) for field in fields if isinstance(data.get(field), (int, float, str))) - return max(values, default=None) + output_cap_values: Final = tuple( + int(raw_value) + for field in fields + for raw_value in (data.get(field),) + if isinstance(raw_value, (int, float, str)) + ) + return max(output_cap_values, default=None) @classmethod def _has_explicit_output_cap(cls, data: object, call_type: str | None) -> bool: @@ -733,18 +742,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _get_output_candidate_count(data: object, call_type: str | None = None) -> int: if not isinstance(data, dict): return 1 - config = ( + config: Final = ( (data.get("config") if "config" in data else data.get("generationConfig")) if call_type in GOOGLE_GENAI_NATIVE_CALL_TYPES else None ) - candidate_values = ( + candidate_values: Final = ( data.get("n"), data.get("best_of"), config.get("candidateCount") if isinstance(config, dict) else None, config.get("candidate_count") if isinstance(config, dict) else None, ) - candidate_count = 1 + candidate_count = 1 # rebind-ok: running maximum across candidate-count aliases for value in candidate_values: try: candidate_count = max(candidate_count, int(value or 1)) @@ -774,31 +783,34 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ if not isinstance(data, dict): return - capped_floor = _PROXY_MaxParallelRequestsHandler_v3._no_max_tokens_output_floor(min_configured_limit) - if call_type in RESPONSES_API_CALL_TYPES: - capped_floor = max(capped_floor, RESPONSES_API_MIN_OUTPUT_TOKENS) - baseline_floor = DEFAULT_MAX_TOKENS_ESTIMATE // _TPM_FLOOR_FRACTION - is_embedding = _PROXY_MaxParallelRequestsHandler_v3._is_embedding_request(data, call_type) + base_capped_floor: Final = _PROXY_MaxParallelRequestsHandler_v3.no_max_tokens_output_floor(min_configured_limit) + capped_floor: Final = ( + max(base_capped_floor, RESPONSES_API_MIN_OUTPUT_TOKENS) + if call_type in RESPONSES_API_CALL_TYPES + else base_capped_floor + ) + baseline_floor: Final = DEFAULT_MAX_TOKENS_ESTIMATE // _TPM_FLOOR_FRACTION + is_embedding: Final = _PROXY_MaxParallelRequestsHandler_v3._is_embedding_request(data, call_type) if ( capped_floor >= baseline_floor or _PROXY_MaxParallelRequestsHandler_v3._has_explicit_output_cap(data, call_type) or is_embedding ): return - effective_cap = max(capped_floor, configured_output_tokens or 0) + effective_cap: Final = max(capped_floor, configured_output_tokens or 0) if call_type in GOOGLE_GENAI_NATIVE_CALL_TYPES: - config_field = "config" if "config" in data or "generationConfig" not in data else "generationConfig" - config = data.get(config_field) + config_field: Final = "config" if "config" in data or "generationConfig" not in data else "generationConfig" + config: Final = data.get(config_field) if config is None or isinstance(config, dict): - data[config_field] = { # mutable-ok: downstream native routing requires a mutable request config + data[config_field] = { # rebind-ok: routed request needs cap # mutable-ok: downstream needs dict **(config or {}), # mutable-ok: downstream native routing requires a mutable request config "maxOutputTokens": effective_cap, } return - cap_field = "max_output_tokens" if call_type in RESPONSES_API_CALL_TYPES else "max_tokens" - existing_cap = data.get(cap_field) + cap_field: Final = "max_output_tokens" if call_type in RESPONSES_API_CALL_TYPES else "max_tokens" + existing_cap: Final = data.get(cap_field) if existing_cap is None or effective_cap < existing_cap: - data[cap_field] = effective_cap + data[cap_field] = effective_cap # rebind-ok: downstream routing requires the bounded output cap def _estimate_tokens_for_request( self, @@ -878,69 +890,57 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return 0, 0 translated_data: Final = self._translate_google_genai_native_request(data, call_type) estimable_data: Final = translated_data if translated_data is not None else data - messages = estimable_data.get("messages") - prompt = estimable_data.get("prompt") - input_text = estimable_data.get("input") + selected_fields: Final[tuple[object | None, object | None, object | None]] = ( + (None, None, estimable_data.get("input")) + if call_type in RESPONSES_API_CALL_TYPES or call_type in EMBEDDING_API_CALL_TYPES + else (None, estimable_data.get("prompt"), None) + if call_type in TEXT_COMPLETION_API_CALL_TYPES + else (estimable_data.get("messages"), None, None) + if call_type + else ( + estimable_data.get("messages"), + estimable_data.get("prompt"), + estimable_data.get("input"), + ) + ) + messages, prompt, input_text = selected_fields - if call_type in RESPONSES_API_CALL_TYPES or call_type in EMBEDDING_API_CALL_TYPES: - messages = None - prompt = None - elif call_type in TEXT_COMPLETION_API_CALL_TYPES: - messages = None - input_text = None - elif call_type: - prompt = None - input_text = None - - match (messages, prompt, input_text): - case (selected_messages, _, _) if selected_messages: - total_chars = len(get_str_from_messages(selected_messages)) - case (_, str() as selected_prompt, _): - total_chars = len(selected_prompt) - case (_, list() as selected_prompt, _): - total_chars = sum(len(str(item)) for item in selected_prompt) - case (_, _, str() as selected_input): - total_chars = len(selected_input) - case (_, _, list() as selected_input): - total_chars = sum(len(str(item)) for item in selected_input) - case _: - total_chars = 0 + total_chars: Final = ( + len(get_str_from_messages(messages)) + if isinstance(messages, list) and messages + else len(prompt) + if isinstance(prompt, str) + else sum(len(str(item)) for item in prompt) + if isinstance(prompt, list) + else len(input_text) + if isinstance(input_text, str) + else sum(len(str(item)) for item in input_text) + if isinstance(input_text, list) + else 0 + ) estimated_input_tokens: Final = max(1, total_chars // DEFAULT_CHARS_PER_TOKEN) if total_chars > 0 else 0 explicit_max_tokens: Final = self._get_explicit_output_cap(data, call_type) is_embedding: Final = self._is_embedding_request(data, call_type) - match (explicit_max_tokens, is_embedding): - case (_, True): - max_tokens_estimate = 0 - case (mt, _) if mt is not None: - max_tokens_estimate = mt - case _ if total_chars == 0 and configured_output_tokens is None: - # Fully contentless request (no messages, prompt, or input). - # Don't apply the conservative output-budget floor here — it - # would over-reserve and could push small TPM limits into a - # false 429. The caller floors at 1 so backpressure still - # applies once the counter is at limit. - max_tokens_estimate = 0 - case _: - # No max_tokens specified — reserve at least the input size with a - # conservative floor so a stream of small concurrent requests can't - # collectively bypass the limit. Cap the floor by a fraction of - # the smallest TPM limit this request will be charged against, - # so a small per-tenant TPM cap can't be tripped by the floor - # alone. - output_floor = self._no_max_tokens_output_floor(min_configured_tpm_limit) - if call_type in RESPONSES_API_CALL_TYPES: - output_floor = max(output_floor, RESPONSES_API_MIN_OUTPUT_TOKENS) - max_tokens_estimate = ( - configured_output_tokens - if configured_output_tokens is not None - else max(estimated_input_tokens, output_floor) - ) + base_output_floor: Final = self.no_max_tokens_output_floor(min_configured_tpm_limit) + output_floor: Final = ( + max(base_output_floor, RESPONSES_API_MIN_OUTPUT_TOKENS) + if call_type in RESPONSES_API_CALL_TYPES + else base_output_floor + ) + max_tokens_estimate: Final = ( + 0 + if is_embedding or (explicit_max_tokens is None and total_chars == 0 and configured_output_tokens is None) + else explicit_max_tokens + if explicit_max_tokens is not None + else configured_output_tokens + if configured_output_tokens is not None + else max(estimated_input_tokens, output_floor) + ) - max_tokens_estimate *= self._get_output_candidate_count(data, call_type) - return estimated_input_tokens, max_tokens_estimate + return estimated_input_tokens, max_tokens_estimate * self._get_output_candidate_count(data, call_type) def _is_redis_cluster(self) -> bool: """ @@ -1755,7 +1755,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): mid-loop, refund applied increments and fall back to in-memory. """ if not descriptor_groups: - return RateLimitResponse(overall_code="OK", statuses=[]) + return RateLimitResponse( + overall_code="OK", + statuses=[], # mutable-ok: response contract requires a status list + ) applied: Final[list[list[AtomicCounterMeta]]] = [] statuses: Final[list[RateLimitStatus]] = [] raw: list[CacheCounterValue] @@ -1946,7 +1949,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ], ) descriptor_state.append( - { + { # mutable-ok: local atomic-counter state is updated during pass two "window_expired": window_expired, "current": current_counter, "window_start": str(now_int if window_expired else int(window_start)), @@ -2096,21 +2099,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): configured, or if the reservation failed), for the caller to stash for post-call reconciliation. """ - itpm_descriptors = [ # mutable-ok: atomic limiter API requires lists + itpm_descriptors: Final = [ # mutable-ok: atomic limiter API requires lists d for d in descriptors if d["key"] == PROJECT_ITPM_DESCRIPTOR_KEY ] - otpm_descriptors = [ # mutable-ok: atomic limiter API requires lists + otpm_descriptors: Final = [ # mutable-ok: atomic limiter API requires lists d for d in descriptors if d["key"] == PROJECT_OTPM_DESCRIPTOR_KEY ] if not itpm_descriptors and not otpm_descriptors: return RateLimitResponse(overall_code="OK", statuses=[]), 0, 0 # mutable-ok: response contract uses a list - itpm_response: RateLimitResponse | None = None - itpm_reserved = 0 - - if itpm_descriptors: - itpm_response = await self.atomic_check_and_increment_by_n( + itpm_response: Final = ( + await self.atomic_check_and_increment_by_n( descriptors=itpm_descriptors, increments=[ # mutable-ok: atomic limiter API requires mutable increment records {"tokens": estimated_input_tokens} # mutable-ok: atomic limiter increment record @@ -2118,12 +2118,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ], parent_otel_span=parent_otel_span, ) - if itpm_response["overall_code"] == "OVER_LIMIT": - return itpm_response, 0, 0 - itpm_reserved = estimated_input_tokens + if itpm_descriptors + else None + ) + if itpm_response is not None and itpm_response["overall_code"] == "OVER_LIMIT": + return itpm_response, 0, 0 + itpm_reserved: Final = estimated_input_tokens if itpm_response is not None else 0 if otpm_descriptors: - otpm_response = await self.atomic_check_and_increment_by_n( + otpm_response: Final = await self.atomic_check_and_increment_by_n( descriptors=otpm_descriptors, increments=[ # mutable-ok: atomic limiter API requires mutable increment records {"tokens": estimated_output_tokens} # mutable-ok: atomic limiter increment record @@ -2142,7 +2145,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=parent_otel_span, ) return otpm_response, 0, 0 - statuses = ( + statuses: Final = ( [ # mutable-ok: response contract uses a list *itpm_response["statuses"], *otpm_response["statuses"], @@ -2172,7 +2175,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async def enforce_project_io_token_quota_for_frame( self, - user_api_key_dict: UserAPIKeyAuth, + user_api_key_dict: UserAPIKeyAuth | None, requested_model: str | None, estimated_input_tokens: int, estimated_output_tokens: int, @@ -2188,8 +2191,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): like the batch rate limiter -- this charges the estimate immediately and never refunds it. """ - descriptors: Final[list[RateLimitDescriptor]] = [] - self._add_project_io_token_rate_limit_descriptors_from_metadata( + if user_api_key_dict is None: + return + descriptors: Final[list[RateLimitDescriptor]] = [] # mutable-ok: descriptor helper appends in place + self.add_project_io_token_rate_limit_descriptors_from_metadata( user_api_key_dict=user_api_key_dict, requested_model=requested_model, descriptors=descriptors, @@ -2911,7 +2916,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) - def _add_project_io_token_rate_limit_descriptors_from_metadata( + def add_project_io_token_rate_limit_descriptors_from_metadata( self, user_api_key_dict: UserAPIKeyAuth, requested_model: str | None, @@ -2926,22 +2931,22 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if requested_model is None or user_api_key_dict.project_id is None: return - itpm_limit_for_project_model = ( + itpm_limit_for_project_model: Final = ( get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_itpm_limit") or {} # mutable-ok: metadata helper returns an optional mapping ) - otpm_limit_for_project_model = ( + otpm_limit_for_project_model: Final = ( get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_otpm_limit") or {} # mutable-ok: metadata helper returns an optional mapping ) - model_itpm_limit = itpm_limit_for_project_model.get(requested_model) - model_otpm_limit = otpm_limit_for_project_model.get(requested_model) + model_itpm_limit: Final = itpm_limit_for_project_model.get(requested_model) + model_otpm_limit: Final = otpm_limit_for_project_model.get(requested_model) if model_itpm_limit is None and model_otpm_limit is None: return - descriptor_value = f"{user_api_key_dict.project_id}:{requested_model}" + descriptor_value: Final = f"{user_api_key_dict.project_id}:{requested_model}" if model_itpm_limit is not None: descriptors.append( RateLimitDescriptor( @@ -3026,10 +3031,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ if not isinstance(block, dict): return DEFAULT_AUDIO_TOKEN_ESTIMATE - input_audio = block.get("input_audio") - b64_data = input_audio.get("data") if isinstance(input_audio, dict) else None + input_audio: Final = block.get("input_audio") + b64_data: Final = input_audio.get("data") if isinstance(input_audio, dict) else None if b64_data and isinstance(b64_data, str): - decoded_bytes = len(b64_data) * 3 // 4 + decoded_bytes: Final = len(b64_data) * 3 // 4 return max(decoded_bytes // _AUDIO_BYTES_PER_TOKEN, DEFAULT_AUDIO_TOKEN_ESTIMATE) return DEFAULT_AUDIO_TOKEN_ESTIMATE @@ -3042,17 +3047,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ if not isinstance(messages, list): return 0 - total = 0 - for message in messages: - content = message.get("content") if isinstance(message, dict) else None - if not isinstance(content, list): - continue - total += sum( - cls._estimate_audio_block_tokens(block) - for block in content - if isinstance(block, dict) and block.get("type") == "input_audio" - ) - return total + return sum( + cls._estimate_audio_block_tokens(block) + for message in messages + if isinstance(message, dict) + for content in (message.get("content"),) + if isinstance(content, list) + for block in content + if isinstance(block, dict) and block.get("type") == "input_audio" + ) @staticmethod def _strip_audio_content_blocks(messages: object) -> object: @@ -3066,7 +3069,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ if not isinstance(messages, list): return messages - sanitized = [] # mutable-ok: token_counter requires a list of message dicts + sanitized: Final[list[object]] = [] # mutable-ok: token_counter requires a list of message dicts for message in messages: if not isinstance(message, dict): sanitized.append(message) @@ -3127,7 +3130,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): @classmethod def _contains_image_content(cls, value: object) -> bool: if isinstance(value, dict): - media_type = value.get("media_type") or value.get("mime_type") + media_type: Final = value.get("media_type") or value.get("mime_type") return ( value.get("type") in ("image", "image_url", "input_image") or (isinstance(media_type, str) and media_type.startswith("image/")) @@ -3159,9 +3162,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): @staticmethod def _rerank_input_to_text(data: Mapping[str, object]) -> str: - documents = data.get("documents") - document_items: Sequence[object] = documents if isinstance(documents, list) else () # pyright: ignore[reportUnknownVariableType] # rerank documents are validated runtime JSON - input_parts: tuple[object, ...] = ( # pyright: ignore[reportUnknownVariableType] # list narrowing preserves unknown JSON element types + documents: Final = data.get("documents") + document_items: Final[Sequence[object]] = documents if isinstance(documents, list) else () # pyright: ignore[reportUnknownVariableType] # rerank documents are validated runtime JSON + input_parts: Final[tuple[object, ...]] = ( # pyright: ignore[reportUnknownVariableType] # list narrowing preserves unknown JSON element types data.get("query"), *document_items, ) @@ -3199,39 +3202,45 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if not isinstance(data, dict): return 0 - selected_text = None - countable_tools = data.get("tools") - countable_tool_choice = data.get("tool_choice") - if call_type in RESPONSES_API_CALL_TYPES: - messages = self._responses_input_to_chat_messages(data) - elif (translated_request := self._translate_google_genai_native_request(data, call_type)) is not None: - messages = translated_request.get("messages") - countable_tools = translated_request.get("tools") - countable_tool_choice = translated_request.get("tool_choice") - elif self._is_embedding_request(data, call_type): - messages = None - selected_text = data.get("input") - pretokenized_input_tokens = self._count_pretokenized_embedding_input(selected_text) - if pretokenized_input_tokens is not None: - return pretokenized_input_tokens - elif call_type in RERANK_API_CALL_TYPES: - messages = None - selected_text = self._rerank_input_to_text(data) # pyright: ignore[reportUnknownArgumentType] # proxy request bodies are runtime-validated JSON - elif call_type in TEXT_COMPLETION_API_CALL_TYPES: - messages = None - selected_text = data.get("prompt") - else: - messages = data.get("messages") - if messages is None: - selected_text = data.get("prompt") - if messages is None and selected_text is None: - selected_text = data.get("input") + is_responses_request: Final = call_type in RESPONSES_API_CALL_TYPES + translated_request: Final = ( + None if is_responses_request else self._translate_google_genai_native_request(data, call_type) + ) + is_embedding_request: Final = self._is_embedding_request(data, call_type) + embedding_text: Final = data.get("input") if is_embedding_request else None + pretokenized_input_tokens: Final = ( + self._count_pretokenized_embedding_input(embedding_text) if is_embedding_request else None + ) + if pretokenized_input_tokens is not None: + return pretokenized_input_tokens - audio_token_estimate = self._estimate_audio_content_tokens(messages) - countable_messages = self._strip_audio_content_blocks(messages) if audio_token_estimate > 0 else messages + prompt: Final = data.get("prompt") + fallback_text: Final = prompt if prompt is not None else data.get("input") + selected_inputs: Final[tuple[object | None, object | None, object | None, object | None]] = ( + (self._responses_input_to_chat_messages(data), None, data.get("tools"), data.get("tool_choice")) + if is_responses_request + else ( + translated_request.get("messages"), + None, + translated_request.get("tools"), + translated_request.get("tool_choice"), + ) + if translated_request is not None + else (None, embedding_text, data.get("tools"), data.get("tool_choice")) + if is_embedding_request + else (None, self._rerank_input_to_text(data), data.get("tools"), data.get("tool_choice")) + if call_type in RERANK_API_CALL_TYPES + else (None, prompt, data.get("tools"), data.get("tool_choice")) + if call_type in TEXT_COMPLETION_API_CALL_TYPES + else (data.get("messages"), fallback_text, data.get("tools"), data.get("tool_choice")) + ) + messages, selected_text, countable_tools, countable_tool_choice = selected_inputs + + audio_token_estimate: Final = self._estimate_audio_content_tokens(messages) + countable_messages: Final = self._strip_audio_content_blocks(messages) if audio_token_estimate > 0 else messages try: - estimate = max( + estimate: Final = max( 0, int( token_counter( @@ -3245,7 +3254,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ), ) return estimate + audio_token_estimate - except Exception: # noqa: BLE001 - any tokenizer/model-resolution/transform failure degrades to the cheap estimate, never a 500 + except Exception: # noqa: BLE001 # tokenizer failures degrade to the cheap estimate if call_type in RERANK_API_CALL_TYPES and isinstance(selected_text, str): return max(0, len(selected_text) // DEFAULT_CHARS_PER_TOKEN) estimated_input_tokens, _ = self._estimate_input_and_output_tokens(data=data, call_type=call_type) @@ -3272,14 +3281,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ if not isinstance(data, dict): return - stash = claim_request_stash_for_data(data) - io_token_descriptors = [ # mutable-ok: reservation API requires descriptor lists + stash: Final = claim_request_stash_for_data(data) + io_token_descriptors: Final = [ # mutable-ok: reservation API requires descriptor lists d for d in descriptors if d["key"] in (PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY) ] if not io_token_descriptors: return - configured_otpm_limits = [ # mutable-ok: min calculation materializes validated limits + configured_otpm_limits: Final = [ # mutable-ok: min calculation materializes validated limits int(v) for d in io_token_descriptors if d["key"] == PROJECT_OTPM_DESCRIPTOR_KEY @@ -3290,8 +3299,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ] if v is not None ] - min_configured_otpm_limit = min(configured_otpm_limits) if configured_otpm_limits else None - configured_itpm_limits = [ # mutable-ok: min calculation materializes validated limits + min_configured_otpm_limit: Final = min(configured_otpm_limits) if configured_otpm_limits else None + configured_itpm_limits: Final = [ # mutable-ok: min calculation materializes validated limits int(v) for d in io_token_descriptors if d["key"] == PROJECT_ITPM_DESCRIPTOR_KEY @@ -3302,14 +3311,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ] if v is not None ] - min_configured_itpm_limit = min(configured_itpm_limits) if configured_itpm_limits else None + min_configured_itpm_limit: Final = min(configured_itpm_limits) if configured_itpm_limits else None - _, estimated_output_tokens = self._estimate_input_and_output_tokens( + _, raw_estimated_output_tokens = self._estimate_input_and_output_tokens( data=data, min_configured_tpm_limit=min_configured_otpm_limit, call_type=call_type, ) - estimated_input_tokens = ( + raw_estimated_input_tokens: Final = ( min_configured_itpm_limit if min_configured_itpm_limit is not None and ( @@ -3319,9 +3328,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) else self._estimate_precise_input_tokens(data=data, model=requested_model, call_type=call_type) ) - estimated_input_tokens = max(estimated_input_tokens, 1) - if not self._has_explicit_output_cap(data, call_type): - estimated_output_tokens = max(estimated_output_tokens, 1) + estimated_input_tokens: Final = max(raw_estimated_input_tokens, 1) + estimated_output_tokens: Final = ( + raw_estimated_output_tokens + if self._has_explicit_output_cap(data, call_type) + else max(raw_estimated_output_tokens, 1) + ) # Hard-cap generation length so an unbounded response can't overshoot # the OTPM budget before post-call reconciliation runs, mirroring the @@ -3353,7 +3365,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=user_api_key_dict.parent_otel_span, ) stash.reservation_released = True - acquisition = stash.parallel_slot + acquisition: Final = stash.parallel_slot if acquisition is not None: await self._release_parallel_request_slots( acquisition=acquisition, @@ -3367,9 +3379,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) if itpm_reserved > 0: - itpm_scopes = [ # mutable-ok: request stash freezes the collected scopes + itpm_scopes: Final = tuple( (d["key"], d["value"]) for d in io_token_descriptors if d["key"] == PROJECT_ITPM_DESCRIPTOR_KEY - ] + ) stash.itpm_reserved_tokens = itpm_reserved stash.itpm_reserved_scopes = frozenset(itpm_scopes) stash.itpm_reserved_window_identities = frozenset( @@ -3378,9 +3390,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if "model_per_project_itpm" in counter_key ) if otpm_reserved > 0: - otpm_scopes = [ # mutable-ok: request stash freezes the collected scopes + otpm_scopes: Final = tuple( (d["key"], d["value"]) for d in io_token_descriptors if d["key"] == PROJECT_OTPM_DESCRIPTOR_KEY - ] + ) stash.otpm_reserved_tokens = otpm_reserved stash.otpm_reserved_scopes = frozenset(otpm_scopes) stash.otpm_reserved_window_identities = frozenset( @@ -3473,7 +3485,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): requested_model=requested_model, descriptors=descriptors, ) - self._add_project_io_token_rate_limit_descriptors_from_metadata( + self.add_project_io_token_rate_limit_descriptors_from_metadata( user_api_key_dict=user_api_key_dict, requested_model=requested_model, descriptors=descriptors, @@ -3546,8 +3558,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # limit. Stays empty/0 whenever no combined-TPM reservation was # made (or it was over limit, in which case execution never # reaches the ITPM/OTPM block -- `_handle_rate_limit_error` raises). - tpm_reservation_scopes: Sequence[tuple[str, str]] = () - tpm_reservation_amount = 0 + tpm_reservation_scopes: Sequence[tuple[str, str]] = () # rebind-ok: set after successful reservation + tpm_reservation_amount = 0 # rebind-ok: set after successful reservation if has_tpm_limits and self.tpm_reservation_enabled: min_configured_tpm_limit: Final = min(configured_tpm_limits) @@ -3633,8 +3645,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) is not None ) - tpm_reservation_scopes = tuple(stash.reserved_scopes) - tpm_reservation_amount = estimated_tokens + tpm_reservation_scopes = tuple( # rebind-ok: record successful reservation scopes + stash.reserved_scopes + ) + tpm_reservation_amount = estimated_tokens # rebind-ok: record successful reservation amount # Merge TPM statuses into the stored rate-limit response # so x-ratelimit-{key}-remaining-tokens / -limit-tokens @@ -3648,7 +3662,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): verbose_proxy_logger.debug( "TPM tokens reserved: %s for model %s", estimated_tokens, requested_model ) - await self._reserve_project_io_tokens_or_raise( descriptors=descriptors, data=data, @@ -3734,7 +3747,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return total_tokens @staticmethod - def _aggregate_only_total_tokens(usage: Usage | dict | None) -> int: + def _aggregate_only_total_tokens(usage: Usage | ResponseAPIUsage | Mapping[str, object] | None) -> int: """Total for usage that carries no input/output split, else 0. A source that can only report one number for the whole request (a @@ -3744,24 +3757,43 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): uncharged, which is how pass-through traffic slips past a TPM limit it is supposed to share. """ - if isinstance(usage, Usage): - prompt_tokens, completion_tokens, total_tokens = ( - usage.prompt_tokens or 0, - usage.completion_tokens or 0, - usage.total_tokens or 0, - ) - elif isinstance(usage, dict): - prompt_tokens, completion_tokens, total_tokens = ( - usage.get("prompt_tokens") or 0, - usage.get("completion_tokens") or 0, + if usage is None: + return 0 + token_counts: Final = ( + (usage.prompt_tokens or 0, usage.completion_tokens or 0, usage.total_tokens or 0) + if isinstance(usage, Usage) + else (usage.input_tokens or 0, usage.output_tokens or 0, usage.total_tokens or 0) + if isinstance(usage, ResponseAPIUsage) + else ( + usage.get("prompt_tokens") or usage.get("input_tokens") or 0, + usage.get("completion_tokens") or usage.get("output_tokens") or 0, usage.get("total_tokens") or 0, ) - else: - return 0 - if prompt_tokens or completion_tokens: + ) + prompt_tokens, completion_tokens, total_tokens = token_counts + if prompt_tokens or completion_tokens or not isinstance(total_tokens, int): return 0 return total_tokens + @staticmethod + def _response_usage( + response_obj: object, + ) -> Usage | ResponseAPIUsage | Mapping[str, object] | None: + if isinstance(response_obj, (Usage, ResponseAPIUsage)): + return response_obj + if isinstance( + response_obj, + (ModelResponse, EmbeddingResponse, TextCompletionResponse, BaseLiteLLMOpenAIResponseObject), + ): + usage: Final = getattr(response_obj, "usage", None) + return usage if isinstance(usage, (Usage, ResponseAPIUsage, dict)) else None + if isinstance(response_obj, dict): + nested_usage: Final = response_obj.get("usage") + if isinstance(nested_usage, (Usage, ResponseAPIUsage, dict)): + return nested_usage + return response_obj + return None + async def _execute_token_increment_script( self, pipeline_operations: list["RedisPipelineIncrementOperation"], @@ -3884,15 +3916,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if self.window_guarded_token_increment_script is not None: try: await self.window_guarded_token_increment_script( - keys=[window_key, operation["key"]], - args=[ + keys=[ # mutable-ok: Redis script interface requires a key list + window_key, + operation["key"], + ], + args=[ # mutable-ok: Redis script interface requires an argument list expected_window_start, operation["increment_value"], operation["ttl"] or 0, ], ) continue - except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the plain increment fallback, never a 500 + except Exception as e: # noqa: BLE001 # Redis failures use the plain increment fallback verbose_proxy_logger.warning( "Window-guarded token adjustment failed for %s: %s", operation["key"], @@ -3978,16 +4013,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if not isinstance(response_obj, RerankResponse) or response_obj.meta is None: return None - rerank_tokens = response_obj.meta.get("tokens") # pyright: ignore[reportUnknownMemberType] # TypedDict's optional generic metadata widens get overloads + rerank_tokens: Final = response_obj.meta.get("tokens") # pyright: ignore[reportUnknownMemberType] # TypedDict's optional generic metadata widens get overloads if rerank_tokens is not None: - input_tokens = rerank_tokens.get("input_tokens") or 0 # pyright: ignore[reportUnknownMemberType] # token fields are typed integers despite the generic get overload - output_tokens = rerank_tokens.get("output_tokens") or 0 # pyright: ignore[reportUnknownMemberType] # token fields are typed integers despite the generic get overload + input_tokens: Final = rerank_tokens.get("input_tokens") or 0 # pyright: ignore[reportUnknownMemberType] # token fields are typed integers despite the generic get overload + output_tokens: Final = rerank_tokens.get("output_tokens") or 0 # pyright: ignore[reportUnknownMemberType] # token fields are typed integers despite the generic get overload if input_tokens or output_tokens: return max(0, input_tokens), max(0, output_tokens), True - billed_units = response_obj.meta.get("billed_units") # pyright: ignore[reportUnknownMemberType] # TypedDict's optional generic metadata widens get overloads + billed_units: Final = response_obj.meta.get("billed_units") # pyright: ignore[reportUnknownMemberType] # TypedDict's optional generic metadata widens get overloads if billed_units is not None: - total_tokens = billed_units.get("total_tokens") or 0 # pyright: ignore[reportUnknownMemberType] # billed total is a typed integer despite the generic get overload + total_tokens: Final = billed_units.get("total_tokens") or 0 # pyright: ignore[reportUnknownMemberType] # billed total is a typed integer despite the generic get overload if total_tokens: return max(0, total_tokens), 0, True return None @@ -4003,68 +4038,57 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): but they're untouched everywhere else (cost/usage logging still sees the full prompt token count). """ - rerank_usage = self._resolve_rerank_token_usage(response_obj) + rerank_usage: Final = self._resolve_rerank_token_usage(response_obj) if rerank_usage is not None: return rerank_usage - usage: object | None = None - if isinstance(response_obj, (Usage, ResponseAPIUsage)): - usage = response_obj - elif isinstance( - response_obj, - (ModelResponse, EmbeddingResponse, TextCompletionResponse, BaseLiteLLMOpenAIResponseObject), - ): - usage = getattr(response_obj, "usage", None) - elif isinstance(response_obj, dict): - usage = response_obj.get("usage") - if usage is None and any( - key in response_obj - for key in ( - "prompt_tokens", - "completion_tokens", - "input_tokens", - "output_tokens", - ) - ): - usage = response_obj + usage: Final = self._response_usage(response_obj) if isinstance(usage, Usage): - prompt_tokens = usage.prompt_tokens or 0 - completion_tokens = usage.completion_tokens or 0 - cached_tokens = 0 - if usage.prompt_tokens_details is not None: - cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0 - elif isinstance(usage, ResponseAPIUsage): - # Responses API usage uses input_tokens/output_tokens instead of - # prompt_tokens/completion_tokens. - prompt_tokens = usage.input_tokens or 0 - completion_tokens = usage.output_tokens or 0 - cached_tokens = 0 - if usage.input_tokens_details is not None: - cached_tokens = usage.input_tokens_details.cached_tokens or 0 - elif isinstance(usage, dict): - prompt_tokens = usage.get("prompt_tokens") or usage.get("input_tokens") or 0 - completion_tokens = usage.get("completion_tokens") or usage.get("output_tokens") or 0 - prompt_details = ( - usage.get("prompt_tokens_details") - or usage.get("input_tokens_details") - or {} # mutable-ok: usage details are optional mappings + prompt_tokens: Final = usage.prompt_tokens or 0 + completion_tokens: Final = usage.completion_tokens or 0 + cached_tokens: Final = ( + getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0 + if usage.prompt_tokens_details is not None + else 0 ) - cached_tokens = ( - (prompt_details.get("cached_tokens", 0) or 0) if isinstance(prompt_details, dict) else 0 - ) or (usage.get("cache_read_input_tokens") or 0) - else: - return 0, 0, False + if prompt_tokens == 0 and completion_tokens == 0: + return 0, 0, False + return max(0, prompt_tokens - cached_tokens), completion_tokens, True - if prompt_tokens == 0 and completion_tokens == 0: - return 0, 0, False - return max(0, prompt_tokens - cached_tokens), completion_tokens, True + if isinstance(usage, ResponseAPIUsage): + response_input_tokens: Final = usage.input_tokens or 0 + response_output_tokens: Final = usage.output_tokens or 0 + response_cached_tokens: Final = ( + usage.input_tokens_details.cached_tokens or 0 if usage.input_tokens_details is not None else 0 + ) + if response_input_tokens == 0 and response_output_tokens == 0: + return 0, 0, False + return max(0, response_input_tokens - response_cached_tokens), response_output_tokens, True + + if isinstance(usage, Mapping): + raw_prompt_tokens: Final = usage.get("prompt_tokens") or usage.get("input_tokens") or 0 + raw_completion_tokens: Final = usage.get("completion_tokens") or usage.get("output_tokens") or 0 + mapped_prompt_tokens: Final = raw_prompt_tokens if isinstance(raw_prompt_tokens, int) else 0 + mapped_completion_tokens: Final = raw_completion_tokens if isinstance(raw_completion_tokens, int) else 0 + prompt_details: Final = usage.get("prompt_tokens_details") or usage.get("input_tokens_details") + raw_cached_tokens: Final = ( + (prompt_details.get("cached_tokens", 0) if isinstance(prompt_details, dict) else 0) + or usage.get("cache_read_input_tokens") + or 0 + ) + mapped_cached_tokens: Final = raw_cached_tokens if isinstance(raw_cached_tokens, int) else 0 + if mapped_prompt_tokens == 0 and mapped_completion_tokens == 0: + return 0, 0, False + return max(0, mapped_prompt_tokens - mapped_cached_tokens), mapped_completion_tokens, True + + return 0, 0, False def _build_io_token_reservation_ops( self, kwargs: object, response_obj: object, - ) -> list[RedisPipelineIncrementOperation] | tuple[ReservationAwareIncrementOperation, ...]: + ) -> Sequence[RedisPipelineIncrementOperation]: """ Reconcile project ITPM/OTPM reservations to actual usage on success: ITPM to billable input tokens, OTPM to actual completion tokens. @@ -4075,25 +4099,33 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ if not isinstance(kwargs, dict): return () - stash = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) + stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) if stash is None: return () - itpm_reserved = stash.itpm_reserved_tokens - otpm_reserved = stash.otpm_reserved_tokens + itpm_reserved: Final = stash.itpm_reserved_tokens + otpm_reserved: Final = stash.otpm_reserved_tokens if itpm_reserved <= 0 and otpm_reserved <= 0: return () - billable_input, completion_tokens, usage_resolved = self._resolve_io_token_reconcile_usage(response_obj) - if not usage_resolved: - billable_input, completion_tokens, usage_resolved = self._resolve_io_token_reconcile_usage( - kwargs.get("combined_usage_object") - ) - if not usage_resolved: - if not stash.reservation_released: - return () - billable_input = itpm_reserved - completion_tokens = otpm_reserved + response_usage: Final = self._resolve_io_token_reconcile_usage(response_obj) + combined_usage: Final = self._resolve_io_token_reconcile_usage(kwargs.get("combined_usage_object")) + aggregate_total: Final = self._aggregate_only_total_tokens( + self._response_usage(response_obj) + ) or self._aggregate_only_total_tokens(self._response_usage(kwargs.get("combined_usage_object"))) + + if not response_usage[2] and not combined_usage[2] and aggregate_total <= 0 and not stash.reservation_released: + return () + resolved_usage: Final = ( + response_usage + if response_usage[2] + else combined_usage + if combined_usage[2] + else (aggregate_total, aggregate_total, True) + if aggregate_total > 0 + else (itpm_reserved, otpm_reserved, False) + ) + billable_input, completion_tokens, _ = resolved_usage if stash.reservation_released or ( not stash.itpm_reserved_window_identities and not stash.otpm_reserved_window_identities @@ -4110,24 +4142,28 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): reserved_tokens=0 if stash.reservation_released else otpm_reserved, ) - itpm_ops: Sequence[ReservationAwareIncrementOperation] = () - if itpm_reserved > 0: - itpm_ops = self._build_project_reservation_ops( + itpm_ops: Final[Sequence[ReservationAwareIncrementOperation]] = ( + self._build_project_reservation_ops( targets=tuple(stash.itpm_reserved_scopes), reserved_scopes=frozenset() if stash.reservation_released else stash.itpm_reserved_scopes, actual_tokens=billable_input, reserved_tokens=itpm_reserved, reservation_window_identities=stash.itpm_reserved_window_identities, ) - otpm_ops: Sequence[ReservationAwareIncrementOperation] = () - if otpm_reserved > 0: - otpm_ops = self._build_project_reservation_ops( + if itpm_reserved > 0 + else () + ) + otpm_ops: Final[Sequence[ReservationAwareIncrementOperation]] = ( + self._build_project_reservation_ops( targets=tuple(stash.otpm_reserved_scopes), reserved_scopes=frozenset() if stash.reservation_released else stash.otpm_reserved_scopes, actual_tokens=completion_tokens, reserved_tokens=otpm_reserved, reservation_window_identities=stash.otpm_reserved_window_identities, ) + if otpm_reserved > 0 + else () + ) return tuple((*itpm_ops, *otpm_ops)) def _collect_tpm_scope_targets( @@ -4522,14 +4558,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # already released it (proxy-level rejection that also bubbles up # here as an LLM-error callback). max_parallel_requests is its # own counter and is always decremented per call. - if stash is None or stash.reservation_released: - reserved_tokens = 0 - itpm_reserved = 0 - otpm_reserved = 0 - else: - reserved_tokens = stash.reserved_tokens - itpm_reserved = stash.itpm_reserved_tokens - otpm_reserved = stash.otpm_reserved_tokens + reserved_tokens, itpm_reserved, otpm_reserved = ( + (0, 0, 0) + if stash is None or stash.reservation_released + else (stash.reserved_tokens, stash.itpm_reserved_tokens, stash.otpm_reserved_tokens) + ) if stash is not None and reserved_tokens > 0: verbose_proxy_logger.debug("Releasing reserved TPM tokens on failure: %s", reserved_tokens) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index af561c9092b..e678fba2852 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -50,6 +50,16 @@ if TYPE_CHECKING: ) +class ProjectQuotaCallback(Protocol): + async def enforce_project_io_token_quota_for_frame( + self, + user_api_key_dict: UserAPIKeyAuth | None, + requested_model: str | None, + estimated_input_tokens: int, + estimated_output_tokens: int, + ) -> None: ... + + @lru_cache(maxsize=1) def _get_openai_response_types(): from litellm.types.llms import openai as openai_types @@ -1345,11 +1355,19 @@ def _extract_frame_quota_estimate_inputs(msg_obj: Mapping[str, object]) -> tuple """ nested: Final = msg_obj.get("response") params: Final[Mapping[str, object]] = ( - nested if _is_json_object(nested) and nested else {k: v for k, v in msg_obj.items() if k != "type"} + nested + if _is_json_object(nested) and nested + else MappingProxyType( # mutable-ok: immediately frozen filtered frame + {k: v for k, v in msg_obj.items() if k != "type"} + ) ) - text_parts: list[str] = [] # mutable-ok: local accumulator built in one pass, not shared - - def _collect_text(value: object) -> None: + text_parts: Final[list[str]] = [] # mutable-ok: local accumulator built in one pass, not shared + pending: Final[list[object]] = [ # mutable-ok: explicit worklist avoids recursion + params.get("input"), + params.get("instructions"), + ] + while pending: + value = pending.pop() if isinstance(value, str): text_parts.append(value) elif _is_json_array(value): @@ -1357,20 +1375,17 @@ def _extract_frame_quota_estimate_inputs(msg_obj: Mapping[str, object]) -> tuple if isinstance(item, str): text_parts.append(item) elif _is_json_object(item): - _collect_text(item.get("content")) - _collect_text(item.get("text")) - - _collect_text(params.get("input")) - _collect_text(params.get("instructions")) + pending.append(item.get("content")) + pending.append(item.get("text")) total_chars: Final = sum(len(part) for part in text_parts) estimated_input_tokens: Final = max(1, total_chars // _FRAME_CHARS_PER_TOKEN_ESTIMATE) if total_chars else 0 - max_output_tokens = params.get("max_output_tokens") + max_output_tokens: Final = params.get("max_output_tokens") return estimated_input_tokens, max_output_tokens if isinstance(max_output_tokens, int) else None async def _enforce_frame_project_quota( - quota_callbacks: Sequence[Any], + quota_callbacks: Sequence[ProjectQuotaCallback], user_api_key_dict: UserAPIKeyAuth | None, model: str | None, raw_message: str, @@ -1387,7 +1402,7 @@ async def _enforce_frame_project_quota( if not _is_json_object(msg_obj) or msg_obj.get("type") != "response.create": return estimated_input_tokens, explicit_max_output_tokens = _extract_frame_quota_estimate_inputs(msg_obj) - estimated_output_tokens = ( + estimated_output_tokens: Final = ( explicit_max_output_tokens if explicit_max_output_tokens is not None else _FRAME_NO_MAX_OUTPUT_TOKENS_FLOOR ) for callback in quota_callbacks: @@ -1433,7 +1448,7 @@ class ResponsesWebSocketStreaming: first_message: str | None = None, guardrail_callbacks: list[Any] | None = None, output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None, - quota_callbacks: list[Any] | None = None, + quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, authorized_model: str | None = None, ): self.websocket = websocket @@ -1446,7 +1461,7 @@ class ResponsesWebSocketStreaming: self.first_message = first_message self.guardrail_callbacks: list[Any] = guardrail_callbacks or [] self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or [] - self.quota_callbacks: list[Any] = quota_callbacks or [] + self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else () # Model name authorized at connection time; enforced on every # response.create frame to prevent deployment-substitution attacks. self.authorized_model: str | None = authorized_model @@ -1870,9 +1885,17 @@ class ResponsesWebSocketStreaming: except RateLimitError as e: try: await self.websocket.send_text( - json.dumps({"type": "error", "error": {"type": "rate_limit_exceeded", "message": str(e)}}) + json.dumps( # mutable-ok: WebSocket wire payload requires JSON objects + { # mutable-ok: WebSocket wire payload requires JSON objects + "type": "error", + "error": { # mutable-ok: nested WebSocket error object + "type": "rate_limit_exceeded", + "message": str(e), + }, + } + ) ) - except Exception: # noqa: BLE001, S110 - best-effort notification, client may already be gone + except Exception: # noqa: BLE001, S110 # client may already be gone pass return False return True @@ -1969,7 +1992,7 @@ class ManagedResponsesWebSocketHandler: timeout: float | None = None, custom_llm_provider: str | None = None, first_message: str | None = None, - quota_callbacks: list[Any] | None = None, + quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, **kwargs: object, ) -> None: self.websocket = websocket @@ -1986,7 +2009,7 @@ class ManagedResponsesWebSocketHandler: self.custom_llm_provider = custom_llm_provider self._connection_provider = self._resolve_provider(model) or custom_llm_provider self.first_message = first_message - self.quota_callbacks: list[Any] = quota_callbacks or [] + self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else () # Carry through safe pass-through kwargs (e.g. extra_headers) self.extra_kwargs: dict[str, object] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS} # In-memory session history: response_id → full accumulated message list. diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py index 66fc84ab1e0..e55185cfa67 100644 --- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py +++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py @@ -3582,18 +3582,24 @@ async def test_streaming_combined_usage_reconciles_project_io_reservations( assert [operation["increment_value"] for operation in otpm_adjustments] == [-45] -def test_aggregate_only_combined_usage_keeps_project_io_reservations(rate_limiter): +def test_aggregate_only_combined_usage_reconciles_project_io_reservations(rate_limiter): handler, _cache = rate_limiter stash = get_or_create_request_stash() stash.itpm_reserved_tokens = 100 stash.itpm_reserved_scopes = frozenset( {(PROJECT_ITPM_DESCRIPTOR_KEY, "project:model")} ) + stash.otpm_reserved_tokens = 80 + stash.otpm_reserved_scopes = frozenset( + {(PROJECT_OTPM_DESCRIPTOR_KEY, "project:model")} + ) kwargs = { "combined_usage_object": Usage(total_tokens=55), } - assert handler._build_io_token_reservation_ops(kwargs, object()) == () + operations = handler._build_io_token_reservation_ops(kwargs, object()) + + assert [operation["increment_value"] for operation in operations] == [-45, -25] def test_raw_split_usage_dict_reconciles_project_io_tokens(rate_limiter): diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 894d99c92e0..21dc0bc2791 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22941 + "limit": 22936 }, "LIT002": { - "limit": 27139 + "limit": 27133 }, "LIT003": { "limit": 269 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16716 + "limit": 16701 }, "LIT011": { - "limit": 5596 + "limit": 5595 }, "LIT012": { "limit": 4519 From 658c67c1523a54d8d2e51ad848af4df3edd40b40 Mon Sep 17 00:00:00 2001 From: Shifat Islam Santo Date: Fri, 14 Aug 2026 14:19:48 -0500 Subject: [PATCH 037/538] fix: preserve prompt cache for mid-conversation system on unflagged Claude models --- .../messages/transformation.py | 56 +++++++++++++------ ...st_messages_mid_conversation_system_e2e.py | 6 +- ...onversation_system_native_providers_e2e.py | 16 +++--- ...azure_anthropic_messages_transformation.py | 20 +++++-- .../test_anthropic_claude3_transformation.py | 55 +++++++++++++----- ...artner_models_anthropic_messages_config.py | 29 +++++++++- 6 files changed, 130 insertions(+), 52 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 4d3354c58b7..1e016cffb0e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -159,8 +159,22 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): def _is_system_role_message(message: Any) -> bool: return isinstance(message, dict) and message.get("role") == "system" + _CONVERTED_SYSTEM_NOTE: Final = ( + "Operator note (not from the user): the following was originally a mid-conversation system-role reminder." + ) + + def _system_role_message_as_user(self, message: dict) -> dict: + return { + **message, + "role": "user", + "content": [ + {"type": "text", "text": self._CONVERTED_SYSTEM_NOTE}, + *self._as_system_content_blocks(message.get("content")), + ], + } + def _normalize_system_role_messages(self, anthropic_messages_request: dict, model: str) -> None: - """Move ``role: "system"`` entries out of ``messages`` per the Anthropic + """Normalize ``role: "system"`` entries in ``messages`` per the Anthropic ``/v1/messages`` contract, which the first-party API, Bedrock Invoke, Vertex, and Azure Foundry all enforce identically. @@ -173,9 +187,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): stay: hoisting one mutates the ``system`` prefix and invalidates the prompt cache for the whole message history. Older Claude models reject the role in every position ("role 'system' is not supported on this model"), - so without the flag every system entry is hoisted to keep the request from - 400-ing. Billing-header system blocks are stripped from the top-level - ``system`` field regardless of whether anything was hoisted. + so without the flag a mid-conversation entry is converted to a user turn + in place (prefixed with an operator note) rather than hoisted: hoisting + would mutate the ``system`` prefix and likewise collapse the cache, while + the in-place conversion keeps everything before it byte-identical. + Billing-header system blocks are stripped from the top-level ``system`` + field regardless of whether anything was hoisted. Subclasses whose upstream rejects the role opt in by calling this from their ``transform_anthropic_messages_request``; the first-party Anthropic @@ -185,21 +202,24 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): messages: Final = anthropic_messages_request.get("messages") if not isinstance(messages, list): return - if _supports_factory( - model=model, - custom_llm_provider=self.custom_llm_provider, - key="supports_mid_conversation_system", - ): - leading_count: Final = next( - (i for i, m in enumerate(messages) if not self._is_system_role_message(m)), - len(messages), + leading_count: Final = next( + (i for i, m in enumerate(messages) if not self._is_system_role_message(m)), + len(messages), + ) + hoisted: Final = messages[:leading_count] + remaining: Final = ( + messages[leading_count:] + if _supports_factory( + model=model, + custom_llm_provider=self.custom_llm_provider, + key="supports_mid_conversation_system", ) - hoisted = messages[:leading_count] - remaining = messages[leading_count:] - else: - hoisted = [m for m in messages if self._is_system_role_message(m)] - remaining = [m for m in messages if not self._is_system_role_message(m)] - if hoisted: + else [ + self._system_role_message_as_user(m) if self._is_system_role_message(m) else m + for m in messages[leading_count:] + ] + ) + if hoisted or remaining != messages: anthropic_messages_request["messages"] = remaining system_content: Final = [ block diff --git a/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py b/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py index fff2109b0cf..e1d836108e0 100644 --- a/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py +++ b/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py @@ -6,7 +6,7 @@ and the 5 family) must keep a mid-conversation system reminder in place inside ``messages`` so the top-level ``system`` prefix stays byte-identical and the prompt cache written on turn one is read back in full on turn two. Models without the flag (Claude 4.7 and older) reject the role inside ``messages`` -outright, so the proxy must hoist the reminder into the top-level ``system`` +outright, so the proxy must convert the reminder to a user turn in place field and the call must still return a completion instead of a provider 400. The conversation shape mirrors what Claude Code sends mid-session: a cached @@ -201,7 +201,7 @@ class TestBedrockInvokeMidConversationSystem: "llm.messages.bedrock_invoke.mid_conversation_system.nonstream.works", exercised_on=[], ) - def test_unflagged_model_hoists_system_reminder_and_succeeds( + def test_unflagged_model_converts_system_reminder_and_succeeds( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: model = _register_invoke_deployment( @@ -227,5 +227,5 @@ class TestBedrockInvokeMidConversationSystem: assert completion.text.strip(), ( f"{model}: conversation with a mid-conversation system reminder " f"returned no text; the reminder was forwarded in place to a model " - f"that rejects role 'system' inside messages instead of being hoisted" + f"that rejects role 'system' inside messages instead of being converted to a user turn" ) diff --git a/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py b/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py index 35ed3dc881a..6eae0d66c59 100644 --- a/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py +++ b/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py @@ -7,13 +7,13 @@ accepted in place on Claude 4.8+/5 (200) but rejected on Claude 4.7 and older ("role 'system' is not supported on this model", 400), and a *leading* system entry is rejected on every model ("messages.0: use the top-level 'system' parameter"). This mirrors Bedrock Invoke (PRs #32578/#32831/#32882); the same -model-gated hoist now runs for these two providers (customer RCA gap #3). +model-gated normalization now runs for these two providers (customer RCA gap #3). Flagged models (``supports_mid_conversation_system`` in the cost map: Claude 4.8+ and the 5 family) must keep the reminder in ``messages`` so the top-level ``system`` prefix stays byte-identical and the prompt cache written on turn one is read back in full on turn two. Unflagged models (Claude 4.7 and older) must -have the reminder hoisted into the top-level ``system`` field so the call +have the reminder converted to a user turn in place so the call returns a completion instead of a provider 400. The conversation shape mirrors what Claude Code sends mid-session: a cached @@ -206,7 +206,7 @@ def _assert_flagged_model_keeps_cache( ) -def _assert_unflagged_model_hoists_and_succeeds( +def _assert_unflagged_model_converts_and_succeeds( client: EndpointsClient, resources: ResourceManager, params: LiteLLMParamsBody ) -> None: model = _register_deployment(client, resources, params) @@ -228,7 +228,7 @@ def _assert_unflagged_model_hoists_and_succeeds( assert completion.text.strip(), ( f"{model}: conversation with a mid-conversation system reminder returned " f"no text; the reminder was forwarded in place to a model that rejects " - f"role 'system' inside messages instead of being hoisted" + f"role 'system' inside messages instead of being converted to a user turn" ) @@ -250,10 +250,10 @@ class TestAzureFoundryMidConversationSystem: "llm.messages.azure_foundry.mid_conversation_system.nonstream.works", exercised_on=[], ) - def test_unflagged_model_hoists_system_reminder_and_succeeds( + def test_unflagged_model_converts_system_reminder_and_succeeds( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: - _assert_unflagged_model_hoists_and_succeeds( + _assert_unflagged_model_converts_and_succeeds( endpoints_client, resources, _azure_params(self.UNFLAGGED_MODEL) ) @@ -276,9 +276,9 @@ class TestVertexMidConversationSystem: "llm.messages.vertex.mid_conversation_system.nonstream.works", exercised_on=[], ) - def test_unflagged_model_hoists_system_reminder_and_succeeds( + def test_unflagged_model_converts_system_reminder_and_succeeds( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: - _assert_unflagged_model_hoists_and_succeeds( + _assert_unflagged_model_converts_and_succeeds( endpoints_client, resources, _vertex_params(self.UNFLAGGED_MODEL) ) diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py index f6446b43fab..add1e9967db 100644 --- a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py @@ -425,7 +425,7 @@ class TestAzureAnthropicMidConversationSystem: {"type": "text", "text": "Cite sources."}, ] - def test_unsupported_model_hoists_mid_conversation_system(self, local_model_cost_map): + def test_unsupported_model_converts_mid_conversation_system_in_place(self, local_model_cost_map): messages = [ {"role": "user", "content": "read the file"}, {"role": "system", "content": "[Truncated: PARTIAL view of big1.txt]"}, @@ -437,13 +437,23 @@ class TestAzureAnthropicMidConversationSystem: ) assert result["messages"] == [ {"role": "user", "content": "read the file"}, + { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + "Operator note (not from the user): the following was " + "originally a mid-conversation system-role reminder." + ), + }, + {"type": "text", "text": "[Truncated: PARTIAL view of big1.txt]"}, + ], + }, {"role": "assistant", "content": "reading"}, {"role": "user", "content": "continue"}, ] - assert result["system"] == [ - {"type": "text", "text": "Base."}, - {"type": "text", "text": "[Truncated: PARTIAL view of big1.txt]"}, - ] + assert result["system"] == [{"type": "text", "text": "Base."}] def test_azure_claude_4_8_plus_cost_map_entries_carry_mid_conversation_system_flag(): diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index fd66667af64..84df706b722 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -2125,13 +2125,14 @@ def test_bedrock_invoke_transform_hoists_only_leading_system_run(local_model_cos ] -def test_bedrock_invoke_transform_hoists_mid_conversation_system_for_older_claude(local_model_cost_map): - """Regression test for Claude Code 400s on pre-Opus-4.8 Bedrock models: - Invoke rejects ``role: "system"`` in every position on Opus 4.7, Sonnet 4.6, - Haiku 4.5, etc. ("role 'system' is not supported on this model"), so on - models without ``supports_mid_conversation_system`` every system entry must - be hoisted into the top-level ``system`` field, mid-conversation ones - included.""" +def test_bedrock_invoke_transform_converts_mid_conversation_system_for_older_claude(local_model_cost_map): + """Invoke rejects ``role: "system"`` in every position on Opus 4.7, Sonnet + 4.6, Haiku 4.5, etc. ("role 'system' is not supported on this model"), but + hoisting a mid-conversation reminder into the top-level ``system`` field + mutates the cached prefix and reprocesses the whole history. On models + without ``supports_mid_conversation_system`` the reminder is converted to a + user turn in place instead: the request stays valid and a cache breakpoint + before the reminder still hits.""" from litellm.types.router import GenericLiteLLMParams cfg = AmazonAnthropicClaudeMessagesConfig() @@ -2156,19 +2157,30 @@ def test_bedrock_invoke_transform_hoists_mid_conversation_system_for_older_claud assert result["messages"] == [ {"role": "user", "content": "read the file"}, + { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + "Operator note (not from the user): the following was " + "originally a mid-conversation system-role reminder." + ), + }, + {"type": "text", "text": "[Truncated: PARTIAL view of big1.txt]"}, + ], + }, {"role": "assistant", "content": "reading"}, {"role": "user", "content": "continue"}, ] - assert result["system"] == [ - {"type": "text", "text": "Base."}, - {"type": "text", "text": "[Truncated: PARTIAL view of big1.txt]"}, - ] + assert result["system"] == [{"type": "text", "text": "Base."}] -def test_bedrock_invoke_transform_hoists_all_system_for_unmapped_model(local_model_cost_map): +def test_bedrock_invoke_transform_converts_system_for_unmapped_model(local_model_cost_map): """A model with no cost-map entry and no fallback-generalization rule gets - the hoist-everything behavior: the safe default is a mutated cache prefix, - never a provider 400 from forwarding a role the model may not accept.""" + the unsupported-model treatment: the safe default converts the reminder to + a user turn in place, never a provider 400 from forwarding a role the model + may not accept, and never a mutated cache prefix.""" from litellm.types.router import GenericLiteLLMParams cfg = AmazonAnthropicClaudeMessagesConfig() @@ -2189,10 +2201,23 @@ def test_bedrock_invoke_transform_hoists_all_system_for_unmapped_model(local_mod assert result["messages"] == [ {"role": "user", "content": "hi"}, + { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + "Operator note (not from the user): the following was " + "originally a mid-conversation system-role reminder." + ), + }, + {"type": "text", "text": "mid-conversation reminder"}, + ], + }, {"role": "assistant", "content": "hello"}, {"role": "user", "content": "continue"}, ] - assert result["system"] == [{"type": "text", "text": "mid-conversation reminder"}] + assert "system" not in result def test_bedrock_invoke_transform_keeps_system_in_place_for_unmapped_future_claude(local_model_cost_map): diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index 292bddf1274..ef7db337a74 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -622,7 +622,7 @@ class TestVertexAnthropicMidConversationSystem: {"type": "text", "text": "Cite sources."}, ] - def test_unsupported_model_hoists_mid_conversation_system(self, local_model_cost_map): + def test_unsupported_model_converts_mid_conversation_system_in_place(self, local_model_cost_map): messages = [ {"role": "user", "content": "read the file"}, {"role": "system", "content": "[Truncated: PARTIAL view of big1.txt]"}, @@ -634,12 +634,35 @@ class TestVertexAnthropicMidConversationSystem: ) assert result["messages"] == [ {"role": "user", "content": "read the file"}, + { + "role": "user", + "content": [ + { + "type": "text", + "text": ( + "Operator note (not from the user): the following was " + "originally a mid-conversation system-role reminder." + ), + }, + {"type": "text", "text": "[Truncated: PARTIAL view of big1.txt]"}, + ], + }, {"role": "assistant", "content": "reading"}, {"role": "user", "content": "continue"}, ] + assert result["system"] == [{"type": "text", "text": "Base."}] + + def test_unsupported_model_still_hoists_leading_system_run(self, local_model_cost_map): + messages = [ + {"role": "system", "content": "You are terse."}, + {"role": "system", "content": "Cite sources."}, + {"role": "user", "content": "hi"}, + ] + result = _vertex_transform("claude-sonnet-4-6", messages) + assert result["messages"] == [{"role": "user", "content": "hi"}] assert result["system"] == [ - {"type": "text", "text": "Base."}, - {"type": "text", "text": "[Truncated: PARTIAL view of big1.txt]"}, + {"type": "text", "text": "You are terse."}, + {"type": "text", "text": "Cite sources."}, ] From 1b5e50727c62cdec0c06ab7c1f54006ad8c8f8a9 Mon Sep 17 00:00:00 2001 From: Shifat Islam Santo Date: Fri, 14 Aug 2026 14:40:10 -0500 Subject: [PATCH 038/538] fix: reuse block builder for lint budget, assert unflagged cache e2e --- .../messages/transformation.py | 6 ++-- ...st_messages_mid_conversation_system_e2e.py | 30 +++++++++++++------ ...onversation_system_native_providers_e2e.py | 28 ++++++++++++----- 3 files changed, 43 insertions(+), 21 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 1e016cffb0e..84fb7a1f45e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -167,10 +167,8 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): return { **message, "role": "user", - "content": [ - {"type": "text", "text": self._CONVERTED_SYSTEM_NOTE}, - *self._as_system_content_blocks(message.get("content")), - ], + "content": self._as_system_content_blocks(self._CONVERTED_SYSTEM_NOTE) + + self._as_system_content_blocks(message.get("content")), } def _normalize_system_role_messages(self, anthropic_messages_request: dict, model: str) -> None: diff --git a/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py b/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py index e1d836108e0..bd59df959c4 100644 --- a/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py +++ b/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py @@ -208,24 +208,36 @@ class TestBedrockInvokeMidConversationSystem: endpoints_client, resources, UNFLAGGED_INVOKE_MODEL ) key = resources.key(models=[model]) + system_block = _cacheable_system_block(unique_marker()) - body = RichMessagesRequest( + primed = _prime_prompt_cache(endpoints_client, key, model, system_block) + + reminder_turn_body = RichMessagesRequest( model=model, - system=[TextBlock(text="You are terse.")], + system=[system_block], messages=[ - _user_turn(f"Say hi. Run {unique_marker()}."), + _user_turn(primed.first_user_text, cached=True), _system_reminder_turn(), - RichMessage(role="assistant", content=[TextBlock(text="Hi.")]), - _user_turn("Say bye."), + RichMessage(role="assistant", content=[TextBlock(text="OK.")]), + _user_turn("Reply with one word again.", cached=True), ], ) - completion = unwrap(_post_messages(endpoints_client, key, body)) + second = unwrap(_post_messages(endpoints_client, key, reminder_turn_body)) - assert completion.role == "assistant", ( - f"{model}: unexpected role {completion.role!r}" + assert second.role == "assistant", ( + f"{model}: unexpected role {second.role!r}" ) - assert completion.text.strip(), ( + assert second.text.strip(), ( f"{model}: conversation with a mid-conversation system reminder " f"returned no text; the reminder was forwarded in place to a model " f"that rejects role 'system' inside messages instead of being converted to a user turn" ) + assert second.usage.cache_read_input_tokens >= primed.full_prefix_tokens, ( + f"{model}: reminder turn read {second.usage.cache_read_input_tokens} " + f"cached tokens, expected at least the {primed.full_prefix_tokens} " + f"cached on turn one ({primed.prefix_read_tokens} system prefix + " + f"{primed.first_turn_creation_tokens} first user turn); the reminder " + f"was hoisted into the top-level system field instead of being " + f"converted to a user turn in place, mutating the cached prefix and " + f"re-billing the conversation at cache-write pricing" + ) diff --git a/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py b/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py index 6eae0d66c59..bf4e662e02e 100644 --- a/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py +++ b/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py @@ -211,25 +211,37 @@ def _assert_unflagged_model_converts_and_succeeds( ) -> None: model = _register_deployment(client, resources, params) key = resources.key(models=[model]) + system_block = _cacheable_system_block(unique_marker()) - body = RichMessagesRequest( + primed = _prime_prompt_cache(client, key, model, system_block) + + reminder_turn_body = RichMessagesRequest( model=model, - system=[TextBlock(text="You are terse.")], + system=[system_block], messages=[ - _user_turn(f"Say hi. Run {unique_marker()}."), + _user_turn(primed.first_user_text, cached=True), _system_reminder_turn(), - RichMessage(role="assistant", content=[TextBlock(text="Hi.")]), - _user_turn("Say bye."), + RichMessage(role="assistant", content=[TextBlock(text="OK.")]), + _user_turn("Reply with one word again.", cached=True), ], ) - completion = unwrap(_post_messages(client, key, body)) + second = unwrap(_post_messages(client, key, reminder_turn_body)) - assert completion.role == "assistant", f"{model}: unexpected role {completion.role!r}" - assert completion.text.strip(), ( + assert second.role == "assistant", f"{model}: unexpected role {second.role!r}" + assert second.text.strip(), ( f"{model}: conversation with a mid-conversation system reminder returned " f"no text; the reminder was forwarded in place to a model that rejects " f"role 'system' inside messages instead of being converted to a user turn" ) + assert second.usage.cache_read_input_tokens >= primed.full_prefix_tokens, ( + f"{model}: reminder turn read {second.usage.cache_read_input_tokens} cached " + f"tokens, expected at least the {primed.full_prefix_tokens} cached on turn " + f"one ({primed.prefix_read_tokens} system prefix + " + f"{primed.first_turn_creation_tokens} first user turn); the reminder was " + f"hoisted into the top-level system field instead of being converted to a " + f"user turn in place, mutating the cached prefix and re-billing the " + f"conversation at cache-write pricing" + ) class TestAzureFoundryMidConversationSystem: From 4cd1c81a2dcd905112a7bc388afb5ed1412b8aa7 Mon Sep 17 00:00:00 2001 From: Shifat Islam Santo Date: Fri, 14 Aug 2026 14:53:11 -0500 Subject: [PATCH 039/538] fix: add supports_mid_conversation_system to bare first-party Claude cost-map keys --- ...odel_prices_and_context_window_backup.json | 5 +++ model_prices_and_context_window.json | 5 +++ ...erimental_pass_through_messages_handler.py | 37 +++++++++++++++++++ 3 files changed, 47 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b288269b0a2..4c1b46f85ab 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -12106,6 +12106,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -12493,6 +12494,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -12528,6 +12530,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -12566,6 +12569,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -47684,6 +47688,7 @@ }, "source": "https://docs.claude.com/en/docs/about-claude/models/overview", "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b288269b0a2..4c1b46f85ab 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -12106,6 +12106,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -12493,6 +12494,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -12528,6 +12530,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -12566,6 +12569,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -47684,6 +47688,7 @@ }, "source": "https://docs.claude.com/en/docs/about-claude/models/overview", "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index f11324ca376..91f5023496a 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -960,3 +960,40 @@ def test_gate_passthrough_skipped_when_only_chat_completions_supported(monkeypat assert result == "translated" assert translation_calls["count"] == 1 assert "config" not in captured + + +def test_first_party_claude_4_8_plus_cost_map_entries_carry_mid_conversation_system_flag(): + """Regional and provider-prefixed Claude 4.8+/5 entries carry + ``supports_mid_conversation_system``, but the bare first-party keys + (``claude-opus-4-8``) that a plain ``custom_llm_provider="anthropic"`` + lookup resolves were missed, so that lookup reports the capability as + unset. Every mapped first-party entry the fallback rule matches must + carry the flag.""" + import json + import os + import re + + import litellm + + cost_map_path = os.path.join( + os.path.dirname(litellm.__file__), "model_prices_and_context_window_backup.json" + ) + with open(cost_map_path) as f: + cost_map = json.load(f) + rules = cost_map["fallback_generalizations"]["rules"] + rule_pattern = next( + (r["pattern"] for r in rules if r["name"] == "claude-mid-conversation-system"), + None, + ) + assert rule_pattern is not None, "claude-mid-conversation-system rule not found in fallback_generalizations" + pattern = re.compile(rule_pattern, re.IGNORECASE) + missing = [ + key + for key, info in cost_map.items() + if isinstance(info, dict) + and info.get("litellm_provider") == "anthropic" + and "claude" in key + and pattern.search(key) + and info.get("supports_mid_conversation_system") is not True + ] + assert missing == [] From ed33687422c544afa7ba6268294744b478026557 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 14 Aug 2026 17:11:05 -0700 Subject: [PATCH 040/538] feat(proxy): auto-suppress the no-Redis banner for confirmed single-worker deployments --- .../migration.sql | 9 ++ .../litellm_proxy_extras/schema.prisma | 11 ++ litellm/proxy/db/proxy_worker_heartbeat.py | 89 +++++++++++++++ .../health_endpoints/_health_endpoints.py | 19 +++- litellm/proxy/proxy_server.py | 32 +++++- litellm/proxy/schema.prisma | 11 ++ schema.prisma | 11 ++ .../proxy/db/test_proxy_worker_heartbeat.py | 81 ++++++++++++++ .../health_endpoints/test_health_endpoints.py | 105 +++++++++++++++--- .../components/NoRedisWarningBanner.test.tsx | 1 + .../src/components/NoRedisWarningBanner.tsx | 8 +- 11 files changed, 349 insertions(+), 28 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260814000000_add_proxy_worker_heartbeat/migration.sql create mode 100644 litellm/proxy/db/proxy_worker_heartbeat.py create mode 100644 tests/test_litellm/proxy/db/test_proxy_worker_heartbeat.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260814000000_add_proxy_worker_heartbeat/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260814000000_add_proxy_worker_heartbeat/migration.sql new file mode 100644 index 00000000000..0a5d9df8aaf --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260814000000_add_proxy_worker_heartbeat/migration.sql @@ -0,0 +1,9 @@ +-- CreateTable +CREATE TABLE "LiteLLM_ProxyWorkerHeartbeat" ( + "worker_id" TEXT NOT NULL, + "hostname" TEXT NOT NULL, + "started_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "last_heartbeat_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "LiteLLM_ProxyWorkerHeartbeat_pkey" PRIMARY KEY ("worker_id") +); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 79d778fb464..09efef813a7 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -945,6 +945,17 @@ model LiteLLM_DailyTagSpend { } +// One row per live proxy worker process. Workers upsert their row on a fixed +// heartbeat; counting rows with a recent heartbeat tells how many workers share +// this database, which lets the Admin UI hide its "no Redis" warning for +// deployments that are provably a single worker. +model LiteLLM_ProxyWorkerHeartbeat { + worker_id String @id + hostname String + started_at DateTime @default(now()) + last_heartbeat_at DateTime @default(now()) +} + // Track the status of cron jobs running. Only allow one pod to run the job at a time model LiteLLM_CronJob { cronjob_id String @id @default(cuid()) // Unique ID for the record diff --git a/litellm/proxy/db/proxy_worker_heartbeat.py b/litellm/proxy/db/proxy_worker_heartbeat.py new file mode 100644 index 00000000000..6a2a4572e43 --- /dev/null +++ b/litellm/proxy/db/proxy_worker_heartbeat.py @@ -0,0 +1,89 @@ +""" +Live proxy worker census, one row per worker process. + +Every uvicorn worker upserts its own row on a fixed heartbeat, so counting +rows with a recent heartbeat answers "how many workers share this database?" +without any coordination. The Admin UI's "no Redis" banner uses that count to +hide itself for deployments that are provably a single worker, where per-worker +rate limits, budgets, and router state are already global. All timestamps are +written and compared with the database's own clock, so pods with skewed clocks +still agree. +""" + +from __future__ import annotations + +import socket +from typing import TYPE_CHECKING, Final + +from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict + +from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + +PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS: Final = 60 +PROXY_WORKER_LIVENESS_WINDOW_SECONDS: Final = 3 * PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS +STALE_ROW_RETENTION_SECONDS: Final = 3600 + +BEAT_SQL: Final = """ +INSERT INTO "LiteLLM_ProxyWorkerHeartbeat" (worker_id, hostname, last_heartbeat_at) +VALUES ($1, $2, NOW()) +ON CONFLICT (worker_id) DO UPDATE SET last_heartbeat_at = NOW() +""" + +PRUNE_SQL: Final = """ +DELETE FROM "LiteLLM_ProxyWorkerHeartbeat" +WHERE last_heartbeat_at < NOW() - make_interval(secs => $1) +""" + +COUNT_SQL: Final = """ +SELECT COUNT(*)::int AS live_workers FROM "LiteLLM_ProxyWorkerHeartbeat" +WHERE last_heartbeat_at > NOW() - make_interval(secs => $1) +""" + +DEREGISTER_SQL: Final = """ +DELETE FROM "LiteLLM_ProxyWorkerHeartbeat" WHERE worker_id = $1 +""" + + +class _LiveWorkerCountRow(TypedDict): + live_workers: ReadOnly[int] + + +_COUNT_ROWS_ADAPTER: Final = TypeAdapter(tuple[_LiveWorkerCountRow, ...]) + + +class ProxyWorkerHeartbeat: + def __init__(self, prisma_client: PrismaClient, worker_id: str | None = None) -> None: + self.prisma_client: Final = prisma_client + self.worker_id: Final[str] = worker_id or str(uuid.uuid4()) + self.hostname: Final = socket.gethostname() + + async def beat(self) -> None: + try: + await self.prisma_client.db.execute_raw(BEAT_SQL, self.worker_id, self.hostname) + await self.prisma_client.db.execute_raw(PRUNE_SQL, STALE_ROW_RETENTION_SECONDS) + except Exception as beat_err: # noqa: BLE001 # a missed heartbeat must never take down the worker + verbose_proxy_logger.debug("Proxy worker heartbeat write failed: %s", beat_err) + + async def deregister(self) -> None: + try: + await self.prisma_client.db.execute_raw(DEREGISTER_SQL, self.worker_id) + except Exception as deregister_err: # noqa: BLE001 # best-effort cleanup; the liveness window ages the row out anyway + verbose_proxy_logger.debug("Proxy worker heartbeat deregister failed: %s", deregister_err) + + +async def count_live_proxy_workers(prisma_client: PrismaClient) -> int | None: + """ + The number of workers with a recent heartbeat, or None when the database + cannot answer. Callers must treat None as "unknown", not as zero. + """ + try: + rows: Final = await prisma_client.db.query_raw(COUNT_SQL, PROXY_WORKER_LIVENESS_WINDOW_SECONDS) + return _COUNT_ROWS_ADAPTER.validate_python(rows)[0]["live_workers"] + except Exception as count_err: # noqa: BLE001 # an unknown count must degrade to "warn", never to a 503 + verbose_proxy_logger.debug("Live proxy worker count unavailable: %s", count_err) + return None diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index e814ec42d26..33894777bc3 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -34,6 +34,7 @@ from litellm.proxy.auth.auth_utils import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.db.proxy_worker_heartbeat import count_live_proxy_workers from litellm.proxy.health_check import ( ADMIN_ONLY_HEALTH_DISPLAY_PARAMS, _clean_endpoint_data, @@ -1451,7 +1452,7 @@ def callback_name(callback): DISABLE_NO_REDIS_WARNING_ENV_VAR: Final = "LITELLM_DISABLE_NO_REDIS_WARNING" -def _show_no_redis_warning() -> bool: +async def _show_no_redis_warning() -> bool: """ Whether the UI should warn that no Redis is configured. @@ -1461,16 +1462,22 @@ def _show_no_redis_warning() -> bool: coordination cache (from a Redis response cache, general_settings. coordination_redis, or the REDIS_* env fallback) and the router's own Redis (router_settings.redis_host), which backs cooldowns and usage-based - routing on its own. Operators who know they run one worker can silence the - warning with LITELLM_DISABLE_NO_REDIS_WARNING=true. + routing on its own. A deployment whose worker-heartbeat census proves it + is exactly one worker needs no cross-worker coordination, so it never + warns; when the census is unavailable or shows more than one worker, the + warning stands unless LITELLM_DISABLE_NO_REDIS_WARNING=true silences it. """ - from litellm.proxy.proxy_server import llm_router, redis_usage_cache + from litellm.proxy.proxy_server import llm_router, prisma_client, redis_usage_cache if redis_usage_cache is not None: return False if llm_router is not None and llm_router.cache.redis_cache is not None: return False - return get_secret_bool(DISABLE_NO_REDIS_WARNING_ENV_VAR, False) is not True + if get_secret_bool(DISABLE_NO_REDIS_WARNING_ENV_VAR, False) is True: + return False + if prisma_client is None: + return True + return await count_live_proxy_workers(prisma_client) != 1 async def _get_health_readiness_details( @@ -1513,7 +1520,7 @@ async def _get_health_readiness_details( # check log level log_level_name: Final = logging.getLevelName(verbose_logger.getEffectiveLevel()) is_detailed_debug: Final = verbose_logger.isEnabledFor(logging.DEBUG) - show_no_redis_warning: Final = _show_no_redis_warning() + show_no_redis_warning: Final = await _show_no_redis_warning() # check DB if prisma_client is not None: # if db passed in, check if it's connected diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 359187f81cb..6ee08a732f2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -379,6 +379,10 @@ from litellm.proxy.db.gateway_request_tracking import ( GatewayRequestAccumulator, flush_gateway_requests, ) +from litellm.proxy.db.proxy_worker_heartbeat import ( + PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS, + ProxyWorkerHeartbeat, +) from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed from litellm.proxy.discovery_endpoints import ui_discovery_endpoints_router from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router @@ -864,9 +868,11 @@ async def _flush_spend_logs_queue_on_shutdown() -> None: verbose_proxy_logger.exception("Error flushing spend logs queue on shutdown: %s", e) -async def proxy_shutdown_event(): +async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = None): global prisma_client, master_key, user_custom_auth, user_custom_key_generate, user_custom_key_update verbose_proxy_logger.info("Shutting down LiteLLM Proxy Server") + if worker_heartbeat is not None and prisma_client: + await worker_heartbeat.deregister() if prisma_client: # Drain the SGR fold first: it lives in memory, so an un-drained interval # is lost, and a write attempted after disconnect raises @@ -1200,7 +1206,7 @@ async def proxy_startup_event(app: FastAPI): ) ### START BATCH WRITING DB + CHECKING NEW MODELS### - if prisma_client is not None: + worker_heartbeat: Final = ( await ProxyStartupEvent.initialize_scheduled_background_jobs( general_settings=general_settings, prisma_client=prisma_client, @@ -1209,7 +1215,10 @@ async def proxy_startup_event(app: FastAPI): proxy_batch_write_at=proxy_batch_write_at, proxy_logging_obj=proxy_logging_obj, ) - + if prisma_client is not None + else None + ) + if prisma_client is not None: await ProxyStartupEvent._update_default_team_member_budget() ## SYNC UI SETTINGS ## @@ -1280,7 +1289,7 @@ async def proxy_startup_event(app: FastAPI): await proxy_config.stop_auth_cache_invalidation_subscriber() - await proxy_shutdown_event() + await proxy_shutdown_event(worker_heartbeat=worker_heartbeat) def _generate_stable_operation_id(route: "APIRoute") -> str: @@ -8665,7 +8674,7 @@ class ProxyStartupEvent: proxy_budget_rescheduler_max_time: int, proxy_batch_write_at: int, proxy_logging_obj: ProxyLogging, - ): + ) -> ProxyWorkerHeartbeat: """Initializes scheduled background jobs""" global store_model_in_db, scheduler @@ -8710,6 +8719,18 @@ class ProxyStartupEvent: # Ensure minimum interval of 30 seconds for batch writing to prevent memory issues batch_writing_interval: Final = proxy_batch_write_at + random.randint(0, 5) + ### PROXY WORKER HEARTBEAT ### + worker_heartbeat: Final = ProxyWorkerHeartbeat(prisma_client=prisma_client) + await worker_heartbeat.beat() + scheduler.add_job( + worker_heartbeat.beat, + "interval", + seconds=PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS, + id="proxy_worker_heartbeat_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + ### RESET BUDGET ### if general_settings.get("disable_reset_budget", False) is False: budget_reset_job: Final = ResetBudgetJob( @@ -9048,6 +9069,7 @@ class ProxyStartupEvent: "APScheduler started with memory leak prevention settings: removed jitter, increased intervals, misfire_grace_time=%s", APSCHEDULER_MISFIRE_GRACE_TIME, ) + return worker_heartbeat @classmethod async def _initialize_spend_tracking_background_jobs(cls, scheduler: AsyncIOScheduler): diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 79d778fb464..09efef813a7 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -945,6 +945,17 @@ model LiteLLM_DailyTagSpend { } +// One row per live proxy worker process. Workers upsert their row on a fixed +// heartbeat; counting rows with a recent heartbeat tells how many workers share +// this database, which lets the Admin UI hide its "no Redis" warning for +// deployments that are provably a single worker. +model LiteLLM_ProxyWorkerHeartbeat { + worker_id String @id + hostname String + started_at DateTime @default(now()) + last_heartbeat_at DateTime @default(now()) +} + // Track the status of cron jobs running. Only allow one pod to run the job at a time model LiteLLM_CronJob { cronjob_id String @id @default(cuid()) // Unique ID for the record diff --git a/schema.prisma b/schema.prisma index 79d778fb464..09efef813a7 100644 --- a/schema.prisma +++ b/schema.prisma @@ -945,6 +945,17 @@ model LiteLLM_DailyTagSpend { } +// One row per live proxy worker process. Workers upsert their row on a fixed +// heartbeat; counting rows with a recent heartbeat tells how many workers share +// this database, which lets the Admin UI hide its "no Redis" warning for +// deployments that are provably a single worker. +model LiteLLM_ProxyWorkerHeartbeat { + worker_id String @id + hostname String + started_at DateTime @default(now()) + last_heartbeat_at DateTime @default(now()) +} + // Track the status of cron jobs running. Only allow one pod to run the job at a time model LiteLLM_CronJob { cronjob_id String @id @default(cuid()) // Unique ID for the record diff --git a/tests/test_litellm/proxy/db/test_proxy_worker_heartbeat.py b/tests/test_litellm/proxy/db/test_proxy_worker_heartbeat.py new file mode 100644 index 00000000000..2209be0dc2e --- /dev/null +++ b/tests/test_litellm/proxy/db/test_proxy_worker_heartbeat.py @@ -0,0 +1,81 @@ +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy.db.proxy_worker_heartbeat import ( + BEAT_SQL, + COUNT_SQL, + DEREGISTER_SQL, + PROXY_WORKER_LIVENESS_WINDOW_SECONDS, + PRUNE_SQL, + STALE_ROW_RETENTION_SECONDS, + ProxyWorkerHeartbeat, + count_live_proxy_workers, +) + + +def _prisma(): + prisma = MagicMock() + prisma.db.execute_raw = AsyncMock() + prisma.db.query_raw = AsyncMock() + return prisma + + +@pytest.mark.asyncio +async def test_beat_upserts_own_row_then_prunes_stale_rows(): + prisma = _prisma() + heartbeat = ProxyWorkerHeartbeat(prisma_client=prisma, worker_id="worker-1") + await heartbeat.beat() + calls = prisma.db.execute_raw.call_args_list + assert calls[0].args == (BEAT_SQL, "worker-1", heartbeat.hostname) + assert calls[1].args == (PRUNE_SQL, STALE_ROW_RETENTION_SECONDS) + + +@pytest.mark.asyncio +async def test_beat_survives_a_database_error(): + prisma = _prisma() + prisma.db.execute_raw = AsyncMock(side_effect=RuntimeError("db down")) + await ProxyWorkerHeartbeat(prisma_client=prisma).beat() + + +def test_each_worker_process_gets_its_own_id(): + prisma = _prisma() + first = ProxyWorkerHeartbeat(prisma_client=prisma) + second = ProxyWorkerHeartbeat(prisma_client=prisma) + assert first.worker_id != second.worker_id + + +@pytest.mark.asyncio +async def test_deregister_deletes_only_its_own_row(): + prisma = _prisma() + await ProxyWorkerHeartbeat(prisma_client=prisma, worker_id="worker-1").deregister() + assert prisma.db.execute_raw.call_args.args == (DEREGISTER_SQL, "worker-1") + + +@pytest.mark.asyncio +async def test_deregister_survives_a_database_error(): + prisma = _prisma() + prisma.db.execute_raw = AsyncMock(side_effect=RuntimeError("db down")) + await ProxyWorkerHeartbeat(prisma_client=prisma, worker_id="worker-1").deregister() + + +@pytest.mark.asyncio +async def test_count_reads_workers_within_the_liveness_window(): + prisma = _prisma() + prisma.db.query_raw.return_value = [{"live_workers": 3}] + assert await count_live_proxy_workers(prisma) == 3 + assert prisma.db.query_raw.call_args.args == (COUNT_SQL, PROXY_WORKER_LIVENESS_WINDOW_SECONDS) + + +@pytest.mark.asyncio +async def test_count_returns_unknown_when_the_query_fails(): + prisma = _prisma() + prisma.db.query_raw.side_effect = RuntimeError("db down") + assert await count_live_proxy_workers(prisma) is None + + +@pytest.mark.asyncio +async def test_count_returns_unknown_for_a_malformed_row(): + prisma = _prisma() + prisma.db.query_raw.return_value = [{"unexpected": "shape"}] + assert await count_live_proxy_workers(prisma) is None diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index e2705bd5fec..831f659051c 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -2467,61 +2467,140 @@ class TestNoRedisWarning: def _router(redis_cache): return SimpleNamespace(cache=SimpleNamespace(redis_cache=redis_cache)) - def test_warns_when_no_redis_is_configured(self, monkeypatch): + @staticmethod + def _prisma_with_workers(live_workers=None, error=None): + prisma = MagicMock() + if error is not None: + prisma.db.query_raw = AsyncMock(side_effect=error) + else: + prisma.db.query_raw = AsyncMock(return_value=[{"live_workers": live_workers}]) + return prisma + + @pytest.mark.asyncio + async def test_warns_when_no_redis_and_no_db_to_count_workers(self, monkeypatch): monkeypatch.delenv("LITELLM_DISABLE_NO_REDIS_WARNING", raising=False) with ( patch("litellm.proxy.proxy_server.redis_usage_cache", None), patch("litellm.proxy.proxy_server.llm_router", self._router(None)), + patch("litellm.proxy.proxy_server.prisma_client", None), ): - assert _show_no_redis_warning() is True + assert await _show_no_redis_warning() is True - def test_warns_when_there_is_no_router_at_all(self, monkeypatch): + @pytest.mark.asyncio + async def test_warns_when_there_is_no_router_at_all(self, monkeypatch): monkeypatch.delenv("LITELLM_DISABLE_NO_REDIS_WARNING", raising=False) with ( patch("litellm.proxy.proxy_server.redis_usage_cache", None), patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), ): - assert _show_no_redis_warning() is True + assert await _show_no_redis_warning() is True - def test_stays_quiet_when_a_coordination_redis_is_configured(self, monkeypatch): + @pytest.mark.asyncio + async def test_stays_quiet_for_a_confirmed_single_worker(self, monkeypatch): + """One live worker needs no cross-worker coordination, so no env var is needed.""" monkeypatch.delenv("LITELLM_DISABLE_NO_REDIS_WARNING", raising=False) + with ( + patch("litellm.proxy.proxy_server.redis_usage_cache", None), + patch("litellm.proxy.proxy_server.llm_router", self._router(None)), + patch("litellm.proxy.proxy_server.prisma_client", self._prisma_with_workers(1)), + ): + assert await _show_no_redis_warning() is False + + @pytest.mark.asyncio + @pytest.mark.parametrize("live_workers", [2, 5]) + async def test_warns_when_multiple_workers_share_the_db(self, monkeypatch, live_workers): + monkeypatch.delenv("LITELLM_DISABLE_NO_REDIS_WARNING", raising=False) + with ( + patch("litellm.proxy.proxy_server.redis_usage_cache", None), + patch("litellm.proxy.proxy_server.llm_router", self._router(None)), + patch("litellm.proxy.proxy_server.prisma_client", self._prisma_with_workers(live_workers)), + ): + assert await _show_no_redis_warning() is True + + @pytest.mark.asyncio + async def test_warns_when_the_worker_census_is_empty(self, monkeypatch): + """Zero rows means the census cannot CONFIRM a single worker, so warn.""" + monkeypatch.delenv("LITELLM_DISABLE_NO_REDIS_WARNING", raising=False) + with ( + patch("litellm.proxy.proxy_server.redis_usage_cache", None), + patch("litellm.proxy.proxy_server.llm_router", self._router(None)), + patch("litellm.proxy.proxy_server.prisma_client", self._prisma_with_workers(0)), + ): + assert await _show_no_redis_warning() is True + + @pytest.mark.asyncio + async def test_warns_when_the_worker_census_query_fails(self, monkeypatch): + monkeypatch.delenv("LITELLM_DISABLE_NO_REDIS_WARNING", raising=False) + with ( + patch("litellm.proxy.proxy_server.redis_usage_cache", None), + patch("litellm.proxy.proxy_server.llm_router", self._router(None)), + patch( + "litellm.proxy.proxy_server.prisma_client", + self._prisma_with_workers(error=RuntimeError("db down")), + ), + ): + assert await _show_no_redis_warning() is True + + @pytest.mark.asyncio + async def test_stays_quiet_when_a_coordination_redis_is_configured(self, monkeypatch): + monkeypatch.delenv("LITELLM_DISABLE_NO_REDIS_WARNING", raising=False) + prisma = self._prisma_with_workers(5) with ( patch("litellm.proxy.proxy_server.redis_usage_cache", MagicMock()), patch("litellm.proxy.proxy_server.llm_router", self._router(None)), + patch("litellm.proxy.proxy_server.prisma_client", prisma), ): - assert _show_no_redis_warning() is False + assert await _show_no_redis_warning() is False + prisma.db.query_raw.assert_not_called() - def test_stays_quiet_when_only_the_router_has_redis(self, monkeypatch): + @pytest.mark.asyncio + async def test_stays_quiet_when_only_the_router_has_redis(self, monkeypatch): """router_settings.redis_host alone backs cooldowns and usage-based routing.""" monkeypatch.delenv("LITELLM_DISABLE_NO_REDIS_WARNING", raising=False) with ( patch("litellm.proxy.proxy_server.redis_usage_cache", None), patch("litellm.proxy.proxy_server.llm_router", self._router(MagicMock())), + patch("litellm.proxy.proxy_server.prisma_client", self._prisma_with_workers(5)), ): - assert _show_no_redis_warning() is False + assert await _show_no_redis_warning() is False + @pytest.mark.asyncio @pytest.mark.parametrize("value", ["true", "True"]) - def test_env_var_suppresses_the_warning(self, monkeypatch, value): + async def test_env_var_suppresses_the_warning_despite_multiple_workers(self, monkeypatch, value): monkeypatch.setenv("LITELLM_DISABLE_NO_REDIS_WARNING", value) with ( patch("litellm.proxy.proxy_server.redis_usage_cache", None), patch("litellm.proxy.proxy_server.llm_router", self._router(None)), + patch("litellm.proxy.proxy_server.prisma_client", self._prisma_with_workers(5)), ): - assert _show_no_redis_warning() is False + assert await _show_no_redis_warning() is False - def test_env_var_set_false_keeps_the_warning(self, monkeypatch): + @pytest.mark.asyncio + async def test_env_var_set_false_keeps_the_warning_for_multiple_workers(self, monkeypatch): monkeypatch.setenv("LITELLM_DISABLE_NO_REDIS_WARNING", "false") with ( patch("litellm.proxy.proxy_server.redis_usage_cache", None), patch("litellm.proxy.proxy_server.llm_router", self._router(None)), + patch("litellm.proxy.proxy_server.prisma_client", self._prisma_with_workers(2)), ): - assert _show_no_redis_warning() is True + assert await _show_no_redis_warning() is True + + @pytest.mark.asyncio + async def test_env_var_set_false_does_not_force_the_warning_for_a_single_worker(self, monkeypatch): + monkeypatch.setenv("LITELLM_DISABLE_NO_REDIS_WARNING", "false") + with ( + patch("litellm.proxy.proxy_server.redis_usage_cache", None), + patch("litellm.proxy.proxy_server.llm_router", self._router(None)), + patch("litellm.proxy.proxy_server.prisma_client", self._prisma_with_workers(1)), + ): + assert await _show_no_redis_warning() is False @pytest.mark.asyncio @pytest.mark.parametrize("has_prisma_client", [True, False]) async def test_readiness_details_carries_the_flag(self, monkeypatch, has_prisma_client): monkeypatch.delenv("LITELLM_DISABLE_NO_REDIS_WARNING", raising=False) - prisma_client = MagicMock() if has_prisma_client else None + prisma_client = self._prisma_with_workers(2) if has_prisma_client else None with ( patch("litellm.proxy.proxy_server.prisma_client", prisma_client), patch("litellm.proxy.proxy_server.redis_usage_cache", None), diff --git a/ui/litellm-dashboard/src/components/NoRedisWarningBanner.test.tsx b/ui/litellm-dashboard/src/components/NoRedisWarningBanner.test.tsx index 8afde8eec94..600315e789e 100644 --- a/ui/litellm-dashboard/src/components/NoRedisWarningBanner.test.tsx +++ b/ui/litellm-dashboard/src/components/NoRedisWarningBanner.test.tsx @@ -20,6 +20,7 @@ describe("NoRedisWarningBanner", () => { renderWithProviders(); expect(screen.getByRole("alert")).toBeInTheDocument(); expect(screen.getByText(/No Redis configured\. Redis is highly recommended/i)).toBeInTheDocument(); + expect(screen.getByText(/more than one worker/i)).toBeInTheDocument(); }); it("should link to the docs page listing what breaks without Redis", () => { diff --git a/ui/litellm-dashboard/src/components/NoRedisWarningBanner.tsx b/ui/litellm-dashboard/src/components/NoRedisWarningBanner.tsx index 93c0f55486d..02433fad521 100644 --- a/ui/litellm-dashboard/src/components/NoRedisWarningBanner.tsx +++ b/ui/litellm-dashboard/src/components/NoRedisWarningBanner.tsx @@ -26,13 +26,13 @@ export const NoRedisWarningBanner: React.FC = ({ acce

No Redis configured. Redis is highly recommended

- Rate limits, budgets, router state, and cache invalidation are per worker without Redis, so limits are - enforced once per worker and spend can overshoot.{" "} + This proxy is running more than one worker (or the worker count could not be verified). Without Redis, rate + limits, budgets, router state, and cache invalidation are per worker, so limits are enforced once per worker + and spend can overshoot.{" "} See everything that does not work without Redis - . If you run a single worker and this is intentional, set{" "} - LITELLM_DISABLE_NO_REDIS_WARNING=true to hide this banner. + . Set LITELLM_DISABLE_NO_REDIS_WARNING=true to hide this banner anyway.

From 3217b8edae27074298717b2a56fd1d5a82b6d517 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 14 Aug 2026 17:23:06 -0700 Subject: [PATCH 041/538] fix(proxy): count worker heartbeats on the primary so replica lag cannot undercount --- litellm/proxy/db/proxy_worker_heartbeat.py | 8 ++++++-- .../proxy/db/test_proxy_worker_heartbeat.py | 13 +++++++++++++ 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/db/proxy_worker_heartbeat.py b/litellm/proxy/db/proxy_worker_heartbeat.py index 6a2a4572e43..990ff48eb18 100644 --- a/litellm/proxy/db/proxy_worker_heartbeat.py +++ b/litellm/proxy/db/proxy_worker_heartbeat.py @@ -20,6 +20,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient @@ -79,10 +80,13 @@ class ProxyWorkerHeartbeat: async def count_live_proxy_workers(prisma_client: PrismaClient) -> int | None: """ The number of workers with a recent heartbeat, or None when the database - cannot answer. Callers must treat None as "unknown", not as zero. + cannot answer. Callers must treat None as "unknown", not as zero. Always + counts on the primary: a lagging read replica must never undercount. """ try: - rows: Final = await prisma_client.db.query_raw(COUNT_SQL, PROXY_WORKER_LIVENESS_WINDOW_SECONDS) + db: Final = prisma_client.db + primary_db: Final = db.writer if isinstance(db, RoutingPrismaWrapper) else db + rows: Final = await primary_db.query_raw(COUNT_SQL, PROXY_WORKER_LIVENESS_WINDOW_SECONDS) return _COUNT_ROWS_ADAPTER.validate_python(rows)[0]["live_workers"] except Exception as count_err: # noqa: BLE001 # an unknown count must degrade to "warn", never to a 503 verbose_proxy_logger.debug("Live proxy worker count unavailable: %s", count_err) diff --git a/tests/test_litellm/proxy/db/test_proxy_worker_heartbeat.py b/tests/test_litellm/proxy/db/test_proxy_worker_heartbeat.py index 2209be0dc2e..33ae6190411 100644 --- a/tests/test_litellm/proxy/db/test_proxy_worker_heartbeat.py +++ b/tests/test_litellm/proxy/db/test_proxy_worker_heartbeat.py @@ -12,6 +12,7 @@ from litellm.proxy.db.proxy_worker_heartbeat import ( ProxyWorkerHeartbeat, count_live_proxy_workers, ) +from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper def _prisma(): @@ -67,6 +68,18 @@ async def test_count_reads_workers_within_the_liveness_window(): assert prisma.db.query_raw.call_args.args == (COUNT_SQL, PROXY_WORKER_LIVENESS_WINDOW_SECONDS) +@pytest.mark.asyncio +async def test_count_reads_from_the_primary_when_reads_route_to_a_replica(): + writer = MagicMock() + writer.query_raw = AsyncMock(return_value=[{"live_workers": 2}]) + reader = MagicMock() + reader.query_raw = AsyncMock(return_value=[{"live_workers": 1}]) + prisma = MagicMock() + prisma.db = RoutingPrismaWrapper(writer=writer, reader=reader) + assert await count_live_proxy_workers(prisma) == 2 + reader.query_raw.assert_not_awaited() + + @pytest.mark.asyncio async def test_count_returns_unknown_when_the_query_fails(): prisma = _prisma() 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 042/538] 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 043/538] 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 a6dc447470f41baea1383b30fcee4c10cf6bdc04 Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Sun, 16 Aug 2026 16:24:23 +0530 Subject: [PATCH 044/538] fix(proxy): stop Responses batch rows from bypassing project OTPM Embeddings rows were identified by body shape (has `input`, no `messages`/`prompt`), which also matches a `/v1/responses` batch row and reserved zero output tokens for it -- letting a project caller run large Responses generations against a quota-limited model without consuming OTPM. Classify embeddings by the row's own `url` instead, and read `max_output_tokens` as a Responses output cap alongside `max_tokens`/`max_completion_tokens`. Co-authored-by: Cursor --- litellm/proxy/hooks/batch_rate_limiter.py | 30 ++++++- .../proxy/hooks/test_batch_file_validation.py | 80 +++++++++++++++++++ 2 files changed, 106 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 3b53d792ce6..4052a9b04fa 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -395,18 +395,40 @@ class _PROXY_BatchRateLimiter(CustomLogger): Batch completion never reconciles actual usage back into the rate limiter, so this pre-call estimate is the only OTPM enforcement a batch gets. Mirrors the real-time no-``max_tokens`` floor so a row - that omits ``max_tokens`` can't be used to bypass OTPM the way an + that omits an output cap can't be used to bypass OTPM the way an unbounded streaming request could. + + Embeddings rows are identified by the row's own ``url`` (the OpenAI + batch schema puts the target route there, e.g. ``/v1/embeddings``), + never by body shape: a `/v1/responses` row also carries `body.input` + with no `messages`/`prompt`, so guessing from body shape alone would + misclassify a token-generating Responses row as a zero-output + embeddings row and let it skip the OTPM reservation entirely. """ + url: Final = entry.get("url") + if isinstance(url, str) and "embeddings" in url: + return 0 # embeddings: no output tokens raw_body: Final = entry.get("body") body: Final[Mapping[str, object]] = ( MappingProxyType(_BATCH_BODY_ADAPTER.validate_python(raw_body)) if isinstance(raw_body, Mapping) else MappingProxyType({}) # mutable-ok: immediately frozen empty fallback ) - if body.get("input") is not None and body.get("messages") is None and body.get("prompt") is None: - return 0 # embeddings: no output tokens - explicit_cap: Final = body.get("max_tokens", body.get("max_completion_tokens")) + # `max_tokens`/`max_completion_tokens` cap chat completions; `/v1/responses` + # rows cap output with `max_output_tokens` instead -- omitting it here + # would fall through to the floor estimate for every capped Responses row. + explicit_cap: Final = next( + ( + v + for v in ( + body.get("max_tokens"), + body.get("max_completion_tokens"), + body.get("max_output_tokens"), + ) + if v is not None + ), + None, + ) if explicit_cap is not None: try: return max(0, int(explicit_cap)) diff --git a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py index ee37cf7d302..e5211973ec3 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -1980,3 +1980,83 @@ async def test_count_input_file_usage_collects_models_after_malformed_line(): ) assert exc.value.status_code == 403 + + +# --------------------------------------------------------------------------- +# VERIA-Low regression: Responses batch rows must not bypass project OTPM +# --------------------------------------------------------------------------- + + +def _output_estimator(): + """A `_PROXY_BatchRateLimiter` whose output-token floor is observable: + the no-`max_tokens` floor mock returns a distinctive sentinel so tests can + tell "floor was used" apart from "an explicit cap was read".""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + limiter = MagicMock() + limiter.no_max_tokens_output_floor.return_value = 999 + return _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=limiter, + ) + + +def test_estimate_entry_output_tokens_zero_for_embeddings_url(): + """A real `/v1/embeddings` row reserves zero output tokens.""" + rate_limiter = _output_estimator() + entry = { + "url": "/v1/embeddings", + "body": {"model": "text-embedding-3-small", "input": "hello world"}, + } + + assert rate_limiter._estimate_entry_output_tokens(entry, None) == 0 + + +def test_estimate_entry_output_tokens_does_not_zero_responses_row_with_input(): + """Pre-fix: a `/v1/responses` row carries `body.input` with no `messages`/ + `prompt`, so the old body-shape heuristic misclassified it as embeddings + and reserved zero output tokens -- a project caller could submit large + Responses generations against a quota-limited model without consuming + OTPM. The row's own `url` (not body shape) must decide this.""" + rate_limiter = _output_estimator() + entry = { + "url": "/v1/responses", + "body": {"model": "gpt-4o", "input": "write me an essay"}, + } + + # No explicit cap on the row, so it must fall back to the no-max-tokens + # floor -- never straight to zero. + assert rate_limiter._estimate_entry_output_tokens(entry, None) == 999 + rate_limiter.parallel_request_limiter.no_max_tokens_output_floor.assert_called_once_with(None) + + +def test_estimate_entry_output_tokens_uses_max_output_tokens_for_responses(): + """`/v1/responses` caps output with `max_output_tokens`, not `max_tokens`/ + `max_completion_tokens`. Pre-fix this field was never inspected, so a + capped Responses row still fell through to the (possibly larger) floor + estimate instead of the caller's own declared cap.""" + rate_limiter = _output_estimator() + entry = { + "url": "/v1/responses", + "body": {"model": "gpt-4o", "input": "hi", "max_output_tokens": 123}, + } + + assert rate_limiter._estimate_entry_output_tokens(entry, None) == 123 + rate_limiter.parallel_request_limiter.no_max_tokens_output_floor.assert_not_called() + + +def test_estimate_entry_output_tokens_prefers_max_tokens_over_max_output_tokens(): + """When a row somehow carries both fields, the chat-style cap wins first -- + `max_output_tokens` is only consulted once the chat-style caps are absent.""" + rate_limiter = _output_estimator() + entry = { + "url": "/v1/chat/completions", + "body": { + "model": "gpt-4o", + "messages": [], + "max_tokens": 50, + "max_output_tokens": 500, + }, + } + + assert rate_limiter._estimate_entry_output_tokens(entry, None) == 50 From ffa37d05b7cf16c9874101f3c737da19b4154aca Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Sun, 16 Aug 2026 14:28:50 -0400 Subject: [PATCH 045/538] 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 046/538] 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 047/538] 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 f86dc8f54e9bbee403c185e02196610e6ef76bd2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 17 Aug 2026 13:09:17 -0700 Subject: [PATCH 048/538] test(proxy): assert non-Bedrock passthrough stream emits no content-type header --- tests/test_litellm/proxy/test_common_request_processing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 7b762457bc8..e567d19dac7 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -4564,7 +4564,7 @@ class TestAllmPassthroughStreamingProviderGate: assert isinstance(result, StreamingResponse) assert result.media_type is None - assert result.headers.get("content-type") != "application/vnd.amazon.eventstream" + assert "content-type" not in result.headers class TestResponseCostHeaderForTypedDictResponses: From d5a4c145778a66df0df5a6a627cae3e40e162645 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 17 Aug 2026 13:21:28 -0700 Subject: [PATCH 049/538] docs(proxy): pre-fix passthrough streams omitted content-type, not octet-stream --- litellm/proxy/common_request_processing.py | 5 +++-- tests/test_litellm/proxy/test_common_request_processing.py | 6 +++--- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 165e58c56f7..376d7785e5a 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2502,8 +2502,9 @@ class ProxyBaseLLMRequestProcessing: buffered (guardrail-rewritten) and the unbuffered relay paths so clients that enforce the event-stream content-type (e.g. Claude Code on Bedrock invoke-with-response-stream) see the correct header instead of - Starlette's application/octet-stream default. Returns None for providers - with no event-stream media type, leaving the response default unchanged. + no content-type at all, which they fall back to reading as + application/octet-stream. Returns None for providers with no + event-stream media type, leaving the response headers unchanged. """ from litellm.llms.pass_through.guardrail_translation.handler import ( LlmPassthroughRouteHandler, diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index e567d19dac7..ddb7b0d510f 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -4516,9 +4516,9 @@ class TestAllmPassthroughStreamingProviderGate: """ Regression for LIT-4561. The unbuffered Bedrock event-stream relay (invoke-with-response-stream, no post-call guardrail rewriting) must set - content-type: application/vnd.amazon.eventstream instead of leaving it to - Starlette's application/octet-stream default, which trips Claude Code's - content-type guard added in 2.1.208 + content-type: application/vnd.amazon.eventstream instead of emitting no + content-type header at all, which trips Claude Code's content-type guard + added in 2.1.208 """ processing_obj = self._build_processing_obj( "bedrock", "model/us.anthropic.claude-sonnet-4-20250514-v1:0/invoke-with-response-stream" From 97b7eaad5cd6d723c4849ab43e22ecdbfb9ce2d5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 17 Aug 2026 14:44:01 -0700 Subject: [PATCH 050/538] fix(mcp): keep issuer-anchored slots on reload, skip discovery for stamped M2M challenge Registry-swap reconciliation used bool(server.url) while registration uses _requires_oauth_discovery, dropping slots for issuer-anchored servers without a url. The preemptive 401 loop awaited discovery before the stamped client_credentials continue, so a deferred discovery failure could 503 requests whose challenge decision never reads metadata --- .../mcp_server/mcp_server_manager.py | 2 +- .../proxy/_experimental/mcp_server/server.py | 6 ++++ .../mcp_server/test_mcp_server.py | 16 ++++++++++ .../mcp_server/test_mcp_server_manager.py | 30 +++++++++++++++++++ 4 files changed, 53 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index cdaf3f2f206..91e5339dda2 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1691,7 +1691,7 @@ class MCPServerManager: def _reconcile_oauth_discovery_slots_for_servers(self, servers: Iterable[MCPServer]) -> None: """Align retry slots after an atomic registry replacement.""" for server in servers: - should_defer = bool(server.url) and _oauth_endpoints_unresolved(server) + should_defer = _requires_oauth_discovery(server.url, server.issuer_is_anchored, server) has_slot = self._oauth_discovery_slot(server.server_id) is not None if should_defer != has_slot: self._set_oauth_discovery_deferred(server.server_id, should_defer) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index e4ac40734dc..2b7f3e96371 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3737,6 +3737,12 @@ if MCP_AVAILABLE: # preemptive challenge and let downstream authorization # return 403. continue + if server is not None and server.auth_type == MCPAuth.oauth2 and server.oauth2_flow == "client_credentials": + # Stamped M2M: the challenge decision below never reads discovered + # metadata, so deferred-discovery failures must not 503 this loop. + # Unstamped rows stay on the discover-first path because filling + # authorization_url/token_url can change their inferred flow. + continue if server is not None: server = await global_mcp_server_manager.ensure_oauth_metadata_discovered(server) if server and server.auth_type == MCPAuth.oauth2: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 0193f2c9152..7b3b8c7a28c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -7969,6 +7969,22 @@ class TestPreemptive401ModeAware: assert manager._oauth_discovery_slot(server.server_id) is None assert exc.value.status_code == 401 + @pytest.mark.asyncio + async def test_stamped_m2m_challenge_skips_deferred_discovery(self): + from litellm.proxy._experimental.mcp_server import server as server_module + + manager = server_module.global_mcp_server_manager + server = _make_oauth2_server("stamped_m2m", oauth2_flow="client_credentials") + + with patch.object( + manager, + "ensure_oauth_metadata_discovered", + new=AsyncMock(side_effect=HTTPException(status_code=503, detail="discovery down")), + ) as discovery: + await self._run(server, None, has_stored_token=False) + + discovery.assert_not_awaited() + @pytest.mark.asyncio async def test_gateway_managed_interactive_no_token_challenges_with_x_litellm_api_key(self): """No stored token, key in x-litellm-api-key (oauth2_headers empty): 401.""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index fe09303b298..76d6a178b9e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -695,6 +695,36 @@ class TestMCPServerManager: assert resolved.token_url == "https://idp.example.com/token" assert manager._oauth_discovery_slot(replacement.server_id) is None + def test_registry_swap_reconcile_keeps_slot_for_issuer_anchored_server_without_url(self): + manager = MCPServerManager() + server = MCPServer( + server_id="anchored-no-url-1", + name="anchored_no_url", + url=None, + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + issuer="https://idp.example.com", + issuer_is_anchored=True, + ) + manager.registry[server.server_id] = server + manager._set_oauth_discovery_deferred(server.server_id, True) + + manager._reconcile_oauth_discovery_slots_for_servers([server]) + + assert manager._oauth_discovery_slot(server.server_id) is not None + + resolved = server.model_copy( + update={ + "authorization_url": "https://idp.example.com/authorize", + "token_url": "https://idp.example.com/token", + } + ) + manager.registry[resolved.server_id] = resolved + manager._reconcile_oauth_discovery_slots_for_servers([resolved]) + + assert manager._oauth_discovery_slot(server.server_id) is None + def _assert_oauth_discovery_state_removed(self, manager, server_id): assert manager._oauth_discovery_slot(server_id) is None From 55e80849d1066f6ba1a4d3bb61597bd27a954a91 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 17 Aug 2026 15:18:53 -0700 Subject: [PATCH 051/538] feat(guardrails): track bedrock guardrail usage units per invocation --- .../migration.sql | 20 ++++ .../litellm_proxy_extras/schema.prisma | 16 +++ .../guardrail_hooks/bedrock_guardrails.py | 7 ++ litellm/proxy/guardrails/usage_endpoints.py | 88 +++++++++++++- litellm/proxy/guardrails/usage_tracking.py | 109 +++++++++++++----- .../proxy/hooks/proxy_track_cost_callback.py | 13 +++ litellm/proxy/schema.prisma | 16 +++ litellm/repositories/table_repositories.py | 4 + litellm/types/utils.py | 8 +- schema.prisma | 16 +++ .../test_bedrock_guardrails.py | 23 ++++ .../proxy/guardrails/test_usage_endpoints.py | 76 ++++++++++++ .../proxy/guardrails/test_usage_tracking.py | 104 +++++++++++++++++ .../hooks/test_proxy_track_cost_callback.py | 67 +++++++++++ .../test_spend_tracking_utils.py | 32 +++++ type-discipline-budget.json | 4 +- 16 files changed, 569 insertions(+), 34 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260817143646_add_daily_guardrail_usage_units/migration.sql create mode 100644 tests/test_litellm/proxy/guardrails/test_usage_tracking.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260817143646_add_daily_guardrail_usage_units/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260817143646_add_daily_guardrail_usage_units/migration.sql new file mode 100644 index 00000000000..6838eb76f3e --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260817143646_add_daily_guardrail_usage_units/migration.sql @@ -0,0 +1,20 @@ +-- CreateTable +CREATE TABLE "LiteLLM_DailyGuardrailUsageUnits" ( + "guardrail_id" TEXT NOT NULL, + "date" TEXT NOT NULL, + "team_id" TEXT NOT NULL, + "api_key" TEXT NOT NULL, + "usage_unit" TEXT NOT NULL, + "units" BIGINT NOT NULL DEFAULT 0, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_DailyGuardrailUsageUnits_pkey" PRIMARY KEY ("guardrail_id","date","team_id","api_key","usage_unit") +); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyGuardrailUsageUnits_date_idx" ON "LiteLLM_DailyGuardrailUsageUnits"("date"); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyGuardrailUsageUnits_guardrail_id_date_idx" ON "LiteLLM_DailyGuardrailUsageUnits"("guardrail_id", "date"); + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 71345d2ccde..d3c277278ff 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1069,6 +1069,22 @@ model LiteLLM_DailyGuardrailMetrics { @@index([guardrail_id]) } +// Daily guardrail billable usage units (one row per guardrail/day/team/key/unit type) +model LiteLLM_DailyGuardrailUsageUnits { + guardrail_id String + date String // YYYY-MM-DD + team_id String // empty string when the request had no team + api_key String // hashed virtual key; empty string when unknown + usage_unit String // provider counter name, e.g. Bedrock's contentPolicyUnits + units BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([guardrail_id, date, team_id, api_key, usage_unit]) + @@index([date]) + @@index([guardrail_id, date]) +} + // Daily policy metrics for usage dashboard (one row per policy per day) model LiteLLM_DailyPolicyMetrics { policy_id String diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index e8c6eba581c..93cbb989e23 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -2053,6 +2053,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): bedrock_action: Final = response.get("action") if isinstance(bedrock_action, str): tracing_detail["guardrail_action"] = bedrock_action + usage: Final = response.get("usage") + if isinstance(usage, dict): + usage_units: Final = { # mutable-ok: json.dumps'd into spend log metadata downstream + key: value for key, value in usage.items() if isinstance(value, int) + } + if usage_units: + tracing_detail["guardrail_usage"] = usage_units return tracing_detail def _extract_violation_category_names(self, response: BedrockGuardrailResponse) -> list[str]: diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index f48e20257db..be324e3d81b 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -4,8 +4,9 @@ GET /guardrails/usage/overview, /guardrails/usage/detail/:id, /guardrails/usage/ """ import json -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Iterable, Mapping, Sequence from datetime import datetime, timedelta, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, overload from fastapi import APIRouter, Depends, Query @@ -16,6 +17,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.repositories.table_repositories import ( DailyGuardrailMetricsRepository, + DailyGuardrailUsageUnitsRepository, DailyPolicyMetricsRepository, GuardrailsRepository, PolicyRepository, @@ -28,6 +30,7 @@ if TYPE_CHECKING: from prisma import types as prisma_types from prisma.actions import ( LiteLLM_DailyGuardrailMetricsActions, + LiteLLM_DailyGuardrailUsageUnitsActions, LiteLLM_DailyPolicyMetricsActions, LiteLLM_GuardrailsTableActions, LiteLLM_PolicyTableActions, @@ -41,6 +44,8 @@ if TYPE_CHECKING: router: Final = APIRouter() +_EMPTY_UNITS: Final[Mapping[str, int]] = MappingProxyType({}) + def _guardrails_table( prisma_client: "PrismaClient", @@ -92,6 +97,38 @@ async def _find_daily_policy_metrics( return await _daily_policy_metrics_table(prisma_client).find_many(where=where) +def _daily_guardrail_usage_units_table( + prisma_client: "PrismaClient", +) -> "LiteLLM_DailyGuardrailUsageUnitsActions[prisma_models.LiteLLM_DailyGuardrailUsageUnits]": + units_table: Final[LiteLLM_DailyGuardrailUsageUnitsActions[prisma_models.LiteLLM_DailyGuardrailUsageUnits]] = ( + DailyGuardrailUsageUnitsRepository(prisma_client).table + ) + return units_table + + +async def _find_daily_guardrail_usage_units( + prisma_client: "PrismaClient", + where: "prisma_types.LiteLLM_DailyGuardrailUsageUnitsWhereInput", +) -> "Sequence[prisma_models.LiteLLM_DailyGuardrailUsageUnits]": + return await _daily_guardrail_usage_units_table(prisma_client).find_many(where=where) + + +def _sum_counter_units(rows: "Iterable[prisma_models.LiteLLM_DailyGuardrailUsageUnits]") -> Mapping[str, int]: + materialized: Final = tuple(rows) + counter_names: Final = frozenset(r.usage_unit for r in materialized) + return MappingProxyType( + {name: sum(int(r.units) for r in materialized if r.usage_unit == name) for name in counter_names} + ) + + +def _units_by( + rows: "Sequence[prisma_models.LiteLLM_DailyGuardrailUsageUnits]", + key_of: "Callable[[prisma_models.LiteLLM_DailyGuardrailUsageUnits], str]", +) -> Mapping[str, Mapping[str, int]]: + keys: Final = frozenset(key_of(r) for r in rows) + return MappingProxyType({key: _sum_counter_units(r for r in rows if key_of(r) == key) for key in keys}) + + # --- Response models --- @@ -140,6 +177,7 @@ class UsageOverviewRow(BaseModel): avgLatency: float | None status: str # healthy | warning | critical trend: str # up | down | stable + usageUnits: Mapping[str, int] # provider counter name -> billable units in range class UsageOverviewResponse(BaseModel): @@ -148,6 +186,12 @@ class UsageOverviewResponse(BaseModel): totalRequests: int totalBlocked: int passRate: float + totalUsageUnits: Mapping[str, int] + + +class UsageUnitsDailyPoint(BaseModel): + date: str + units: Mapping[str, int] class UsageDetailResponse(BaseModel): @@ -163,6 +207,10 @@ class UsageDetailResponse(BaseModel): trend: str description: str | None time_series: list[UsageChartPoint] + usage_units: Mapping[str, int] + usage_units_daily: Sequence[UsageUnitsDailyPoint] + usage_units_by_team: Mapping[str, Mapping[str, int]] # team_id ("" = no team) -> counter -> units + usage_units_by_key: Mapping[str, Mapping[str, int]] # hashed api key ("" = unknown) -> counter -> units class UsageLogEntry(BaseModel): @@ -278,6 +326,7 @@ def _guardrail_overview_rows( guardrails: "Sequence[_DbOrConfigGuardrail]", agg: Mapping[str, _MetricTotals], prev_agg: Mapping[str, float], + units_agg: Mapping[str, Mapping[str, int]], ) -> list[UsageOverviewRow]: rows: Final[list[UsageOverviewRow]] = [] covered_keys: Final[set[str]] = set() @@ -303,6 +352,7 @@ def _guardrail_overview_rows( prev_fail = float(prev_agg.get(k, 0.0) or 0.0) break trend = _trend_from_comparison(fail_rate, prev_fail) + row_units: Mapping[str, int] = next((units_agg[k] for k in lookup_keys if k in units_agg), _EMPTY_UNITS) rows.append( UsageOverviewRow( id=gid, @@ -315,6 +365,7 @@ def _guardrail_overview_rows( avgLatency=None, status=_status_from_fail_rate(fail_rate), trend=trend, + usageUnits=row_units, ) ) # Add rows for guardrails with metrics but not in guardrails table (e.g. MCP, config) @@ -337,6 +388,7 @@ def _guardrail_overview_rows( avgLatency=None, status=_status_from_fail_rate(fail_rate), trend=trend, + usageUnits=units_agg.get(agg_key, _EMPTY_UNITS), ) ) return rows @@ -366,6 +418,7 @@ def _policy_overview_rows( avgLatency=None, status=_status_from_fail_rate(fail_rate), trend=trend, + usageUnits=_EMPTY_UNITS, ) ) return rows @@ -386,7 +439,9 @@ async def guardrails_usage_overview( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - return UsageOverviewResponse(rows=[], chart=[], totalRequests=0, totalBlocked=0, passRate=100.0) + return UsageOverviewResponse( + rows=[], chart=[], totalRequests=0, totalBlocked=0, passRate=100.0, totalUsageUnits=_EMPTY_UNITS + ) now: Final = datetime.now(timezone.utc) end: Final = end_date or now.strftime("%Y-%m-%d") @@ -413,19 +468,28 @@ async def guardrails_usage_overview( prisma_client, where={"date": {"gte": start_prev, "lt": start}} ) + units_where: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsWhereInput] = { + "date": {"gte": start, "lte": end} + } + units_rows: Final[ + Sequence[prisma_models.LiteLLM_DailyGuardrailUsageUnits] + ] = await _find_daily_guardrail_usage_units(prisma_client, where=units_where) + agg: Final = _aggregate_daily_metrics(metrics, "guardrail_id") prev_agg: Final = _prev_fail_rates(metrics_prev, "guardrail_id") + units_agg: Final = _units_by(units_rows, lambda r: r.guardrail_id) chart: Final = _chart_from_metrics(metrics) total_requests: Final = sum(a["requests"] for a in agg.values()) total_blocked: Final = sum(a["blocked"] for a in agg.values()) pass_rate: Final = (100.0 * (total_requests - total_blocked) / total_requests) if total_requests else 100.0 - rows: Final = _guardrail_overview_rows(guardrails, agg, prev_agg) + rows: Final = _guardrail_overview_rows(guardrails, agg, prev_agg, units_agg) return UsageOverviewResponse( rows=rows, chart=chart, totalRequests=total_requests, totalBlocked=total_blocked, passRate=round(pass_rate, 1), + totalUsageUnits=_sum_counter_units(units_rows), ) except Exception as e: from litellm.proxy.utils import handle_exception_on_proxy @@ -485,6 +549,13 @@ async def guardrails_usage_detail( "date": {"lt": start}, }, ) + units_where: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsWhereInput] = { + "guardrail_id": {"in": metric_ids}, + "date": {"gte": start, "lte": end}, + } + units_rows: Final[ + Sequence[prisma_models.LiteLLM_DailyGuardrailUsageUnits] + ] = await _find_daily_guardrail_usage_units(prisma_client, where=units_where) requests: Final = sum(int(m.requests_evaluated or 0) for m in metrics) blocked: Final = sum(int(m.blocked_count or 0) for m in metrics) @@ -510,6 +581,8 @@ async def guardrails_usage_detail( litellm_params: Final = _to_dict(_get_guardrail_field(guardrail, "litellm_params")) guardrail_info: Final = _to_dict(_get_guardrail_field(guardrail, "guardrail_info")) _guardrail_name: Final = _get_guardrail_field(guardrail, "guardrail_name") + daily_unit_sums: Final = sorted(_units_by(units_rows, lambda r: r.date).items()) + units_daily: Final = tuple(UsageUnitsDailyPoint(date=d, units=units) for d, units in daily_unit_sums) return UsageDetailResponse( guardrail_id=guardrail_id, @@ -524,6 +597,10 @@ async def guardrails_usage_detail( trend=trend, description=guardrail_info.get("description"), time_series=time_series, + usage_units=_sum_counter_units(units_rows), + usage_units_daily=units_daily, + usage_units_by_team=_units_by(units_rows, lambda r: r.team_id), + usage_units_by_key=_units_by(units_rows, lambda r: r.api_key), ) @@ -743,7 +820,9 @@ async def policies_usage_overview( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - return UsageOverviewResponse(rows=[], chart=[], totalRequests=0, totalBlocked=0, passRate=100.0) + return UsageOverviewResponse( + rows=[], chart=[], totalRequests=0, totalBlocked=0, passRate=100.0, totalUsageUnits=_EMPTY_UNITS + ) now: Final = datetime.now(timezone.utc) end: Final = end_date or now.strftime("%Y-%m-%d") @@ -776,6 +855,7 @@ async def policies_usage_overview( totalRequests=total_requests, totalBlocked=total_blocked, passRate=round(pass_rate, 1), + totalUsageUnits=_EMPTY_UNITS, ) except Exception as e: from litellm.proxy.utils import handle_exception_on_proxy diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index 54dfe8eece1..cebf329b601 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -5,16 +5,25 @@ insert into SpendLogGuardrailIndex when spend logs are written. import json from collections import defaultdict +from collections.abc import Iterator, Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Final +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_proxy_logger from litellm.proxy.utils import PrismaClient from litellm.repositories.table_repositories import ( DailyGuardrailMetricsRepository, + DailyGuardrailUsageUnitsRepository, SpendLogGuardrailIndexRepository, ) +if TYPE_CHECKING: + from prisma import types as prisma_types + +_UsageUnitKey = tuple[str, str, str, str, str] +"""(guardrail_id, date, team_id, api_key, usage_unit)""" + def _guardrail_status_to_action(status: str | None) -> str: """Map StandardLogging guardrail_status to blocked/passed/flagged.""" @@ -28,7 +37,7 @@ def _guardrail_status_to_action(status: str | None) -> str: return "passed" -def _parse_guardrail_info_from_payload(payload: dict[str, Any]) -> list[dict[str, Any]]: +def _parse_guardrail_info_from_payload(payload: Mapping[str, Any]) -> Sequence[Mapping[str, Any]]: """Extract guardrail_information from spend log payload metadata.""" meta = payload.get("metadata") if not meta: @@ -53,6 +62,68 @@ def _date_str(dt: datetime) -> str: return dt.astimezone(timezone.utc).strftime("%Y-%m-%d") +def _parse_payload_start_time(payload: Mapping[str, Any]) -> datetime | None: + start_time: Final = payload.get("startTime") + if isinstance(start_time, datetime): + return start_time + if not isinstance(start_time, str): + return None + try: + return datetime.fromisoformat(start_time.replace("Z", "+00:00")) + except (ValueError, TypeError): + return None + + +def _iter_usage_unit_increments(logs_to_process: Sequence[Mapping[str, Any]]) -> Iterator[tuple[_UsageUnitKey, int]]: + for payload in logs_to_process: + start_time = _parse_payload_start_time(payload) + if start_time is None: + continue + date_key = _date_str(start_time) + team_id = str(payload.get("team_id") or "") + api_key = str(payload.get("api_key") or "") + for entry in _parse_guardrail_info_from_payload(payload): + guardrail_id = entry.get("guardrail_id") or entry.get("guardrail_name") or "" + usage = entry.get("guardrail_usage") + if not guardrail_id or not isinstance(usage, dict): + continue + for unit_name, units in usage.items(): + if isinstance(units, int) and not isinstance(units, bool) and units > 0: + yield (guardrail_id, date_key, team_id, api_key, unit_name), units + + +def _sum_usage_unit_increments(logs_to_process: Sequence[Mapping[str, Any]]) -> Mapping[_UsageUnitKey, int]: + increments: Final = tuple(_iter_usage_unit_increments(logs_to_process)) + keys: Final = frozenset(k for k, _ in increments) + return MappingProxyType({key: sum(u for k, u in increments if k == key) for key in keys}) + + +async def _upsert_usage_unit_row(prisma_client: PrismaClient, key: _UsageUnitKey, units: int) -> None: + guardrail_id, date_key, team_id, api_key, usage_unit = key + row: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsCreateInput] = { + "guardrail_id": guardrail_id, + "date": date_key, + "team_id": team_id, + "api_key": api_key, + "usage_unit": usage_unit, + "units": units, + } + where: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsWhereUniqueInput] = { + "guardrail_id_date_team_id_api_key_usage_unit": { + "guardrail_id": guardrail_id, + "date": date_key, + "team_id": team_id, + "api_key": api_key, + "usage_unit": usage_unit, + } + } + data: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsUpsertInput] = { + "create": row, + "update": {"units": {"increment": units}}, + } + await DailyGuardrailUsageUnitsRepository(prisma_client).table.upsert(where=where, data=data) + + async def process_spend_logs_guardrail_usage( prisma_client: PrismaClient, logs_to_process: list[dict[str, Any]], @@ -76,14 +147,9 @@ async def process_spend_logs_guardrail_usage( for payload in logs_to_process: request_id = payload.get("request_id") - start_time = payload.get("startTime") - if not request_id or not start_time: + start_time = _parse_payload_start_time(payload) + if not request_id or start_time is None: continue - if isinstance(start_time, str): - try: - start_time = datetime.fromisoformat(start_time.replace("Z", "+00:00")) - except (ValueError, TypeError): - continue date_key = _date_str(start_time) for entry in _parse_guardrail_info_from_payload(payload): @@ -109,31 +175,17 @@ async def process_spend_logs_guardrail_usage( } ) - if not daily_guardrail and not index_rows: + usage_unit_totals: Final = _sum_usage_unit_increments(logs_to_process) + + if not daily_guardrail and not index_rows and not usage_unit_totals: return try: # Insert index rows (skip duplicates by request_id + guardrail_id) if index_rows: - index_data: Final = [] - for r in index_rows: - st = r["start_time"] - if isinstance(st, str): - try: - st = datetime.fromisoformat(st.replace("Z", "+00:00")) - except (ValueError, TypeError): - continue - index_data.append( - { - "request_id": r["request_id"], - "guardrail_id": r["guardrail_id"], - "policy_id": r.get("policy_id"), - "start_time": st, - } - ) try: await SpendLogGuardrailIndexRepository(prisma_client).table.create_many( - data=index_data, + data=index_rows, skip_duplicates=True, ) except Exception as e: @@ -168,5 +220,8 @@ async def process_spend_logs_guardrail_usage( }, }, ) + + for unit_key, units in usage_unit_totals.items(): + await _upsert_usage_unit_row(prisma_client, unit_key, units) except Exception as e: verbose_proxy_logger.warning("Guardrail usage tracking failed (non-fatal): %s", e) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 4551680e1b4..da44c06040d 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -125,6 +125,19 @@ class _ProxyDBLogger(CustomLogger): existing_metadata: Final[dict] = request_data.get("metadata", None) or {} existing_metadata.update(_metadata) + # Guardrail hooks write standard_logging_guardrail_information into the + # request's litellm_metadata bucket when one exists (get_or_create_metadata_bucket + # prefers it). Failure rows are serialized from the metadata bucket lifted below, + # so carry the guardrail info over or blocked invocations lose it in spend logs. + litellm_metadata_bucket: Final = request_data.get("litellm_metadata") + if ( + isinstance(litellm_metadata_bucket, dict) + and "standard_logging_guardrail_information" not in existing_metadata + ): + guardrail_info: Final = litellm_metadata_bucket.get("standard_logging_guardrail_information") + if guardrail_info is not None: + existing_metadata["standard_logging_guardrail_information"] = guardrail_info + if "litellm_params" not in request_data: request_data["litellm_params"] = {} diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 71345d2ccde..d3c277278ff 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1069,6 +1069,22 @@ model LiteLLM_DailyGuardrailMetrics { @@index([guardrail_id]) } +// Daily guardrail billable usage units (one row per guardrail/day/team/key/unit type) +model LiteLLM_DailyGuardrailUsageUnits { + guardrail_id String + date String // YYYY-MM-DD + team_id String // empty string when the request had no team + api_key String // hashed virtual key; empty string when unknown + usage_unit String // provider counter name, e.g. Bedrock's contentPolicyUnits + units BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([guardrail_id, date, team_id, api_key, usage_unit]) + @@index([date]) + @@index([guardrail_id, date]) +} + // Daily policy metrics for usage dashboard (one row per policy per day) model LiteLLM_DailyPolicyMetrics { policy_id String diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index be19f290ba6..131f4d377ef 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -158,6 +158,10 @@ class DailyGuardrailMetricsRepository(PrismaTableRepository): table_name = "litellm_dailyguardrailmetrics" +class DailyGuardrailUsageUnitsRepository(PrismaTableRepository): + table_name = "litellm_dailyguardrailusageunits" + + class PolicyAttachmentRepository(PrismaTableRepository): table_name = "litellm_policyattachmenttable" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 272fbabf807..9e68a5de9e5 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -39,7 +39,7 @@ from pydantic import ( field_serializer, field_validator, ) -from typing_extensions import Required, TypedDict +from typing_extensions import ReadOnly, Required, TypedDict from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -3007,6 +3007,11 @@ class StandardLoggingGuardrailInformation(TypedDict, total=False): surface it as a queryable span attribute without parsing the raw guardrail_response blob.""" + guardrail_usage: ReadOnly[Mapping[str, int] | None] + """Provider-reported billable usage counters for this invocation, keyed by the + provider's counter name (e.g. Bedrock's ``contentPolicyUnits``). Kept as a + sibling of guardrail_response so spend-log prompt redaction never drops it.""" + class EvalVerdict(TypedDict, total=False): criterion_name: str @@ -3050,6 +3055,7 @@ class GuardrailTracingDetail(TypedDict, total=False): risk_score: float | None violation_categories: list[str] | None guardrail_action: str | None + guardrail_usage: ReadOnly[Mapping[str, int] | None] StandardLoggingPayloadStatus = Literal["success", "failure"] diff --git a/schema.prisma b/schema.prisma index 71345d2ccde..d3c277278ff 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1069,6 +1069,22 @@ model LiteLLM_DailyGuardrailMetrics { @@index([guardrail_id]) } +// Daily guardrail billable usage units (one row per guardrail/day/team/key/unit type) +model LiteLLM_DailyGuardrailUsageUnits { + guardrail_id String + date String // YYYY-MM-DD + team_id String // empty string when the request had no team + api_key String // hashed virtual key; empty string when unknown + usage_unit String // provider counter name, e.g. Bedrock's contentPolicyUnits + units BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([guardrail_id, date, team_id, api_key, usage_unit]) + @@index([date]) + @@index([guardrail_id, date]) +} + // Daily policy metrics for usage dashboard (one row per policy per day) model LiteLLM_DailyPolicyMetrics { policy_id String diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index f4f4003d5ee..53921e7e74a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5077,3 +5077,26 @@ async def test_apply_guardrail_failure_logs_a_dict_not_a_bare_string(): logged = mock_log.call_args.kwargs["guardrail_json_response"] assert isinstance(logged, dict), f"expected a dict, got {type(logged).__name__}" assert "error" in logged + + +def test_build_tracing_detail_surfaces_usage_counters(): + """LIT-5650: the billable usage block Bedrock returns per ApplyGuardrail call must + land on the tracing detail as guardrail_usage so it reaches spend logs as a + sibling of guardrail_response (which default redaction replaces wholesale).""" + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") + + detail = guardrail._build_tracing_detail( + { + "action": "GUARDRAIL_INTERVENED", + "usage": {"topicPolicyUnits": 1, "contentPolicyUnits": 2, "wordPolicyUnits": 0, "oddball": "not-an-int"}, + } + ) + + assert detail["guardrail_usage"] == {"topicPolicyUnits": 1, "contentPolicyUnits": 2, "wordPolicyUnits": 0} + + +def test_build_tracing_detail_omits_guardrail_usage_when_bedrock_reports_none(): + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") + + assert "guardrail_usage" not in guardrail._build_tracing_detail({"action": "NONE"}) + assert "guardrail_usage" not in guardrail._build_tracing_detail({"action": "NONE", "usage": {}}) diff --git a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py index bf7b1b3b238..23c7554c286 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py @@ -79,18 +79,38 @@ def _metric(guardrail_id: str, date: str = "2026-04-25", requests: int = 10, pas return m +def _units_row( + guardrail_id: str, + date: str = "2026-04-25", + team_id: str = "", + api_key: str = "", + usage_unit: str = "contentPolicyUnits", + units: int = 1, +) -> Any: + r = MagicMock() + r.guardrail_id = guardrail_id + r.date = date + r.team_id = team_id + r.api_key = api_key + r.usage_unit = usage_unit + r.units = units + return r + + def _prisma( *, find_many=None, find_unique=None, metrics=None, index_find_many=None, + units=None, ) -> MagicMock: client = MagicMock() db = client.db db.litellm_guardrailstable.find_many = AsyncMock(return_value=find_many or []) db.litellm_guardrailstable.find_unique = AsyncMock(return_value=find_unique) db.litellm_dailyguardrailmetrics.find_many = AsyncMock(return_value=metrics or []) + db.litellm_dailyguardrailusageunits.find_many = AsyncMock(return_value=units or []) db.litellm_spendlogguardrailindex.find_many = AsyncMock(return_value=index_find_many or []) db.litellm_spendlogguardrailindex.count = AsyncMock(return_value=0) db.litellm_spendlogs.find_many = AsyncMock(return_value=[]) @@ -215,6 +235,62 @@ async def test_overview_excludes_db_sourced_in_memory_entry(): assert "stale" not in ids +@pytest.mark.asyncio +async def test_overview_reports_usage_units_per_row_and_total(): + """LIT-5650: billable units must surface per guardrail row (matched by + logical name like the daily metrics) and as a response-level total.""" + prisma = _prisma( + find_many=[], + metrics=[_metric("yaml-pii", requests=4, passed=3, blocked=1)], + units=[ + _units_row("yaml-pii", usage_unit="topicPolicyUnits", units=4), + _units_row("yaml-pii", usage_unit="contentPolicyUnits", units=3), + _units_row("yaml-pii", team_id="team-a", usage_unit="contentPolicyUnits", units=2), + _units_row("other-guard", usage_unit="topicPolicyUnits", units=7), + ], + ) + handler = _config_handler(_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii")) + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN) + row = next(r for r in resp.rows if r.id == "yaml-uuid") + assert row.usageUnits == {"topicPolicyUnits": 4, "contentPolicyUnits": 5} + assert resp.totalUsageUnits == {"topicPolicyUnits": 11, "contentPolicyUnits": 5} + + +@pytest.mark.asyncio +async def test_detail_breaks_units_down_by_day_team_and_key(): + prisma = _prisma( + find_unique=None, + units=[ + _units_row("yaml-pii", date="2026-04-25", team_id="team-a", api_key="hash-1", units=2), + _units_row("yaml-pii", date="2026-04-25", team_id="", api_key="hash-2", units=1), + _units_row( + "yaml-pii", date="2026-04-24", team_id="team-a", api_key="hash-1", usage_unit="topicPolicyUnits" + ), + ], + ) + handler = _config_handler(_yaml_guardrail()) + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_detail( + guardrail_id="yaml-1", start_date=START, end_date=END, user_api_key_dict=ADMIN + ) + assert resp.usage_units == {"contentPolicyUnits": 3, "topicPolicyUnits": 1} + assert [p.model_dump() for p in resp.usage_units_daily] == [ + {"date": "2026-04-24", "units": {"topicPolicyUnits": 1}}, + {"date": "2026-04-25", "units": {"contentPolicyUnits": 3}}, + ] + assert resp.usage_units_by_team == { + "team-a": {"contentPolicyUnits": 2, "topicPolicyUnits": 1}, + "": {"contentPolicyUnits": 1}, + } + assert resp.usage_units_by_key == { + "hash-1": {"contentPolicyUnits": 2, "topicPolicyUnits": 1}, + "hash-2": {"contentPolicyUnits": 1}, + } + + # ---- logs ------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py new file mode 100644 index 00000000000..7adc1dd314d --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -0,0 +1,104 @@ +import json +from datetime import datetime, timezone +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy.guardrails.usage_tracking import process_spend_logs_guardrail_usage + + +def _prisma() -> MagicMock: + client = MagicMock() + db = client.db + db.litellm_dailyguardrailmetrics.upsert = AsyncMock() + db.litellm_dailyguardrailusageunits.upsert = AsyncMock() + db.litellm_spendlogguardrailindex.create_many = AsyncMock() + return client + + +def _payload( + request_id: str, + *, + team_id: str | None = "team-a", + api_key: str = "hashed-key-1", + usage: dict[str, Any] | None = None, + guardrail_status: str = "success", +) -> dict[str, Any]: + entry: dict[str, Any] = { + "guardrail_id": "bedrock-guard", + "guardrail_status": guardrail_status, + } + if usage is not None: + entry["guardrail_usage"] = usage + return { + "request_id": request_id, + "startTime": datetime(2026, 8, 17, 12, 0, tzinfo=timezone.utc), + "team_id": team_id, + "api_key": api_key, + "metadata": json.dumps({"guardrail_information": [entry]}), + } + + +def _units_upserts(prisma: MagicMock) -> dict[tuple, int]: + calls = prisma.db.litellm_dailyguardrailusageunits.upsert.call_args_list + out: dict[tuple, int] = {} + for c in calls: + where = c.kwargs["where"]["guardrail_id_date_team_id_api_key_usage_unit"] + create = c.kwargs["data"]["create"] + assert create["units"] == c.kwargs["data"]["update"]["units"]["increment"] + assert {k: create[k] for k in where} == where + out[tuple(where[k] for k in ("guardrail_id", "date", "team_id", "api_key", "usage_unit"))] = create["units"] + return out + + +@pytest.mark.asyncio +async def test_usage_units_rolled_up_by_guardrail_team_key_and_date(): + """ + LIT-5650: billable units must aggregate per (guardrail, date, team, key, + counter): same-key payloads sum into one upsert, a team-less payload gets + its own empty-string-team row, and blocked invocations (which Bedrock + still bills for) count exactly like passed ones. + """ + prisma = _prisma() + logs = [ + _payload("r1", usage={"topicPolicyUnits": 1, "contentPolicyUnits": 1}), + _payload( + "r2", + usage={"topicPolicyUnits": 1, "contentPolicyUnits": 2}, + guardrail_status="guardrail_intervened", + ), + _payload("r3", team_id=None, api_key="hashed-key-2", usage={"topicPolicyUnits": 1}), + ] + + await process_spend_logs_guardrail_usage(prisma, logs) + + assert _units_upserts(prisma) == { + ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 2, + ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 3, + ("bedrock-guard", "2026-08-17", "", "hashed-key-2", "topicPolicyUnits"): 1, + } + + +@pytest.mark.asyncio +async def test_zero_and_non_int_usage_counters_are_skipped(): + prisma = _prisma() + logs = [ + _payload( + "r1", + usage={ + "topicPolicyUnits": 1, + "wordPolicyUnits": 0, + "contentPolicyImageUnits": 0, + "oddball": "not-an-int", + "boolish": True, + }, + ), + _payload("r2", usage=None), + ] + + await process_spend_logs_guardrail_usage(prisma, logs) + + assert _units_upserts(prisma) == { + ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 1, + } 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 2b162774aea..bca8210baa6 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 @@ -85,6 +85,73 @@ async def test_async_post_call_failure_hook(): assert metadata["original_key"] == "original_value" +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_carries_guardrail_info_from_litellm_metadata(): + """ + LIT-5650 regression: on a pre_call guardrail block the unified guardrail + layer seeds request_data["litellm_metadata"], so the guardrail hook writes + standard_logging_guardrail_information there, while the failure spend log + is serialized from request_data["metadata"]. Blocked invocations still + consume provider usage units, so the info must be carried over or the + failure row logs guardrail_information: null. + """ + logger = _ProxyDBLogger() + guardrail_info = [ + { + "guardrail_name": "bedrock-guard", + "guardrail_status": "guardrail_intervened", + "guardrail_usage": {"topicPolicyUnits": 1, "contentPolicyUnits": 1}, + } + ] + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {"original_key": "original_value"}, + "litellm_metadata": {"standard_logging_guardrail_information": guardrail_info}, + "proxy_server_request": {"request_id": "test_request_id"}, + } + + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database: + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("Violated guardrail policy"), + user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"), + ) + + metadata = mock_update_database.call_args[1]["kwargs"]["litellm_params"]["metadata"] + assert metadata["standard_logging_guardrail_information"] == guardrail_info + assert metadata["original_key"] == "original_value" + + +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_does_not_clobber_guardrail_info_in_metadata(): + logger = _ProxyDBLogger() + metadata_bucket_info = [{"guardrail_name": "from-metadata-bucket"}] + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {"standard_logging_guardrail_information": metadata_bucket_info}, + "litellm_metadata": {"standard_logging_guardrail_information": [{"guardrail_name": "from-litellm-bucket"}]}, + "proxy_server_request": {"request_id": "test_request_id"}, + } + + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database: + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("Test exception"), + user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"), + ) + + metadata = mock_update_database.call_args[1]["kwargs"]["litellm_params"]["metadata"] + assert metadata["standard_logging_guardrail_information"] == metadata_bucket_info + + @pytest.mark.asyncio async def test_async_post_call_failure_hook_non_llm_route(): # Setup 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 0f6ac3f9b4f..33d835652cd 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 @@ -1565,6 +1565,38 @@ def test_sanitize_guardrail_information_redacts_prompt_fields_when_flag_false( } +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") +def test_sanitize_guardrail_information_preserves_guardrail_usage_when_flag_false( + mock_should_store, +): + """ + LIT-5650 regression: provider-reported billable usage counters live in + guardrail_usage, a sibling of guardrail_response, precisely so the + default spend-log redaction cannot drop them. The response blob (which + also embeds a usage copy) must still be redacted wholesale. + """ + mock_should_store.return_value = False + guardrail_info = [ + { + "guardrail_name": "bedrock-guard", + "guardrail_status": "guardrail_intervened", + "guardrail_response": { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "Sorry, the model cannot answer this question."}], + "usage": {"topicPolicyUnits": 1, "contentPolicyUnits": 1}, + }, + "guardrail_usage": {"topicPolicyUnits": 1, "contentPolicyUnits": 1, "wordPolicyUnits": 0}, + } + ] + + result = _sanitize_guardrail_information_for_spend_logs(guardrail_info) + + assert result is not None + entry = result[0] + assert entry["guardrail_response"] == REDACTED_BY_LITELM_STRING + assert entry["guardrail_usage"] == {"topicPolicyUnits": 1, "contentPolicyUnits": 1, "wordPolicyUnits": 0} + + @patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_sanitize_guardrail_information_passthrough_when_flag_true( mock_should_store, diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 6d70a6aa5f4..a37bb7ef49f 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22909 + "limit": 22906 }, "LIT002": { - "limit": 26898 + "limit": 26896 }, "LIT003": { "limit": 269 From b7593a99c74a8f842c9d65e44a41725320a6c5d2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 17 Aug 2026 16:01:53 -0700 Subject: [PATCH 052/538] fix(guardrails): keep remaining usage upserts when one write fails Per-row guards in the daily metrics and usage unit flush so a single DB error no longer drops the rest of the batch, plus removal of narrating comments flagged in review --- litellm/proxy/guardrails/usage_endpoints.py | 6 +- litellm/proxy/guardrails/usage_tracking.py | 61 ++++++++++++------- .../proxy/hooks/proxy_track_cost_callback.py | 4 -- .../proxy/guardrails/test_usage_tracking.py | 20 ++++++ 4 files changed, 61 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index be324e3d81b..da29f27458d 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -177,7 +177,7 @@ class UsageOverviewRow(BaseModel): avgLatency: float | None status: str # healthy | warning | critical trend: str # up | down | stable - usageUnits: Mapping[str, int] # provider counter name -> billable units in range + usageUnits: Mapping[str, int] class UsageOverviewResponse(BaseModel): @@ -209,8 +209,8 @@ class UsageDetailResponse(BaseModel): time_series: list[UsageChartPoint] usage_units: Mapping[str, int] usage_units_daily: Sequence[UsageUnitsDailyPoint] - usage_units_by_team: Mapping[str, Mapping[str, int]] # team_id ("" = no team) -> counter -> units - usage_units_by_key: Mapping[str, Mapping[str, int]] # hashed api key ("" = unknown) -> counter -> units + usage_units_by_team: Mapping[str, Mapping[str, int]] + usage_units_by_key: Mapping[str, Mapping[str, int]] class UsageLogEntry(BaseModel): diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index cebf329b601..22c98868598 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -196,32 +196,47 @@ async def process_spend_logs_guardrail_usage( n = int(agg["requests_evaluated"]) if n == 0: continue - await DailyGuardrailMetricsRepository(prisma_client).table.upsert( - where={ - "guardrail_id_date": { - "guardrail_id": guardrail_id, - "date": date_key, - } - }, - data={ - "create": { - "guardrail_id": guardrail_id, - "date": date_key, - "requests_evaluated": n, - "passed_count": int(agg["passed_count"]), - "blocked_count": int(agg["blocked_count"]), - "flagged_count": int(agg["flagged_count"]), + try: + await DailyGuardrailMetricsRepository(prisma_client).table.upsert( + where={ + "guardrail_id_date": { + "guardrail_id": guardrail_id, + "date": date_key, + } }, - "update": { - "requests_evaluated": {"increment": n}, - "passed_count": {"increment": int(agg["passed_count"])}, - "blocked_count": {"increment": int(agg["blocked_count"])}, - "flagged_count": {"increment": int(agg["flagged_count"])}, + data={ + "create": { + "guardrail_id": guardrail_id, + "date": date_key, + "requests_evaluated": n, + "passed_count": int(agg["passed_count"]), + "blocked_count": int(agg["blocked_count"]), + "flagged_count": int(agg["flagged_count"]), + }, + "update": { + "requests_evaluated": {"increment": n}, + "passed_count": {"increment": int(agg["passed_count"])}, + "blocked_count": {"increment": int(agg["blocked_count"])}, + "flagged_count": {"increment": int(agg["flagged_count"])}, + }, }, - }, - ) + ) + except Exception as metrics_error: + verbose_proxy_logger.warning( + "Guardrail usage tracking: daily metrics upsert failed for %s on %s (non-fatal): %s", + guardrail_id, + date_key, + metrics_error, + ) for unit_key, units in usage_unit_totals.items(): - await _upsert_usage_unit_row(prisma_client, unit_key, units) + try: + await _upsert_usage_unit_row(prisma_client, unit_key, units) + except Exception as unit_error: + verbose_proxy_logger.warning( + "Guardrail usage tracking: usage unit upsert failed for %s (non-fatal): %s", + unit_key, + unit_error, + ) except Exception as e: verbose_proxy_logger.warning("Guardrail usage tracking failed (non-fatal): %s", e) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index da44c06040d..5dc92d82bda 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -125,10 +125,6 @@ class _ProxyDBLogger(CustomLogger): existing_metadata: Final[dict] = request_data.get("metadata", None) or {} existing_metadata.update(_metadata) - # Guardrail hooks write standard_logging_guardrail_information into the - # request's litellm_metadata bucket when one exists (get_or_create_metadata_bucket - # prefers it). Failure rows are serialized from the metadata bucket lifted below, - # so carry the guardrail info over or blocked invocations lose it in spend logs. litellm_metadata_bucket: Final = request_data.get("litellm_metadata") if ( isinstance(litellm_metadata_bucket, dict) diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index 7adc1dd314d..67953c66720 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -80,6 +80,26 @@ async def test_usage_units_rolled_up_by_guardrail_team_key_and_date(): } +@pytest.mark.asyncio +async def test_one_failing_upsert_does_not_drop_remaining_writes(): + """ + A DB error on one daily-metrics or usage-unit upsert must not cancel the + remaining upserts in the flushed batch, or the usage endpoints would + permanently under-report billable counters (batches are never retried). + """ + prisma = _prisma() + prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = RuntimeError("db down") + prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [RuntimeError("db down"), None] + logs = [ + _payload("r1", usage={"topicPolicyUnits": 1}), + _payload("r2", team_id=None, api_key="hashed-key-2", usage={"topicPolicyUnits": 1}), + ] + + await process_spend_logs_guardrail_usage(prisma, logs) + + assert prisma.db.litellm_dailyguardrailusageunits.upsert.call_count == 2 + + @pytest.mark.asyncio async def test_zero_and_non_int_usage_counters_are_skipped(): prisma = _prisma() From 5437139b947259057b253b16d350b09ea89d086d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 17 Aug 2026 16:43:58 -0700 Subject: [PATCH 053/538] chore(guardrails): sync lazy OpenAPI snapshot and dashboard types for usage units Regenerates the guardrails and policy_engine fragments of the lazy OpenAPI snapshot for the usage-unit fields, regenerates schema.d.ts from it, and exports DailyGuardrailUsageUnitsRepository next to its sibling repositories --- litellm/proxy/_lazy_openapi_snapshot.json | 141 ++++++++++++++++-- litellm/repositories/__init__.py | 2 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 43 +++++- 3 files changed, 171 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 7fe02c6d8bc..026a02d6b1d 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -11349,6 +11349,40 @@ "title": "UpdateGuardrailRequest", "type": "object" }, + "UsageChartPoint": { + "properties": { + "blocked": { + "title": "Blocked", + "type": "integer" + }, + "date": { + "title": "Date", + "type": "string" + }, + "passed": { + "title": "Passed", + "type": "integer" + }, + "score": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Score" + } + }, + "required": [ + "date", + "passed", + "blocked" + ], + "title": "UsageChartPoint", + "type": "object" + }, "UsageDetailResponse": { "properties": { "avgLatency": { @@ -11410,8 +11444,7 @@ }, "time_series": { "items": { - "additionalProperties": true, - "type": "object" + "$ref": "#/components/schemas/UsageChartPoint" }, "title": "Time Series", "type": "array" @@ -11423,6 +11456,40 @@ "type": { "title": "Type", "type": "string" + }, + "usage_units": { + "additionalProperties": { + "type": "integer" + }, + "title": "Usage Units", + "type": "object" + }, + "usage_units_by_key": { + "additionalProperties": { + "additionalProperties": { + "type": "integer" + }, + "type": "object" + }, + "title": "Usage Units By Key", + "type": "object" + }, + "usage_units_by_team": { + "additionalProperties": { + "additionalProperties": { + "type": "integer" + }, + "type": "object" + }, + "title": "Usage Units By Team", + "type": "object" + }, + "usage_units_daily": { + "items": { + "$ref": "#/components/schemas/UsageUnitsDailyPoint" + }, + "title": "Usage Units Daily", + "type": "array" } }, "required": [ @@ -11437,7 +11504,11 @@ "status", "trend", "description", - "time_series" + "time_series", + "usage_units", + "usage_units_daily", + "usage_units_by_team", + "usage_units_by_key" ], "title": "UsageDetailResponse", "type": "object" @@ -11572,8 +11643,7 @@ "properties": { "chart": { "items": { - "additionalProperties": true, - "type": "object" + "$ref": "#/components/schemas/UsageChartPoint" }, "title": "Chart", "type": "array" @@ -11596,6 +11666,13 @@ "totalRequests": { "title": "Totalrequests", "type": "integer" + }, + "totalUsageUnits": { + "additionalProperties": { + "type": "integer" + }, + "title": "Totalusageunits", + "type": "object" } }, "required": [ @@ -11603,7 +11680,8 @@ "chart", "totalRequests", "totalBlocked", - "passRate" + "passRate", + "totalUsageUnits" ], "title": "UsageOverviewResponse", "type": "object" @@ -11663,6 +11741,13 @@ "type": { "title": "Type", "type": "string" + }, + "usageUnits": { + "additionalProperties": { + "type": "integer" + }, + "title": "Usageunits", + "type": "object" } }, "required": [ @@ -11675,11 +11760,33 @@ "avgScore", "avgLatency", "status", - "trend" + "trend", + "usageUnits" ], "title": "UsageOverviewRow", "type": "object" }, + "UsageUnitsDailyPoint": { + "properties": { + "date": { + "title": "Date", + "type": "string" + }, + "units": { + "additionalProperties": { + "type": "integer" + }, + "title": "Units", + "type": "object" + } + }, + "required": [ + "date", + "units" + ], + "title": "UsageUnitsDailyPoint", + "type": "object" + }, "ValidationError": { "properties": { "loc": { @@ -21477,6 +21584,13 @@ "totalRequests": { "title": "Totalrequests", "type": "integer" + }, + "totalUsageUnits": { + "additionalProperties": { + "type": "integer" + }, + "title": "Totalusageunits", + "type": "object" } }, "required": [ @@ -21484,7 +21598,8 @@ "chart", "totalRequests", "totalBlocked", - "passRate" + "passRate", + "totalUsageUnits" ], "title": "UsageOverviewResponse", "type": "object" @@ -21544,6 +21659,13 @@ "type": { "title": "Type", "type": "string" + }, + "usageUnits": { + "additionalProperties": { + "type": "integer" + }, + "title": "Usageunits", + "type": "object" } }, "required": [ @@ -21556,7 +21678,8 @@ "avgScore", "avgLatency", "status", - "trend" + "trend", + "usageUnits" ], "title": "UsageOverviewRow", "type": "object" diff --git a/litellm/repositories/__init__.py b/litellm/repositories/__init__.py index e2e7f1fac73..881f7a66cea 100644 --- a/litellm/repositories/__init__.py +++ b/litellm/repositories/__init__.py @@ -28,6 +28,7 @@ from litellm.repositories.table_repositories import ( ClaudeCodePluginRepository, ConfigOverridesRepository, DailyGuardrailMetricsRepository, + DailyGuardrailUsageUnitsRepository, DailyPolicyMetricsRepository, DailyTagSpendRepository, DailyToolSpendRepository, @@ -101,6 +102,7 @@ __all__ = [ "ConfigRepository", "CredentialsRepository", "DailyGuardrailMetricsRepository", + "DailyGuardrailUsageUnitsRepository", "DailyPolicyMetricsRepository", "DailyTagSpendRepository", "DailyToolSpendRepository", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 2bc21735248..318f02750ce 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35089,13 +35089,29 @@ export interface components { /** Status */ status: string; /** Time Series */ - time_series: { - [key: string]: unknown; - }[]; + time_series: components["schemas"]["UsageChartPoint"][]; /** Trend */ trend: string; /** Type */ type: string; + /** Usage Units */ + usage_units: { + [key: string]: number; + }; + /** Usage Units By Key */ + usage_units_by_key: { + [key: string]: { + [key: string]: number; + }; + }; + /** Usage Units By Team */ + usage_units_by_team: { + [key: string]: { + [key: string]: number; + }; + }; + /** Usage Units Daily */ + usage_units_daily: components["schemas"]["UsageUnitsDailyPoint"][]; }; /** UsageLogEntry */ UsageLogEntry: { @@ -35132,9 +35148,7 @@ export interface components { /** UsageOverviewResponse */ UsageOverviewResponse: { /** Chart */ - chart: { - [key: string]: unknown; - }[]; + chart: components["schemas"]["UsageChartPoint"][]; /** Passrate */ passRate: number; /** Rows */ @@ -35143,6 +35157,10 @@ export interface components { totalBlocked: number; /** Totalrequests */ totalRequests: number; + /** Totalusageunits */ + totalUsageUnits: { + [key: string]: number; + }; }; /** UsageOverviewRow */ UsageOverviewRow: { @@ -35166,6 +35184,19 @@ export interface components { trend: string; /** Type */ type: string; + /** Usageunits */ + usageUnits: { + [key: string]: number; + }; + }; + /** UsageUnitsDailyPoint */ + UsageUnitsDailyPoint: { + /** Date */ + date: string; + /** Units */ + units: { + [key: string]: number; + }; }; /** * UserAPIKeyAuth From 8ba2263d4c9d8518513f48ed3522fa18e39fd16d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 17 Aug 2026 17:19:05 -0700 Subject: [PATCH 054/538] perf(guardrails): aggregate usage units in one sorted pass The flush and the usage endpoints summed units with a scan per distinct key, quadratic in rows times keys; group sorted rows instead. Skip payloads without a request_id like the metrics path, type the flush key as a NamedTuple, and drop the (guardrail_id, date) index that the primary key already covers --- .../migration.sql | 4 -- .../litellm_proxy_extras/schema.prisma | 1 - litellm/proxy/guardrails/usage_endpoints.py | 14 ++++-- litellm/proxy/guardrails/usage_tracking.py | 47 +++++++++++-------- litellm/proxy/schema.prisma | 1 - schema.prisma | 1 - .../proxy/guardrails/test_usage_endpoints.py | 4 ++ .../proxy/guardrails/test_usage_tracking.py | 16 +++++++ type-discipline-budget.json | 4 +- 9 files changed, 58 insertions(+), 34 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260817143646_add_daily_guardrail_usage_units/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260817143646_add_daily_guardrail_usage_units/migration.sql index 6838eb76f3e..7244312c6b0 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260817143646_add_daily_guardrail_usage_units/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260817143646_add_daily_guardrail_usage_units/migration.sql @@ -14,7 +14,3 @@ CREATE TABLE "LiteLLM_DailyGuardrailUsageUnits" ( -- CreateIndex CREATE INDEX "LiteLLM_DailyGuardrailUsageUnits_date_idx" ON "LiteLLM_DailyGuardrailUsageUnits"("date"); - --- CreateIndex -CREATE INDEX "LiteLLM_DailyGuardrailUsageUnits_guardrail_id_date_idx" ON "LiteLLM_DailyGuardrailUsageUnits"("guardrail_id", "date"); - diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index d3c277278ff..24c0f1f11cc 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1082,7 +1082,6 @@ model LiteLLM_DailyGuardrailUsageUnits { @@id([guardrail_id, date, team_id, api_key, usage_unit]) @@index([date]) - @@index([guardrail_id, date]) } // Daily policy metrics for usage dashboard (one row per policy per day) diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index da29f27458d..a73efed30ad 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -6,6 +6,7 @@ GET /guardrails/usage/overview, /guardrails/usage/detail/:id, /guardrails/usage/ import json from collections.abc import Callable, Iterable, Mapping, Sequence from datetime import datetime, timedelta, timezone +from itertools import groupby from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, overload @@ -113,11 +114,14 @@ async def _find_daily_guardrail_usage_units( return await _daily_guardrail_usage_units_table(prisma_client).find_many(where=where) +def _counter_name(row: "prisma_models.LiteLLM_DailyGuardrailUsageUnits") -> str: + return row.usage_unit + + def _sum_counter_units(rows: "Iterable[prisma_models.LiteLLM_DailyGuardrailUsageUnits]") -> Mapping[str, int]: - materialized: Final = tuple(rows) - counter_names: Final = frozenset(r.usage_unit for r in materialized) + ordered: Final = sorted(rows, key=_counter_name) return MappingProxyType( - {name: sum(int(r.units) for r in materialized if r.usage_unit == name) for name in counter_names} + {name: sum(int(r.units) for r in group) for name, group in groupby(ordered, key=_counter_name)} ) @@ -125,8 +129,8 @@ def _units_by( rows: "Sequence[prisma_models.LiteLLM_DailyGuardrailUsageUnits]", key_of: "Callable[[prisma_models.LiteLLM_DailyGuardrailUsageUnits], str]", ) -> Mapping[str, Mapping[str, int]]: - keys: Final = frozenset(key_of(r) for r in rows) - return MappingProxyType({key: _sum_counter_units(r for r in rows if key_of(r) == key) for key in keys}) + ordered: Final = sorted(rows, key=key_of) + return MappingProxyType({key: _sum_counter_units(group) for key, group in groupby(ordered, key=key_of)}) # --- Response models --- diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index 22c98868598..c25cf07b4d4 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -7,8 +7,10 @@ import json from collections import defaultdict from collections.abc import Iterator, Mapping, Sequence from datetime import datetime, timezone +from itertools import groupby +from operator import itemgetter from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, NamedTuple from litellm._logging import verbose_proxy_logger from litellm.proxy.utils import PrismaClient @@ -21,8 +23,13 @@ from litellm.repositories.table_repositories import ( if TYPE_CHECKING: from prisma import types as prisma_types -_UsageUnitKey = tuple[str, str, str, str, str] -"""(guardrail_id, date, team_id, api_key, usage_unit)""" + +class _UsageUnitKey(NamedTuple): + guardrail_id: str + date: str + team_id: str + api_key: str + usage_unit: str def _guardrail_status_to_action(status: str | None) -> str: @@ -77,44 +84,44 @@ def _parse_payload_start_time(payload: Mapping[str, Any]) -> datetime | None: def _iter_usage_unit_increments(logs_to_process: Sequence[Mapping[str, Any]]) -> Iterator[tuple[_UsageUnitKey, int]]: for payload in logs_to_process: start_time = _parse_payload_start_time(payload) - if start_time is None: + if not payload.get("request_id") or start_time is None: continue date_key = _date_str(start_time) team_id = str(payload.get("team_id") or "") api_key = str(payload.get("api_key") or "") for entry in _parse_guardrail_info_from_payload(payload): - guardrail_id = entry.get("guardrail_id") or entry.get("guardrail_name") or "" + guardrail_id = str(entry.get("guardrail_id") or entry.get("guardrail_name") or "") usage = entry.get("guardrail_usage") if not guardrail_id or not isinstance(usage, dict): continue for unit_name, units in usage.items(): if isinstance(units, int) and not isinstance(units, bool) and units > 0: - yield (guardrail_id, date_key, team_id, api_key, unit_name), units + yield _UsageUnitKey(guardrail_id, date_key, team_id, api_key, str(unit_name)), units def _sum_usage_unit_increments(logs_to_process: Sequence[Mapping[str, Any]]) -> Mapping[_UsageUnitKey, int]: - increments: Final = tuple(_iter_usage_unit_increments(logs_to_process)) - keys: Final = frozenset(k for k, _ in increments) - return MappingProxyType({key: sum(u for k, u in increments if k == key) for key in keys}) + ordered: Final = sorted(_iter_usage_unit_increments(logs_to_process), key=itemgetter(0)) + return MappingProxyType( + {key: sum(units for _, units in group) for key, group in groupby(ordered, key=itemgetter(0))} + ) async def _upsert_usage_unit_row(prisma_client: PrismaClient, key: _UsageUnitKey, units: int) -> None: - guardrail_id, date_key, team_id, api_key, usage_unit = key row: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsCreateInput] = { - "guardrail_id": guardrail_id, - "date": date_key, - "team_id": team_id, - "api_key": api_key, - "usage_unit": usage_unit, + "guardrail_id": key.guardrail_id, + "date": key.date, + "team_id": key.team_id, + "api_key": key.api_key, + "usage_unit": key.usage_unit, "units": units, } where: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsWhereUniqueInput] = { "guardrail_id_date_team_id_api_key_usage_unit": { - "guardrail_id": guardrail_id, - "date": date_key, - "team_id": team_id, - "api_key": api_key, - "usage_unit": usage_unit, + "guardrail_id": key.guardrail_id, + "date": key.date, + "team_id": key.team_id, + "api_key": key.api_key, + "usage_unit": key.usage_unit, } } data: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsUpsertInput] = { diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index d3c277278ff..24c0f1f11cc 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1082,7 +1082,6 @@ model LiteLLM_DailyGuardrailUsageUnits { @@id([guardrail_id, date, team_id, api_key, usage_unit]) @@index([date]) - @@index([guardrail_id, date]) } // Daily policy metrics for usage dashboard (one row per policy per day) diff --git a/schema.prisma b/schema.prisma index d3c277278ff..24c0f1f11cc 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1082,7 +1082,6 @@ model LiteLLM_DailyGuardrailUsageUnits { @@id([guardrail_id, date, team_id, api_key, usage_unit]) @@index([date]) - @@index([guardrail_id, date]) } // Daily policy metrics for usage dashboard (one row per policy per day) diff --git a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py index 23c7554c286..b0f98b81c05 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py @@ -256,6 +256,8 @@ async def test_overview_reports_usage_units_per_row_and_total(): row = next(r for r in resp.rows if r.id == "yaml-uuid") assert row.usageUnits == {"topicPolicyUnits": 4, "contentPolicyUnits": 5} assert resp.totalUsageUnits == {"topicPolicyUnits": 11, "contentPolicyUnits": 5} + units_where = prisma.db.litellm_dailyguardrailusageunits.find_many.call_args.kwargs["where"] + assert units_where == {"date": {"gte": START, "lte": END}} @pytest.mark.asyncio @@ -289,6 +291,8 @@ async def test_detail_breaks_units_down_by_day_team_and_key(): "hash-1": {"contentPolicyUnits": 2, "topicPolicyUnits": 1}, "hash-2": {"contentPolicyUnits": 1}, } + units_where = prisma.db.litellm_dailyguardrailusageunits.find_many.call_args.kwargs["where"] + assert units_where == {"guardrail_id": {"in": ["yaml-pii", "yaml-1"]}, "date": {"gte": START, "lte": END}} # ---- logs ------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index 67953c66720..4f606d489d5 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -122,3 +122,19 @@ async def test_zero_and_non_int_usage_counters_are_skipped(): assert _units_upserts(prisma) == { ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 1, } + + +@pytest.mark.asyncio +async def test_payload_without_request_id_is_skipped_like_the_metrics_path(): + prisma = _prisma() + logs = [ + {**_payload("ignored", usage={"topicPolicyUnits": 5}), "request_id": None}, + _payload("r2", usage={"topicPolicyUnits": 1}), + ] + + await process_spend_logs_guardrail_usage(prisma, logs) + + assert _units_upserts(prisma) == { + ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 1, + } + assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_args.kwargs["data"]["create"]["requests_evaluated"] == 1 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index a37bb7ef49f..d0ec219b85c 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22906 + "limit": 22903 }, "LIT002": { - "limit": 26896 + "limit": 26894 }, "LIT003": { "limit": 269 From ae23bf85d257d0d48cb2e03ae7361957b8d50210 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 17 Aug 2026 17:41:57 -0700 Subject: [PATCH 055/538] fix(guardrails): retry failed daily metrics and usage unit upserts with backoff A transient DB error during the spend log flush dropped that batch's guardrail metrics and usage unit rows for good. Retry only the rows that failed, up to 3 times with 1s/2s/4s backoff, mirroring the daily spend writer, and inject the sleep so tests stay fast. Lowers the lint budgets the refactor freed up --- basedpyright-code-budget.json | 2 +- litellm/proxy/guardrails/usage_tracking.py | 133 +++++++++++------- ruff-strict-budget.json | 2 +- .../proxy/guardrails/test_usage_tracking.py | 54 ++++++- type-discipline-budget.json | 4 +- 5 files changed, 137 insertions(+), 58 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index ba131ecac4c..1a9086ccee7 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5681 }, "reportMissingTypeArgument": { - "limit": 15609 + "limit": 15608 }, "reportMissingTypeStubs": { "limit": 40 diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index c25cf07b4d4..727f767469f 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -3,14 +3,16 @@ Track guardrail and policy usage for the dashboard: upsert daily metrics and insert into SpendLogGuardrailIndex when spend logs are written. """ +import asyncio import json from collections import defaultdict -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from datetime import datetime, timezone +from functools import partial from itertools import groupby from operator import itemgetter from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, NamedTuple +from typing import TYPE_CHECKING, Any, Final, NamedTuple, TypeVar from litellm._logging import verbose_proxy_logger from litellm.proxy.utils import PrismaClient @@ -24,6 +26,12 @@ if TYPE_CHECKING: from prisma import types as prisma_types +_UPSERT_RETRY_TIMES: Final = 3 + +_RowKey = TypeVar("_RowKey") +_RowValue = TypeVar("_RowValue") + + class _UsageUnitKey(NamedTuple): guardrail_id: str date: str @@ -32,6 +40,46 @@ class _UsageUnitKey(NamedTuple): usage_unit: str +class _MetricsKey(NamedTuple): + guardrail_id: str + date: str + + +async def _attempt_upsert( + upsert_row: Callable[[_RowKey, _RowValue], Awaitable[None]], key: _RowKey, value: _RowValue +) -> Exception | None: + try: + await upsert_row(key, value) + except Exception as error: + return error + return None + + +async def _upsert_rows_with_retry( + rows: Mapping[_RowKey, _RowValue], + upsert_row: Callable[[_RowKey, _RowValue], Awaitable[None]], + label: str, + sleep: Callable[[float], Awaitable[None]], + retries_left: int = _UPSERT_RETRY_TIMES, +) -> None: + outcomes: Final = {key: await _attempt_upsert(upsert_row, key, value) for key, value in rows.items()} + failed: Final = MappingProxyType({key: rows[key] for key, error in outcomes.items() if error is not None}) + if not failed: + return + if retries_left == 0: + for key in failed: + verbose_proxy_logger.warning( + "Guardrail usage tracking: %s upsert failed for %s after %d retries (non-fatal): %s", + label, + key, + _UPSERT_RETRY_TIMES, + outcomes[key], + ) + return + await sleep(2 ** (_UPSERT_RETRY_TIMES - retries_left)) + await _upsert_rows_with_retry(failed, upsert_row, label, sleep, retries_left - 1) + + def _guardrail_status_to_action(status: str | None) -> str: """Map StandardLogging guardrail_status to blocked/passed/flagged.""" if not status: @@ -131,9 +179,33 @@ async def _upsert_usage_unit_row(prisma_client: PrismaClient, key: _UsageUnitKey await DailyGuardrailUsageUnitsRepository(prisma_client).table.upsert(where=where, data=data) +async def _upsert_metrics_row(prisma_client: PrismaClient, key: _MetricsKey, agg: Mapping[str, int]) -> None: + n: Final = int(agg["requests_evaluated"]) + await DailyGuardrailMetricsRepository(prisma_client).table.upsert( + where={"guardrail_id_date": {"guardrail_id": key.guardrail_id, "date": key.date}}, + data={ + "create": { + "guardrail_id": key.guardrail_id, + "date": key.date, + "requests_evaluated": n, + "passed_count": int(agg["passed_count"]), + "blocked_count": int(agg["blocked_count"]), + "flagged_count": int(agg["flagged_count"]), + }, + "update": { + "requests_evaluated": {"increment": n}, + "passed_count": {"increment": int(agg["passed_count"])}, + "blocked_count": {"increment": int(agg["blocked_count"])}, + "flagged_count": {"increment": int(agg["flagged_count"])}, + }, + }, + ) + + async def process_spend_logs_guardrail_usage( prisma_client: PrismaClient, logs_to_process: list[dict[str, Any]], + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, ) -> None: """ After spend logs are written: update DailyGuardrailMetrics and insert @@ -142,7 +214,7 @@ async def process_spend_logs_guardrail_usage( if not logs_to_process: return # Aggregate daily metrics by (guardrail_id, date). Latency/score metrics dropped. - daily_guardrail: Final[dict[tuple, dict[str, Any]]] = defaultdict( + daily_guardrail: Final[dict[_MetricsKey, dict[str, Any]]] = defaultdict( lambda: { "requests_evaluated": 0, "passed_count": 0, @@ -163,7 +235,7 @@ async def process_spend_logs_guardrail_usage( guardrail_id = entry.get("guardrail_id") or entry.get("guardrail_name") or "" if not guardrail_id: continue - key = (guardrail_id, date_key) + key = _MetricsKey(guardrail_id, date_key) daily_guardrail[key]["requests_evaluated"] += 1 action = _guardrail_status_to_action(entry.get("guardrail_status")) if action == "passed": @@ -199,51 +271,12 @@ async def process_spend_logs_guardrail_usage( verbose_proxy_logger.debug("Guardrail usage tracking: index create_many skipped: %s", e) # Upsert daily guardrail metrics (counts only; latency/score dropped) - for (guardrail_id, date_key), agg in daily_guardrail.items(): - n = int(agg["requests_evaluated"]) - if n == 0: - continue - try: - await DailyGuardrailMetricsRepository(prisma_client).table.upsert( - where={ - "guardrail_id_date": { - "guardrail_id": guardrail_id, - "date": date_key, - } - }, - data={ - "create": { - "guardrail_id": guardrail_id, - "date": date_key, - "requests_evaluated": n, - "passed_count": int(agg["passed_count"]), - "blocked_count": int(agg["blocked_count"]), - "flagged_count": int(agg["flagged_count"]), - }, - "update": { - "requests_evaluated": {"increment": n}, - "passed_count": {"increment": int(agg["passed_count"])}, - "blocked_count": {"increment": int(agg["blocked_count"])}, - "flagged_count": {"increment": int(agg["flagged_count"])}, - }, - }, - ) - except Exception as metrics_error: - verbose_proxy_logger.warning( - "Guardrail usage tracking: daily metrics upsert failed for %s on %s (non-fatal): %s", - guardrail_id, - date_key, - metrics_error, - ) - - for unit_key, units in usage_unit_totals.items(): - try: - await _upsert_usage_unit_row(prisma_client, unit_key, units) - except Exception as unit_error: - verbose_proxy_logger.warning( - "Guardrail usage tracking: usage unit upsert failed for %s (non-fatal): %s", - unit_key, - unit_error, - ) + metrics_rows: Final = MappingProxyType( + {key: agg for key, agg in daily_guardrail.items() if int(agg["requests_evaluated"]) > 0} + ) + await _upsert_rows_with_retry(metrics_rows, partial(_upsert_metrics_row, prisma_client), "daily metrics", sleep) + await _upsert_rows_with_retry( + usage_unit_totals, partial(_upsert_usage_unit_row, prisma_client), "usage unit", sleep + ) except Exception as e: verbose_proxy_logger.warning("Guardrail usage tracking failed (non-fatal): %s", e) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 15396a95632..ffbb8dc2332 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -78,7 +78,7 @@ "limit": 1 }, "C901": { - "limit": 313 + "limit": 312 }, "D419": { "limit": 6 diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index 4f606d489d5..62a4bbbbe6e 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -80,24 +80,70 @@ async def test_usage_units_rolled_up_by_guardrail_team_key_and_date(): } +def _fake_sleep() -> tuple[AsyncMock, list[float]]: + delays: list[float] = [] + sleep = AsyncMock(side_effect=lambda delay: delays.append(delay)) + return sleep, delays + + @pytest.mark.asyncio async def test_one_failing_upsert_does_not_drop_remaining_writes(): """ A DB error on one daily-metrics or usage-unit upsert must not cancel the remaining upserts in the flushed batch, or the usage endpoints would - permanently under-report billable counters (batches are never retried). + permanently under-report billable counters. """ prisma = _prisma() prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = RuntimeError("db down") - prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [RuntimeError("db down"), None] + prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [RuntimeError("db down"), None, None] + sleep, _ = _fake_sleep() logs = [ _payload("r1", usage={"topicPolicyUnits": 1}), _payload("r2", team_id=None, api_key="hashed-key-2", usage={"topicPolicyUnits": 1}), ] - await process_spend_logs_guardrail_usage(prisma, logs) + await process_spend_logs_guardrail_usage(prisma, logs, sleep=sleep) - assert prisma.db.litellm_dailyguardrailusageunits.upsert.call_count == 2 + assert _units_upserts(prisma) == { + ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 1, + ("bedrock-guard", "2026-08-17", "", "hashed-key-2", "topicPolicyUnits"): 1, + } + + +@pytest.mark.asyncio +async def test_transient_upsert_failure_is_retried_with_backoff_for_failed_rows_only(): + """ + A transient DB error must not permanently drop billed units from the + aggregates: only the rows that failed are re-sent, after exponential + backoff, and the batch ends once every row has landed. + """ + prisma = _prisma() + prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [RuntimeError("blip"), None, None] + sleep, delays = _fake_sleep() + logs = [ + _payload("r1", usage={"topicPolicyUnits": 1}), + _payload("r2", team_id=None, api_key="hashed-key-2", usage={"topicPolicyUnits": 1}), + ] + + await process_spend_logs_guardrail_usage(prisma, logs, sleep=sleep) + + calls = prisma.db.litellm_dailyguardrailusageunits.upsert.call_args_list + assert len(calls) == 3 + assert calls[2].kwargs["where"] == calls[0].kwargs["where"] + assert delays == [1] + + +@pytest.mark.asyncio +async def test_persistent_upsert_failure_stops_after_three_retries(): + prisma = _prisma() + prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = RuntimeError("db down") + sleep, delays = _fake_sleep() + + await process_spend_logs_guardrail_usage(prisma, [_payload("r1", usage={"topicPolicyUnits": 1})], sleep=sleep) + + assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_count == 4 + assert delays == [1, 2, 4] + assert prisma.db.litellm_dailyguardrailusageunits.upsert.call_count == 1 @pytest.mark.asyncio diff --git a/type-discipline-budget.json b/type-discipline-budget.json index d0ec219b85c..3d2923c39ac 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22903 + "limit": 22900 }, "LIT002": { - "limit": 26894 + "limit": 26893 }, "LIT003": { "limit": 269 From dc54b16d3c84c87835d365e193c6f029dbe6aa81 Mon Sep 17 00:00:00 2001 From: Bruno Felthes Date: Mon, 17 Aug 2026 21:43:24 -0300 Subject: [PATCH 056/538] fix(fireworks): skip accounts/ rewrite for FW-* Foundry deployment ids resolve_fireworks_resource_name prefixes bare names with accounts/fireworks/models/ (or routers/ for *-fast). Azure AI Foundry hosts Fireworks models under deployment ids like FW-Kimi-K3; rewriting those yields 404 DeploymentNotFound. Leave names that already start with FW- unchanged. Native Fireworks short names still get the accounts/ path. Co-authored-by: Cursor --- litellm/llms/fireworks_ai/common_utils.py | 5 +++++ .../llms/fireworks_ai/test_fireworks_ai_common_utils.py | 3 +++ 2 files changed, 8 insertions(+) diff --git a/litellm/llms/fireworks_ai/common_utils.py b/litellm/llms/fireworks_ai/common_utils.py index e07e7a26f9e..522c955770c 100644 --- a/litellm/llms/fireworks_ai/common_utils.py +++ b/litellm/llms/fireworks_ai/common_utils.py @@ -33,6 +33,11 @@ def resolve_fireworks_resource_name(model: str) -> str: stripped: Final = model.removeprefix("fireworks_ai/") if stripped.startswith("accounts/") or "#" in stripped: return stripped + # Azure AI Foundry (and similar OpenAI-compat hosts) expose Fireworks + # deployments as ids like ``FW-Kimi-K3``. Rewriting those to + # ``accounts/fireworks/models/FW-…`` yields DeploymentNotFound. + if stripped.startswith("FW-"): + return stripped if stripped.startswith(("routers/", "models/")): return f"accounts/fireworks/{stripped}" if stripped.endswith("-fast"): diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py index 4af395baf41..b52c910d5a6 100644 --- a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py +++ b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py @@ -39,6 +39,9 @@ from litellm.llms.fireworks_ai.common_utils import resolve_fireworks_resource_na "glm-4p6#accounts/gitlab/deployments/2fb7764c", "glm-4p6#accounts/gitlab/deployments/2fb7764c", ), + ("FW-Kimi-K3", "FW-Kimi-K3"), + ("fireworks_ai/FW-Kimi-K3", "FW-Kimi-K3"), + ("FW-GLM-5.2-Fast", "FW-GLM-5.2-Fast"), ], ) def test_resolve_fireworks_resource_name(model, expected): From 35176fa64aec30de01a9b650bb3fc88da71b9244 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 17 Aug 2026 19:05:50 -0700 Subject: [PATCH 057/538] fix(guardrails): retry usage upserts only on connection errors The daily guardrail metrics and usage-unit upserts are non-idempotent increments, but the retry loop re-sent every failed row on any exception. An ambiguous post-send failure such as a read timeout after the write had already committed therefore stacked a second increment and inflated the billable unit totals served by the guardrail usage endpoints. Retry only DB_RETRY_SAFE_ERROR_TYPES (httpx.ConnectError), the same rule the spend writer and autorouter rollup use for increment upserts, and log any other failure once as terminal for that row while the rest of the batch still lands. Follows up #37225 --- litellm/proxy/guardrails/usage_tracking.py | 19 +++-- .../proxy/guardrails/test_usage_tracking.py | 71 +++++++++++++++++-- 2 files changed, 79 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index 727f767469f..0e2d37c8d57 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -15,6 +15,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, NamedTuple, TypeVar from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES from litellm.proxy.utils import PrismaClient from litellm.repositories.table_repositories import ( DailyGuardrailMetricsRepository, @@ -63,11 +64,21 @@ async def _upsert_rows_with_retry( retries_left: int = _UPSERT_RETRY_TIMES, ) -> None: outcomes: Final = {key: await _attempt_upsert(upsert_row, key, value) for key, value in rows.items()} - failed: Final = MappingProxyType({key: rows[key] for key, error in outcomes.items() if error is not None}) - if not failed: + for key, error in outcomes.items(): + if error is not None and not isinstance(error, DB_RETRY_SAFE_ERROR_TYPES): + verbose_proxy_logger.warning( + "Guardrail usage tracking: %s upsert failed for %s and is not safe to retry (non-fatal): %s", + label, + key, + error, + ) + retryable: Final = MappingProxyType( + {key: rows[key] for key, error in outcomes.items() if isinstance(error, DB_RETRY_SAFE_ERROR_TYPES)} + ) + if not retryable: return if retries_left == 0: - for key in failed: + for key in retryable: verbose_proxy_logger.warning( "Guardrail usage tracking: %s upsert failed for %s after %d retries (non-fatal): %s", label, @@ -77,7 +88,7 @@ async def _upsert_rows_with_retry( ) return await sleep(2 ** (_UPSERT_RETRY_TIMES - retries_left)) - await _upsert_rows_with_retry(failed, upsert_row, label, sleep, retries_left - 1) + await _upsert_rows_with_retry(retryable, upsert_row, label, sleep, retries_left - 1) def _guardrail_status_to_action(status: str | None) -> str: diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index 62a4bbbbe6e..693a73b152d 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -3,6 +3,7 @@ from datetime import datetime, timezone from typing import Any from unittest.mock import AsyncMock, MagicMock +import httpx import pytest from litellm.proxy.guardrails.usage_tracking import process_spend_logs_guardrail_usage @@ -94,8 +95,8 @@ async def test_one_failing_upsert_does_not_drop_remaining_writes(): permanently under-report billable counters. """ prisma = _prisma() - prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = RuntimeError("db down") - prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [RuntimeError("db down"), None, None] + prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = httpx.ConnectError("db down") + prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [httpx.ConnectError("db down"), None, None] sleep, _ = _fake_sleep() logs = [ _payload("r1", usage={"topicPolicyUnits": 1}), @@ -113,12 +114,13 @@ async def test_one_failing_upsert_does_not_drop_remaining_writes(): @pytest.mark.asyncio async def test_transient_upsert_failure_is_retried_with_backoff_for_failed_rows_only(): """ - A transient DB error must not permanently drop billed units from the - aggregates: only the rows that failed are re-sent, after exponential - backoff, and the batch ends once every row has landed. + A connection error (the write provably never reached the database) must + not permanently drop billed units from the aggregates: only the rows that + failed are re-sent, after exponential backoff, and the batch ends once + every row has landed. """ prisma = _prisma() - prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [RuntimeError("blip"), None, None] + prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [httpx.ConnectError("blip"), None, None] sleep, delays = _fake_sleep() logs = [ _payload("r1", usage={"topicPolicyUnits": 1}), @@ -136,7 +138,7 @@ async def test_transient_upsert_failure_is_retried_with_backoff_for_failed_rows_ @pytest.mark.asyncio async def test_persistent_upsert_failure_stops_after_three_retries(): prisma = _prisma() - prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = RuntimeError("db down") + prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = httpx.ConnectError("db down") sleep, delays = _fake_sleep() await process_spend_logs_guardrail_usage(prisma, [_payload("r1", usage={"topicPolicyUnits": 1})], sleep=sleep) @@ -146,6 +148,61 @@ async def test_persistent_upsert_failure_stops_after_three_retries(): assert prisma.db.litellm_dailyguardrailusageunits.upsert.call_count == 1 +def _units_upsert_wheres(prisma: MagicMock) -> list[tuple]: + return [ + tuple( + c.kwargs["where"]["guardrail_id_date_team_id_api_key_usage_unit"][k] + for k in ("guardrail_id", "date", "team_id", "api_key", "usage_unit") + ) + for c in prisma.db.litellm_dailyguardrailusageunits.upsert.call_args_list + ] + + +@pytest.mark.asyncio +async def test_post_send_failure_is_never_retried_so_increments_cannot_double_count(): + """ + Follow-up to #37225: the units upsert is a non-idempotent increment, so an + ambiguous post-send failure (read timeout after the statement may have + committed) must be attempted exactly once. Re-sending it stacks a second + increment and inflates billable unit totals. Only a connection error proves + the write never reached the database and may be retried; the other rows in + the batch still land either way. + """ + prisma = _prisma() + prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [ + httpx.ReadTimeout("read timed out"), + httpx.ConnectError("refused"), + None, + ] + sleep, delays = _fake_sleep() + logs = [ + _payload("r1", usage={"topicPolicyUnits": 1}), + _payload("r2", team_id=None, api_key="hashed-key-2", usage={"topicPolicyUnits": 1}), + ] + + await process_spend_logs_guardrail_usage(prisma, logs, sleep=sleep) + + timed_out_row = ("bedrock-guard", "2026-08-17", "", "hashed-key-2", "topicPolicyUnits") + refused_row = ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits") + assert _units_upsert_wheres(prisma) == [timed_out_row, refused_row, refused_row] + assert delays == [1] + + +@pytest.mark.asyncio +async def test_generic_upsert_exception_is_terminal_for_that_row_only(): + prisma = _prisma() + prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = RuntimeError("constraint violation") + sleep, delays = _fake_sleep() + + await process_spend_logs_guardrail_usage(prisma, [_payload("r1", usage={"topicPolicyUnits": 1})], sleep=sleep) + + assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_count == 1 + assert delays == [] + assert _units_upserts(prisma) == { + ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 1, + } + + @pytest.mark.asyncio async def test_zero_and_non_int_usage_counters_are_skipped(): prisma = _prisma() From 15823b1be345daf7e0ffaa2191ffe524feaa80f7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 17 Aug 2026 19:38:43 -0700 Subject: [PATCH 058/538] fix(guardrails): degrade usage units to empty when the units table is missing GET /guardrails/usage/overview and GET /guardrails/usage/detail/{id} 500ed on a database that has not applied 20260817143646_add_daily_guardrail_usage_units yet (pip installs on litellm-proxy-extras 0.4.86 with DISABLE_SCHEMA_UPDATE=true). Both endpoints now return their metrics with empty units and log one warning until the migration lands. --- litellm/proxy/guardrails/usage_endpoints.py | 12 +++++- .../proxy/guardrails/test_usage_endpoints.py | 39 +++++++++++++++++++ 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index a73efed30ad..ca89c7587ba 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -14,6 +14,7 @@ from fastapi import APIRouter, Depends, Query from pydantic import BaseModel from typing_extensions import NotRequired, ReadOnly, TypedDict +from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.repositories.table_repositories import ( @@ -111,7 +112,16 @@ async def _find_daily_guardrail_usage_units( prisma_client: "PrismaClient", where: "prisma_types.LiteLLM_DailyGuardrailUsageUnitsWhereInput", ) -> "Sequence[prisma_models.LiteLLM_DailyGuardrailUsageUnits]": - return await _daily_guardrail_usage_units_table(prisma_client).find_many(where=where) + from prisma.errors import TableNotFoundError + + try: + return await _daily_guardrail_usage_units_table(prisma_client).find_many(where=where) + except TableNotFoundError as e: + verbose_proxy_logger.warning( + "Guardrail usage units are unavailable until the LiteLLM_DailyGuardrailUsageUnits migration is applied: %s", + e, + ) + return () def _counter_name(row: "prisma_models.LiteLLM_DailyGuardrailUsageUnits") -> str: diff --git a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py index b0f98b81c05..f63c08a2c39 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py @@ -19,6 +19,7 @@ import pytest sys.path.insert(0, os.path.abspath("../../..")) from fastapi import HTTPException +from prisma.errors import TableNotFoundError from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler @@ -295,6 +296,44 @@ async def test_detail_breaks_units_down_by_day_team_and_key(): assert units_where == {"guardrail_id": {"in": ["yaml-pii", "yaml-1"]}, "date": {"gte": START, "lte": END}} +def _units_table_missing() -> TableNotFoundError: + return TableNotFoundError( + data={"user_facing_error": {"meta": {"table": "public.LiteLLM_DailyGuardrailUsageUnits"}}} + ) + + +@pytest.mark.asyncio +async def test_overview_degrades_units_to_empty_when_units_table_is_missing(): + prisma = _prisma(metrics=[_metric("yaml-pii", requests=4, passed=3, blocked=1)]) + prisma.db.litellm_dailyguardrailusageunits.find_many = AsyncMock(side_effect=_units_table_missing()) + handler = _config_handler(_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii")) + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN) + row = next(r for r in resp.rows if r.id == "yaml-uuid") + assert (row.requestsEvaluated, row.usageUnits) == (4, {}) + assert (resp.totalRequests, resp.totalBlocked, resp.totalUsageUnits) == (4, 1, {}) + + +@pytest.mark.asyncio +async def test_detail_degrades_units_to_empty_when_units_table_is_missing(): + prisma = _prisma(metrics=[_metric("yaml-pii", requests=4, passed=3, blocked=1)]) + prisma.db.litellm_dailyguardrailusageunits.find_many = AsyncMock(side_effect=_units_table_missing()) + handler = _config_handler(_yaml_guardrail()) + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_detail( + guardrail_id="yaml-1", start_date=START, end_date=END, user_api_key_dict=ADMIN + ) + assert (resp.requestsEvaluated, resp.failRate) == (4, 25.0) + assert (resp.usage_units, list(resp.usage_units_daily), resp.usage_units_by_team, resp.usage_units_by_key) == ( + {}, + [], + {}, + {}, + ) + + # ---- logs ------------------------------------------------------------------- From a738c45fc7223d08646e7d86e5617d2449e6fbf9 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 17 Aug 2026 19:50:05 -0700 Subject: [PATCH 059/538] fix(proxy): strip callback credentials from the auth object stamped into request metadata (#37233) * fix(proxy): strip callback credentials from the auth object stamped into request metadata * style(proxy): drop the restating half of the stamp-site comment * test(proxy): pin that the stamped auth copy carries header-derived identity --- litellm/proxy/common_utils/callback_utils.py | 8 +- litellm/proxy/litellm_pre_call_utils.py | 20 ++- .../proxy/common_utils/test_callback_utils.py | 2 + .../proxy/test_litellm_pre_call_utils.py | 160 ++++++++++++++++++ 4 files changed, 180 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 4afd7c76a35..9379a8577a3 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -39,10 +39,10 @@ _EXTRA_SENSITIVE_CALLBACK_KEYS: Final = {"gcs_path_service_account"} # already-encrypted input cheaply (no decrypt-attempt round trip) and # avoid double-encrypting if `LITELLM_SALT_KEY` is rotated between writes. _CALLBACK_VAR_ENCRYPTED_PREFIX: Final = "litellm_enc::" -# Metadata slots that hold operator-configured callback setup (and therefore -# integration credentials). Resolved from UserAPIKeyAuth during pre-call setup, -# never read back off the copies stamped into request metadata. -_CALLBACK_CONFIG_SLOTS: Final = frozenset({"logging", "callback_settings"}) +# Metadata slots that hold operator-configured callback and secret-manager setup +# (and therefore integration credentials). Resolved from UserAPIKeyAuth during +# pre-call setup, never read back off the copies stamped into request metadata. +_CALLBACK_CONFIG_SLOTS: Final = frozenset({"logging", "callback_settings", "secret_manager_settings"}) blue_color_code: Final = "\033[94m" reset_color_code: Final = "\033[0m" diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 6172f3a9158..ef6e590ba26 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1303,8 +1303,15 @@ class LiteLLMProxyRequestSetup: ) if user_api_key_dict.budget_reservation is not None: data[_metadata_variable_name]["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation - # Add the full UserAPIKeyAuth object for MCP server access control - data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict + # UserAPIKeyAuth object for MCP server access control + data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict.model_copy( + update={ + "metadata": strip_callback_config(user_api_key_dict.metadata), + "team_metadata": strip_callback_config(user_api_key_dict.team_metadata), + "project_metadata": strip_callback_config(user_api_key_dict.project_metadata), + "organization_metadata": strip_callback_config(user_api_key_dict.organization_metadata), + } + ) return data @staticmethod @@ -1326,10 +1333,11 @@ class LiteLLMProxyRequestSetup: ) # ignore any special fields - added_metadata: Final = {} - for k, v in management_endpoint_metadata.items(): - if k not in (LiteLLM_ManagementEndpoint_MetadataFields_Premium + LiteLLM_ManagementEndpoint_MetadataFields): - added_metadata[k] = v + added_metadata: Final = { + k: v + for k, v in (strip_callback_config(management_endpoint_metadata) or {}).items() + if k not in (LiteLLM_ManagementEndpoint_MetadataFields_Premium + LiteLLM_ManagementEndpoint_MetadataFields) + } if data[_metadata_variable_name].get("user_api_key_auth_metadata") is None: data[_metadata_variable_name]["user_api_key_auth_metadata"] = {} data[_metadata_variable_name]["user_api_key_auth_metadata"].update(added_metadata) diff --git a/tests/test_litellm/proxy/common_utils/test_callback_utils.py b/tests/test_litellm/proxy/common_utils/test_callback_utils.py index 59963bd3707..515a7b27c7b 100644 --- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -492,6 +492,7 @@ def test_strip_callback_config_drops_credential_bearing_slots(): } ], "callback_settings": {"callback_vars": {"langfuse_secret_key": "litellm_enc::other"}}, + "secret_manager_settings": {"vault_token": "vt-secret"}, "priority": "high", "guardrails": ["presidio"], "langsmith_provisioning": {"api_key_id": "prov-1"}, @@ -501,6 +502,7 @@ def test_strip_callback_config_drops_credential_bearing_slots(): assert "logging" not in stripped assert "callback_settings" not in stripped + assert "secret_manager_settings" not in stripped assert stripped["priority"] == "high" assert stripped["guardrails"] == ["presidio"] assert stripped["langsmith_provisioning"] == {"api_key_id": "prov-1"} diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index d3b1e089489..0171e83be08 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -229,6 +229,43 @@ async def test_add_litellm_data_to_request_parses_string_metadata(): assert updated_data["metadata"]["generation_name"] == "gen123" +@pytest.mark.asyncio +async def test_stamped_auth_object_reflects_header_derived_identity(): + """ + Regression (LIT-5487): the stamped object is a copy taken partway through request setup, + so it only carries header-derived identity if the stamp still runs after those fields are + resolved. Moving the stamp earlier would silently misattribute spend. + """ + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + request_mock = MagicMock(spec=Request) + request_mock.url = MagicMock() + request_mock.url.path = "/v1/chat/completions" + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json", "user": "end-user-from-header"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + user_api_key_dict = UserAPIKeyAuth(api_key="hashed-key", metadata={}, team_metadata={}) + + updated_data = await add_litellm_data_to_request( + data={"model": "gpt-3.5-turbo"}, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={"user_header_name": "user"}, + version="test-version", + ) + + # precondition: the header was actually resolved onto the live object + assert user_api_key_dict.end_user_id == "end-user-from-header" + + stamped = updated_data["metadata"]["user_api_key_auth"] + assert stamped.end_user_id == "end-user-from-header" + + @pytest.mark.asyncio async def test_add_litellm_data_to_request_strips_admin_injection_slots(): """User-supplied user_api_key_metadata / user_api_key_team_metadata / @@ -2306,6 +2343,129 @@ def test_add_user_api_key_auth_to_request_metadata(): assert result["messages"] == [{"role": "user", "content": "Hello"}] +def _auth_with_callback_credentials() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="hashed-test-key-123", + key_alias="test-key-alias", + team_id="test-team-789", + team_alias="test-team-alias", + metadata={ + "logging": [{"callback_name": "langfuse", "callback_vars": {"langfuse_secret_key": "sk-KEY-CANARY"}}], + "rpm_limit_type": "guaranteed_throughput", + }, + team_metadata={ + "callback_settings": {"langfuse": {"callback_vars": {"langfuse_secret_key": "sk-TEAM-CANARY"}}}, + "secret_manager_settings": {"vault_token": "vt-TEAM-CANARY"}, + "model_rpm_limit": {"gpt-4": 10}, + }, + project_metadata={ + "logging": [{"callback_vars": {"langfuse_secret_key": "sk-PROJECT-CANARY"}}], + "project_tier": "gold", + }, + organization_metadata={ + "secret_manager_settings": {"vault_token": "vt-ORG-CANARY"}, + "org_tier": "platinum", + }, + ) + + +def test_stamped_auth_object_carries_no_callback_credentials(): + """ + Regression (LIT-5487): the UserAPIKeyAuth stamped into request metadata reaches every + raw-metadata logging integration, so it must not carry team/key callback credentials. + """ + user_api_key_dict = _auth_with_callback_credentials() + otel_span = object() + user_api_key_dict.parent_otel_span = otel_span + user_api_key_dict.budget_reservation = {"amount": 1.0} + user_api_key_dict.via_virtual_key = True + data = {"litellm_metadata": {}} + + result = LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( + data=data, + user_api_key_dict=user_api_key_dict, + _metadata_variable_name="litellm_metadata", + ) + + stamped = result["litellm_metadata"]["user_api_key_auth"] + emitted = json.dumps( + { + "metadata": stamped.metadata, + "team_metadata": stamped.team_metadata, + "project_metadata": stamped.project_metadata, + "organization_metadata": stamped.organization_metadata, + }, + default=str, + ) + assert "sk-KEY-CANARY" not in emitted + assert "sk-TEAM-CANARY" not in emitted + assert "vt-TEAM-CANARY" not in emitted + assert "sk-PROJECT-CANARY" not in emitted + assert "vt-ORG-CANARY" not in emitted + + # consumers keep the type and the non-credential slots they read + assert isinstance(stamped, UserAPIKeyAuth) + assert stamped.key_alias == "test-key-alias" + assert stamped.team_id == "test-team-789" + assert stamped.team_alias == "test-team-alias" + assert stamped.api_key == "hashed-test-key-123" + assert stamped.metadata["rpm_limit_type"] == "guaranteed_throughput" + assert stamped.team_metadata["model_rpm_limit"] == {"gpt-4": 10} + assert stamped.project_metadata["project_tier"] == "gold" + assert stamped.organization_metadata["org_tier"] == "platinum" + + # server-only markers are excluded from model_dump, so rebuilding the object + # instead of copying it would silently drop them + assert stamped.via_virtual_key is True + assert stamped.budget_reservation == {"amount": 1.0} + assert stamped.parent_otel_span is otel_span + + +def test_stamping_does_not_mutate_the_cached_auth_object(): + """ + Regression (LIT-5487): UserAPIKeyAuth is cached and model_copy is shallow, so stripping + in place would poison the shared dicts and silently kill team callbacks fleet-wide. + """ + user_api_key_dict = _auth_with_callback_credentials() + metadata_before = copy.deepcopy(user_api_key_dict.metadata) + team_metadata_before = copy.deepcopy(user_api_key_dict.team_metadata) + + LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( + data={"litellm_metadata": {}}, + user_api_key_dict=user_api_key_dict, + _metadata_variable_name="litellm_metadata", + ) + + assert user_api_key_dict.metadata == metadata_before + assert user_api_key_dict.team_metadata == team_metadata_before + + +def test_management_endpoint_metadata_drops_callback_credentials(): + """ + Regression (LIT-5487): user_api_key_auth_metadata is part of StandardLoggingPayload, so a + callback_settings-shaped team must not push credentials into it. + """ + data = {"litellm_metadata": {}} + + result = LiteLLMProxyRequestSetup.add_management_endpoint_metadata_to_request_metadata( + data=data, + management_endpoint_metadata={ + "callback_settings": {"langfuse": {"callback_vars": {"langfuse_secret_key": "sk-TEAM-CANARY"}}}, + "secret_manager_settings": {"vault_token": "vt-TEAM-CANARY"}, + "logging": [{"callback_vars": {"langfuse_secret_key": "sk-LOGGING-CANARY"}}], + "other_field": "value", + }, + _metadata_variable_name="litellm_metadata", + ) + + auth_metadata = result["litellm_metadata"]["user_api_key_auth_metadata"] + emitted = json.dumps(auth_metadata, default=str) + assert "sk-TEAM-CANARY" not in emitted + assert "vt-TEAM-CANARY" not in emitted + assert "sk-LOGGING-CANARY" not in emitted + assert auth_metadata["other_field"] == "value" + + @pytest.mark.parametrize( "data, model_group_settings, expected_headers_added", [ From 23bbb1d2429e74369a97f5c1f6524e20e1a9f3b7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 17 Aug 2026 20:56:36 -0700 Subject: [PATCH 060/538] refactor(mcp): pass materialized server tuples to oauth discovery helpers Annotate _prime_oauth_metadata_discovery_for_servers and _reconcile_oauth_discovery_slots_for_servers with Sequence and snapshot registry views with tuple() at the call sites. This drops the Iterable addition to the collections.abc import, restoring that line to its base spelling so the branch merges cleanly with litellm_internal_staging, which adds Mapping on the same line --- .../_experimental/mcp_server/mcp_server_manager.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 91e5339dda2..f963fa7a106 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -13,7 +13,7 @@ import json import os import re import time -from collections.abc import AsyncIterator, Callable, Iterable, Sequence +from collections.abc import AsyncIterator, Callable, Sequence from contextlib import asynccontextmanager from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypedDict, cast @@ -1684,11 +1684,11 @@ class MCPServerManager: """ self._get_or_start_oauth_discovery_task(server) - def _prime_oauth_metadata_discovery_for_servers(self, servers: Iterable[MCPServer]) -> None: + def _prime_oauth_metadata_discovery_for_servers(self, servers: Sequence[MCPServer]) -> None: for server in servers: self.prime_oauth_metadata_discovery(server) - def _reconcile_oauth_discovery_slots_for_servers(self, servers: Iterable[MCPServer]) -> None: + def _reconcile_oauth_discovery_slots_for_servers(self, servers: Sequence[MCPServer]) -> None: """Align retry slots after an atomic registry replacement.""" for server in servers: should_defer = _requires_oauth_discovery(server.url, server.issuer_is_anchored, server) @@ -2105,7 +2105,7 @@ class MCPServerManager: await self._hydrate_config_servers_dcr_clients() - self._prime_oauth_metadata_discovery_for_servers(self.config_mcp_servers.values()) + self._prime_oauth_metadata_discovery_for_servers(tuple(self.config_mcp_servers.values())) self.initialize_tool_name_to_mcp_server_name_mapping() @@ -5862,8 +5862,9 @@ class MCPServerManager: # this replacement was being staged. Reconcile every published entry # synchronously after the swap so a lost publication cannot also leave # the replacement unresolved with no retry slot. - self._reconcile_oauth_discovery_slots_for_servers(registered_registry.values()) - self._prime_oauth_metadata_discovery_for_servers(registered_registry.values()) + registered_servers: Final = tuple(registered_registry.values()) + self._reconcile_oauth_discovery_slots_for_servers(registered_servers) + self._prime_oauth_metadata_discovery_for_servers(registered_servers) if registered_openapi_tools: self.initialize_tool_name_to_mcp_server_name_mapping() From 3d523d6d81816d692927e00a44ad998665731b82 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 18 Aug 2026 13:14:42 +0000 Subject: [PATCH 061/538] fix(model_prices): add provider-announced deprecation_date to 205 registry entries Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 205 ++++++++++++++++++ model_prices_and_context_window.json | 205 ++++++++++++++++++ 2 files changed, 410 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 78b53cefc53..1a130c4ac0a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -54,6 +54,7 @@ "output_cost_per_image": 0.04 }, "1024-x-1024/dall-e-2": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 1.9e-08, "litellm_provider": "openai", "mode": "image_generation", @@ -67,6 +68,7 @@ "output_cost_per_image": 0.08 }, "256-x-256/dall-e-2": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 2.4414e-07, "litellm_provider": "openai", "mode": "image_generation", @@ -80,6 +82,7 @@ "output_cost_per_image": 0.018 }, "512-x-512/dall-e-2": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 6.86e-08, "litellm_provider": "openai", "mode": "image_generation", @@ -2887,6 +2890,7 @@ "supports_function_calling": true }, "azure_ai/claude-haiku-4-5": { + "deprecation_date": "2026-10-19", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -2908,6 +2912,7 @@ "supports_vision": true }, "azure_ai/claude-opus-4-5": { + "deprecation_date": "2026-10-19", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -2930,6 +2935,7 @@ "supports_output_config": true }, "azure_ai/claude-opus-4-6": { + "deprecation_date": "2027-02-02", "supports_adaptive_thinking": true, "input_cost_per_token": 5e-06, "output_cost_per_token": 2.5e-05, @@ -2959,6 +2965,7 @@ "supports_max_reasoning_effort": true }, "azure_ai/claude-opus-4-7": { + "deprecation_date": "2027-04-06", "supports_adaptive_thinking": true, "input_cost_per_token": 5e-06, "output_cost_per_token": 2.5e-05, @@ -3083,6 +3090,7 @@ "supports_max_reasoning_effort": true }, "azure_ai/claude-opus-4-1": { + "deprecation_date": "2026-08-05", "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, @@ -3104,6 +3112,7 @@ "supports_vision": true }, "azure_ai/claude-sonnet-4-5": { + "deprecation_date": "2026-10-19", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -3156,6 +3165,7 @@ "supports_max_reasoning_effort": true }, "azure_ai/claude-sonnet-4-6": { + "deprecation_date": "2027-02-10", "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, @@ -3226,6 +3236,7 @@ "supports_tool_choice": true }, "azure_ai/gpt-5.5": { + "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -3318,6 +3329,7 @@ "supports_minimal_reasoning_effort": false }, "azure_ai/gpt-5.4": { + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens": 5e-07, "cache_read_input_token_cost_priority": 5e-07, @@ -3364,6 +3376,7 @@ "supports_minimal_reasoning_effort": true }, "azure_ai/gpt-5.4-2026-03-05": { + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens": 5e-07, "cache_read_input_token_cost_priority": 5e-07, @@ -3410,6 +3423,7 @@ "supports_minimal_reasoning_effort": true }, "azure_ai/gpt-5.4-pro": { + "deprecation_date": "2027-09-07", "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "cache_read_input_token_cost_priority": 6e-06, @@ -3455,6 +3469,7 @@ "supports_minimal_reasoning_effort": true }, "azure_ai/gpt-5.4-pro-2026-03-05": { + "deprecation_date": "2027-09-07", "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "cache_read_input_token_cost_priority": 6e-06, @@ -3500,6 +3515,7 @@ "supports_minimal_reasoning_effort": true }, "azure_ai/gpt-5.4-mini": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_priority": 1.5e-07, "input_cost_per_token": 7.5e-07, @@ -3540,6 +3556,7 @@ "supports_minimal_reasoning_effort": false }, "azure_ai/gpt-5.4-mini-2026-03-17": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_priority": 1.5e-07, "input_cost_per_token": 7.5e-07, @@ -3580,6 +3597,7 @@ "supports_minimal_reasoning_effort": false }, "azure_ai/gpt-5.4-nano": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_priority": 4e-08, "input_cost_per_token": 2e-07, @@ -3620,6 +3638,7 @@ "supports_minimal_reasoning_effort": false }, "azure_ai/gpt-5.4-nano-2026-03-17": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_priority": 4e-08, "input_cost_per_token": 2e-07, @@ -3849,6 +3868,7 @@ "supports_vision": true }, "azure/eu/gpt-5.1": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.38e-06, "litellm_provider": "azure", @@ -3918,6 +3938,7 @@ "supports_none_reasoning_effort": true }, "azure/eu/gpt-5.1-codex": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.38e-06, "litellm_provider": "azure", @@ -3948,6 +3969,7 @@ "supports_vision": true }, "azure/eu/gpt-5.1-codex-mini": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 2.75e-07, "litellm_provider": "azure", @@ -4107,6 +4129,7 @@ "supports_vision": true }, "azure/global-standard/gpt-4o-mini": { + "deprecation_date": "2027-04-14", "input_cost_per_token": 1.5e-07, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -4155,6 +4178,7 @@ "supports_vision": true }, "azure/global/gpt-5.1": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -4224,6 +4248,7 @@ "supports_none_reasoning_effort": true }, "azure/global/gpt-5.1-codex": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -4254,6 +4279,7 @@ "supports_vision": true }, "azure/global/gpt-5.1-codex-mini": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 2.5e-07, "litellm_provider": "azure", @@ -4492,6 +4518,7 @@ "supports_vision": true }, "azure/gpt-4.1": { + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -4559,6 +4586,7 @@ "supports_web_search": false }, "azure/gpt-4.1-mini": { + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 4e-07, "input_cost_per_token_batches": 2e-07, @@ -4626,6 +4654,7 @@ "supports_web_search": false }, "azure/gpt-4.1-nano": { + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, @@ -4902,6 +4931,7 @@ "supports_vision": false }, "azure/gpt-4o-mini": { + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 1.65e-07, "litellm_provider": "azure", @@ -5344,6 +5374,7 @@ "supports_vision": true }, "azure/gpt-5": { + "deprecation_date": "2027-02-09", "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -5507,6 +5538,7 @@ "supports_vision": true }, "azure/gpt-5-mini": { + "deprecation_date": "2027-02-09", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 2.5e-07, "litellm_provider": "azure", @@ -5572,6 +5604,7 @@ "supports_vision": true }, "azure/gpt-5-nano": { + "deprecation_date": "2027-02-09", "cache_read_input_token_cost": 5e-09, "input_cost_per_token": 5e-08, "litellm_provider": "azure", @@ -5667,6 +5700,7 @@ "supports_vision": true }, "azure/gpt-5.1": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -5736,6 +5770,7 @@ "supports_none_reasoning_effort": true }, "azure/gpt-5.1-codex": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -5797,6 +5832,7 @@ "supports_vision": true }, "azure/gpt-5.1-codex-mini": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 2.5e-07, "litellm_provider": "azure", @@ -5827,6 +5863,7 @@ "supports_vision": true }, "azure/gpt-5.2": { + "deprecation_date": "2027-06-08", "cache_read_input_token_cost": 1.75e-07, "input_cost_per_token": 1.75e-06, "litellm_provider": "azure", @@ -6136,6 +6173,7 @@ "supports_web_search": true }, "azure/gpt-5.4": { + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens": 5e-07, "cache_read_input_token_cost_priority": 5e-07, @@ -6180,6 +6218,7 @@ "supports_minimal_reasoning_effort": true }, "azure/us/gpt-5.4": { + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.8e-07, "cache_read_input_token_cost_priority": 5.5e-07, "input_cost_per_token": 2.75e-06, @@ -6218,6 +6257,7 @@ "supports_minimal_reasoning_effort": true }, "azure/eu/gpt-5.4": { + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.8e-07, "cache_read_input_token_cost_priority": 5.5e-07, "input_cost_per_token": 2.75e-06, @@ -6379,6 +6419,7 @@ "supports_minimal_reasoning_effort": true }, "azure/gpt-5.4-pro": { + "deprecation_date": "2027-09-07", "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -7045,6 +7086,7 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.5": { + "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -7095,6 +7137,7 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.5": { + "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -7142,6 +7185,7 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.5": { + "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -7408,6 +7452,7 @@ "supports_web_search": true }, "azure/gpt-5.4-mini": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 7.5e-07, "litellm_provider": "azure", @@ -7489,6 +7534,7 @@ "supports_xhigh_reasoning_effort": true }, "azure/gpt-5.4-nano": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2e-07, "litellm_provider": "azure", @@ -7601,6 +7647,7 @@ "output_cost_per_token": 0.0 }, "azure/high/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.59263611e-07, "litellm_provider": "azure", "mode": "image_generation", @@ -7610,6 +7657,7 @@ ] }, "azure/high/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.58945719e-07, "litellm_provider": "azure", "mode": "image_generation", @@ -7619,6 +7667,7 @@ ] }, "azure/high/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.58945719e-07, "litellm_provider": "azure", "mode": "image_generation", @@ -7628,6 +7677,7 @@ ] }, "azure/low/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.0490417e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7637,6 +7687,7 @@ ] }, "azure/low/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.0172526e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7646,6 +7697,7 @@ ] }, "azure/low/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.0172526e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7655,6 +7707,7 @@ ] }, "azure/medium/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7664,6 +7717,7 @@ ] }, "azure/medium/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7673,6 +7727,7 @@ ] }, "azure/medium/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7695,6 +7750,7 @@ ] }, "azure/gpt-image-1.5": { + "deprecation_date": "2027-06-16", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 5e-06, "input_cost_per_image_token": 8e-06, @@ -7720,6 +7776,7 @@ ] }, "azure/gpt-image-2": { + "deprecation_date": "2027-10-21", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 5e-06, "input_cost_per_image_token": 8e-06, @@ -7751,6 +7808,7 @@ "supports_pdf_input": true }, "azure/low/1024-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 2.0751953125e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7760,6 +7818,7 @@ ] }, "azure/low/1024-x-1536/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 2.0751953125e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7769,6 +7828,7 @@ ] }, "azure/low/1536-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 2.0345052083e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7778,6 +7838,7 @@ ] }, "azure/medium/1024-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 8.056640625e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7787,6 +7848,7 @@ ] }, "azure/medium/1024-x-1536/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 8.056640625e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7796,6 +7858,7 @@ ] }, "azure/medium/1536-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 7.9752604167e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7805,6 +7868,7 @@ ] }, "azure/high/1024-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 3.173828125e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7814,6 +7878,7 @@ ] }, "azure/high/1024-x-1536/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 3.173828125e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7823,6 +7888,7 @@ ] }, "azure/high/1536-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 3.1575520833e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7850,6 +7916,7 @@ "supports_function_calling": true }, "azure/o1": { + "deprecation_date": "2026-10-21", "cache_read_input_token_cost": 7.5e-06, "input_cost_per_token": 1.5e-05, "litellm_provider": "azure", @@ -7944,6 +8011,7 @@ "supports_vision": false }, "azure/o3": { + "deprecation_date": "2026-10-21", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "azure", @@ -8041,6 +8109,7 @@ "supports_web_search": true }, "azure/o3-mini": { + "deprecation_date": "2026-10-01", "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", @@ -8071,6 +8140,7 @@ "supports_vision": false }, "azure/o3-pro": { + "deprecation_date": "2026-12-17", "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "azure", @@ -8132,6 +8202,7 @@ "supports_vision": true }, "azure/o4-mini": { + "deprecation_date": "2026-10-16", "cache_read_input_token_cost": 2.75e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", @@ -8580,6 +8651,7 @@ "supports_vision": true }, "azure/us/gpt-5.1": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.38e-06, "litellm_provider": "azure", @@ -8649,6 +8721,7 @@ "supports_none_reasoning_effort": true }, "azure/us/gpt-5.1-codex": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.38e-06, "litellm_provider": "azure", @@ -8679,6 +8752,7 @@ "supports_vision": true }, "azure/us/gpt-5.1-codex-mini": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 2.75e-07, "litellm_provider": "azure", @@ -8876,6 +8950,7 @@ ] }, "azure_ai/FW-DeepSeek-V3.2": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 3.1e-07, "input_cost_per_token": 6.2e-07, "litellm_provider": "azure_ai", @@ -8906,6 +8981,7 @@ "supports_tool_choice": true }, "azure_ai/FW-GLM-5": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 2.2e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "azure_ai", @@ -8921,6 +8997,7 @@ "supports_tool_choice": true }, "azure_ai/FW-GLM-5.1": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 2.86e-07, "input_cost_per_token": 1.54e-06, "litellm_provider": "azure_ai", @@ -8987,6 +9064,7 @@ "supports_tool_choice": true }, "azure_ai/FW-Kimi-K2.5": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 1.1e-07, "input_cost_per_token": 6.6e-07, "litellm_provider": "azure_ai", @@ -9079,6 +9157,7 @@ "supports_vision": true }, "azure_ai/FW-MiniMax-M2.5": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 3.3e-08, "input_cost_per_token": 3.3e-07, "litellm_provider": "azure_ai", @@ -9164,6 +9243,7 @@ ] }, "azure_ai/MAI-Image-2e": { + "deprecation_date": "2026-08-15", "input_cost_per_token": 5e-06, "litellm_provider": "azure_ai", "mode": "image_generation", @@ -9175,6 +9255,7 @@ ] }, "azure_ai/Llama-3.2-11B-Vision-Instruct": { + "deprecation_date": "2026-06-13", "input_cost_per_token": 3.7e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9188,6 +9269,7 @@ "supports_vision": true }, "azure_ai/Llama-3.2-90B-Vision-Instruct": { + "deprecation_date": "2026-06-13", "input_cost_per_token": 2.04e-06, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9249,6 +9331,7 @@ "supports_tool_choice": true }, "azure_ai/Meta-Llama-3.1-405B-Instruct": { + "deprecation_date": "2026-06-13", "input_cost_per_token": 5.33e-06, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9271,6 +9354,7 @@ "supports_tool_choice": true }, "azure_ai/Meta-Llama-3.1-8B-Instruct": { + "deprecation_date": "2026-06-13", "input_cost_per_token": 3e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9452,6 +9536,7 @@ "supports_reasoning": true }, "azure_ai/mistral-document-ai-2505": { + "deprecation_date": "2026-07-20", "litellm_provider": "azure_ai", "ocr_cost_per_page": 0.003, "mode": "ocr", @@ -9529,6 +9614,7 @@ "output_cost_per_token": 0.0 }, "azure_ai/cohere-rerank-v3.5": { + "deprecation_date": "2026-05-14", "input_cost_per_query": 0.002, "input_cost_per_token": 0.0, "litellm_provider": "azure_ai", @@ -9591,6 +9677,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-r1": { + "deprecation_date": "2026-08-13", "input_cost_per_token": 1.35e-06, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9614,6 +9701,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-v3-0324": { + "deprecation_date": "2026-07-13", "input_cost_per_token": 1.14e-06, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9626,6 +9714,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-v3.1": { + "deprecation_date": "2026-07-13", "input_cost_per_token": 1.23e-06, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -9639,6 +9728,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-v4-pro": { + "deprecation_date": "2028-02-20", "input_cost_per_token": 1.74e-06, "litellm_provider": "azure_ai", "max_input_tokens": 1000000, @@ -9652,6 +9742,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-v4-flash": { + "deprecation_date": "2028-02-20", "input_cost_per_token": 1.9e-07, "litellm_provider": "azure_ai", "max_input_tokens": 1000000, @@ -9683,6 +9774,7 @@ "supports_embedding_image_input": true }, "azure_ai/global/grok-3": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 3e-06, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -9697,6 +9789,7 @@ "supports_web_search": true }, "azure_ai/global/grok-3-mini": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 2.5e-07, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -9712,6 +9805,7 @@ "supports_web_search": true }, "azure_ai/grok-3": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 3e-06, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -9726,6 +9820,7 @@ "supports_web_search": true }, "azure_ai/grok-3-mini": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 2.5e-07, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -9773,6 +9868,7 @@ "supports_web_search": true }, "azure_ai/grok-4-fast-non-reasoning": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", @@ -9786,6 +9882,7 @@ "supports_web_search": true }, "azure_ai/grok-4-fast-reasoning": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", @@ -9863,6 +9960,7 @@ "supports_tool_choice": true }, "azure_ai/kimi-k2.5": { + "deprecation_date": "2027-01-26", "input_cost_per_token": 6e-07, "litellm_provider": "azure_ai", "max_input_tokens": 262144, @@ -9877,6 +9975,7 @@ "supports_vision": true }, "azure_ai/kimi-k2.6": { + "deprecation_date": "2027-04-16", "input_cost_per_token": 9.5e-07, "litellm_provider": "azure_ai", "max_input_tokens": 262144, @@ -10004,6 +10103,7 @@ "supports_vision": true }, "babbage-002": { + "deprecation_date": "2026-09-28", "input_cost_per_token": 4e-07, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, @@ -11999,6 +12099,7 @@ ] }, "claude-haiku-4-5-20251001": { + "deprecation_date": "2026-10-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -12022,6 +12123,7 @@ "prompt_cache_min_tokens": 4096 }, "claude-haiku-4-5": { + "deprecation_date": "2026-10-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -12170,6 +12272,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-sonnet-4-5": { + "deprecation_date": "2026-09-29", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, @@ -12203,6 +12306,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-sonnet-4-5-20250929": { + "deprecation_date": "2026-09-29", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, @@ -12237,6 +12341,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-sonnet-5": { + "deprecation_date": "2027-06-30", "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -12273,6 +12378,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-sonnet-4-6": { + "deprecation_date": "2027-02-17", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -12419,6 +12525,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-opus-4-5-20251101": { + "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12448,6 +12555,7 @@ "prompt_cache_min_tokens": 4096 }, "claude-opus-4-5": { + "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12477,6 +12585,7 @@ "prompt_cache_min_tokens": 4096 }, "claude-opus-4-6": { + "deprecation_date": "2027-02-05", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12513,6 +12622,7 @@ "prompt_cache_min_tokens": 4096 }, "claude-opus-4-6-20260205": { + "deprecation_date": "2027-02-05", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12549,6 +12659,7 @@ "prompt_cache_min_tokens": 4096 }, "claude-opus-4-7": { + "deprecation_date": "2027-04-16", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12587,6 +12698,7 @@ "prompt_cache_min_tokens": 2048 }, "claude-opus-4-7-20260416": { + "deprecation_date": "2027-04-16", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12625,6 +12737,7 @@ "prompt_cache_min_tokens": 2048 }, "claude-fable-5": { + "deprecation_date": "2027-06-09", "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, "cache_read_input_token_cost": 1e-06, @@ -12660,6 +12773,7 @@ "prompt_cache_min_tokens": 512 }, "claude-opus-5": { + "deprecation_date": "2027-07-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12698,6 +12812,7 @@ "prompt_cache_min_tokens": 512 }, "claude-opus-4-8": { + "deprecation_date": "2027-05-28", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -14801,6 +14916,7 @@ "mode": "search" }, "davinci-002": { + "deprecation_date": "2026-09-28", "input_cost_per_token": 2e-06, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, @@ -18353,6 +18469,7 @@ } }, "gemini-2.5-flash": { + "deprecation_date": "2026-10-20", "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -18398,6 +18515,7 @@ "supports_image_size": false }, "gemini-2.5-flash-image": { + "deprecation_date": "2026-10-02", "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -18442,6 +18560,7 @@ "supports_image_size": false }, "gemini-3-pro-image": { + "deprecation_date": "2027-05-28", "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -18522,6 +18641,7 @@ "web_search_billing_unit": "per_query" }, "gemini-3.1-flash-image": { + "deprecation_date": "2027-05-28", "input_cost_per_image": 0.00056, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", @@ -18646,6 +18766,7 @@ "web_search_billing_unit": "per_query" }, "gemini-3.1-flash-lite": { + "deprecation_date": "2027-05-07", "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, @@ -18702,6 +18823,7 @@ "web_search_billing_unit": "per_query" }, "gemini-3.5-flash-lite": { + "deprecation_date": "2027-07-21", "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_flex": 2e-08, "cache_read_input_token_cost_priority": 5e-08, @@ -18791,6 +18913,7 @@ "supports_web_search": true }, "gemini-2.5-flash-lite": { + "deprecation_date": "2026-10-20", "cache_read_input_token_cost": 1e-08, "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, @@ -19062,6 +19185,7 @@ "supports_image_size": false }, "gemini-2.5-pro": { + "deprecation_date": "2026-10-20", "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, @@ -19373,6 +19497,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.5-flash": { + "deprecation_date": "2027-05-19", "cache_read_input_token_cost": 1.5e-07, "input_cost_per_token": 1.5e-06, "input_cost_per_audio_token": 1e-06, @@ -19809,6 +19934,7 @@ "web_search_billing_unit": "per_query" }, "gemini/gemini-robotics-er-1.6-preview": { + "deprecation_date": "2026-08-31", "input_cost_per_audio_token": 2e-06, "input_cost_per_token": 1e-06, "litellm_provider": "gemini", @@ -19879,6 +20005,7 @@ "supports_vision": true }, "gemini-embedding-001": { + "deprecation_date": "2028-05-20", "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 2048, @@ -21492,6 +21619,7 @@ "supports_vision": true }, "gemini-3.5-flash": { + "deprecation_date": "2027-05-19", "cache_read_input_token_cost": 1.5e-07, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 1.5e-06, @@ -23004,6 +23132,7 @@ "supports_tool_choice": true }, "gpt-3.5-turbo-instruct": { + "deprecation_date": "2026-09-28", "input_cost_per_token": 1.5e-06, "litellm_provider": "text-completion-openai", "max_input_tokens": 8192, @@ -24135,6 +24264,7 @@ "supports_pdf_input": true }, "low/1024-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.009, "litellm_provider": "openai", "mode": "image_generation", @@ -24146,6 +24276,7 @@ "supports_pdf_input": true }, "low/1024-x-1536/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -24157,6 +24288,7 @@ "supports_pdf_input": true }, "low/1536-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -24168,6 +24300,7 @@ "supports_pdf_input": true }, "medium/1024-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.034, "litellm_provider": "openai", "mode": "image_generation", @@ -24179,6 +24312,7 @@ "supports_pdf_input": true }, "medium/1024-x-1536/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.05, "litellm_provider": "openai", "mode": "image_generation", @@ -24190,6 +24324,7 @@ "supports_pdf_input": true }, "medium/1536-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.05, "litellm_provider": "openai", "mode": "image_generation", @@ -24201,6 +24336,7 @@ "supports_pdf_input": true }, "high/1024-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.133, "litellm_provider": "openai", "mode": "image_generation", @@ -24212,6 +24348,7 @@ "supports_pdf_input": true }, "high/1024-x-1536/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.2, "litellm_provider": "openai", "mode": "image_generation", @@ -24223,6 +24360,7 @@ "supports_pdf_input": true }, "high/1536-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.2, "litellm_provider": "openai", "mode": "image_generation", @@ -24234,6 +24372,7 @@ "supports_pdf_input": true }, "standard/1024-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.009, "litellm_provider": "openai", "mode": "image_generation", @@ -24245,6 +24384,7 @@ "supports_pdf_input": true }, "standard/1024-x-1536/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -24256,6 +24396,7 @@ "supports_pdf_input": true }, "standard/1536-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -24267,6 +24408,7 @@ "supports_pdf_input": true }, "1024-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.009, "litellm_provider": "openai", "mode": "image_generation", @@ -24278,6 +24420,7 @@ "supports_pdf_input": true }, "1024-x-1536/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -24289,6 +24432,7 @@ "supports_pdf_input": true }, "1536-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -27202,18 +27346,21 @@ "output_cost_per_second": 0.0 }, "hd/1024-x-1024/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 7.629e-08, "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_pixel": 0.0 }, "hd/1024-x-1792/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 6.539e-08, "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_pixel": 0.0 }, "hd/1792-x-1024/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 6.539e-08, "litellm_provider": "openai", "mode": "image_generation", @@ -27260,6 +27407,7 @@ "max_output_tokens": 8192 }, "high/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.167, "input_cost_per_pixel": 1.59263611e-07, "litellm_provider": "openai", @@ -27270,6 +27418,7 @@ ] }, "high/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.25, "input_cost_per_pixel": 1.58945719e-07, "litellm_provider": "openai", @@ -27280,6 +27429,7 @@ ] }, "high/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.25, "input_cost_per_pixel": 1.58945719e-07, "litellm_provider": "openai", @@ -28067,6 +28217,7 @@ "supports_tool_choice": true }, "low/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.011, "input_cost_per_pixel": 1.0490417e-08, "litellm_provider": "openai", @@ -28077,6 +28228,7 @@ ] }, "low/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.016, "input_cost_per_pixel": 1.0172526e-08, "litellm_provider": "openai", @@ -28087,6 +28239,7 @@ ] }, "low/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.016, "input_cost_per_pixel": 1.0172526e-08, "litellm_provider": "openai", @@ -28111,6 +28264,7 @@ "output_cost_per_image": 0.072 }, "medium/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.042, "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "openai", @@ -28121,6 +28275,7 @@ ] }, "medium/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.063, "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "openai", @@ -28131,6 +28286,7 @@ ] }, "medium/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.063, "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "openai", @@ -28141,6 +28297,7 @@ ] }, "low/1024-x-1024/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.005, "litellm_provider": "openai", "mode": "image_generation", @@ -28149,6 +28306,7 @@ ] }, "low/1024-x-1536/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.006, "litellm_provider": "openai", "mode": "image_generation", @@ -28157,6 +28315,7 @@ ] }, "low/1536-x-1024/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.006, "litellm_provider": "openai", "mode": "image_generation", @@ -28165,6 +28324,7 @@ ] }, "medium/1024-x-1024/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.011, "litellm_provider": "openai", "mode": "image_generation", @@ -28173,6 +28333,7 @@ ] }, "medium/1024-x-1536/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.015, "litellm_provider": "openai", "mode": "image_generation", @@ -28181,6 +28342,7 @@ ] }, "medium/1536-x-1024/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.015, "litellm_provider": "openai", "mode": "image_generation", @@ -30074,6 +30236,7 @@ ] }, "multimodalembedding@001": { + "deprecation_date": "2027-04-01", "input_cost_per_character": 2e-07, "input_cost_per_image": 0.0001, "input_cost_per_token": 8e-07, @@ -35772,18 +35935,21 @@ "output_cost_per_image": 0.14 }, "standard/1024-x-1024/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 3.81469e-08, "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_pixel": 0.0 }, "standard/1024-x-1792/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 4.359e-08, "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_pixel": 0.0 }, "standard/1792-x-1024/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 4.359e-08, "litellm_provider": "openai", "mode": "image_generation", @@ -35847,6 +36013,7 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models" }, "text-embedding-005": { + "deprecation_date": "2027-04-01", "input_cost_per_character": 2.5e-08, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-embedding-models", @@ -35920,6 +36087,7 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, "text-moderation-007": { + "deprecation_date": "2025-10-27", "input_cost_per_token": 0.0, "litellm_provider": "openai", "max_input_tokens": 32768, @@ -35929,6 +36097,7 @@ "output_cost_per_token": 0.0 }, "text-moderation-latest": { + "deprecation_date": "2025-10-27", "input_cost_per_token": 0.0, "litellm_provider": "openai", "max_input_tokens": 32768, @@ -35938,6 +36107,7 @@ "output_cost_per_token": 0.0 }, "text-moderation-stable": { + "deprecation_date": "2025-10-27", "input_cost_per_token": 0.0, "litellm_provider": "openai", "max_input_tokens": 32768, @@ -35947,6 +36117,7 @@ "output_cost_per_token": 0.0 }, "text-multilingual-embedding-002": { + "deprecation_date": "2027-04-01", "input_cost_per_character": 2.5e-08, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-embedding-models", @@ -38434,6 +38605,7 @@ "supports_tool_choice": true }, "vertex_ai/claude-haiku-4-5": { + "deprecation_date": "2026-10-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -38457,6 +38629,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-haiku-4-5@20251001": { + "deprecation_date": "2026-10-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -38609,6 +38782,7 @@ "supports_vision": true }, "vertex_ai/claude-opus-4": { + "deprecation_date": "2026-05-14", "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, @@ -38636,6 +38810,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-opus-4-1": { + "deprecation_date": "2026-08-05", "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, @@ -38654,6 +38829,7 @@ "supports_vision": true }, "vertex_ai/claude-opus-4-1@20250805": { + "deprecation_date": "2026-08-05", "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, @@ -38672,6 +38848,7 @@ "supports_vision": true }, "vertex_ai/claude-opus-4-5": { + "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -38700,6 +38877,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-5@20251101": { + "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -38729,6 +38907,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-6": { + "deprecation_date": "2027-02-05", "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -38759,6 +38938,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-6@default": { + "deprecation_date": "2027-02-05", "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -38789,6 +38969,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-7": { + "deprecation_date": "2027-04-16", "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -38820,6 +39001,7 @@ "prompt_cache_min_tokens": 2048 }, "vertex_ai/claude-opus-4-7@default": { + "deprecation_date": "2027-04-16", "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -38851,6 +39033,7 @@ "prompt_cache_min_tokens": 2048 }, "vertex_ai/claude-fable-5": { + "deprecation_date": "2027-06-08", "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, @@ -38882,6 +39065,7 @@ "supports_max_reasoning_effort": true }, "vertex_ai/claude-fable-5@default": { + "deprecation_date": "2027-06-08", "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, @@ -38913,6 +39097,7 @@ "supports_max_reasoning_effort": true }, "vertex_ai/claude-opus-5": { + "deprecation_date": "2027-01-24", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -38945,6 +39130,7 @@ "prompt_cache_min_tokens": 512 }, "vertex_ai/claude-opus-5@default": { + "deprecation_date": "2027-01-24", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -38977,6 +39163,7 @@ "prompt_cache_min_tokens": 512 }, "vertex_ai/claude-opus-4-8": { + "deprecation_date": "2027-05-28", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -39009,6 +39196,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-opus-4-8@default": { + "deprecation_date": "2027-05-28", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -39041,6 +39229,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-4-5": { + "deprecation_date": "2026-09-29", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -39069,6 +39258,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-5": { + "deprecation_date": "2026-12-24", "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, @@ -39131,6 +39321,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-4-5@20250929": { + "deprecation_date": "2026-09-29", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -39160,6 +39351,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-opus-4@20250514": { + "deprecation_date": "2026-05-14", "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, @@ -39187,6 +39379,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-4": { + "deprecation_date": "2026-05-14", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -39218,6 +39411,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-4@20250514": { + "deprecation_date": "2026-05-14", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -39382,6 +39576,7 @@ "supports_tool_choice": true }, "vertex_ai/gemini-2.5-flash-image": { + "deprecation_date": "2026-10-02", "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -39427,6 +39622,7 @@ "supports_image_size": false }, "vertex_ai/gemini-3-pro-image": { + "deprecation_date": "2027-05-28", "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -39459,6 +39655,7 @@ "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" }, "vertex_ai/gemini-3.1-flash-image": { + "deprecation_date": "2027-05-28", "input_cost_per_image": 0.00056, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", @@ -39535,6 +39732,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.1-flash-lite": { + "deprecation_date": "2027-05-07", "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, @@ -39591,6 +39789,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.5-flash-lite": { + "deprecation_date": "2027-07-21", "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_flex": 2e-08, "cache_read_input_token_cost_priority": 5e-08, @@ -40308,6 +40507,7 @@ "supports_tool_choice": true }, "vertex_ai/veo-2.0-generate-001": { + "deprecation_date": "2026-06-30", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -40322,6 +40522,7 @@ ] }, "vertex_ai/veo-3.0-fast-generate-001": { + "deprecation_date": "2026-06-30", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -40336,6 +40537,7 @@ ] }, "vertex_ai/veo-3.0-generate-001": { + "deprecation_date": "2026-06-30", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -40378,6 +40580,7 @@ ] }, "vertex_ai/veo-3.1-generate-001": { + "deprecation_date": "2026-11-17", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -40392,6 +40595,7 @@ ] }, "vertex_ai/veo-3.1-fast-generate-001": { + "deprecation_date": "2026-11-17", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -46773,6 +46977,7 @@ } }, "vertex_ai/claude-sonnet-5@default": { + "deprecation_date": "2026-12-24", "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 78b53cefc53..1a130c4ac0a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -54,6 +54,7 @@ "output_cost_per_image": 0.04 }, "1024-x-1024/dall-e-2": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 1.9e-08, "litellm_provider": "openai", "mode": "image_generation", @@ -67,6 +68,7 @@ "output_cost_per_image": 0.08 }, "256-x-256/dall-e-2": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 2.4414e-07, "litellm_provider": "openai", "mode": "image_generation", @@ -80,6 +82,7 @@ "output_cost_per_image": 0.018 }, "512-x-512/dall-e-2": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 6.86e-08, "litellm_provider": "openai", "mode": "image_generation", @@ -2887,6 +2890,7 @@ "supports_function_calling": true }, "azure_ai/claude-haiku-4-5": { + "deprecation_date": "2026-10-19", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -2908,6 +2912,7 @@ "supports_vision": true }, "azure_ai/claude-opus-4-5": { + "deprecation_date": "2026-10-19", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -2930,6 +2935,7 @@ "supports_output_config": true }, "azure_ai/claude-opus-4-6": { + "deprecation_date": "2027-02-02", "supports_adaptive_thinking": true, "input_cost_per_token": 5e-06, "output_cost_per_token": 2.5e-05, @@ -2959,6 +2965,7 @@ "supports_max_reasoning_effort": true }, "azure_ai/claude-opus-4-7": { + "deprecation_date": "2027-04-06", "supports_adaptive_thinking": true, "input_cost_per_token": 5e-06, "output_cost_per_token": 2.5e-05, @@ -3083,6 +3090,7 @@ "supports_max_reasoning_effort": true }, "azure_ai/claude-opus-4-1": { + "deprecation_date": "2026-08-05", "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, @@ -3104,6 +3112,7 @@ "supports_vision": true }, "azure_ai/claude-sonnet-4-5": { + "deprecation_date": "2026-10-19", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -3156,6 +3165,7 @@ "supports_max_reasoning_effort": true }, "azure_ai/claude-sonnet-4-6": { + "deprecation_date": "2027-02-10", "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, @@ -3226,6 +3236,7 @@ "supports_tool_choice": true }, "azure_ai/gpt-5.5": { + "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -3318,6 +3329,7 @@ "supports_minimal_reasoning_effort": false }, "azure_ai/gpt-5.4": { + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens": 5e-07, "cache_read_input_token_cost_priority": 5e-07, @@ -3364,6 +3376,7 @@ "supports_minimal_reasoning_effort": true }, "azure_ai/gpt-5.4-2026-03-05": { + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens": 5e-07, "cache_read_input_token_cost_priority": 5e-07, @@ -3410,6 +3423,7 @@ "supports_minimal_reasoning_effort": true }, "azure_ai/gpt-5.4-pro": { + "deprecation_date": "2027-09-07", "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "cache_read_input_token_cost_priority": 6e-06, @@ -3455,6 +3469,7 @@ "supports_minimal_reasoning_effort": true }, "azure_ai/gpt-5.4-pro-2026-03-05": { + "deprecation_date": "2027-09-07", "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "cache_read_input_token_cost_priority": 6e-06, @@ -3500,6 +3515,7 @@ "supports_minimal_reasoning_effort": true }, "azure_ai/gpt-5.4-mini": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_priority": 1.5e-07, "input_cost_per_token": 7.5e-07, @@ -3540,6 +3556,7 @@ "supports_minimal_reasoning_effort": false }, "azure_ai/gpt-5.4-mini-2026-03-17": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_priority": 1.5e-07, "input_cost_per_token": 7.5e-07, @@ -3580,6 +3597,7 @@ "supports_minimal_reasoning_effort": false }, "azure_ai/gpt-5.4-nano": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_priority": 4e-08, "input_cost_per_token": 2e-07, @@ -3620,6 +3638,7 @@ "supports_minimal_reasoning_effort": false }, "azure_ai/gpt-5.4-nano-2026-03-17": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_priority": 4e-08, "input_cost_per_token": 2e-07, @@ -3849,6 +3868,7 @@ "supports_vision": true }, "azure/eu/gpt-5.1": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.38e-06, "litellm_provider": "azure", @@ -3918,6 +3938,7 @@ "supports_none_reasoning_effort": true }, "azure/eu/gpt-5.1-codex": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.38e-06, "litellm_provider": "azure", @@ -3948,6 +3969,7 @@ "supports_vision": true }, "azure/eu/gpt-5.1-codex-mini": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 2.75e-07, "litellm_provider": "azure", @@ -4107,6 +4129,7 @@ "supports_vision": true }, "azure/global-standard/gpt-4o-mini": { + "deprecation_date": "2027-04-14", "input_cost_per_token": 1.5e-07, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -4155,6 +4178,7 @@ "supports_vision": true }, "azure/global/gpt-5.1": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -4224,6 +4248,7 @@ "supports_none_reasoning_effort": true }, "azure/global/gpt-5.1-codex": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -4254,6 +4279,7 @@ "supports_vision": true }, "azure/global/gpt-5.1-codex-mini": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 2.5e-07, "litellm_provider": "azure", @@ -4492,6 +4518,7 @@ "supports_vision": true }, "azure/gpt-4.1": { + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -4559,6 +4586,7 @@ "supports_web_search": false }, "azure/gpt-4.1-mini": { + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 4e-07, "input_cost_per_token_batches": 2e-07, @@ -4626,6 +4654,7 @@ "supports_web_search": false }, "azure/gpt-4.1-nano": { + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, @@ -4902,6 +4931,7 @@ "supports_vision": false }, "azure/gpt-4o-mini": { + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 1.65e-07, "litellm_provider": "azure", @@ -5344,6 +5374,7 @@ "supports_vision": true }, "azure/gpt-5": { + "deprecation_date": "2027-02-09", "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -5507,6 +5538,7 @@ "supports_vision": true }, "azure/gpt-5-mini": { + "deprecation_date": "2027-02-09", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 2.5e-07, "litellm_provider": "azure", @@ -5572,6 +5604,7 @@ "supports_vision": true }, "azure/gpt-5-nano": { + "deprecation_date": "2027-02-09", "cache_read_input_token_cost": 5e-09, "input_cost_per_token": 5e-08, "litellm_provider": "azure", @@ -5667,6 +5700,7 @@ "supports_vision": true }, "azure/gpt-5.1": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -5736,6 +5770,7 @@ "supports_none_reasoning_effort": true }, "azure/gpt-5.1-codex": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -5797,6 +5832,7 @@ "supports_vision": true }, "azure/gpt-5.1-codex-mini": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 2.5e-07, "litellm_provider": "azure", @@ -5827,6 +5863,7 @@ "supports_vision": true }, "azure/gpt-5.2": { + "deprecation_date": "2027-06-08", "cache_read_input_token_cost": 1.75e-07, "input_cost_per_token": 1.75e-06, "litellm_provider": "azure", @@ -6136,6 +6173,7 @@ "supports_web_search": true }, "azure/gpt-5.4": { + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens": 5e-07, "cache_read_input_token_cost_priority": 5e-07, @@ -6180,6 +6218,7 @@ "supports_minimal_reasoning_effort": true }, "azure/us/gpt-5.4": { + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.8e-07, "cache_read_input_token_cost_priority": 5.5e-07, "input_cost_per_token": 2.75e-06, @@ -6218,6 +6257,7 @@ "supports_minimal_reasoning_effort": true }, "azure/eu/gpt-5.4": { + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.8e-07, "cache_read_input_token_cost_priority": 5.5e-07, "input_cost_per_token": 2.75e-06, @@ -6379,6 +6419,7 @@ "supports_minimal_reasoning_effort": true }, "azure/gpt-5.4-pro": { + "deprecation_date": "2027-09-07", "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -7045,6 +7086,7 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.5": { + "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -7095,6 +7137,7 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.5": { + "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -7142,6 +7185,7 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.5": { + "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -7408,6 +7452,7 @@ "supports_web_search": true }, "azure/gpt-5.4-mini": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 7.5e-07, "litellm_provider": "azure", @@ -7489,6 +7534,7 @@ "supports_xhigh_reasoning_effort": true }, "azure/gpt-5.4-nano": { + "deprecation_date": "2027-09-21", "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2e-07, "litellm_provider": "azure", @@ -7601,6 +7647,7 @@ "output_cost_per_token": 0.0 }, "azure/high/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.59263611e-07, "litellm_provider": "azure", "mode": "image_generation", @@ -7610,6 +7657,7 @@ ] }, "azure/high/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.58945719e-07, "litellm_provider": "azure", "mode": "image_generation", @@ -7619,6 +7667,7 @@ ] }, "azure/high/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.58945719e-07, "litellm_provider": "azure", "mode": "image_generation", @@ -7628,6 +7677,7 @@ ] }, "azure/low/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.0490417e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7637,6 +7687,7 @@ ] }, "azure/low/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.0172526e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7646,6 +7697,7 @@ ] }, "azure/low/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 1.0172526e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7655,6 +7707,7 @@ ] }, "azure/medium/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7664,6 +7717,7 @@ ] }, "azure/medium/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7673,6 +7727,7 @@ ] }, "azure/medium/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7695,6 +7750,7 @@ ] }, "azure/gpt-image-1.5": { + "deprecation_date": "2027-06-16", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 5e-06, "input_cost_per_image_token": 8e-06, @@ -7720,6 +7776,7 @@ ] }, "azure/gpt-image-2": { + "deprecation_date": "2027-10-21", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 5e-06, "input_cost_per_image_token": 8e-06, @@ -7751,6 +7808,7 @@ "supports_pdf_input": true }, "azure/low/1024-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 2.0751953125e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7760,6 +7818,7 @@ ] }, "azure/low/1024-x-1536/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 2.0751953125e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7769,6 +7828,7 @@ ] }, "azure/low/1536-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 2.0345052083e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7778,6 +7838,7 @@ ] }, "azure/medium/1024-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 8.056640625e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7787,6 +7848,7 @@ ] }, "azure/medium/1024-x-1536/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 8.056640625e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7796,6 +7858,7 @@ ] }, "azure/medium/1536-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 7.9752604167e-09, "litellm_provider": "azure", "mode": "image_generation", @@ -7805,6 +7868,7 @@ ] }, "azure/high/1024-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 3.173828125e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7814,6 +7878,7 @@ ] }, "azure/high/1024-x-1536/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 3.173828125e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7823,6 +7888,7 @@ ] }, "azure/high/1536-x-1024/gpt-image-1-mini": { + "deprecation_date": "2027-04-07", "input_cost_per_pixel": 3.1575520833e-08, "litellm_provider": "azure", "mode": "image_generation", @@ -7850,6 +7916,7 @@ "supports_function_calling": true }, "azure/o1": { + "deprecation_date": "2026-10-21", "cache_read_input_token_cost": 7.5e-06, "input_cost_per_token": 1.5e-05, "litellm_provider": "azure", @@ -7944,6 +8011,7 @@ "supports_vision": false }, "azure/o3": { + "deprecation_date": "2026-10-21", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "azure", @@ -8041,6 +8109,7 @@ "supports_web_search": true }, "azure/o3-mini": { + "deprecation_date": "2026-10-01", "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", @@ -8071,6 +8140,7 @@ "supports_vision": false }, "azure/o3-pro": { + "deprecation_date": "2026-12-17", "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "azure", @@ -8132,6 +8202,7 @@ "supports_vision": true }, "azure/o4-mini": { + "deprecation_date": "2026-10-16", "cache_read_input_token_cost": 2.75e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", @@ -8580,6 +8651,7 @@ "supports_vision": true }, "azure/us/gpt-5.1": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.38e-06, "litellm_provider": "azure", @@ -8649,6 +8721,7 @@ "supports_none_reasoning_effort": true }, "azure/us/gpt-5.1-codex": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.38e-06, "litellm_provider": "azure", @@ -8679,6 +8752,7 @@ "supports_vision": true }, "azure/us/gpt-5.1-codex-mini": { + "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 2.75e-07, "litellm_provider": "azure", @@ -8876,6 +8950,7 @@ ] }, "azure_ai/FW-DeepSeek-V3.2": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 3.1e-07, "input_cost_per_token": 6.2e-07, "litellm_provider": "azure_ai", @@ -8906,6 +8981,7 @@ "supports_tool_choice": true }, "azure_ai/FW-GLM-5": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 2.2e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "azure_ai", @@ -8921,6 +8997,7 @@ "supports_tool_choice": true }, "azure_ai/FW-GLM-5.1": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 2.86e-07, "input_cost_per_token": 1.54e-06, "litellm_provider": "azure_ai", @@ -8987,6 +9064,7 @@ "supports_tool_choice": true }, "azure_ai/FW-Kimi-K2.5": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 1.1e-07, "input_cost_per_token": 6.6e-07, "litellm_provider": "azure_ai", @@ -9079,6 +9157,7 @@ "supports_vision": true }, "azure_ai/FW-MiniMax-M2.5": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 3.3e-08, "input_cost_per_token": 3.3e-07, "litellm_provider": "azure_ai", @@ -9164,6 +9243,7 @@ ] }, "azure_ai/MAI-Image-2e": { + "deprecation_date": "2026-08-15", "input_cost_per_token": 5e-06, "litellm_provider": "azure_ai", "mode": "image_generation", @@ -9175,6 +9255,7 @@ ] }, "azure_ai/Llama-3.2-11B-Vision-Instruct": { + "deprecation_date": "2026-06-13", "input_cost_per_token": 3.7e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9188,6 +9269,7 @@ "supports_vision": true }, "azure_ai/Llama-3.2-90B-Vision-Instruct": { + "deprecation_date": "2026-06-13", "input_cost_per_token": 2.04e-06, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9249,6 +9331,7 @@ "supports_tool_choice": true }, "azure_ai/Meta-Llama-3.1-405B-Instruct": { + "deprecation_date": "2026-06-13", "input_cost_per_token": 5.33e-06, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9271,6 +9354,7 @@ "supports_tool_choice": true }, "azure_ai/Meta-Llama-3.1-8B-Instruct": { + "deprecation_date": "2026-06-13", "input_cost_per_token": 3e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9452,6 +9536,7 @@ "supports_reasoning": true }, "azure_ai/mistral-document-ai-2505": { + "deprecation_date": "2026-07-20", "litellm_provider": "azure_ai", "ocr_cost_per_page": 0.003, "mode": "ocr", @@ -9529,6 +9614,7 @@ "output_cost_per_token": 0.0 }, "azure_ai/cohere-rerank-v3.5": { + "deprecation_date": "2026-05-14", "input_cost_per_query": 0.002, "input_cost_per_token": 0.0, "litellm_provider": "azure_ai", @@ -9591,6 +9677,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-r1": { + "deprecation_date": "2026-08-13", "input_cost_per_token": 1.35e-06, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9614,6 +9701,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-v3-0324": { + "deprecation_date": "2026-07-13", "input_cost_per_token": 1.14e-06, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -9626,6 +9714,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-v3.1": { + "deprecation_date": "2026-07-13", "input_cost_per_token": 1.23e-06, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -9639,6 +9728,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-v4-pro": { + "deprecation_date": "2028-02-20", "input_cost_per_token": 1.74e-06, "litellm_provider": "azure_ai", "max_input_tokens": 1000000, @@ -9652,6 +9742,7 @@ "supports_tool_choice": true }, "azure_ai/deepseek-v4-flash": { + "deprecation_date": "2028-02-20", "input_cost_per_token": 1.9e-07, "litellm_provider": "azure_ai", "max_input_tokens": 1000000, @@ -9683,6 +9774,7 @@ "supports_embedding_image_input": true }, "azure_ai/global/grok-3": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 3e-06, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -9697,6 +9789,7 @@ "supports_web_search": true }, "azure_ai/global/grok-3-mini": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 2.5e-07, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -9712,6 +9805,7 @@ "supports_web_search": true }, "azure_ai/grok-3": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 3e-06, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -9726,6 +9820,7 @@ "supports_web_search": true }, "azure_ai/grok-3-mini": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 2.5e-07, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -9773,6 +9868,7 @@ "supports_web_search": true }, "azure_ai/grok-4-fast-non-reasoning": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", @@ -9786,6 +9882,7 @@ "supports_web_search": true }, "azure_ai/grok-4-fast-reasoning": { + "deprecation_date": "2026-05-01", "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", @@ -9863,6 +9960,7 @@ "supports_tool_choice": true }, "azure_ai/kimi-k2.5": { + "deprecation_date": "2027-01-26", "input_cost_per_token": 6e-07, "litellm_provider": "azure_ai", "max_input_tokens": 262144, @@ -9877,6 +9975,7 @@ "supports_vision": true }, "azure_ai/kimi-k2.6": { + "deprecation_date": "2027-04-16", "input_cost_per_token": 9.5e-07, "litellm_provider": "azure_ai", "max_input_tokens": 262144, @@ -10004,6 +10103,7 @@ "supports_vision": true }, "babbage-002": { + "deprecation_date": "2026-09-28", "input_cost_per_token": 4e-07, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, @@ -11999,6 +12099,7 @@ ] }, "claude-haiku-4-5-20251001": { + "deprecation_date": "2026-10-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -12022,6 +12123,7 @@ "prompt_cache_min_tokens": 4096 }, "claude-haiku-4-5": { + "deprecation_date": "2026-10-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -12170,6 +12272,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-sonnet-4-5": { + "deprecation_date": "2026-09-29", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, @@ -12203,6 +12306,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-sonnet-4-5-20250929": { + "deprecation_date": "2026-09-29", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, @@ -12237,6 +12341,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-sonnet-5": { + "deprecation_date": "2027-06-30", "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -12273,6 +12378,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-sonnet-4-6": { + "deprecation_date": "2027-02-17", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -12419,6 +12525,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-opus-4-5-20251101": { + "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12448,6 +12555,7 @@ "prompt_cache_min_tokens": 4096 }, "claude-opus-4-5": { + "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12477,6 +12585,7 @@ "prompt_cache_min_tokens": 4096 }, "claude-opus-4-6": { + "deprecation_date": "2027-02-05", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12513,6 +12622,7 @@ "prompt_cache_min_tokens": 4096 }, "claude-opus-4-6-20260205": { + "deprecation_date": "2027-02-05", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12549,6 +12659,7 @@ "prompt_cache_min_tokens": 4096 }, "claude-opus-4-7": { + "deprecation_date": "2027-04-16", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12587,6 +12698,7 @@ "prompt_cache_min_tokens": 2048 }, "claude-opus-4-7-20260416": { + "deprecation_date": "2027-04-16", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12625,6 +12737,7 @@ "prompt_cache_min_tokens": 2048 }, "claude-fable-5": { + "deprecation_date": "2027-06-09", "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, "cache_read_input_token_cost": 1e-06, @@ -12660,6 +12773,7 @@ "prompt_cache_min_tokens": 512 }, "claude-opus-5": { + "deprecation_date": "2027-07-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -12698,6 +12812,7 @@ "prompt_cache_min_tokens": 512 }, "claude-opus-4-8": { + "deprecation_date": "2027-05-28", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -14801,6 +14916,7 @@ "mode": "search" }, "davinci-002": { + "deprecation_date": "2026-09-28", "input_cost_per_token": 2e-06, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, @@ -18353,6 +18469,7 @@ } }, "gemini-2.5-flash": { + "deprecation_date": "2026-10-20", "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -18398,6 +18515,7 @@ "supports_image_size": false }, "gemini-2.5-flash-image": { + "deprecation_date": "2026-10-02", "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -18442,6 +18560,7 @@ "supports_image_size": false }, "gemini-3-pro-image": { + "deprecation_date": "2027-05-28", "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -18522,6 +18641,7 @@ "web_search_billing_unit": "per_query" }, "gemini-3.1-flash-image": { + "deprecation_date": "2027-05-28", "input_cost_per_image": 0.00056, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", @@ -18646,6 +18766,7 @@ "web_search_billing_unit": "per_query" }, "gemini-3.1-flash-lite": { + "deprecation_date": "2027-05-07", "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, @@ -18702,6 +18823,7 @@ "web_search_billing_unit": "per_query" }, "gemini-3.5-flash-lite": { + "deprecation_date": "2027-07-21", "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_flex": 2e-08, "cache_read_input_token_cost_priority": 5e-08, @@ -18791,6 +18913,7 @@ "supports_web_search": true }, "gemini-2.5-flash-lite": { + "deprecation_date": "2026-10-20", "cache_read_input_token_cost": 1e-08, "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, @@ -19062,6 +19185,7 @@ "supports_image_size": false }, "gemini-2.5-pro": { + "deprecation_date": "2026-10-20", "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, @@ -19373,6 +19497,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.5-flash": { + "deprecation_date": "2027-05-19", "cache_read_input_token_cost": 1.5e-07, "input_cost_per_token": 1.5e-06, "input_cost_per_audio_token": 1e-06, @@ -19809,6 +19934,7 @@ "web_search_billing_unit": "per_query" }, "gemini/gemini-robotics-er-1.6-preview": { + "deprecation_date": "2026-08-31", "input_cost_per_audio_token": 2e-06, "input_cost_per_token": 1e-06, "litellm_provider": "gemini", @@ -19879,6 +20005,7 @@ "supports_vision": true }, "gemini-embedding-001": { + "deprecation_date": "2028-05-20", "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 2048, @@ -21492,6 +21619,7 @@ "supports_vision": true }, "gemini-3.5-flash": { + "deprecation_date": "2027-05-19", "cache_read_input_token_cost": 1.5e-07, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 1.5e-06, @@ -23004,6 +23132,7 @@ "supports_tool_choice": true }, "gpt-3.5-turbo-instruct": { + "deprecation_date": "2026-09-28", "input_cost_per_token": 1.5e-06, "litellm_provider": "text-completion-openai", "max_input_tokens": 8192, @@ -24135,6 +24264,7 @@ "supports_pdf_input": true }, "low/1024-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.009, "litellm_provider": "openai", "mode": "image_generation", @@ -24146,6 +24276,7 @@ "supports_pdf_input": true }, "low/1024-x-1536/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -24157,6 +24288,7 @@ "supports_pdf_input": true }, "low/1536-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -24168,6 +24300,7 @@ "supports_pdf_input": true }, "medium/1024-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.034, "litellm_provider": "openai", "mode": "image_generation", @@ -24179,6 +24312,7 @@ "supports_pdf_input": true }, "medium/1024-x-1536/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.05, "litellm_provider": "openai", "mode": "image_generation", @@ -24190,6 +24324,7 @@ "supports_pdf_input": true }, "medium/1536-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.05, "litellm_provider": "openai", "mode": "image_generation", @@ -24201,6 +24336,7 @@ "supports_pdf_input": true }, "high/1024-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.133, "litellm_provider": "openai", "mode": "image_generation", @@ -24212,6 +24348,7 @@ "supports_pdf_input": true }, "high/1024-x-1536/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.2, "litellm_provider": "openai", "mode": "image_generation", @@ -24223,6 +24360,7 @@ "supports_pdf_input": true }, "high/1536-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.2, "litellm_provider": "openai", "mode": "image_generation", @@ -24234,6 +24372,7 @@ "supports_pdf_input": true }, "standard/1024-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.009, "litellm_provider": "openai", "mode": "image_generation", @@ -24245,6 +24384,7 @@ "supports_pdf_input": true }, "standard/1024-x-1536/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -24256,6 +24396,7 @@ "supports_pdf_input": true }, "standard/1536-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -24267,6 +24408,7 @@ "supports_pdf_input": true }, "1024-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.009, "litellm_provider": "openai", "mode": "image_generation", @@ -24278,6 +24420,7 @@ "supports_pdf_input": true }, "1024-x-1536/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -24289,6 +24432,7 @@ "supports_pdf_input": true }, "1536-x-1024/gpt-image-1.5": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.013, "litellm_provider": "openai", "mode": "image_generation", @@ -27202,18 +27346,21 @@ "output_cost_per_second": 0.0 }, "hd/1024-x-1024/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 7.629e-08, "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_pixel": 0.0 }, "hd/1024-x-1792/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 6.539e-08, "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_pixel": 0.0 }, "hd/1792-x-1024/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 6.539e-08, "litellm_provider": "openai", "mode": "image_generation", @@ -27260,6 +27407,7 @@ "max_output_tokens": 8192 }, "high/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.167, "input_cost_per_pixel": 1.59263611e-07, "litellm_provider": "openai", @@ -27270,6 +27418,7 @@ ] }, "high/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.25, "input_cost_per_pixel": 1.58945719e-07, "litellm_provider": "openai", @@ -27280,6 +27429,7 @@ ] }, "high/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.25, "input_cost_per_pixel": 1.58945719e-07, "litellm_provider": "openai", @@ -28067,6 +28217,7 @@ "supports_tool_choice": true }, "low/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.011, "input_cost_per_pixel": 1.0490417e-08, "litellm_provider": "openai", @@ -28077,6 +28228,7 @@ ] }, "low/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.016, "input_cost_per_pixel": 1.0172526e-08, "litellm_provider": "openai", @@ -28087,6 +28239,7 @@ ] }, "low/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.016, "input_cost_per_pixel": 1.0172526e-08, "litellm_provider": "openai", @@ -28111,6 +28264,7 @@ "output_cost_per_image": 0.072 }, "medium/1024-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.042, "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "openai", @@ -28121,6 +28275,7 @@ ] }, "medium/1024-x-1536/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.063, "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "openai", @@ -28131,6 +28286,7 @@ ] }, "medium/1536-x-1024/gpt-image-1": { + "deprecation_date": "2026-10-23", "input_cost_per_image": 0.063, "input_cost_per_pixel": 4.0054321e-08, "litellm_provider": "openai", @@ -28141,6 +28297,7 @@ ] }, "low/1024-x-1024/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.005, "litellm_provider": "openai", "mode": "image_generation", @@ -28149,6 +28306,7 @@ ] }, "low/1024-x-1536/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.006, "litellm_provider": "openai", "mode": "image_generation", @@ -28157,6 +28315,7 @@ ] }, "low/1536-x-1024/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.006, "litellm_provider": "openai", "mode": "image_generation", @@ -28165,6 +28324,7 @@ ] }, "medium/1024-x-1024/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.011, "litellm_provider": "openai", "mode": "image_generation", @@ -28173,6 +28333,7 @@ ] }, "medium/1024-x-1536/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.015, "litellm_provider": "openai", "mode": "image_generation", @@ -28181,6 +28342,7 @@ ] }, "medium/1536-x-1024/gpt-image-1-mini": { + "deprecation_date": "2026-12-01", "input_cost_per_image": 0.015, "litellm_provider": "openai", "mode": "image_generation", @@ -30074,6 +30236,7 @@ ] }, "multimodalembedding@001": { + "deprecation_date": "2027-04-01", "input_cost_per_character": 2e-07, "input_cost_per_image": 0.0001, "input_cost_per_token": 8e-07, @@ -35772,18 +35935,21 @@ "output_cost_per_image": 0.14 }, "standard/1024-x-1024/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 3.81469e-08, "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_pixel": 0.0 }, "standard/1024-x-1792/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 4.359e-08, "litellm_provider": "openai", "mode": "image_generation", "output_cost_per_pixel": 0.0 }, "standard/1792-x-1024/dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_pixel": 4.359e-08, "litellm_provider": "openai", "mode": "image_generation", @@ -35847,6 +36013,7 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models" }, "text-embedding-005": { + "deprecation_date": "2027-04-01", "input_cost_per_character": 2.5e-08, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-embedding-models", @@ -35920,6 +36087,7 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, "text-moderation-007": { + "deprecation_date": "2025-10-27", "input_cost_per_token": 0.0, "litellm_provider": "openai", "max_input_tokens": 32768, @@ -35929,6 +36097,7 @@ "output_cost_per_token": 0.0 }, "text-moderation-latest": { + "deprecation_date": "2025-10-27", "input_cost_per_token": 0.0, "litellm_provider": "openai", "max_input_tokens": 32768, @@ -35938,6 +36107,7 @@ "output_cost_per_token": 0.0 }, "text-moderation-stable": { + "deprecation_date": "2025-10-27", "input_cost_per_token": 0.0, "litellm_provider": "openai", "max_input_tokens": 32768, @@ -35947,6 +36117,7 @@ "output_cost_per_token": 0.0 }, "text-multilingual-embedding-002": { + "deprecation_date": "2027-04-01", "input_cost_per_character": 2.5e-08, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-embedding-models", @@ -38434,6 +38605,7 @@ "supports_tool_choice": true }, "vertex_ai/claude-haiku-4-5": { + "deprecation_date": "2026-10-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -38457,6 +38629,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-haiku-4-5@20251001": { + "deprecation_date": "2026-10-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -38609,6 +38782,7 @@ "supports_vision": true }, "vertex_ai/claude-opus-4": { + "deprecation_date": "2026-05-14", "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, @@ -38636,6 +38810,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-opus-4-1": { + "deprecation_date": "2026-08-05", "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, @@ -38654,6 +38829,7 @@ "supports_vision": true }, "vertex_ai/claude-opus-4-1@20250805": { + "deprecation_date": "2026-08-05", "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, @@ -38672,6 +38848,7 @@ "supports_vision": true }, "vertex_ai/claude-opus-4-5": { + "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -38700,6 +38877,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-5@20251101": { + "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -38729,6 +38907,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-6": { + "deprecation_date": "2027-02-05", "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -38759,6 +38938,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-6@default": { + "deprecation_date": "2027-02-05", "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -38789,6 +38969,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-opus-4-7": { + "deprecation_date": "2027-04-16", "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -38820,6 +39001,7 @@ "prompt_cache_min_tokens": 2048 }, "vertex_ai/claude-opus-4-7@default": { + "deprecation_date": "2027-04-16", "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -38851,6 +39033,7 @@ "prompt_cache_min_tokens": 2048 }, "vertex_ai/claude-fable-5": { + "deprecation_date": "2027-06-08", "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, @@ -38882,6 +39065,7 @@ "supports_max_reasoning_effort": true }, "vertex_ai/claude-fable-5@default": { + "deprecation_date": "2027-06-08", "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, @@ -38913,6 +39097,7 @@ "supports_max_reasoning_effort": true }, "vertex_ai/claude-opus-5": { + "deprecation_date": "2027-01-24", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -38945,6 +39130,7 @@ "prompt_cache_min_tokens": 512 }, "vertex_ai/claude-opus-5@default": { + "deprecation_date": "2027-01-24", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -38977,6 +39163,7 @@ "prompt_cache_min_tokens": 512 }, "vertex_ai/claude-opus-4-8": { + "deprecation_date": "2027-05-28", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -39009,6 +39196,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-opus-4-8@default": { + "deprecation_date": "2027-05-28", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "cache_creation_input_token_cost": 6.25e-06, @@ -39041,6 +39229,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-4-5": { + "deprecation_date": "2026-09-29", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -39069,6 +39258,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-5": { + "deprecation_date": "2026-12-24", "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, @@ -39131,6 +39321,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-4-5@20250929": { + "deprecation_date": "2026-09-29", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -39160,6 +39351,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-opus-4@20250514": { + "deprecation_date": "2026-05-14", "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, @@ -39187,6 +39379,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-4": { + "deprecation_date": "2026-05-14", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -39218,6 +39411,7 @@ "prompt_cache_min_tokens": 1024 }, "vertex_ai/claude-sonnet-4@20250514": { + "deprecation_date": "2026-05-14", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -39382,6 +39576,7 @@ "supports_tool_choice": true }, "vertex_ai/gemini-2.5-flash-image": { + "deprecation_date": "2026-10-02", "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -39427,6 +39622,7 @@ "supports_image_size": false }, "vertex_ai/gemini-3-pro-image": { + "deprecation_date": "2027-05-28", "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -39459,6 +39655,7 @@ "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" }, "vertex_ai/gemini-3.1-flash-image": { + "deprecation_date": "2027-05-28", "input_cost_per_image": 0.00056, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", @@ -39535,6 +39732,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.1-flash-lite": { + "deprecation_date": "2027-05-07", "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, @@ -39591,6 +39789,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.5-flash-lite": { + "deprecation_date": "2027-07-21", "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_flex": 2e-08, "cache_read_input_token_cost_priority": 5e-08, @@ -40308,6 +40507,7 @@ "supports_tool_choice": true }, "vertex_ai/veo-2.0-generate-001": { + "deprecation_date": "2026-06-30", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -40322,6 +40522,7 @@ ] }, "vertex_ai/veo-3.0-fast-generate-001": { + "deprecation_date": "2026-06-30", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -40336,6 +40537,7 @@ ] }, "vertex_ai/veo-3.0-generate-001": { + "deprecation_date": "2026-06-30", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -40378,6 +40580,7 @@ ] }, "vertex_ai/veo-3.1-generate-001": { + "deprecation_date": "2026-11-17", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -40392,6 +40595,7 @@ ] }, "vertex_ai/veo-3.1-fast-generate-001": { + "deprecation_date": "2026-11-17", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -46773,6 +46977,7 @@ } }, "vertex_ai/claude-sonnet-5@default": { + "deprecation_date": "2026-12-24", "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, From 0896015927caf2e38988c7add242fb3b6f96d3e2 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 18 Aug 2026 08:43:25 -0700 Subject: [PATCH 062/538] test(ui): await the playground model combobox before clicking it (#36850) The ChatUI playground test helper looked up the model combobox with a synchronous getByPlaceholderText. That control renders its placeholder from the model-loading flag, so the text is "Loading models..." until the mocked fetch resolves, and the element the helper wants does not exist yet. Under a loaded full-suite run the query could land inside that window and fail with "Unable to find an element with the placeholder text of: Select a Model", while the same test passed in isolation every time. Switch the helper to findByPlaceholderText so it waits for the control to come back after loading. --- .../(dashboard)/playground/components/chat_ui/ChatUI.test.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx index 24c56b276af..9d9c5dc86ab 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx @@ -35,7 +35,7 @@ const STREAMING_ENABLED_ARG_INDEX = 25; async function openComboboxByPlaceholder(placeholder: string) { const user = userEvent.setup(); - const combobox = screen.getByPlaceholderText(placeholder); + const combobox = await screen.findByPlaceholderText(placeholder); await user.click(combobox); return combobox; } From b08032c5f743b5e5e009f5e1b4086177c2cf1403 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 18 Aug 2026 08:51:22 -0700 Subject: [PATCH 063/538] refactor(ui): migrate budget and skill forms to react-hook-form and shadcn (#37262) * refactor(ui): migrate budget and skill forms to react-hook-form and shadcn Moves three dashboard forms off antd Form onto react-hook-form plus the shared shadcn field kit, and onto semantic color tokens so they render correctly in dark mode. The submitted request bodies are unchanged. Budget create and edit previously relied on antd InputNumber precision={2}, which rounds the submitted value rather than only the display. That rounding is now an explicit shared helper so the wire payload stays identical, and the helper carries unit tests covering key presence, null passthrough, negatives and non-finite input. Characterization tests for both budget modals were written against the antd implementation first and pass unchanged against the migrated components, which is what pins the payload. They are named .integration.test.tsx per the dashboard test tiers, with the pure rounding logic unit tested separately. * fix(ui): keep collapsed Optional Settings values so reopening does not lose them react-hook-form shouldUnregister deletes a field's value when its section unmounts, so typing a budget, collapsing Optional Settings and reopening it submitted the seeded default instead of what was typed. antd reported only mounted fields in onFinish but preserved their values in its store, so the two behaviours have to be reproduced separately. Drop shouldUnregister, drive the section from controlled state, and blank the section's fields at submit while it is closed. Seed the edit form from the five form fields rather than the whole budget record, which shouldUnregister had been masking. --- .../_components/budgetPrecision.test.ts | 45 ++ .../budgets/_components/budgetPrecision.ts | 18 + .../budget_modal.integration.test.tsx | 139 ++++++ .../budgets/_components/budget_modal.tsx | 166 +++++-- .../edit_budget_modal.integration.test.tsx | 112 +++++ .../budgets/_components/edit_budget_modal.tsx | 169 +++++-- .../skills/_components/add_plugin_form.tsx | 420 +++++++++++------- 7 files changed, 806 insertions(+), 263 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budgetPrecision.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budgetPrecision.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.integration.test.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budgetPrecision.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budgetPrecision.test.ts new file mode 100644 index 00000000000..358ea2980a6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budgetPrecision.test.ts @@ -0,0 +1,45 @@ +import { describe, expect, it } from "vitest"; + +import { applyBudgetPrecision } from "./budgetPrecision"; + +describe("applyBudgetPrecision", () => { + it("rounds each precision field to two decimals, matching antd InputNumber precision={2}", () => { + const typed = { budget_id: "b", tpm_limit: 500.567, rpm_limit: 7.005, max_budget: 42.567 }; + const rounded = { budget_id: "b", tpm_limit: 500.57, rpm_limit: 7.01, max_budget: 42.57 }; + + expect(applyBudgetPrecision(typed)).toEqual(rounded); + }); + + it("leaves non-precision fields untouched even when numeric", () => { + expect(applyBudgetPrecision({ soft_budget: 1.239, budget_duration: "30d" })).toEqual({ + soft_budget: 1.239, + budget_duration: "30d", + }); + }); + + it("preserves key presence exactly, so an omitted field is not reintroduced as undefined", () => { + expect(Object.keys(applyBudgetPrecision({ budget_id: "b", tpm_limit: 1 }))).toEqual(["budget_id", "tpm_limit"]); + }); + + it("passes null and undefined through without coercing them to a number", () => { + expect(applyBudgetPrecision({ tpm_limit: null, rpm_limit: undefined, max_budget: 1.005 })).toEqual({ + tpm_limit: null, + rpm_limit: undefined, + max_budget: 1.01, + }); + }); + + it("rounds negatives away from zero the way antd does", () => { + expect(applyBudgetPrecision({ max_budget: -1.005 })).toEqual({ max_budget: -1.01 }); + }); + + it("returns non-finite values unchanged rather than emitting NaN", () => { + expect(applyBudgetPrecision({ max_budget: Number.POSITIVE_INFINITY })).toEqual({ + max_budget: Number.POSITIVE_INFINITY, + }); + }); + + it("does not disturb a value that already has two or fewer decimals", () => { + expect(applyBudgetPrecision({ max_budget: 42.5, tpm_limit: 500 })).toEqual({ max_budget: 42.5, tpm_limit: 500 }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budgetPrecision.ts b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budgetPrecision.ts new file mode 100644 index 00000000000..51930a6f604 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budgetPrecision.ts @@ -0,0 +1,18 @@ +const PRECISION_FIELDS: ReadonlySet = new Set(["tpm_limit", "rpm_limit", "max_budget"]); + +const roundToPrecision = (value: number): number => { + const shifted = Number(`${Math.abs(value)}e2`); + if (!Number.isFinite(shifted)) { + return value; + } + const rounded = Number(`${Math.round(shifted)}e-2`); + return value < 0 ? -rounded : rounded; +}; + +export const applyBudgetPrecision = >(formValues: TValues): TValues => + Object.fromEntries( + Object.entries(formValues).map(([key, value]) => [ + key, + PRECISION_FIELDS.has(key) && typeof value === "number" ? roundToPrecision(value) : value, + ]), + ) as TValues; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.integration.test.tsx new file mode 100644 index 00000000000..419da23af3a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.integration.test.tsx @@ -0,0 +1,139 @@ +import { render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import React from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import BudgetModal from "./budget_modal"; + +const { createMock } = vi.hoisted(() => ({ createMock: vi.fn() })); + +vi.mock("@/app/(dashboard)/hooks/budgets/useBudgets", () => ({ + useCreateBudget: () => ({ mutateAsync: createMock }), +})); + +const FULL_PAYLOAD = { + budget_id: "budget-alpha", + tpm_limit: 500.57, + rpm_limit: 7, + max_budget: 42.57, + budget_duration: "30d", +}; + +const renderModal = () => render(); + +const create = async (user: ReturnType) => + user.click(screen.getByRole("button", { name: "Create Budget" })); + +const openOptionalSettings = async (user: ReturnType) => { + await user.click(screen.getByText("Optional Settings")); + await screen.findByLabelText("Max Budget (USD)"); +}; + +describe("BudgetModal", () => { + beforeEach(() => { + vi.clearAllMocks(); + createMock.mockResolvedValue(undefined); + }); + + it("submits only the mounted fields when Optional Settings stays collapsed", async () => { + const user = userEvent.setup(); + renderModal(); + + await user.type(screen.getByLabelText("Budget ID"), "budget-alpha"); + await user.type(screen.getByLabelText("Max Tokens per minute"), "500.567"); + await user.type(screen.getByLabelText("Max Requests per minute"), "7"); + await create(user); + + await waitFor(() => expect(createMock).toHaveBeenCalledTimes(1)); + expect(createMock.mock.calls[0][0]).toEqual({ + budget_id: "budget-alpha", + tpm_limit: 500.57, + rpm_limit: 7, + }); + }); + + it("submits every field once Optional Settings is expanded", async () => { + const user = userEvent.setup(); + renderModal(); + + await user.type(screen.getByLabelText("Budget ID"), "budget-alpha"); + await user.type(screen.getByLabelText("Max Tokens per minute"), "500.567"); + await user.type(screen.getByLabelText("Max Requests per minute"), "7"); + + await openOptionalSettings(user); + await user.type(screen.getByLabelText("Max Budget (USD)"), "42.567"); + + await user.click(screen.getByRole("combobox")); + await user.click(await screen.findByText("monthly")); + + await create(user); + + await waitFor(() => expect(createMock).toHaveBeenCalledTimes(1)); + expect(createMock.mock.calls[0][0]).toEqual(FULL_PAYLOAD); + }); + + it("drops Optional Settings values again when the section is collapsed before submit", async () => { + const user = userEvent.setup(); + renderModal(); + + await user.type(screen.getByLabelText("Budget ID"), "budget-alpha"); + + await openOptionalSettings(user); + await user.type(screen.getByLabelText("Max Budget (USD)"), "42.567"); + await user.click(screen.getByRole("combobox")); + await user.click(await screen.findByText("monthly")); + + await user.click(screen.getByText("Optional Settings")); + await waitFor(() => expect(screen.queryByLabelText("Max Budget (USD)")).not.toBeInTheDocument()); + await create(user); + + await waitFor(() => expect(createMock).toHaveBeenCalledTimes(1)); + expect(createMock.mock.calls[0][0]).toEqual({ budget_id: "budget-alpha" }); + }); + + it("submits a cleared number field as null", async () => { + const user = userEvent.setup(); + renderModal(); + + await user.type(screen.getByLabelText("Budget ID"), "budget-alpha"); + await user.type(screen.getByLabelText("Max Tokens per minute"), "5"); + await user.clear(screen.getByLabelText("Max Tokens per minute")); + await create(user); + + await waitFor(() => expect(createMock).toHaveBeenCalledTimes(1)); + expect(createMock.mock.calls[0][0]).toEqual({ + budget_id: "budget-alpha", + tpm_limit: null, + }); + }); + + it("blocks submit while Budget ID is empty", async () => { + const user = userEvent.setup(); + renderModal(); + + await user.type(screen.getByLabelText("Max Tokens per minute"), "5"); + await create(user); + + await waitFor(() => expect(screen.getByLabelText("Budget ID")).toHaveAttribute("aria-invalid", "true")); + expect(createMock).not.toHaveBeenCalled(); + }); + + it("keeps a typed Optional Setting when the section is collapsed and reopened, as antd's store did", async () => { + const user = userEvent.setup(); + renderModal(); + await user.type(screen.getByLabelText("Budget ID"), "probe-budget"); + + await openOptionalSettings(user); + await user.type(screen.getByLabelText("Max Budget (USD)"), "42.5"); + + await user.click(screen.getByText("Optional Settings")); + await user.click(screen.getByText("Optional Settings")); + + expect(await screen.findByLabelText("Max Budget (USD)")).toHaveValue(42.5); + + await create(user); + + await waitFor(() => expect(createMock).toHaveBeenCalledTimes(1)); + expect(createMock.mock.calls[0][0]).toMatchObject({ budget_id: "probe-budget", max_budget: 42.5 }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx index b4658aa9991..a0ca8bc7ae9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx @@ -1,33 +1,65 @@ +import { ChevronRight } from "lucide-react"; import React from "react"; -import { TextInput, Accordion, AccordionHeader, AccordionBody } from "@tremor/react"; -import { Button as Button2, Modal, Form, InputNumber, Select } from "antd"; +import { Modal } from "antd"; +import { z } from "zod/v4"; import { useCreateBudget } from "@/app/(dashboard)/hooks/budgets/useBudgets"; +import { applyBudgetPrecision } from "./budgetPrecision"; import NotificationsManager from "@/components/molecules/notifications_manager"; +import { FieldGroup } from "@/components/shared/form/field"; +import { FormField } from "@/components/shared/form/FormField"; +import { Button } from "@/components/ui/button"; +import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; +import { Input } from "@/components/ui/input"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { useZodForm } from "@/lib/forms/useZodForm"; + +const budgetShape = { + budget_id: z.string().min(1, "Please input a human-friendly name for the budget"), + tpm_limit: z.number().nullish(), + rpm_limit: z.number().nullish(), + max_budget: z.number().nullish(), + budget_duration: z.string().nullish(), +}; + +const budgetSchema = z.object(budgetShape); + +type BudgetFormValues = z.output; + +const BUDGET_DURATION_OPTIONS = [ + { value: "24h", label: "daily" }, + { value: "7d", label: "weekly" }, + { value: "30d", label: "monthly" }, +]; interface BudgetModalProps { isModalVisible: boolean; setIsModalVisible: React.Dispatch>; } const BudgetModal: React.FC = ({ isModalVisible, setIsModalVisible }) => { - const [form] = Form.useForm(); + const [optionalSettingsOpen, setOptionalSettingsOpen] = React.useState(false); + const form = useZodForm(budgetSchema, { defaultValues: { budget_id: "" } }); const createBudget = useCreateBudget(); const handleOk = () => { setIsModalVisible(false); - form.resetFields(); + form.reset(); }; const handleCancel = () => { setIsModalVisible(false); - form.resetFields(); + form.reset(); }; - const handleCreate = async (formValues: Record) => { + const handleCreate = async (formValues: BudgetFormValues) => { try { NotificationsManager.info("Making API Call"); - await createBudget.mutateAsync(formValues); + await createBudget.mutateAsync( + applyBudgetPrecision( + optionalSettingsOpen ? formValues : { ...formValues, max_budget: undefined, budget_duration: undefined }, + ), + ); NotificationsManager.success("Budget Created"); - form.resetFields(); + form.reset(); setIsModalVisible(false); } catch (error) { console.error("Error creating the budget:", error); @@ -44,51 +76,93 @@ const BudgetModal: React.FC = ({ isModalVisible, setIsModalVis onOk={handleOk} onCancel={handleCancel} > -
- <> - + + - - - - - - - - + {({ ref, ...field }) => } + + + {({ ref, value, onChange, ...field }) => ( + onChange(event.target.value === "" ? null : event.target.valueAsNumber)} + /> + )} + + + {({ ref, value, onChange, ...field }) => ( + onChange(event.target.value === "" ? null : event.target.valueAsNumber)} + /> + )} + - - + + Optional Settings - - - - - - - - - - - + + + + + {({ ref, value, onChange, ...field }) => ( + onChange(event.target.value === "" ? null : event.target.valueAsNumber)} + /> + )} + + + {({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => ( + + )} + + + +
- Create Budget +
- + ); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.integration.test.tsx new file mode 100644 index 00000000000..207f789e7f0 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.integration.test.tsx @@ -0,0 +1,112 @@ +import { render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import React from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import type { components } from "@/lib/http/schema"; + +import EditBudgetModal from "./edit_budget_modal"; + +const { updateMock } = vi.hoisted(() => ({ updateMock: vi.fn() })); + +vi.mock("@/app/(dashboard)/hooks/budgets/useBudgets", () => ({ + useUpdateBudget: () => ({ mutateAsync: updateMock }), +})); + +type BudgetItem = components["schemas"]["BudgetListItem"]; + +const EXISTING_BUDGET: BudgetItem = { + budget_id: "budget-alpha", + max_budget: 100, + budget_duration: "7d", + tpm_limit: 1000, + rpm_limit: 10, + soft_budget: 25, + budget_reset_at: "2026-02-01T00:00:00Z", + created_at: "2026-01-01T00:00:00Z", + updated_at: "2026-01-02T00:00:00Z", +}; + +const renderModal = () => + render(); + +const save = async (user: ReturnType) => + user.click(screen.getByRole("button", { name: "Save" })); + +const openOptionalSettings = async (user: ReturnType) => { + await user.click(screen.getByText("Optional Settings")); + await screen.findByLabelText("Max Budget (USD)"); +}; + +describe("EditBudgetModal", () => { + beforeEach(() => { + vi.clearAllMocks(); + updateMock.mockResolvedValue(undefined); + }); + + it("submits only the mounted fields when Optional Settings stays collapsed", async () => { + const user = userEvent.setup(); + renderModal(); + + await user.clear(screen.getByLabelText("Max Tokens per minute")); + await user.type(screen.getByLabelText("Max Tokens per minute"), "500.567"); + await save(user); + + await waitFor(() => expect(updateMock).toHaveBeenCalledTimes(1)); + expect(updateMock.mock.calls[0][0]).toEqual({ + budget_id: "budget-alpha", + tpm_limit: 500.57, + rpm_limit: 10, + }); + }); + + it("submits every field once Optional Settings is expanded", async () => { + const user = userEvent.setup(); + renderModal(); + + await user.clear(screen.getByLabelText("Max Tokens per minute")); + await user.type(screen.getByLabelText("Max Tokens per minute"), "500.567"); + await user.clear(screen.getByLabelText("Max Requests per minute")); + await user.type(screen.getByLabelText("Max Requests per minute"), "7"); + + await openOptionalSettings(user); + await user.clear(screen.getByLabelText("Max Budget (USD)")); + await user.type(screen.getByLabelText("Max Budget (USD)"), "42.567"); + + await user.click(screen.getByRole("combobox")); + await user.click(await screen.findByText("monthly")); + + await save(user); + + await waitFor(() => expect(updateMock).toHaveBeenCalledTimes(1)); + const expected = { + budget_id: "budget-alpha", + tpm_limit: 500.57, + rpm_limit: 7, + max_budget: 42.57, + budget_duration: "30d", + }; + + expect(updateMock.mock.calls[0][0]).toEqual(expected); + }); + + it("keeps a typed Optional Setting when the section is collapsed and reopened, as antd's store did", async () => { + const user = userEvent.setup(); + renderModal(); + + await openOptionalSettings(user); + const maxBudget = screen.getByLabelText("Max Budget (USD)"); + await user.clear(maxBudget); + await user.type(maxBudget, "99.25"); + + await user.click(screen.getByText("Optional Settings")); + await user.click(screen.getByText("Optional Settings")); + + expect(await screen.findByLabelText("Max Budget (USD)")).toHaveValue(99.25); + + await save(user); + + await waitFor(() => expect(updateMock).toHaveBeenCalledTimes(1)); + expect(updateMock.mock.calls[0][0]).toMatchObject({ max_budget: 99.25 }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx index 98a6996948f..1ed4e6b4f21 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx @@ -1,9 +1,36 @@ +import { ChevronRight } from "lucide-react"; import React, { useEffect } from "react"; -import { TextInput, Accordion, AccordionHeader, AccordionBody } from "@tremor/react"; -import { Button as Button2, Modal, Form, InputNumber, Select } from "antd"; +import { Modal } from "antd"; +import { useForm } from "react-hook-form"; import { useUpdateBudget } from "@/app/(dashboard)/hooks/budgets/useBudgets"; import { budgetItem } from "@/app/(dashboard)/hooks/budgets/useBudgets"; +import { applyBudgetPrecision } from "./budgetPrecision"; import NotificationsManager from "@/components/molecules/notifications_manager"; +import { FieldGroup } from "@/components/shared/form/field"; +import { FormField } from "@/components/shared/form/FormField"; +import { Button } from "@/components/ui/button"; +import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; +import { Input } from "@/components/ui/input"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; + +type EditBudgetFormValues = Pick< + budgetItem, + "budget_id" | "tpm_limit" | "rpm_limit" | "max_budget" | "budget_duration" +>; + +const toFormValues = (budget: budgetItem): EditBudgetFormValues => ({ + budget_id: budget.budget_id, + tpm_limit: budget.tpm_limit, + rpm_limit: budget.rpm_limit, + max_budget: budget.max_budget, + budget_duration: budget.budget_duration, +}); + +const BUDGET_DURATION_OPTIONS = [ + { value: "24h", label: "daily" }, + { value: "7d", label: "weekly" }, + { value: "30d", label: "monthly" }, +]; interface EditBudgetModalProps { isModalVisible: boolean; @@ -11,29 +38,34 @@ interface EditBudgetModalProps { existingBudget: budgetItem; } const EditBudgetModal: React.FC = ({ isModalVisible, setIsModalVisible, existingBudget }) => { - const [form] = Form.useForm(); + const [optionalSettingsOpen, setOptionalSettingsOpen] = React.useState(false); + const form = useForm({ defaultValues: toFormValues(existingBudget) }); const updateBudget = useUpdateBudget(); useEffect(() => { - form.setFieldsValue(existingBudget); + form.reset(toFormValues(existingBudget)); }, [existingBudget, form]); const handleOk = () => { setIsModalVisible(false); - form.resetFields(); + form.reset(); }; const handleCancel = () => { setIsModalVisible(false); - form.resetFields(); + form.reset(); }; - const handleUpdate = async (formValues: Record) => { + const handleUpdate = async (formValues: EditBudgetFormValues) => { try { NotificationsManager.info("Making API Call"); - await updateBudget.mutateAsync(formValues); + await updateBudget.mutateAsync( + applyBudgetPrecision( + optionalSettingsOpen ? formValues : { ...formValues, max_budget: undefined, budget_duration: undefined }, + ), + ); NotificationsManager.success("Budget Updated"); - form.resetFields(); + form.reset(); setIsModalVisible(false); } catch (error) { console.error("Error updating the budget:", error); @@ -43,48 +75,93 @@ const EditBudgetModal: React.FC = ({ isModalVisible, setIs return ( -
- <> - - - - - - - - - + + + + {({ ref, ...field }) => } + + + {({ ref, value, onChange, ...field }) => ( + onChange(event.target.value === "" ? null : event.target.valueAsNumber)} + /> + )} + + + {({ ref, value, onChange, ...field }) => ( + onChange(event.target.value === "" ? null : event.target.valueAsNumber)} + /> + )} + - - + + Optional Settings - - - - - - - - - - - + + + + + {({ ref, value, onChange, ...field }) => ( + onChange(event.target.value === "" ? null : event.target.valueAsNumber)} + /> + )} + + + {({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => ( + + )} + + + +
- Save +
-
+ ); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.tsx index 686f7130024..3be4a0fa085 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.tsx @@ -1,13 +1,29 @@ import React, { useState } from "react"; -import { Modal, Form, Input, Select } from "antd"; +import { Modal } from "antd"; +import { CircleHelp } from "lucide-react"; +import { z } from "zod/v4"; import MessageManager from "@/components/molecules/message_manager"; -import { Button } from "@tremor/react"; import { registerClaudeCodePlugin } from "@/components/networking"; +import { FieldGroup } from "@/components/shared/form/field"; +import { FormField } from "@/components/shared/form/FormField"; +import { Button } from "@/components/ui/button"; +import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; +import { + Combobox, + ComboboxContent, + ComboboxEmpty, + ComboboxInput, + ComboboxItem, + ComboboxList, +} from "@/components/ui/combobox"; +import { Input } from "@/components/ui/input"; +import { Textarea } from "@/components/ui/textarea"; +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; +import { useZodForm } from "@/lib/forms/useZodForm"; import { validatePluginName, isValidSemanticVersion, isValidEmail, - isValidUrl, parseKeywords, parseSkillSource, isValidSubPath, @@ -15,9 +31,6 @@ import { } from "@/components/claude_code_plugins/helpers"; import { PluginAuthor, PluginSource, SkillRegisterRequest } from "@/components/claude_code_plugins/types"; -const { TextArea } = Input; -const { Option } = Select; - interface AddPluginFormProps { visible: boolean; onClose: () => void; @@ -25,24 +38,51 @@ interface AddPluginFormProps { onSuccess: () => void; } -interface AddPluginFormValues { - name: string; - skillUrl?: string; - subPath?: string; - version?: string; - description?: string; - authorName?: string; - authorEmail?: string; - homepage?: string; - category?: string; - keywords?: string; - domain?: string; - namespace?: string; -} +const addPluginShape = { + skillUrl: z.string().min(1, "Please enter a repository URL"), + subPath: z + .string() + .refine( + (value) => !value || isValidSubPath(value), + "Subfolder must be a relative path like plugins/my-skill (letters, numbers, dots, hyphens, underscores)", + ), + name: z + .string() + .min(1, "Please enter skill name") + .regex(/^[a-z0-9-]+$/, "Name must be kebab-case (lowercase, numbers, hyphens only)"), + domain: z.string(), + namespace: z.string(), + description: z.string(), + category: z.string(), + keywords: z.string(), + version: z.string(), + authorName: z.string(), + authorEmail: z + .string() + .refine((value) => value === "" || z.email().safeParse(value).success, "Please enter a valid email"), +}; + +const addPluginSchema = z.object(addPluginShape); + +type AddPluginFormValues = z.infer; + +const EMPTY_VALUES: AddPluginFormValues = { + skillUrl: "", + subPath: "", + name: "", + domain: "", + namespace: "", + description: "", + category: "", + keywords: "", + version: "", + authorName: "", + authorEmail: "", +}; const buildAuthor = (values: AddPluginFormValues): PluginAuthor | undefined => { - const name = values.authorName?.trim(); - const email = values.authorEmail?.trim(); + const name = values.authorName.trim(); + const email = values.authorEmail.trim(); if (!name) { return undefined; } @@ -57,7 +97,6 @@ const buildRegisterRequest = (values: AddPluginFormValues, source: PluginSource) ...(values.version ? { version: values.version.trim() } : {}), ...(values.description ? { description: values.description.trim() } : {}), ...(author ? { author } : {}), - ...(values.homepage ? { homepage: values.homepage.trim() } : {}), ...(values.category ? { category: values.category } : {}), ...(values.keywords ? { keywords: parseKeywords(values.keywords) } : {}), ...(values.domain ? { domain: values.domain.trim() } : {}), @@ -76,8 +115,18 @@ const PREDEFINED_CATEGORIES = [ "Documentation", ]; +const labelWithHint = (label: string, hint: string): React.ReactNode => ( + <> + {label} + + } /> + {hint} + + +); + const AddPluginForm: React.FC = ({ visible, onClose, accessToken, onSuccess }) => { - const [form] = Form.useForm(); + const form = useZodForm(addPluginSchema, { defaultValues: EMPTY_VALUES }); const [isSubmitting, setIsSubmitting] = useState(false); const [urlPreview, setUrlPreview] = useState(null); const [urlEncodesSubdir, setUrlEncodesSubdir] = useState(false); @@ -85,24 +134,16 @@ const AddPluginForm: React.FC = ({ visible, onClose, accessT const recomputePreview = (skillUrl: string, subPath: string) => { const encodesSubdir = parseSkillSource(skillUrl)?.parsed.source === "git-subdir"; setUrlEncodesSubdir(encodesSubdir); - if (encodesSubdir && form.getFieldValue("subPath")) { - form.setFieldsValue({ subPath: "" }); + if (encodesSubdir && form.getValues("subPath")) { + form.setValue("subPath", ""); } const preview = parseSkillSource(skillUrl, encodesSubdir ? undefined : subPath); setUrlPreview(preview); - if (preview && !form.getFieldValue("name")) { - form.setFieldsValue({ name: preview.suggestedName }); + if (preview && !form.getValues("name")) { + form.setValue("name", preview.suggestedName); } }; - const handleUrlChange = (e: React.ChangeEvent) => { - recomputePreview(e.target.value, form.getFieldValue("subPath") ?? ""); - }; - - const handleSubPathChange = (e: React.ChangeEvent) => { - recomputePreview(form.getFieldValue("skillUrl") ?? "", e.target.value); - }; - const handleSubmit = async (values: AddPluginFormValues) => { if (!accessToken) { MessageManager.error("No access token available"); @@ -129,16 +170,11 @@ const AddPluginForm: React.FC = ({ visible, onClose, accessT return; } - if (values.homepage && !isValidUrl(values.homepage)) { - MessageManager.error("Invalid homepage URL format"); - return; - } - setIsSubmitting(true); try { await registerClaudeCodePlugin(accessToken, buildRegisterRequest(values, urlPreview.parsed)); MessageManager.success("Skill registered successfully"); - form.resetFields(); + form.reset(EMPTY_VALUES); setUrlPreview(null); setUrlEncodesSubdir(false); onSuccess(); @@ -152,7 +188,7 @@ const AddPluginForm: React.FC = ({ visible, onClose, accessT }; const handleCancel = () => { - form.resetFields(); + form.reset(EMPTY_VALUES); setUrlPreview(null); setUrlEncodesSubdir(false); onClose(); @@ -160,150 +196,192 @@ const AddPluginForm: React.FC = ({ visible, onClose, accessT return ( -
- {/* Smart URL Input */} - - - + + + + + {({ ref, onChange, ...field }) => ( + { + onChange(event); + recomputePreview(event.target.value, form.getValues("subPath")); + }} + /> + )} + - {/* Optional subfolder for monorepos */} - - !value || isValidSubPath(value) - ? Promise.resolve() - : Promise.reject( - new Error( - "Subfolder must be a relative path like plugins/my-skill (letters, numbers, dots, hyphens, underscores)", - ), - ), - }, - ]} - tooltip="Path within the repository where the skill lives (e.g., plugins/my-skill). Leave empty if the skill is at the repo root." - extra={urlEncodesSubdir ? "The URL already points to a subfolder, so this field is disabled" : undefined} - > - - + + {({ ref, onChange, ...field }) => ( + { + onChange(event); + recomputePreview(form.getValues("skillUrl"), event.target.value); + }} + disabled={urlEncodesSubdir} + /> + )} + - {/* Parsed preview */} - {urlPreview && ( -
- Detected: {urlPreview.label} -
- )} + {urlPreview && ( +
+ Detected: {urlPreview.label} +
+ )} - {/* Skill Name */} - - - + + {({ ref, ...field }) => } + - {/* Domain and Namespace — side by side */} -
- - - - - - -
+
+ + {({ ref, ...field }) => ( + + )} + + + {({ ref, ...field }) => } + +
- {/* Description */} - -